fix: environment variables specified in Model or Estimator are not passed through to SageMaker ModelStep (#160)

1. fix: environment variables are overwritten and not passed through to SageMaker ModelStep

The environment variables are not passed from the model to ModelStep parameters (they are overwritten here) for models without instance type (models other than sagemake.model.FrameworkModel)
With this change, the Model env parameters are passed to the ModelStep.

2. fix: env variables defined in the Estimator are not translated to Model when calling TrainingStep.get_expected_model()

Fixed by passing the env from the estimator to the expected Model
Added a new estimator (pca_estimator_with_env) with defined env parameters to use in test and confirm that the env variables are passed through to the Model With this change, the env variables defined in the estimator are passed to the expected Model when calling TrainingStep.get_expected_model()

Author: ca-nguyen

src/stepfunctions/steps/sagemaker.py

@@ -159,6 +159,8 @@ def get_expected_model(self, model_name=None):
             model.name = model_name
         else:
             model.name = self.job_name
+        if self.estimator.environment:
+            model.env = self.estimator.environment
         model.model_data = self.output()["ModelArtifacts"]["S3ModelArtifacts"]
         return model
 
@@ -295,7 +297,7 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No
                 'ExecutionRoleArn': model.role,
                 'ModelName': model_name or model.name,
                 'PrimaryContainer': {
-                    'Environment': {},
+                    'Environment': model.env,
                     'Image': model.image_uri,
                     'ModelDataUrl': model.model_data
                 }

tests/unit/test_sagemaker_steps.py

@@ -68,6 +68,41 @@ def pca_estimator():
     return pca
 
 
+@pytest.fixture
+def pca_estimator_with_env():
+    s3_output_location = 's3://sagemaker/models'
+
+    pca = sagemaker.estimator.Estimator(
+        PCA_IMAGE,
+        role=EXECUTION_ROLE,
+        instance_count=1,
+        instance_type='ml.c4.xlarge',
+        output_path=s3_output_location,
+        environment={
+            'JobName': "job_name",
+            'ModelName': "model_name"
+        },
+        subnets=[
+            'subnet-00000000000000000',
+            'subnet-00000000000000001'
+        ]
+    )
+
+    pca.set_hyperparameters(
+        feature_dim=50000,
+        num_components=10,
+        subtract_mean=True,
+        algorithm_mode='randomized',
+        mini_batch_size=200
+    )
+
+    pca.sagemaker_session = MagicMock()
+    pca.sagemaker_session.boto_region_name = 'us-east-1'
+    pca.sagemaker_session._default_bucket = 'sagemaker'
+
+    return pca
+
+
 @pytest.fixture
 def pca_estimator_with_debug_hook():
     s3_output_location = 's3://sagemaker/models'
@@ -156,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(
@@ -537,6 +597,63 @@ def test_training_step_creation_with_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_training_step_creation_with_model_with_env(pca_estimator_with_env):
+    training_step = TrainingStep('Training', estimator=pca_estimator_with_env, job_name='TrainingJob')
+    model_step = ModelStep('Training - Save Model', training_step.get_expected_model(model_name=training_step.output()['TrainingJobName']))
+    training_step.next(model_step)
+    assert training_step.to_dict() == {
+        'Type': 'Task',
+        'Parameters': {
+            'AlgorithmSpecification': {
+                'TrainingImage': PCA_IMAGE,
+                'TrainingInputMode': 'File'
+            },
+            'OutputDataConfig': {
+                'S3OutputPath': 's3://sagemaker/models'
+            },
+            'StoppingCondition': {
+                'MaxRuntimeInSeconds': 86400
+            },
+            'ResourceConfig': {
+                'InstanceCount': 1,
+                'InstanceType': 'ml.c4.xlarge',
+                'VolumeSizeInGB': 30
+            },
+            'RoleArn': EXECUTION_ROLE,
+            'HyperParameters': {
+                'feature_dim': '50000',
+                'num_components': '10',
+                'subtract_mean': 'True',
+                'algorithm_mode': 'randomized',
+                'mini_batch_size': '200'
+            },
+            'TrainingJobName': 'TrainingJob'
+        },
+        'Resource': 'arn:aws:states:::sagemaker:createTrainingJob.sync',
+        'Next': 'Training - Save Model'
+    }
+
+    assert model_step.to_dict() == {
+        'Type': 'Task',
+        'Resource': 'arn:aws:states:::sagemaker:createModel',
+        'Parameters': {
+            'ExecutionRoleArn': EXECUTION_ROLE,
+            'ModelName.$': "$['TrainingJobName']",
+            'PrimaryContainer': {
+                'Environment': {
+                    'JobName': 'job_name',
+                    'ModelName': 'model_name'
+                },
+                'Image': PCA_IMAGE,
+                'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']"
+            }
+        },
+        '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_training_step_creation_with_framework(tensorflow_estimator):
@@ -806,6 +923,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):
@@ -859,6 +1001,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(