From fc3565e9db9d6ba2d741d98c3699472c51440d22 Mon Sep 17 00:00:00 2001
From: Taleb Zeghmi <talebz@zillowgroup.com>
Date: Thu, 7 Dec 2023 15:50:30 -0800
Subject: [PATCH] AIP_JOIN_METAFLOW_S3OP_NUM_WORKERS in aip_constants.py

---
 metaflow/plugins/aip/aip_constants.py     | 2 ++
 metaflow/plugins/aip/aip_metaflow_step.py | 5 ++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/metaflow/plugins/aip/aip_constants.py b/metaflow/plugins/aip/aip_constants.py
index 49e578b0db0..be0f3c4fd10 100644
--- 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
index ea49a61fd33..96ef88ac44f 100644
--- 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,