Skip to content

Commit 1a9bfa4

Browse files
rauldiazRaul Diaz Garciaaviruthenmujtaba1747
authored
feature: allow custom runproc.sh in FrameworkProcessor (#5742)
Co-authored-by: Raul Diaz Garcia <raul.diaz2.contractor@bbva.com> Co-authored-by: aviruthen <91846056+aviruthen@users.noreply.github.com> Co-authored-by: Syed Mujtaba <42322958+mujtaba1747@users.noreply.github.com>
1 parent 2f6be9c commit 1a9bfa4

2 files changed

Lines changed: 118 additions & 5 deletions

File tree

‎sagemaker-core/src/sagemaker/core/processing.py‎

Lines changed: 53 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1197,6 +1197,7 @@ def run(
11971197
job_name: Optional[str] = None,
11981198
experiment_config: Optional[Dict[str, str]] = None,
11991199
kms_key: Optional[str] = None,
1200+
entry_point: Optional[str] = None
12001201
):
12011202
"""Runs a processing job.
12021203
@@ -1224,6 +1225,9 @@ def run(
12241225
experiment_config (dict[str, str]): Experiment management configuration.
12251226
kms_key (str): The ARN of the KMS key that is used to encrypt the
12261227
user code file (default: None).
1228+
entry_point (str): Path (absolute or relative) to a custom entrypoint script file
1229+
(e.g., runproc.sh). The python script call is appended automatically.
1230+
12271231
Returns:
12281232
None or pipeline step arguments in case the Processor instance is built with
12291233
:class:`~sagemaker.workflow.pipeline_context.PipelineSession`
@@ -1235,6 +1239,7 @@ def run(
12351239
job_name,
12361240
inputs,
12371241
kms_key,
1242+
entry_point
12381243
)
12391244

12401245
# Submit a processing job.
@@ -1258,6 +1263,7 @@ def _pack_and_upload_code(
12581263
job_name,
12591264
inputs,
12601265
kms_key=None,
1266+
entry_point=None
12611267
):
12621268
"""Pack local code bundle and upload to Amazon S3."""
12631269
if code.startswith("s3://"):
@@ -1294,7 +1300,7 @@ def _pack_and_upload_code(
12941300
script = os.path.basename(code)
12951301
evaluated_kms_key = kms_key if kms_key else self.output_kms_key
12961302
s3_runproc_sh = self._create_and_upload_runproc(
1297-
script, evaluated_kms_key, entrypoint_s3_uri
1303+
script, evaluated_kms_key, entrypoint_s3_uri, entry_point, source_dir
12981304
)
12991305

13001306
return s3_runproc_sh, inputs, job_name
@@ -1332,12 +1338,12 @@ def _set_entrypoint(self, command, user_script_name):
13321338
)
13331339
self.entrypoint = self.framework_entrypoint_command + [user_script_location]
13341340

1335-
def _create_and_upload_runproc(self, user_script, kms_key, entrypoint_s3_uri):
1341+
def _create_and_upload_runproc(self, user_script, kms_key, entrypoint_s3_uri, entry_point=None, source_dir=None):
13361342
"""Create runproc shell script and upload to S3 bucket."""
13371343
from sagemaker.core.workflow.utilities import _pipeline_config, hash_object
13381344

13391345
if _pipeline_config and _pipeline_config.pipeline_name:
1340-
runproc_file_str = self._generate_framework_script(user_script)
1346+
runproc_file_str = self._generate_framework_script(user_script, entry_point, source_dir)
13411347
runproc_file_hash = hash_object(runproc_file_str)
13421348
s3_uri = s3.s3_path_join(
13431349
self._s3_code_prefix(),
@@ -1354,16 +1360,19 @@ def _create_and_upload_runproc(self, user_script, kms_key, entrypoint_s3_uri):
13541360
)
13551361
else:
13561362
s3_runproc_sh = s3.S3Uploader.upload_string_as_file_body(
1357-
self._generate_framework_script(user_script),
1363+
self._generate_framework_script(user_script, entry_point, source_dir),
13581364
desired_s3_uri=entrypoint_s3_uri,
13591365
kms_key=kms_key,
13601366
sagemaker_session=self.sagemaker_session,
13611367
)
13621368

13631369
return s3_runproc_sh
13641370

1365-
def _generate_framework_script(self, user_script: str) -> str:
1371+
def _generate_framework_script(self, user_script: str, entry_point: str = None, source_dir: str = None) -> str:
13661372
"""Generate the framework entrypoint file (as text) for a processing job."""
1373+
if entry_point:
1374+
return self._generate_custom_framework_script(user_script, entry_point, source_dir)
1375+
13671376
return dedent(
13681377
"""\
13691378
#!/bin/bash
@@ -1401,6 +1410,45 @@ def _generate_framework_script(self, user_script: str) -> str:
14011410
entry_point=user_script,
14021411
)
14031412

