Create job base class

Author: Dan Choi

src/sagemaker/job.py

@@ -0,0 +1,109 @@
+# Copyright 2017 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 abc import abstractmethod
+from six import string_types
+
+from sagemaker.session import s3_input
+
+
+class _Job(object):
+    def __init__(self, sagemaker_session, training_job_name):
+        self.sagemaker_session = sagemaker_session
+        self.job_name = training_job_name
+
+    @abstractmethod
+    def start_new(cls, estimator, inputs):
+        """Create a new Amazon SageMaker job from the estimator.
+
+        Args:
+            estimator (sagemaker.estimator.EstimatorBase): Estimator object created by the user.
+            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.
+        """
+
+        pass
+
+    @staticmethod
+    def _load_config(inputs, estimator):
+        input_config = _Job._format_inputs_to_input_config(inputs)
+        role = estimator.sagemaker_session.expand_role(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,
+                                                                estimator.train_volume_size)
+        stop_condition = _Job._prepare_stopping_condition(estimator.train_max_run)
+
+        return input_config, role, output_config, resource_config, 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)
+        elif isinstance(inputs, s3_input):
+            input_dict['training'] = inputs
+        elif isinstance(inputs, dict):
+            for k, v in inputs.items():
+                input_dict[k] = _Job._format_s3_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_s3_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):
+            return input
+        else:
+            raise ValueError('Cannot format input {}. Expecting one of str or s3_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)
+        else:
+            self.sagemaker_session.wait_for_job(self.job_name)