Add airflow module that could export training config (#480)

yangaws · web-flow · commit 4153e4124d3f · 2018-11-14T20:18:23.000-08:00
diff --git a/src/sagemaker/job.py b/src/sagemaker/job.py
@@ -51,9 +51,9 @@ def wait(self):
         pass
 
     @staticmethod
-    def _load_config(inputs, estimator):
-        input_config = _Job._format_inputs_to_input_config(inputs)
-        role = estimator.sagemaker_session.expand_role(estimator.role)
+    def _load_config(inputs, estimator, expand_role=True, validate_uri=True):
+        input_config = _Job._format_inputs_to_input_config(inputs, validate_uri)
+        role = estimator.sagemaker_session.expand_role(estimator.role) if expand_role else estimator.role
         output_config = _Job._prepare_output_config(estimator.output_path, estimator.output_kms_key)
         resource_config = _Job._prepare_resource_config(estimator.train_instance_count,
                                                         estimator.train_instance_type,
@@ -62,7 +62,8 @@ def _load_config(inputs, estimator):
         stop_condition = _Job._prepare_stop_condition(estimator.train_max_run)
         vpc_config = estimator.get_vpc_config()
 
-        model_channel = _Job._prepare_model_channel(input_config, estimator.model_uri, estimator.model_channel_name)
+        model_channel = _Job._prepare_model_channel(input_config, estimator.model_uri, estimator.model_channel_name,
+                                                    validate_uri)
         if model_channel:
             input_config = [] if input_config is None else input_config
             input_config.append(model_channel)
@@ -75,7 +76,7 @@ def _load_config(inputs, estimator):
                 'vpc_config': vpc_config}
 
     @staticmethod
-    def _format_inputs_to_input_config(inputs):
+    def _format_inputs_to_input_config(inputs, validate_uri=True):
         if inputs is None:
             return None
 
@@ -86,14 +87,14 @@ def _format_inputs_to_input_config(inputs):
 
         input_dict = {}
         if isinstance(inputs, string_types):
-            input_dict['training'] = _Job._format_string_uri_input(inputs)
+            input_dict['training'] = _Job._format_string_uri_input(inputs, validate_uri)
         elif isinstance(inputs, s3_input):
             input_dict['training'] = inputs
         elif isinstance(inputs, file_input):
             input_dict['training'] = inputs
         elif isinstance(inputs, dict):
             for k, v in inputs.items():
-                input_dict[k] = _Job._format_string_uri_input(v)
+                input_dict[k] = _Job._format_string_uri_input(v, validate_uri)
         elif isinstance(inputs, list):
             input_dict = _Job._format_record_set_list_input(inputs)
         else:
@@ -111,15 +112,16 @@ def _convert_input_to_channel(channel_name, channel_s3_input):
         return channel_config
 
     @staticmethod
-    def _format_string_uri_input(uri_input):
-        if isinstance(uri_input, str):
-            if uri_input.startswith('s3://'):
-                return s3_input(uri_input)
-            elif uri_input.startswith('file://'):
-                return file_input(uri_input)
-            else:
-                raise ValueError('Training input data must be a valid S3 or FILE URI: must start with "s3://" or '
-                                 '"file://"')
+    def _format_string_uri_input(uri_input, validate_uri=True):
+        if isinstance(uri_input, str) and validate_uri and uri_input.startswith('s3://'):
+            return s3_input(uri_input)
+        elif isinstance(uri_input, str) and validate_uri and uri_input.startswith('file://'):
+            return file_input(uri_input)
+        elif isinstance(uri_input, str) and validate_uri:
+            raise ValueError('Training input data must be a valid S3 or FILE URI: must start with "s3://" or '
+                             '"file://"')
+        elif isinstance(uri_input, str):
+            return s3_input(uri_input)
         elif isinstance(uri_input, s3_input):
             return uri_input
         elif isinstance(uri_input, file_input):
@@ -128,7 +130,7 @@ def _format_string_uri_input(uri_input):
             raise ValueError('Cannot format input {}. Expecting one of str, s3_input, or file_input'.format(uri_input))
 
     @staticmethod
-    def _prepare_model_channel(input_config, model_uri=None, model_channel_name=None):
+    def _prepare_model_channel(input_config, model_uri=None, model_channel_name=None, validate_uri=True):
         if not model_uri:
             return
         elif not model_channel_name:
@@ -139,22 +141,24 @@ def _prepare_model_channel(input_config, model_uri=None, model_channel_name=None
                 if channel['ChannelName'] == model_channel_name:
                     raise ValueError('Duplicate channels not allowed.')
 
-        model_input = _Job._format_model_uri_input(model_uri)
+        model_input = _Job._format_model_uri_input(model_uri, validate_uri)
         model_channel = _Job._convert_input_to_channel(model_channel_name, model_input)
 
         return model_channel
 
     @staticmethod
-    def _format_model_uri_input(model_uri):
-        if isinstance(model_uri, string_types):
-            if model_uri.startswith('s3://'):
-                return s3_input(model_uri, input_mode='File', distribution='FullyReplicated',
-                                content_type='application/x-sagemaker-model')
-            elif model_uri.startswith('file://'):
-                return file_input(model_uri)
-            else:
-                raise ValueError('Model URI must be a valid S3 or FILE URI: must start with "s3://" or '
-                                 '"file://')
+    def _format_model_uri_input(model_uri, validate_uri=True):
+        if isinstance(model_uri, string_types)and validate_uri and model_uri.startswith('s3://'):
+            return s3_input(model_uri, input_mode='File', distribution='FullyReplicated',
+                            content_type='application/x-sagemaker-model')
+        elif isinstance(model_uri, string_types) and validate_uri and model_uri.startswith('file://'):
+            return file_input(model_uri)
+        elif isinstance(model_uri, string_types) and validate_uri:
+            raise ValueError('Model URI must be a valid S3 or FILE URI: must start with "s3://" or '
+                             '"file://')
+        elif isinstance(model_uri, string_types):
+            return s3_input(model_uri, input_mode='File', distribution='FullyReplicated',
+                            content_type='application/x-sagemaker-model')
         else:
             raise ValueError('Cannot format model URI {}. Expecting str'.format(model_uri))
 
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
@@ -169,11 +169,11 @@ def default_bucket(self):
         if self._default_bucket:
             return self._default_bucket
 
-        s3 = self.boto_session.resource('s3')
         account = self.boto_session.client('sts').get_caller_identity()['Account']
         region = self.boto_session.region_name
         default_bucket = 'sagemaker-{}-{}'.format(region, account)
 
+        s3 = self.boto_session.resource('s3')
         try:
             # 'us-east-1' cannot be specified because it is the default region:
             # https://github.com/boto/boto3/issues/125
diff --git a/src/sagemaker/utils.py b/src/sagemaker/utils.py
@@ -26,6 +26,9 @@
 import six
 
 
+AIRFLOW_TIME_MACRO = "{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}"
+
+
 # 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.
@@ -58,6 +61,21 @@ def name_from_base(base, max_length=63, short=False):
     return '{}-{}'.format(trimmed_base, timestamp)
 
 
+def airflow_name_from_base(base):
+    """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.
+
+    Returns:
+        str: Input parameter with appended macro.
+    """
+
+    return "{}-{}".format(base, AIRFLOW_TIME_MACRO)
+
+
 def base_name_from_image(image):
     """Extract the base name of the image to use as the 'algorithm name' for the job.
 
diff --git a/src/sagemaker/workflow/__init__.py b/src/sagemaker/workflow/__init__.py
@@ -0,0 +1,12 @@
+# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
diff --git a/src/sagemaker/workflow/airflow.py b/src/sagemaker/workflow/airflow.py
@@ -0,0 +1,136 @@
+# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import print_function, absolute_import
+
+import os
+
+import sagemaker
+from sagemaker import job, utils, model
+from sagemaker.amazon import amazon_estimator
+
+
+def prepare_framework(estimator, s3_operations):
+    """Prepare S3 operations (specify where to upload source_dir) and environment variables
+    related to framework.
+
+    Args:
+        estimator (sagemaker.estimator.Estimator): The framework estimator to get information from and update.
+        s3_operations (dict): The dict to specify s3 operations (upload source_dir).
+    """
+    bucket = estimator.code_location if estimator.code_location else estimator.sagemaker_session._default_bucket
+    key = '{}/source/sourcedir.tar.gz'.format(estimator._current_job_name)
+    script = os.path.basename(estimator.entry_point)
+    if estimator.source_dir and estimator.source_dir.lower().startswith('s3://'):
+        code_dir = estimator.source_dir
+    else:
+        code_dir = 's3://{}/{}'.format(bucket, key)
+        s3_operations['S3Upload'] = [{
+            'Path': estimator.source_dir or script,
+            'Bucket': bucket,
+            'Key': key,
+            'Tar': True
+        }]
+    estimator._hyperparameters[model.DIR_PARAM_NAME] = code_dir
+    estimator._hyperparameters[model.SCRIPT_PARAM_NAME] = script
+    estimator._hyperparameters[model.CLOUDWATCH_METRICS_PARAM_NAME] = estimator.enable_cloudwatch_metrics
+    estimator._hyperparameters[model.CONTAINER_LOG_LEVEL_PARAM_NAME] = estimator.container_log_level
+    estimator._hyperparameters[model.JOB_NAME_PARAM_NAME] = estimator._current_job_name
+    estimator._hyperparameters[model.SAGEMAKER_REGION_PARAM_NAME] = estimator.sagemaker_session.boto_region_name
+
+
+def prepare_amazon_algorithm_estimator(estimator, inputs):
+    """ Set up amazon algorithm estimator, adding the required `feature_dim` hyperparameter from training data.
+
+    Args:
+        estimator (sagemaker.amazon.amazon_estimator.AmazonAlgorithmEstimatorBase):
+            An estimator for a built-in Amazon algorithm to get information from and update.
+        inputs (single or list of sagemaker.amazon.amazon_estimator.RecordSet):
+            The training data, must be in RecordSet format.
+    """
+    if isinstance(inputs, list):
+        for record in inputs:
+            if isinstance(record, amazon_estimator.RecordSet) and record.channel == 'train':
+                estimator.feature_dim = record.feature_dim
+                break
+    elif isinstance(inputs, amazon_estimator.RecordSet):
+        estimator.feature_dim = inputs.feature_dim
+    else:
+        raise TypeError('Training data must be represented in RecordSet or list of RecordSets')
+
+
+def training_config(estimator, inputs=None, job_name=None):  # noqa: C901 - suppress complexity warning for this method
+    """Export Airflow training config from an estimator
+
+    Args:
+        estimator (sagemaker.estimator.EstimatroBase):
+            The estimator to export training config from. Can be a BYO estimator,
+            Framework estimator or Amazon algorithm estimator.
+        inputs (str, dict, single or list of sagemaker.amazon.amazon_estimator.RecordSet):
+            The training data.
+        job_name (str): Specify a training job name if needed.
+
+    Returns:
+        A dict of training config that can be directly used by SageMakerTrainingOperator
+            in Airflow.
+    """
+    default_bucket = estimator.sagemaker_session.default_bucket()
+    s3_operations = {}
+
+    if job_name is not None:
+        estimator._current_job_name = job_name
+    else:
+        base_name = estimator.base_job_name or utils.base_name_from_image(estimator.train_image())
+        estimator._current_job_name = utils.airflow_name_from_base(base_name)
+
+    if estimator.output_path is None:
+        estimator.output_path = 's3://{}/'.format(default_bucket)
+
+    if isinstance(estimator, sagemaker.estimator.Framework):
+        prepare_framework(estimator, s3_operations)
+
+    elif isinstance(estimator, amazon_estimator.AmazonAlgorithmEstimatorBase):
+        prepare_amazon_algorithm_estimator(estimator, inputs)
+
+    job_config = job._Job._load_config(inputs, estimator, expand_role=False, validate_uri=False)
+
+    train_config = {
+        'AlgorithmSpecification': {
+            'TrainingImage': estimator.train_image(),
+            'TrainingInputMode': estimator.input_mode
+        },
+        'OutputDataConfig': job_config['output_config'],
+        'TrainingJobName': estimator._current_job_name,
+        'StoppingCondition': job_config['stop_condition'],
+        'ResourceConfig': job_config['resource_config'],
+        'RoleArn': job_config['role'],
+    }
+
+    if job_config['input_config'] is not None:
+        train_config['InputDataConfig'] = job_config['input_config']
+
+    if job_config['vpc_config'] is not None:
+        train_config['VpcConfig'] = job_config['vpc_config']
+
+    if estimator.hyperparameters() is not None:
+        hyperparameters = {str(k): str(v) for (k, v) in estimator.hyperparameters().items()}
+
+    if hyperparameters and len(hyperparameters) > 0:
+        train_config['HyperParameters'] = hyperparameters
+
+    if estimator.tags is not None:
+        train_config['Tags'] = estimator.tags
+
+    if s3_operations:
+        train_config['S3Operations'] = s3_operations
+
+    return train_config
diff --git a/tests/unit/test_airflow.py b/tests/unit/test_airflow.py
