refactor nproc_per_node for backwards compatibility

brunopistone · brunopistone · commit 00eb637eb93a · 2025-01-15T21:57:59.000Z
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(

Original file line number	Diff line number	Diff line change
`@@ -1081,8 +1081,7 @@ def _prepare_and_upload_runtime_scripts(`
`1081`	`1081`
`1082`	`1082`	`if nproc_per_node is not None and nproc_per_node > 0:`
`1083`	`1083`	`entry_point_script = entry_point_script.replace(`
`1084`		`- "$SM_NPROC_PER_NODE",`
`1085`		`- str(nproc_per_node)`
	`1084`	`+ "$SM_NPROC_PER_NODE", str(nproc_per_node)`
`1086`	`1085`	`)`
`1087`	`1086`
`1088`	`1087`	`with open(entrypoint_script_path, "w", newline="\n") as file:`
Original file line number	Diff line number	Diff line change
`@@ -908,6 +908,7 @@ def test_remote_decorator_fields_consistency(get_execution_role, session):`
`908`	`908`	`"max_wait_time_in_seconds",`
`909`	`909`	`"custom_file_filter",`
`910`	`910`	`"use_torchrun",`
	`911`	`+ "nproc_per_node",`
`911`	`912`	`}`
`912`	`913`
`913`	`914`	`job_settings = _JobSettings(`