refactor nproc_per_node for backwards compatibility

aws · Jan 15, 2025 · 00eb637 · 00eb637
1 parent adcc38e
commit 00eb637
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/sagemaker/remote_function/job.py b/src/sagemaker/remote_function/job.py
@@ -1081,8 +1081,7 @@ def _prepare_and_upload_runtime_scripts(
 
             if nproc_per_node is not None and nproc_per_node > 0:
                 entry_point_script = entry_point_script.replace(
-                    "$SM_NPROC_PER_NODE",
-                    str(nproc_per_node)
+                    "$SM_NPROC_PER_NODE", str(nproc_per_node)
                 )
 
         with open(entrypoint_script_path, "w", newline="\n") as file:

diff --git a/tests/unit/sagemaker/feature_store/feature_processor/test_feature_scheduler.py b/tests/unit/sagemaker/feature_store/feature_processor/test_feature_scheduler.py
@@ -908,6 +908,7 @@ def test_remote_decorator_fields_consistency(get_execution_role, session):
         "max_wait_time_in_seconds",
         "custom_file_filter",
         "use_torchrun",
+        "nproc_per_node",
     }
 
     job_settings = _JobSettings(