Update env param directly and add test

Author: ca-nguyen

Diff for: src/stepfunctions/steps/sagemaker.py

@@ -286,10 +286,20 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No
             instance_type (str, optional): The EC2 instance type to deploy this Model to. For example, 'ml.p2.xlarge'.
             tags (list[dict], optional): `List to tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
         """
-        if isinstance(model, Model):
+        if isinstance(model, FrameworkModel):
             parameters = model_config(model=model, instance_type=instance_type, role=model.role, image_uri=model.image_uri)
             if model_name:
                 parameters['ModelName'] = model_name
+        elif isinstance(model, Model):
+            parameters = {
+                'ExecutionRoleArn': model.role,
+                'ModelName': model_name or model.name,
+                'PrimaryContainer': {
+                    'Environment': model.env,
+                    'Image': model.image_uri,
+                    'ModelDataUrl': model.model_data
+                }
+            }
         else:
             raise ValueError("Expected 'model' parameter to be of type 'sagemaker.model.Model', but received type '{}'".format(type(model).__name__))
 

Diff for: tests/unit/test_sagemaker_steps.py

@@ -81,7 +81,11 @@ def pca_estimator_with_env():
         environment={
             'JobName': "job_name",
             'ModelName': "model_name"
-        }
+        },
+        subnets=[
+            'subnet-00000000000000000',
+            'subnet-00000000000000001'
+        ]
     )
 
     pca.set_hyperparameters(
@@ -187,6 +191,31 @@ def pca_model():
     )
 
 
+@pytest.fixture
+def pca_model_with_env():
+    model_data = 's3://sagemaker/models/pca.tar.gz'
+    return Model(
+        model_data=model_data,
+        image_uri=PCA_IMAGE,
+        role=EXECUTION_ROLE,
+        name='pca-model',
+        env={
+            'JobName': "job_name",
+            'ModelName': "model_name"
+        },
+        vpc_config={
+            "SecurityGroupIds": ["sg-00000000000000000"],
+            "Subnets": ["subnet-00000000000000000", "subnet-00000000000000001"]
+        },
+        image_config={
+            "RepositoryAccessMode": "Vpc",
+            "RepositoryAuthConfig": {
+                "RepositoryCredentialsProviderArn": "arn"
+            }
+        }
+    )
+
+
 @pytest.fixture
 def pca_transformer(pca_model):
     return Transformer(
@@ -855,6 +884,31 @@ def test_get_expected_model(pca_estimator):
     }
 
 
+@patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call)
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_get_expected_model_with_env(pca_estimator_with_env):
+    training_step = TrainingStep('Training', estimator=pca_estimator_with_env, job_name='TrainingJob')
+    expected_model = training_step.get_expected_model()
+    model_step = ModelStep('Create model', model=expected_model, model_name='pca-model')
+    assert model_step.to_dict() == {
+        'Type': 'Task',
+        'Parameters': {
+            'ExecutionRoleArn': EXECUTION_ROLE,
+            'ModelName': 'pca-model',
+            'PrimaryContainer': {
+                'Environment': {
+                    'JobName': 'job_name',
+                    'ModelName': 'model_name'
+                },
+                'Image': expected_model.image_uri,
+                'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']"
+            }
+        },
+        'Resource': 'arn:aws:states:::sagemaker:createModel',
+        'End': True
+    }
+
+
 @patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call)
 @patch.object(boto3.session.Session, 'region_name', 'us-east-1')
 def test_get_expected_model_with_framework_estimator(tensorflow_estimator):
@@ -908,6 +962,29 @@ def test_model_step_creation(pca_model):
     }
 
 
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_model_step_creation_with_env(pca_model_with_env):
+    step = ModelStep('Create model', model=pca_model_with_env, model_name='pca-model', tags=DEFAULT_TAGS)
+    assert step.to_dict() == {
+        'Type': 'Task',
+        'Parameters': {
+            'ExecutionRoleArn': EXECUTION_ROLE,
+            'ModelName': 'pca-model',
+            'PrimaryContainer': {
+                'Environment': {
+                    'JobName': 'job_name',
+                    'ModelName': 'model_name'
+                },
+                'Image': pca_model_with_env.image_uri,
+                'ModelDataUrl': pca_model_with_env.model_data
+            },
+            'Tags': DEFAULT_TAGS_LIST
+        },
+        'Resource': 'arn:aws:states:::sagemaker:createModel',
+        'End': True
+    }
+
+
 @patch.object(boto3.session.Session, 'region_name', 'us-east-1')
 def test_endpoint_config_step_creation(pca_model):
     data_capture_config = DataCaptureConfig(