Feature: Adding placeholder output_path parameter to TrainingStep to …

…propagate to Estimator

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