Skip to content

Commit 5d1daa5

Browse files
author
svia3
committed
fix: add warning message for job-prefixed pipeline steps when no job name is provided
1 parent d083396 commit 5d1daa5

File tree

2 files changed

+38
-2
lines changed

2 files changed

+38
-2
lines changed

src/sagemaker/workflow/utilities.py

+10
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,14 @@
4141
"if desired."
4242
)
4343

44+
JOB_KEY_NONE_WARN_MSG_TEMPLATE = (
45+
"Invalid input: use_custom_job_prefix flag is set but the name field [{}] has not been "
46+
"specified. Please refer to the AWS Docs to identify which field should be set to enable the "
47+
"custom-prefixing feature for jobs created via a pipeline execution. "
48+
"https://docs.aws.amazon.com/sagemaker/latest/dg/"
49+
"build-and-manage-access.html#build-and-manage-step-permissions-prefix"
50+
)
51+
4452
if TYPE_CHECKING:
4553
from sagemaker.workflow.step_collections import StepCollection
4654

@@ -458,6 +466,8 @@ def trim_request_dict(request_dict, job_key, config):
458466
request_dict.pop(job_key, None) # safely return null in case of KeyError
459467
else:
460468
if job_key in request_dict:
469+
if request_dict[job_key] is None or len(request_dict[job_key]) == 0:
470+
raise ValueError(JOB_KEY_NONE_WARN_MSG_TEMPLATE.format(job_key))
461471
request_dict[job_key] = base_from_name(request_dict[job_key]) # trim timestamp
462472

463473
return request_dict

tests/unit/sagemaker/workflow/test_model_step.py

+28-2
Original file line numberDiff line numberDiff line change
@@ -1096,7 +1096,6 @@ def test_pass_in_wrong_type_of_retry_policies(pipeline_session, model):
10961096

10971097

10981098
def test_register_model_step_using_custom_model_package_name(pipeline_session):
1099-
11001099
custom_model_prefix = "custom-model-package-prefix"
11011100
model = Model(
11021101
name="MyModel",
@@ -1139,7 +1138,6 @@ def test_register_model_step_using_custom_model_package_name(pipeline_session):
11391138

11401139

11411140
def test_create_model_step_using_custom_model_name(pipeline_session):
1142-
11431141
custom_model_prefix = "custom-model-prefix"
11441142
model = Model(
11451143
name=custom_model_prefix,
@@ -1171,3 +1169,31 @@ def test_create_model_step_using_custom_model_name(pipeline_session):
11711169
steps = json.loads(pipeline.definition())["Steps"]
11721170
assert len(steps) == 1
11731171
assert "ModelName" not in steps[0]["Arguments"]
1172+
1173+
1174+
def test_create_model_step_using_custom_model_name_set_to_none(pipeline_session):
1175+
# Name of the model not specified, will resolve to None.
1176+
model = Model(
1177+
image_uri="my-image",
1178+
sagemaker_session=pipeline_session,
1179+
model_data="s3://",
1180+
role=ROLE,
1181+
)
1182+
step_create_model = ModelStep(name="MyModelStep", step_args=model.create())
1183+
1184+
# 1. Toggle on custom-prefixing model package name popped
1185+
config = PipelineDefinitionConfig(use_custom_job_prefix=True)
1186+
1187+
with pytest.raises(ValueError) as error:
1188+
pipeline = Pipeline(
1189+
name="MyPipeline",
1190+
steps=[step_create_model],
1191+
sagemaker_session=pipeline_session,
1192+
pipeline_definition_config=config,
1193+
)
1194+
pipeline.definition()
1195+
1196+
assert (
1197+
"Invalid input: use_custom_job_prefix flag is set but the name field "
1198+
"[ModelName] has not been specified." in str(error.value)
1199+
)

0 commit comments

Comments
 (0)