Refactor training job class (aws#10)

Author: ChoiByungWook

src/sagemaker/estimator.py

@@ -17,15 +17,14 @@
 import os
 from abc import ABCMeta
 from abc import abstractmethod
-from six import with_metaclass, string_types
+from six import with_metaclass
 
 from sagemaker.fw_utils import tar_and_upload_dir, parse_s3_url, UploadedCode, validate_source_dir
-from sagemaker.local import LocalSession, file_input
-
+from sagemaker.job import _Job
+from sagemaker.local import LocalSession
 from sagemaker.model import Model
 from sagemaker.model import (SCRIPT_PARAM_NAME, DIR_PARAM_NAME, CLOUDWATCH_METRICS_PARAM_NAME,
                              CONTAINER_LOG_LEVEL_PARAM_NAME, JOB_NAME_PARAM_NAME, SAGEMAKER_REGION_PARAM_NAME)
-
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.session import Session
 from sagemaker.session import s3_input
@@ -310,10 +309,9 @@ def delete_endpoint(self):
         self.sagemaker_session.delete_endpoint(self.latest_training_job.name)
 
 
-class _TrainingJob(object):
+class _TrainingJob(_Job):
     def __init__(self, sagemaker_session, training_job_name):
-        self.sagemaker_session = sagemaker_session
-        self.job_name = training_job_name
+        super(_TrainingJob, self).__init__(sagemaker_session, training_job_name)
 
     @classmethod
     def start_new(cls, estimator, inputs):
@@ -324,7 +322,8 @@ def start_new(cls, estimator, inputs):
             inputs (str): Parameters used when called  :meth:`~sagemaker.estimator.EstimatorBase.fit`.
 
         Returns:
-            sagemaker.estimator.Framework: Constructed object that captures all information about the started job.
+            sagemaker.estimator._TrainingJob: Constructed object that captures all information about the started
+            training job.
         """
 
         local_mode = estimator.sagemaker_session.local_mode
@@ -334,86 +333,19 @@ def start_new(cls, estimator, inputs):
             if not local_mode:
                 raise ValueError('File URIs are supported in local mode only. Please use a S3 URI instead.')
 
-        input_config = _TrainingJob._format_inputs_to_input_config(inputs)
-        role = estimator.sagemaker_session.expand_role(estimator.role)
-        output_config = _TrainingJob._prepare_output_config(estimator.output_path, estimator.output_kms_key)
-        resource_config = _TrainingJob._prepare_resource_config(estimator.train_instance_count,
-                                                                estimator.train_instance_type,
-                                                                estimator.train_volume_size)
-        stop_condition = _TrainingJob._prepare_stopping_condition(estimator.train_max_run)
+        config = _Job._load_config(inputs, estimator)
 
         if estimator.hyperparameters() is not None:
             hyperparameters = {str(k): str(v) for (k, v) in estimator.hyperparameters().items()}
 
         estimator.sagemaker_session.train(image=estimator.train_image(), input_mode=estimator.input_mode,
-                                          input_config=input_config, role=role, job_name=estimator._current_job_name,
-                                          output_config=output_config, resource_config=resource_config,
-                                          hyperparameters=hyperparameters, stop_condition=stop_condition)
+                                          input_config=config['input_config'], role=config['role'],
+                                          job_name=estimator._current_job_name, output_config=config['output_config'],
+                                          resource_config=config['resource_config'], hyperparameters=hyperparameters,
+                                          stop_condition=config['stop_condition'])
 
         return cls(estimator.sagemaker_session, estimator._current_job_name)
 
-    @staticmethod
-    def _format_inputs_to_input_config(inputs):
-        input_dict = {}
-        if isinstance(inputs, string_types):
-            input_dict['training'] = _TrainingJob._format_string_uri_input(inputs)
-        elif isinstance(inputs, s3_input):
-            input_dict['training'] = inputs
-        elif isinstance(input, file_input):
-            input_dict['training'] = inputs
-        elif isinstance(inputs, dict):
-            for k, v in inputs.items():
-                input_dict[k] = _TrainingJob._format_string_uri_input(v)
-        else:
-            raise ValueError('Cannot format input {}. Expecting one of str, dict or s3_input'.format(inputs))
-
-        channels = []
-        for channel_name, channel_s3_input in input_dict.items():
-            channel_config = channel_s3_input.config.copy()
-            channel_config['ChannelName'] = channel_name
-            channels.append(channel_config)
-        return channels
-
-    @staticmethod
-    def _format_string_uri_input(input):
-        if isinstance(input, str):
-            if input.startswith('s3://'):
-                return s3_input(input)
-            elif input.startswith('file://'):
-                return file_input(input)
-            else:
-                raise ValueError('Training input data must be a valid S3 or FILE URI: must start with "s3://" or '
-                                 '"file://"')
-        elif isinstance(input, s3_input):
-            return input
-        elif isinstance(input, file_input):
-            return input
-        else:
-            raise ValueError('Cannot format input {}. Expecting one of str, s3_input, or file_input'.format(input))
-
-    @staticmethod
-    def _prepare_output_config(s3_path, kms_key_id):
-        config = {'S3OutputPath': s3_path}
-        if kms_key_id is not None:
-            config['KmsKeyId'] = kms_key_id
-        return config
-
-    @staticmethod
-    def _prepare_resource_config(instance_count, instance_type, volume_size):
-        resource_config = {'InstanceCount': instance_count,
-                           'InstanceType': instance_type,
-                           'VolumeSizeInGB': volume_size}
-        return resource_config
-
-    @staticmethod
-    def _prepare_stopping_condition(max_run):
-        stop_condition = {'MaxRuntimeInSeconds': max_run}
-        return stop_condition
-
-    @property
-    def name(self):
-        return self.job_name
-
     def wait(self, logs=True):
         if logs:
             self.sagemaker_session.logs_for_job(self.job_name, wait=True)
@@ -474,8 +406,7 @@ def train_image(self):
         """
         Returns the docker image to use for training.
 
-        The fit() method, that does the model training, calls this method to find the image to use
-        for model training.
+        The fit() method, that does the model training, calls this method to find the image to use for model training.
         """
         return self.image_name
 

src/sagemaker/hpo.py

@@ -98,7 +98,7 @@ def start_new(cls, tuner, inputs):
         resource_config = _TrainingJob._prepare_resource_config(tuner.estimator.train_instance_count,
                                                                 tuner.estimator.train_instance_type,
                                                                 tuner.estimator.train_volume_size)
-        stop_condition = _TrainingJob._prepare_stopping_condition(tuner.estimator.train_max_run)
+        stop_condition = _TrainingJob._prepare_stop_condition(tuner.estimator.train_max_run)
 
         if tuner.estimator.hyperparameters() is None:
             raise ValueError('Cannot tune estimator without hyperparameters')

src/sagemaker/job.py

@@ -10,9 +10,12 @@
 # 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 absolute_import
+
 from abc import abstractmethod
 from six import string_types
 
+from sagemaker.local import file_input
 from sagemaker.session import s3_input
 
 
@@ -55,22 +58,29 @@ def _load_config(inputs, estimator):
         resource_config = _Job._prepare_resource_config(estimator.train_instance_count,
                                                         estimator.train_instance_type,
                                                         estimator.train_volume_size)
-        stopping_condition = _Job._prepare_stopping_condition(estimator.train_max_run)
+        stop_condition = _Job._prepare_stop_condition(estimator.train_max_run)
 
-        return input_config, role, output_config, resource_config, stopping_condition
+        return {'input_config': input_config,
+                'role': role,
+                'output_config': output_config,
+                'resource_config': resource_config,
+                'stop_condition': stop_condition}
 
     @staticmethod
     def _format_inputs_to_input_config(inputs):
         input_dict = {}
         if isinstance(inputs, string_types):
-            input_dict['training'] = _Job._format_s3_uri_input(inputs)
+            input_dict['training'] = _Job._format_string_uri_input(inputs)
         elif isinstance(inputs, s3_input):
             input_dict['training'] = inputs
+        elif isinstance(input, file_input):
+            input_dict['training'] = inputs
         elif isinstance(inputs, dict):
             for k, v in inputs.items():
-                input_dict[k] = _Job._format_s3_uri_input(v)
+                input_dict[k] = _Job._format_string_uri_input(v)
         else:
-            raise ValueError('Cannot format input {}. Expecting one of str, dict or s3_input'.format(inputs))
+            raise ValueError(
+                'Cannot format input {}. Expecting one of str, dict or s3_input'.format(inputs))
 
         channels = []
         for channel_name, channel_s3_input in input_dict.items():
@@ -80,15 +90,24 @@ def _format_inputs_to_input_config(inputs):
         return channels
 
     @staticmethod
-    def _format_s3_uri_input(input):
+    def _format_string_uri_input(input):
         if isinstance(input, str):
-            if not input.startswith('s3://'):
-                raise ValueError('Training input data must be a valid S3 URI and must start with "s3://"')
-            return s3_input(input)
-        if isinstance(input, s3_input):
+            if input.startswith('s3://'):
+                return s3_input(input)
+            elif input.startswith('file://'):
+                return file_input(input)
+            else:
+                raise ValueError(
+                    'Training input data must be a valid S3 or FILE URI: must start with "s3://" or '
+                    '"file://"')
+        elif isinstance(input, s3_input):
+            return input
+        elif isinstance(input, file_input):
             return input
         else:
-            raise ValueError('Cannot format input {}. Expecting one of str or s3_input'.format(input))
+            raise ValueError(
+                'Cannot format input {}. Expecting one of str, s3_input, or file_input'.format(
+                    input))
 
     @staticmethod
     def _prepare_output_config(s3_path, kms_key_id):
@@ -104,7 +123,7 @@ def _prepare_resource_config(instance_count, instance_type, volume_size):
                 'VolumeSizeInGB': volume_size}
 
     @staticmethod
-    def _prepare_stopping_condition(max_run):
+    def _prepare_stop_condition(max_run):
         return {'MaxRuntimeInSeconds': max_run}
 
     @property