Skip to content

fix: rename RegisterModel inner steps to prevent duplicate step names #3240

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 14, 2022
Merged

fix: rename RegisterModel inner steps to prevent duplicate step names #3240

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions src/sagemaker/workflow/step_collections.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,9 @@ def properties(self):
class RegisterModel(StepCollection): # pragma: no cover
"""Register Model step collection for workflow."""

_REGISTER_MODEL_NAME_BASE = "RegisterModel"
_REPACK_MODEL_NAME_BASE = "RepackModel"

def __init__(
self,
name: str,
Expand Down Expand Up @@ -168,7 +171,7 @@ def __init__(
kwargs = dict(**kwargs, output_kms_key=kwargs.pop("model_kms_key", None))

repack_model_step = _RepackModelStep(
name=f"{name}RepackModel",
name="{}-{}".format(self.name, self._REPACK_MODEL_NAME_BASE),
depends_on=depends_on,
retry_policies=repack_model_step_retry_policies,
sagemaker_session=estimator.sagemaker_session,
Expand Down Expand Up @@ -212,7 +215,7 @@ def __init__(
model_name = model_entity.name or model_entity._framework_name

repack_model_step = _RepackModelStep(
name=f"{model_name}RepackModel",
name="{}-{}".format(model_name, self._REPACK_MODEL_NAME_BASE),
depends_on=depends_on,
retry_policies=repack_model_step_retry_policies,
sagemaker_session=sagemaker_session,
Expand Down Expand Up @@ -256,7 +259,7 @@ def __init__(
)

register_model_step = _RegisterModelStep(
name=name,
name="{}-{}".format(self.name, self._REGISTER_MODEL_NAME_BASE),
estimator=estimator,
model_data=model_data,
content_types=content_types,
Expand Down
8 changes: 4 additions & 4 deletions tests/integ/sagemaker/workflow/test_model_create_and_registration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,8 @@ def test_conditional_pytorch_training_model_registration(
model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
content_types=["*"],
response_types=["*"],
inference_instances=["*"],
transform_instances=["*"],
inference_instances=["ml.m5.xlarge"],
transform_instances=["ml.m5.xlarge"],
description="test-description",
sample_payload_url=sample_payload_url,
task=task,
Expand Down Expand Up @@ -234,7 +234,7 @@ def test_mxnet_model_registration(
content_types=["*"],
response_types=["*"],
inference_instances=["ml.m5.xlarge"],
transform_instances=["*"],
transform_instances=["ml.m5.xlarge"],
description="test-description",
sample_payload_url=sample_payload_url,
task=task,
Expand Down Expand Up @@ -670,7 +670,7 @@ def test_model_registration_with_drift_check_baselines(
)
continue
assert execution_steps[0]["StepStatus"] == "Succeeded"
assert execution_steps[0]["StepName"] == "MyRegisterModelStep"
assert execution_steps[0]["StepName"] == "MyRegisterModelStep-RegisterModel"

response = sagemaker_session.sagemaker_client.describe_model_package(
ModelPackageName=execution_steps[0]["Metadata"]["RegisterModel"]["Arn"]
Expand Down
43 changes: 27 additions & 16 deletions tests/unit/sagemaker/workflow/test_step_collections.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
_REPACK_MODEL_NAME_BASE,
)
from sagemaker.workflow.parameters import ParameterString
from sagemaker.workflow.pipeline import Pipeline
from sagemaker.workflow.pipeline import Pipeline, PipelineGraph
from sagemaker.workflow.pipeline_context import PipelineSession
from sagemaker.workflow.utilities import list_to_request
from tests.unit import DATA_DIR
Expand Down Expand Up @@ -268,7 +268,7 @@ def test_step_collection_properties(pipeline_session, sagemaker_session):
steps = register_model.steps
assert len(steps) == 1
assert register_model.properties.ModelPackageName.expr == {
"Get": f"Steps.{register_model_step_name}.ModelPackageName"
"Get": f"Steps.{register_model_step_name}-RegisterModel.ModelPackageName"
}

# Custom StepCollection
Expand Down Expand Up @@ -330,10 +330,9 @@ def test_step_collection_is_depended_on(pipeline_session, sagemaker_session):
step_list = json.loads(pipeline.definition())["Steps"]
assert len(step_list) == 7
for step in step_list:
if step["Name"] not in ["MyStep2", "MyStep3", f"{model_name}RepackModel"]:
if step["Name"] not in ["MyStep2", "MyStep3", f"{model_name}-RepackModel"]:
assert "DependsOn" not in step
continue
if step["Name"] == f"{model_name}RepackModel":
elif step["Name"] == f"{model_name}-RepackModel":
assert set(step["DependsOn"]) == {
"MyStep1",
f"{model_step_name}-{_REPACK_MODEL_NAME_BASE}-{model_name}",
Expand All @@ -344,9 +343,21 @@ def test_step_collection_is_depended_on(pipeline_session, sagemaker_session):
"MyStep1",
f"{model_step_name}-{_REPACK_MODEL_NAME_BASE}-{model_name}",
f"{model_step_name}-{_CREATE_MODEL_NAME_BASE}",
f"{model_name}RepackModel",
register_model_name,
f"{model_name}-RepackModel",
f"{register_model_name}-RegisterModel",
}
adjacency_list = PipelineGraph.from_pipeline(pipeline).adjacency_list
assert ordered(adjacency_list) == ordered(
{
"MyStep1": ["MyStep2", "MyStep3", "MyModel-RepackModel"],
"MyStep2": [],
"MyStep3": [],
"MyModelStep-RepackModel-MyModel": ["MyModelStep-CreateModel"],
"MyModelStep-CreateModel": ["MyStep2", "MyStep3", "MyModel-RepackModel"],
"MyModel-RepackModel": [],
"RegisterModelStep-RegisterModel": ["MyStep2", "MyStep3"],
}
)


def test_register_model(estimator, model_metrics, drift_check_baselines):
Expand Down Expand Up @@ -378,7 +389,7 @@ def test_register_model(estimator, model_metrics, drift_check_baselines):
assert ordered(register_model.request_dicts()) == ordered(
[
{
"Name": "RegisterModelStep",
"Name": "RegisterModelStep-RegisterModel",
"Type": "RegisterModel",
"DependsOn": ["TestStep"],
"DisplayName": "RegisterModelStep",
Expand Down Expand Up @@ -450,7 +461,7 @@ def test_register_model_tf(estimator_tf, model_metrics, drift_check_baselines):
assert ordered(register_model.request_dicts()) == ordered(
[
{
"Name": "RegisterModelStep",
"Name": "RegisterModelStep-RegisterModel",
"Type": "RegisterModel",
"Description": "description",
"Arguments": {
Expand Down Expand Up @@ -526,7 +537,7 @@ def test_register_model_sip(estimator, model_metrics, drift_check_baselines):
assert ordered(register_model.request_dicts()) == ordered(
[
{
"Name": "RegisterModelStep",
"Name": "RegisterModelStep-RegisterModel",
"Type": "RegisterModel",
"Description": "description",
"DependsOn": ["TestStep"],
Expand Down Expand Up @@ -618,7 +629,7 @@ def test_register_model_with_model_repack_with_estimator(

for request_dict in request_dicts:
if request_dict["Type"] == "Training":
assert request_dict["Name"] == "RegisterModelStepRepackModel"
assert request_dict["Name"] == "RegisterModelStep-RepackModel"
assert len(request_dict["DependsOn"]) == 1
assert request_dict["DependsOn"][0] == "TestStep"
arguments = request_dict["Arguments"]
Expand Down Expand Up @@ -671,7 +682,7 @@ def test_register_model_with_model_repack_with_estimator(
}
)
elif request_dict["Type"] == "RegisterModel":
assert request_dict["Name"] == "RegisterModelStep"
assert request_dict["Name"] == "RegisterModelStep-RegisterModel"
assert "DependsOn" not in request_dict
arguments = request_dict["Arguments"]
assert len(arguments["InferenceSpecification"]["Containers"]) == 1
Expand Down Expand Up @@ -745,7 +756,7 @@ def test_register_model_with_model_repack_with_model(model, model_metrics, drift

for request_dict in request_dicts:
if request_dict["Type"] == "Training":
assert request_dict["Name"] == "modelNameRepackModel"
assert request_dict["Name"] == "modelName-RepackModel"
assert len(request_dict["DependsOn"]) == 1
assert request_dict["DependsOn"][0] == "TestStep"
arguments = request_dict["Arguments"]
Expand Down Expand Up @@ -798,7 +809,7 @@ def test_register_model_with_model_repack_with_model(model, model_metrics, drift
}
)
elif request_dict["Type"] == "RegisterModel":
assert request_dict["Name"] == "RegisterModelStep"
assert request_dict["Name"] == "RegisterModelStep-RegisterModel"
assert "DependsOn" not in request_dict
arguments = request_dict["Arguments"]
assert len(arguments["InferenceSpecification"]["Containers"]) == 1
Expand Down Expand Up @@ -874,7 +885,7 @@ def test_register_model_with_model_repack_with_pipeline_model(

for request_dict in request_dicts:
if request_dict["Type"] == "Training":
assert request_dict["Name"] == "modelNameRepackModel"
assert request_dict["Name"] == "modelName-RepackModel"
assert len(request_dict["DependsOn"]) == 1
assert request_dict["DependsOn"][0] == "TestStep"
arguments = request_dict["Arguments"]
Expand Down Expand Up @@ -927,7 +938,7 @@ def test_register_model_with_model_repack_with_pipeline_model(
}
)
elif request_dict["Type"] == "RegisterModel":
assert request_dict["Name"] == "RegisterModelStep"
assert request_dict["Name"] == "RegisterModelStep-RegisterModel"
assert "DependsOn" not in request_dict
arguments = request_dict["Arguments"]
assert len(arguments["InferenceSpecification"]["Containers"]) == 1
Expand Down