From 00eb637eb93ac3486cb9cbf3dd90e967b339a407 Mon Sep 17 00:00:00 2001
From: brunopistone <brn.pistone@gmail.com>
Date: Wed, 15 Jan 2025 21:57:59 +0000
Subject: [PATCH] refactor nproc_per_node for backwards compatibility

---
 src/sagemaker/remote_function/job.py                           | 3 +--
 .../feature_store/feature_processor/test_feature_scheduler.py  | 1 +
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/sagemaker/remote_function/job.py b/src/sagemaker/remote_function/job.py
index 86574c0bdf..8b4df8e2e4 100644
--- 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
index 35dea8b204..57f4a54f78 100644
--- 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(