1413+
def _generate_custom_framework_script(
1414+
self, user_script: str, entry_point: str, source_dir: str = None
1415+
) -> str:
1416+
"""
1417+
Generate a custom framework script with a user-provided entrypoint embedded.
1418+
1419+
Reads the entry_point file and embeds its content in the script,
1420+
then appends the command to execute the user script.
1421+
1422+
Args:
1423+
user_script (str): Relative path to the user script in the source bundle
1424+
entry_point (str): Path to the custom entrypoint script file
1425+
source_dir (str): Path to the source directory. If provided and entry_point
1426+
is relative, it will be combined with source_dir.
1427+
1428+
Returns:
1429+
str: The generated script content
1430+
"""
1431+
# Resolve the full path to the entry_point file
1432+
if source_dir and not os.path.isabs(entry_point):
1433+
full_entry_point_path = os.path.join(source_dir, entry_point)
1434+
else:
1435+
full_entry_point_path = entry_point
1436+
1437+
# Read the entry_point file content
1438+
with open(full_entry_point_path, "r", encoding="utf-8") as f:
1439+
entry_point_content = f.read()
1440+
1441+
# Generate the script with embedded entry_point content
1442+
return dedent("""\
1443+
{entry_point_content}
1444+
1445+
{entry_point_command} {entry_point} "$@"
1446+
""").format(
1447+
entry_point_content=entry_point_content,
1448+
entry_point_command=" ".join(self.command),
1449+
entry_point=user_script,
1450+
)
1451+
14041452

14051453
class FeatureStoreOutput(ApiObject):
14061454
"""Configuration for processing job outputs in Amazon SageMaker Feature Store."""

‎sagemaker-core/tests/unit/test_processing.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -865,6 +865,71 @@ def test_create_and_upload_runproc_without_pipeline(self, mock_session):
865865
)
866866
assert result == "s3://bucket/runproc.sh"
867867

868+
def test_generate_framework_script_with_custom_entry_point(self, mock_session):
869+
processor = FrameworkProcessor(
870+
role="arn:aws:iam::123456789012:role/SageMakerRole",
871+
image_uri="test-image:latest",
872+
command=["python3"],
873+
instance_count=1,
874+
instance_type="ml.m5.xlarge",
875+
sagemaker_session=mock_session,
876+
)
877+
878+
custom_script_content = "#!/bin/bash\necho 'THIS IS THE CUSTOM runproc.sh'\nset -e\n"
879+
880+
with tempfile.NamedTemporaryFile(mode="w", suffix=".sh", delete=False) as f:
881+
f.write(custom_script_content)
882+
entry_point_path = f.name
883+
884+
try:
885+
script = processor._generate_framework_script(
886+
"train.py", entry_point=entry_point_path
887+
)
888+
assert custom_script_content in script
889+
assert "python3 train.py" in script
890+
assert "tar -xzf sourcedir.tar.gz" not in script
891+
finally:
892+
os.unlink(entry_point_path)
893+
894+
def test_generate_framework_script_with_custom_entry_point_and_source_dir(self, mock_session):
895+
processor = FrameworkProcessor(
896+
role="arn:aws:iam::123456789012:role/SageMakerRole",
897+
image_uri="test-image:latest",
898+
command=["python3"],
899+
instance_count=1,
900+
instance_type="ml.m5.xlarge",
901+
sagemaker_session=mock_session,
902+
)
903+
904+
with tempfile.TemporaryDirectory() as tmpdir:
905+
custom_script_content = "#!/bin/bash\necho 'custom from source_dir'\n"
906+
script_path = os.path.join(tmpdir, "custom_runproc.sh")
907+
with open(script_path, "w") as f:
908+
f.write(custom_script_content)
909+
910+
script = processor._generate_framework_script(
911+
"train.py",
912+
entry_point="custom_runproc.sh",
913+
source_dir=tmpdir,
914+
)
915+
assert custom_script_content in script
916+
assert "python3 train.py" in script
917+
918+
def test_generate_framework_script_with_default_entry_point(self, mock_session):
919+
processor = FrameworkProcessor(
920+
role="arn:aws:iam::123456789012:role/SageMakerRole",
921+
image_uri="test-image:latest",
922+
command=["python3"],
923+
instance_count=1,
924+
instance_type="ml.m5.xlarge",
925+
sagemaker_session=mock_session,
926+
)
927+
928+
script = processor._generate_framework_script("train.py")
929+
assert "#!/bin/bash" in script
930+
assert "tar -xzf sourcedir.tar.gz" in script
931+
assert "python3 train.py" in script
932+
868933

869934
class TestHelperFunctions:
870935
def test_processing_input_to_request_dict(self):

0 commit comments

Comments
 (0)