AIP_JOIN_METAFLOW_S3OP_NUM_WORKERS in aip_constants.py

zillow · Dec 7, 2023 · fc3565e · fc3565e
1 parent fe315c9
commit fc3565e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/metaflow/plugins/aip/aip_constants.py b/metaflow/plugins/aip/aip_constants.py
@@ -28,3 +28,5 @@
 STDERR_PATH = os.path.join(LOGS_DIR, STDERR_FILE)
 
 AIP_CLI_DEFAULT_RETRY = 3
+
+AIP_JOIN_METAFLOW_S3OP_NUM_WORKERS = 5
diff --git a/metaflow/plugins/aip/aip_metaflow_step.py b/metaflow/plugins/aip/aip_metaflow_step.py
@@ -22,6 +22,7 @@
     STDOUT_PATH,
     STEP_ENVIRONMENT_VARIABLES,
     TASK_ID_ENV_NAME,
+    AIP_JOIN_METAFLOW_S3OP_NUM_WORKERS,
 )
 
 from ... import R
@@ -377,7 +378,9 @@ def aip_metaflow_step(
 
     if is_join_step and "METAFLOW_S3OP_NUM_WORKERS" not in os.environ:
         # AIP-7487: Metaflow joins steps require lots of memory
-        os.environ["METAFLOW_S3OP_NUM_WORKERS"] = "5"
+        os.environ["METAFLOW_S3OP_NUM_WORKERS"] = str(
+            AIP_JOIN_METAFLOW_S3OP_NUM_WORKERS
+        )
 
     env: Dict[str, str] = {
         **os.environ,
Original file line number	Diff line number	Diff line change
Expand Up		@@ -28,3 +28,5 @@
		STDERR_PATH = os.path.join(LOGS_DIR, STDERR_FILE)

		AIP_CLI_DEFAULT_RETRY = 3

		AIP_JOIN_METAFLOW_S3OP_NUM_WORKERS = 5