Specify task of training/tuning job in Airflow transform/deploy operator (#590)

Author: yangaws

CHANGELOG.rst

@@ -2,10 +2,15 @@
 CHANGELOG
 =========
 
+1.17.1.dev
+==========
+
+* enhancement: Workflow: Specify tasks from which training/tuning operator to transform/deploy in related operators
+
 1.17.0
 ======
 
-* bug-fix: Revert appending Airflow retry id to default job name
+* bug-fix: Workflow: Revert appending Airflow retry id to default job name
 * bug-fix: Session: don't allow get_execution_role() to return an ARN that's not a role but has "role" in the name
 * bug-fix: Remove ``__all__`` from ``__init__.py`` files
 * doc-fix: Add TFRecord split type to docs
@@ -22,7 +27,7 @@ CHANGELOG
 ======
 
 * bug-fix: Local Mode: Allow support for SSH in local mode
-* bug-fix: Append retry id to default Airflow job name to avoid name collisions in retry
+* bug-fix: Workflow: Append retry id to default Airflow job name to avoid name collisions in retry
 * bug-fix: Local Mode: No longer requires s3 permissions to run local entry point file
 * feature: Estimators: add support for PyTorch 1.0.0
 * bug-fix: Local Mode: Move dependency on sagemaker_s3_output from rl.estimator to model

src/sagemaker/utils.py

@@ -27,12 +27,6 @@
 import six
 
 
-AIRFLOW_TIME_MACRO = "{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}"
-AIRFLOW_TIME_MACRO_LEN = 19
-AIRFLOW_TIME_MACRO_SHORT = "{{ execution_date.strftime('%y%m%d-%H%M') }}"
-AIRFLOW_TIME_MACRO_SHORT_LEN = 11
-
-
 # Use the base name of the image as the job name if the user doesn't give us one
 def name_from_image(image):
     """Create a training job name based on the image name and a timestamp.
@@ -73,25 +67,6 @@ def unique_name_from_base(base, max_length=63):
     return '{}-{}-{}'.format(trimmed, ts, unique)
 
 
-def airflow_name_from_base(base, max_length=63, short=False):
-    """Append airflow execution_date macro (https://airflow.apache.org/code.html?#macros)
-    to the provided string. The macro will beevaluated in Airflow operator runtime.
-    This guarantees that different operators will have same name returned by this function.
-
-    Args:
-        base (str): String used as prefix to generate the unique name.
-        max_length (int): Maximum length for the resulting string.
-        short (bool): Whether or not to use a truncated timestamp.
-
-    Returns:
-        str: Input parameter with appended macro.
-    """
-    macro = AIRFLOW_TIME_MACRO_SHORT if short else AIRFLOW_TIME_MACRO
-    length = AIRFLOW_TIME_MACRO_SHORT_LEN if short else AIRFLOW_TIME_MACRO_LEN
-    trimmed_base = base[:max_length - length - 1]
-    return "{}-{}".format(trimmed_base, macro)
-
-
 def base_name_from_image(image):
     """Extract the base name of the image to use as the 'algorithm name' for the job.
 

src/sagemaker/workflow/README.rst

@@ -57,7 +57,11 @@ dictionary that can be generated by the SageMaker Python SDK. For example:
                                    inputs=your_training_data_s3_uri)
 
     # trans_config specifies SageMaker batch transform configuration
+    # task_id specifies which operator the training job associatd with; task_type specifies whether the operator is a
+    # training operator or tuning operator
     trans_config = transform_config_from_estimator(estimator=estimator,
+                                                   task_id='tf_training',
+                                                   task_type='training',
                                                    instance_count=1,
                                                    instance_type='ml.m4.xlarge',
                                                    data=your_transform_data_s3_uri,
@@ -82,13 +86,13 @@ Now you can pass these configurations to the corresponding SageMaker operators a
               schedule_interval='@once')
 
     train_op = SageMakerTrainingOperator(
-        task_id='training',
+        task_id='tf_training',
         config=train_config,
         wait_for_completion=True,
         dag=dag)
 
     transform_op = SageMakerTransformOperator(
-        task_id='transform',
+        task_id='tf_transform',
         config=trans_config,
         wait_for_completion=True,
         dag=dag)