feat: Adds support for Placeholders in TrainingStep to set S3 location for InputDataConfig and OutputDataConfig

src/stepfunctions/steps/sagemaker.py

@@ -43,7 +43,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, experiment_config=None, wait_for_completion=True, tags=None, **kwargs):
+    def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=None, mini_batch_size=None, experiment_config=None, wait_for_completion=True, tags=None, output_path=None, **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.
@@ -69,6 +69,13 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
             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)
             tags (list[dict], optional): `List to tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
+            output_path (Placeholder, optional): S3 location for saving the training result (model
+                artifacts and output files) to propagate to estimator. If not specified, results are
+                stored to a default bucket. If the bucket with the specific name
+                does not exist, the estimator creates the bucket during the
+                :meth:`~sagemaker.estimator.EstimatorBase.fit` method execution.
+                file:// urls are used for local mode. For example: 'file://model/'
+                will save to the model folder in the current directory.
         """
         self.estimator = estimator
         self.job_name = job_name
@@ -103,6 +110,10 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
         if isinstance(job_name, (ExecutionInput, StepInput)):
             parameters['TrainingJobName'] = job_name
 
+        if output_path is not None:
+            if isinstance(output_path, (ExecutionInput, StepInput)):
+                parameters['OutputDataConfig']['S3OutputPath'] = output_path
+
         if hyperparameters is not None:
             parameters['HyperParameters'] = hyperparameters