fix: change CreateModel to Model in step type (#536)

Author: metrizable

src/sagemaker/workflow/steps.py

@@ -49,7 +49,7 @@ class StepTypeEnum(Enum, metaclass=DefaultEnumMeta):
     """Enum of step types."""
 
     CONDITION = "Condition"
-    CREATE_MODEL = "CreateModel"
+    CREATE_MODEL = "Model"
     FAIL = "Fail"
     PROCESSING = "Processing"
     REGISTER_MODEL = "RegisterModel"

tests/integ/test_workflow.py

@@ -22,7 +22,8 @@
 
 from botocore.config import Config
 from botocore.exceptions import WaiterError
-from sagemaker.inputs import TrainingInput
+from sagemaker.inputs import CreateModelInput, TrainingInput
+from sagemaker.model import Model
 from sagemaker.processing import ProcessingInput, ProcessingOutput
 from sagemaker.pytorch.estimator import PyTorch
 from sagemaker.session import get_execution_role, Session
@@ -35,6 +36,7 @@
     ParameterString,
 )
 from sagemaker.workflow.steps import (
+    CreateModelStep,
     ProcessingStep,
     TrainingStep,
 )
@@ -95,7 +97,7 @@ def pipeline_name():
     return f"my-pipeline-{int(time.time() * 10**7)}"
 
 
-def test_two_step_definition(
+def test_three_step_definition(
     sagemaker_session, workflow_session, region_name, role, script_dir, pipeline_name
 ):
     framework_version = "0.20.0"
@@ -140,10 +142,26 @@ def test_two_step_definition(
         ),
     )
 
+    model = Model(
+        image_uri=sklearn_train.image_uri,
+        model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
+        sagemaker_session=sagemaker_session,
+        role=role,
+    )
+    model_inputs = CreateModelInput(
+        instance_type="ml.m5.large",
+        accelerator_type="ml.eia1.medium",
+    )
+    step_model = CreateModelStep(
+        name="my-model",
+        model=model,
+        inputs=model_inputs,
+    )
+
     pipeline = Pipeline(
         name=pipeline_name,
         parameters=[instance_type, instance_count],
-        steps=[step_process, step_train],
+        steps=[step_process, step_train, step_model],
         sagemaker_session=workflow_session,
     )
 
@@ -160,7 +178,7 @@ def test_two_step_definition(
     )
 
     steps = definition["Steps"]
-    assert len(steps) == 2
+    assert len(steps) == 3
 
     names_and_types = []
     processing_args = {}
@@ -171,11 +189,14 @@ def test_two_step_definition(
             processing_args = step["Arguments"]
         if step["Type"] == "Training":
             training_args = step["Arguments"]
+        if step["Type"] == "Model":
+            model_args = step["Arguments"]
 
     assert set(names_and_types) == set(
         [
             ("my-process", "Processing"),
             ("my-train", "Training"),
+            ("my-model", "Model"),
         ]
     )
 
@@ -193,6 +214,9 @@ def test_two_step_definition(
     assert training_args["InputDataConfig"][0]["DataSource"]["S3DataSource"]["S3Uri"] == {
         "Get": "Steps.my-process.ProcessingOutputConfig.Outputs['train_data'].S3Output.S3Uri"
     }
+    assert model_args["PrimaryContainer"]["ModelDataUrl"] == {
+        "Get": "Steps.my-train.ModelArtifacts.S3ModelArtifacts"
+    }
 
 
 # TODO-reinvent-2020: Modify use of the workflow client
@@ -324,11 +348,27 @@ def test_conditional_pytorch_training_model_registration(
         transform_instances=["*"],
     )
 
+    model = Model(
+        image_uri=pytorch_estimator.training_image_uri(),
+        model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
+        sagemaker_session=sagemaker_session,
+        role=role,
+    )
+    model_inputs = CreateModelInput(
+        instance_type="ml.m5.large",
+        accelerator_type="ml.eia1.medium",
+    )
+    step_model = CreateModelStep(
+        name="pytorch-model",
+        model=model,
+        inputs=model_inputs,
+    )
+
     step_cond = ConditionStep(
         name="cond-good-enough",
         conditions=[ConditionGreaterThanOrEqualTo(left=good_enough_input, right=1)],
         if_steps=[step_train, step_register],
-        else_steps=[],
+        else_steps=[step_model],
     )
 
     pipeline = Pipeline(

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -194,10 +194,10 @@ def test_estimator_transformer(estimator):
     request_dicts = estimator_transformer.request_dicts()
     assert len(request_dicts) == 2
     for request_dict in request_dicts:
-        if request_dict["Type"] == "CreateModel":
+        if request_dict["Type"] == "Model":
             assert request_dict == {
                 "Name": "EstimatorTransformerStepCreateModelStep",
-                "Type": "CreateModel",
+                "Type": "Model",
                 "Arguments": {
                     "ExecutionRoleArn": "DummyRole",
                     "PrimaryContainer": {

tests/unit/sagemaker/workflow/test_steps.py

@@ -220,7 +220,7 @@ def test_create_model_step(sagemaker_session):
 
     assert step.to_request() == {
         "Name": "MyCreateModelStep",
-        "Type": "CreateModel",
+        "Type": "Model",
         "Arguments": {
             "ExecutionRoleArn": "DummyRole",
             "PrimaryContainer": {"Environment": {}, "Image": "fakeimage"},