Merge pull request aws#9 from aws/refactor_job

ChoiByungWook · web-flow · commit 0a655e87eadc · 2018-05-16T16:42:36.000-07:00
Add job base class
diff --git a/src/sagemaker/job.py b/src/sagemaker/job.py
@@ -0,0 +1,112 @@
+# 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):
+    """Handle creating, starting and waiting for Amazon SageMaker jobs to finish.
+
+    This class shouldn't be directly instantiated.
+
+    Subclasses must define a way to create, start and wait for an Amazon SageMaker job.
+    """
+
+    def __init__(self, sagemaker_session, job_name):
+        self.sagemaker_session = sagemaker_session
+        self.job_name = 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.job: Constructed object that captures all information about the started job.
+        """
+        pass
+
+    @abstractmethod
+    def wait(self):
+        """Wait for the Amazon SageMaker job to finish.
+        """
+        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)
+        stopping_condition = _Job._prepare_stopping_condition(estimator.train_max_run)
+
+        return input_config, role, output_config, resource_config, stopping_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):
+        return {'InstanceCount': instance_count,
+                'InstanceType': instance_type,
+                'VolumeSizeInGB': volume_size}
+
+    @staticmethod
+    def _prepare_stopping_condition(max_run):
+        return {'MaxRuntimeInSeconds': max_run}
+
+    @property
+    def name(self):
+        return self.job_name
diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py
@@ -0,0 +1,166 @@
+# 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.
+import pytest
+from mock import Mock
+
+from sagemaker.estimator import Estimator
+from sagemaker.job import _Job
+from sagemaker.session import s3_input
+
+BUCKET_NAME = 's3://mybucket/train'
+S3_OUTPUT_PATH = 's3://bucket/prefix'
+INSTANCE_COUNT = 1
+INSTANCE_TYPE = 'c4.4xlarge'
+VOLUME_SIZE = 1
+MAX_RUNTIME = 1
+ROLE = 'DummyRole'
+IMAGE_NAME = 'fakeimage'
+JOB_NAME = 'fakejob'
+
+
+@pytest.fixture()
+def estimator(sagemaker_session):
+    return Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, VOLUME_SIZE, MAX_RUNTIME,
+                     output_path=S3_OUTPUT_PATH, sagemaker_session=sagemaker_session)
+
+
+@pytest.fixture()
+def job(sagemaker_session):
+    return _Job(sagemaker_session, JOB_NAME)
+
+
+@pytest.fixture()
+def sagemaker_session():
+    boto_mock = Mock(name='boto_session')
+    mock_session = Mock(name='sagemaker_session', boto_session=boto_mock)
+    mock_session.expand_role = Mock(name='expand_role', return_value=ROLE)
+
+    return mock_session
+
+
+def test_load_config(job, estimator):
+    inputs = s3_input(BUCKET_NAME)
+
+    input_config, role, output_config, resource_config, stopping_condition = \
+        job._load_config(inputs, estimator)
+
+    assert input_config[0]['DataSource']['S3DataSource']['S3Uri'] == BUCKET_NAME
+    assert role == ROLE
+    assert output_config['S3OutputPath'] == S3_OUTPUT_PATH
+    assert 'KmsKeyId' not in output_config
+    assert resource_config['InstanceCount'] == INSTANCE_COUNT
+    assert resource_config['InstanceType'] == INSTANCE_TYPE
+    assert resource_config['VolumeSizeInGB'] == VOLUME_SIZE
+    assert stopping_condition['MaxRuntimeInSeconds'] == MAX_RUNTIME
+
+
+def test_format_inputs_to_input_config_string(job):
+    inputs = BUCKET_NAME
+
+    channels = job._format_inputs_to_input_config(inputs)
+
+    assert channels[0]['DataSource']['S3DataSource']['S3Uri'] == inputs
+
+
+def test_format_inputs_to_input_config_s3_input(job):
+    inputs = s3_input(BUCKET_NAME)
+
+    channels = job._format_inputs_to_input_config(inputs)
+
+    assert channels[0]['DataSource']['S3DataSource']['S3Uri'] == inputs.config['DataSource'][
+        'S3DataSource']['S3Uri']
+
+
+def test_format_inputs_to_input_config_dict(job):
+    inputs = {'train': BUCKET_NAME}
+
+    channels = job._format_inputs_to_input_config(inputs)
+
+    assert channels[0]['DataSource']['S3DataSource']['S3Uri'] == inputs['train']
+
+
+def test_format_inputs_to_input_config_exception(job):
+    inputs = 1
+
+    with pytest.raises(ValueError):
+        job._format_inputs_to_input_config(inputs)
+
+
+def test_format_s3_uri_input_string(job):
+    inputs = BUCKET_NAME
+
+    s3_uri_input = job._format_s3_uri_input(inputs)
+
+    assert s3_uri_input.config['DataSource']['S3DataSource']['S3Uri'] == inputs
+
+
+def test_format_s3_uri_input_string_exception(job):
+    inputs = 'mybucket/train'
+
+    with pytest.raises(ValueError):
+        job._format_s3_uri_input(inputs)
+
+
+def test_format_s3_uri_input(job):
+    inputs = s3_input(BUCKET_NAME)
+
+    s3_uri_input = job._format_s3_uri_input(inputs)
+
+    assert s3_uri_input.config['DataSource']['S3DataSource']['S3Uri'] == inputs.config[
+        'DataSource']['S3DataSource']['S3Uri']
+
+
+def test_format_s3_uri_input_exception(job):
+    inputs = 1
+
+    with pytest.raises(ValueError):
+        job._format_s3_uri_input(inputs)
+
+
+def test_prepare_output_config(job):
+    kms_key_id = 'kms_key'
+
+    config = job._prepare_output_config(BUCKET_NAME, kms_key_id)
+
+    assert config['S3OutputPath'] == BUCKET_NAME
+    assert config['KmsKeyId'] == kms_key_id
+
+
+def test_prepare_output_config_kms_key_none(job):
+    s3_path = BUCKET_NAME
+    kms_key_id = None
+
+    config = job._prepare_output_config(s3_path, kms_key_id)
+
+    assert config['S3OutputPath'] == s3_path
+    assert 'KmsKeyId' not in config
+
+
+def test_prepare_resource_config(job):
+    resource_config = job._prepare_resource_config(INSTANCE_COUNT, INSTANCE_TYPE, VOLUME_SIZE)
+
+    assert resource_config['InstanceCount'] == INSTANCE_COUNT
+    assert resource_config['InstanceType'] == INSTANCE_TYPE
+    assert resource_config['VolumeSizeInGB'] == VOLUME_SIZE
+
+
+def test_prepare_stopping_condition(job):
+    max_run = 1
+
+    stopping_condition = job._prepare_stopping_condition(max_run)
+
+    assert stopping_condition['MaxRuntimeInSeconds'] == max_run
+
+
+def test_name(job):
+    assert job.name == JOB_NAME
