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(