Merge pull request #276 from zillow/tz/AIP-7487-join-memory

AIP-7487 join memory fix w/ METAFLOW_S3OP_NUM_WORKERS=5

Author: talebzeghmi

metaflow/datatools/s3op.py

@@ -35,7 +35,7 @@
 from metaflow.multicore_utils import parallel_map
 from metaflow.datatools.s3util import aws_retry, read_in_chunks, get_timestamp
 
-NUM_WORKERS_DEFAULT = 64
+NUM_WORKERS_DEFAULT = int(os.environ.get("METAFLOW_S3OP_NUM_WORKERS", 64))
 
 DOWNLOAD_FILE_THRESHOLD = 2 * TransferConfig().multipart_threshold
 DOWNLOAD_MAX_CHUNK = 2 * 1024 * 1024 * 1024 - 1

metaflow/plugins/aip/aip.py

@@ -1488,6 +1488,8 @@ def _create_metaflow_step_op(
             metaflow_execution_cmd += f" --namespace {flow_variables.namespace}"
         if step_variables.is_split_index:
             metaflow_execution_cmd += " --is_split_index"
+        if node.type == "join":
+            metaflow_execution_cmd += " --is-join-step"
 
         metaflow_execution_cmd += ' --preceding_component_outputs_dict "'
         for key in preceding_component_outputs_dict:

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

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
@@ -282,6 +283,7 @@ def _command(
 @click.option("--user_code_retries", type=int)
 @click.option("--workflow_name")
 @click.option("--is-interruptible/--not-interruptible", default=False)
+@click.option("--is-join-step", is_flag=True, default=False)
 def aip_metaflow_step(
     volume_dir: str,
     environment: str,
@@ -306,6 +308,7 @@ def aip_metaflow_step(
     user_code_retries: int,
     workflow_name: str,
     is_interruptible: bool,
+    is_join_step: bool,
 ) -> None:
     """
     (1) Renders and runs the Metaflow package_commands and Metaflow step
@@ -373,6 +376,12 @@ def aip_metaflow_step(
     ):
         metaflow_configs_new["METAFLOW_USER"] = "aip-user"
 
+    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"] = str(
+            AIP_JOIN_METAFLOW_S3OP_NUM_WORKERS
+        )
+
     env: Dict[str, str] = {
         **os.environ,
         **metaflow_configs_new,