feature: adding support for ExperimentConfig in training and transform steps (#23)

Author: brightsparc

src/stepfunctions/steps/sagemaker.py

@@ -27,7 +27,7 @@ class TrainingStep(Task):
     Creates a Task State to execute a `SageMaker Training Job <https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateTrainingJob.html>`_. The TrainingStep will also create a model by default, and the model shares the same name as the training job.
     """
 
-    def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=None, mini_batch_size=None, wait_for_completion=True, **kwargs):
+    def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=None, mini_batch_size=None, experiment_config=None, wait_for_completion=True, **kwargs):
         """
         Args:
             state_id (str): State name whose length **must be** less than or equal to 128 unicode characters. State names **must be** unique within the scope of the whole state machine.
@@ -50,6 +50,7 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
                     where each instance is a different channel of training data.
             hyperparameters (dict, optional): Specify the hyper parameters for the training. (Default: None)
             mini_batch_size (int): Specify this argument only when estimator is a built-in estimator of an Amazon algorithm. For other estimators, batch size should be specified in the estimator.
+            experiment_config (dict, optional): Specify the experiment config for the training. (Default: None)
             wait_for_completion (bool, optional): Boolean value set to `True` if the Task state should wait for the training job to complete before proceeding to the next step in the workflow. Set to `False` if the Task state should submit the training job and proceed to the next step. (default: True)
         """
         self.estimator = estimator
@@ -71,6 +72,9 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
         if hyperparameters is not None:
             parameters['HyperParameters'] = hyperparameters
 
+        if experiment_config is not None:
+            parameters['ExperimentConfig'] = experiment_config
+
         if 'S3Operations' in parameters:
             del parameters['S3Operations']
 
@@ -101,7 +105,7 @@ class TransformStep(Task):
     Creates a Task State to execute a `SageMaker Transform Job <https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateTransformJob.html>`_.
     """
 
-    def __init__(self, state_id, transformer, job_name, model_name, data, data_type='S3Prefix', content_type=None, compression_type=None, split_type=None, wait_for_completion=True, **kwargs):
+    def __init__(self, state_id, transformer, job_name, model_name, data, data_type='S3Prefix', content_type=None, compression_type=None, split_type=None, experiment_config=None, wait_for_completion=True, **kwargs):
         """
         Args:
             state_id (str): State name whose length **must be** less than or equal to 128 unicode characters. State names **must be** unique within the scope of the whole state machine.
@@ -119,6 +123,7 @@ def __init__(self, state_id, transformer, job_name, model_name, data, data_type=
             content_type (str): MIME type of the input data (default: None).
             compression_type (str): Compression type of the input data, if compressed (default: None). Valid values: 'Gzip', None.
             split_type (str): The record delimiter for the input object (default: 'None'). Valid values: 'None', 'Line', 'RecordIO', and 'TFRecord'.
+            experiment_config (dict, optional): Specify the experiment config for the transform. (Default: None)
             wait_for_completion(bool, optional): Boolean value set to `True` if the Task state should wait for the transform job to complete before proceeding to the next step in the workflow. Set to `False` if the Task state should submit the transform job and proceed to the next step. (default: True)
         """
         if wait_for_completion:
@@ -151,6 +156,9 @@ def __init__(self, state_id, transformer, job_name, model_name, data, data_type=
 
         parameters['ModelName'] = model_name
 
+        if experiment_config is not None:
+            parameters['ExperimentConfig'] = experiment_config
+
         kwargs[Field.Parameters.value] = parameters
         super(TransformStep, self).__init__(state_id, **kwargs)
 

tests/unit/test_sagemaker_steps.py

@@ -98,7 +98,14 @@ def tensorflow_estimator():
 
 @patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call)
 def test_training_step_creation(pca_estimator):
-    step = TrainingStep('Training', estimator=pca_estimator, job_name='TrainingJob')
+    step = TrainingStep('Training', 
+        estimator=pca_estimator, 
+        job_name='TrainingJob', 
+        experiment_config={
+            'ExperimentName': 'pca_experiment',
+            'TrialName': 'pca_trial',
+            'TrialComponentDisplayName': 'Training'
+        })
     assert step.to_dict() == {
         'Type': 'Task',
         'Parameters': {
@@ -125,6 +132,11 @@ def test_training_step_creation(pca_estimator):
                 'algorithm_mode': 'randomized',
                 'mini_batch_size': '200'
             },
+            'ExperimentConfig': {
+                'ExperimentName': 'pca_experiment',
+                'TrialName': 'pca_trial',
+                'TrialComponentDisplayName': 'Training'                
+            },
             'TrainingJobName': 'TrainingJob'
         },
         'Resource': 'arn:aws:states:::sagemaker:createTrainingJob.sync',
@@ -243,7 +255,12 @@ def test_transform_step_creation(pca_transformer):
         transformer=pca_transformer,
         data='s3://sagemaker/inference',
         job_name='transform-job',
-        model_name='pca-model'
+        model_name='pca-model',
+        experiment_config={
+            'ExperimentName': 'pca_experiment',
+            'TrialName': 'pca_trial',
+            'TrialComponentDisplayName': 'Transform'
+        }
     )
     assert step.to_dict() == {
         'Type': 'Task',
@@ -264,6 +281,11 @@ def test_transform_step_creation(pca_transformer):
             'TransformResources': {
                 'InstanceCount': 1,
                 'InstanceType': 'ml.c4.xlarge'
+            },
+            'ExperimentConfig': {
+                'ExperimentName': 'pca_experiment',
+                'TrialName': 'pca_trial',
+                'TrialComponentDisplayName': 'Transform'                
             }
         },
         'Resource': 'arn:aws:states:::sagemaker:createTransformJob.sync',