Fix and add new Unit Tests, Flake 8

Also fix some TODOs here and there.

Author: Ignacio Quintero

src/sagemaker/estimator.py

@@ -12,7 +12,6 @@
 # language governing permissions and limitations under the License.
 from __future__ import print_function, absolute_import
 
-import os
 import json
 import logging
 from abc import ABCMeta
@@ -22,18 +21,15 @@
 from sagemaker.fw_utils import tar_and_upload_dir
 from sagemaker.fw_utils import parse_s3_url
 from sagemaker.fw_utils import UploadedCode
-
 from sagemaker.local.local_session import LocalSession, file_input
 
 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
-
 from sagemaker.utils import base_name_from_image, name_from_base
 
 
@@ -382,7 +378,8 @@ def _format_string_uri_input(input):
             elif input.startswith('file://'):
                 return file_input(input)
             else:
-                raise ValueError('Training input data must be a valid S3 or FILE URI and must start with "s3://" or "file://"')
+                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):

src/sagemaker/local/image.py

@@ -93,7 +93,6 @@ def train(self, input_data_config, hyperparameters):
         # mount the local directory to the container. For S3 Data we will download the S3 data
         # first.
         for channel in input_data_config:
-
             if channel['DataSource'] and 'S3DataSource' in channel['DataSource']:
                 uri = channel['DataSource']['S3DataSource']['S3Uri']
             elif channel['DataSource'] and 'FileDataSource' in channel['DataSource']:
@@ -112,9 +111,7 @@ def train(self, input_data_config, hyperparameters):
                 bucket_name = parsed_uri.netloc
                 self._download_folder(bucket_name, key, channel_dir)
             elif parsed_uri.scheme == 'file':
-                # TODO Check why this is file:/... and not file:///...
-                # TODO use the parsed_uri.xxx and use os.path.join
-                path = uri.lstrip('file:')
+                path = parsed_uri.path
                 volumes.append(_Volume(path, channel=channel_name))
             else:
                 raise ValueError('Unknown URI scheme {}'.format(parsed_uri.scheme))

src/sagemaker/local/local_session.py

@@ -179,42 +179,25 @@ def logs_for_job(self, job_name, wait=False, poll=5):
         # on local mode.
         pass
 
-# TODO Naming consistent with session.s3_input. May want to change both
-# (e.g. S3Input and FileInput)
+
 class file_input(object):
     """Amazon SageMaker channel configuration for FILE data sources, used in local mode.
 
     Attributes:
-        config (dict[str, dict]): A SageMaker ``DataSource`` referencing a SageMaker ``S3DataSource``.
+        config (dict[str, dict]): A SageMaker ``DataSource`` referencing a SageMaker ``FileDataSource``.
     """
 
-    def __init__(self, fileUri):
+    def __init__(self, fileUri, content_type=None):
         """Create a definition for input data used by an SageMaker training job in local mode.
         """
-
-        """TODO Keeping this consistent with s3_input data structure. May be
-        better to have a Type key under DataSource, but that really would mess
-        with the standard implementation....
-        """
-
         self.config = {
             'DataSource': {
                 'FileDataSource': {
-                    # TODO Ok to hardcode this here or allow input?
                     'FileDataDistributionType': 'FullyReplicated',
                     'FileUri': fileUri
                 }
             }
         }
 
-        # As per docs, leave unset in FILE mode
-        # if compression is not None:
-        #     self.config['CompressionType'] = compression
-
-        # if content_type is not None:
-        #     self.config['ContentType'] = content_type
-
-        # As per docs, leave unset in FILE mode
-        # if record_wrapping is not None:
-        #     self.config['RecordWrapperType'] = record_wrapping
-
+        if content_type is not None:
+            self.config['ContentType'] = content_type

tests/unit/test_estimator.py

@@ -101,7 +101,7 @@ def test_sagemaker_s3_uri_invalid(sagemaker_session):
         t = DummyFramework(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                            train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE)
         t.fit('thisdoesntstartwiths3')
-    assert 'must be a valid S3 URI' in str(error)
+    assert 'must be a valid S3 or FILE URI' in str(error)
 
 
 @patch('time.strftime', return_value=TIMESTAMP)
@@ -427,9 +427,8 @@ def test_unsupported_type():
 
 
 def test_unsupported_type_in_dict():
-    with pytest.raises(ValueError) as error:
+    with pytest.raises(ValueError):
         _TrainingJob._format_inputs_to_input_config({'a': 66})
-    assert 'Expecting one of str or s3_input' in str(error)
 
 
 #################################################################################

tests/unit/test_image.py

@@ -27,10 +27,9 @@
     {
         'ChannelName': 'a',
         'DataSource': {
-            'S3DataSource': {
-                'S3DataDistributionType': 'FullyReplicated',
-                'S3DataType': 'S3Prefix',
-                'S3Uri': '/tmp/source1'
+            'FileDataSource': {
+                'FileDataDistributionType': 'FullyReplicated',
+                'FileUri': 'file:///tmp/source1'
             }
         }
     },

tests/unit/test_local_session.py

@@ -35,7 +35,15 @@ def test_create_training_job(train, LocalSession):
     image = "my-docker-image:1.0"
 
     algo_spec = {'TrainingImage': image}
-    input_data_config = {}
+    input_data_config = [{
+        'ChannelName': 'a',
+        'DataSource': {
+            'S3DataSource': {
+                'S3DataDistributionType': 'FullyReplicated',
+                'S3Uri': 's3://my_bucket/tmp/source1'
+            }
+        }
+    }]
     output_data_config = {}
     resource_config = {'InstanceType': 'local', 'InstanceCount': instance_count}
     hyperparameters = {'a': 1, 'b': 'bee'}
@@ -61,6 +69,67 @@ def test_create_training_job(train, LocalSession):
     assert response['ModelArtifacts']['S3ModelArtifacts'] == expected['ModelArtifacts']['S3ModelArtifacts']
 
 
+@patch('sagemaker.local.image._SageMakerContainer.train', return_value="/some/path/to/model")
+@patch('sagemaker.local.local_session.LocalSession')
+def test_create_training_job_invalid_data_source(train, LocalSession):
+    local_sagemaker_client = sagemaker.local.local_session.LocalSagemakerClient()
+
+    instance_count = 2
+    image = "my-docker-image:1.0"
+
+    algo_spec = {'TrainingImage': image}
+
+    # InvalidDataSource is not supported. S3DataSource and FileDataSource are currently the only
+    # valid Data Sources. We expect a ValueError if we pass this input data config.
+    input_data_config = [{
+        'ChannelName': 'a',
+        'DataSource': {
+            'InvalidDataSource': {
+                'FileDataDistributionType': 'FullyReplicated',
+                'FileUri': 'ftp://myserver.com/tmp/source1'
+            }
+        }
+    }]
+
+    output_data_config = {}
+    resource_config = {'InstanceType': 'local', 'InstanceCount': instance_count}
+    hyperparameters = {'a': 1, 'b': 'bee'}
+
+    with pytest.raises(ValueError):
+        local_sagemaker_client.create_training_job("my-training-job", algo_spec, 'arn:my-role', input_data_config,
+                                                   output_data_config, resource_config, None, hyperparameters)
+
+
+@patch('sagemaker.local.image._SageMakerContainer.train', return_value="/some/path/to/model")
+@patch('sagemaker.local.local_session.LocalSession')
+def test_create_training_job_not_fully_replicated(train, LocalSession):
+    local_sagemaker_client = sagemaker.local.local_session.LocalSagemakerClient()
+
+    instance_count = 2
+    image = "my-docker-image:1.0"
+
+    algo_spec = {'TrainingImage': image}
+
+    # Local Mode only supports FullyReplicated as Data Distribution type.
+    input_data_config = [{
+        'ChannelName': 'a',
+        'DataSource': {
+            'S3DataSource': {
+                'S3DataDistributionType': 'ShardedByS3Key',
+                'S3Uri': 's3://my_bucket/tmp/source1'
+            }
+        }
+    }]
+
+    output_data_config = {}
+    resource_config = {'InstanceType': 'local', 'InstanceCount': instance_count}
+    hyperparameters = {'a': 1, 'b': 'bee'}
+
+    with pytest.raises(RuntimeError):
+        local_sagemaker_client.create_training_job("my-training-job", algo_spec, 'arn:my-role', input_data_config,
+                                                   output_data_config, resource_config, None, hyperparameters)
+
+
 @patch('sagemaker.local.local_session.LocalSession')
 def test_create_model(LocalSession):
     local_sagemaker_client = sagemaker.local.local_session.LocalSagemakerClient()
@@ -130,3 +199,34 @@ def test_create_endpoint_fails(serve, request, LocalSession):
 
     with pytest.raises(RuntimeError):
         local_sagemaker_client.create_endpoint('my-endpoint', 'some-endpoint-config')
+
+
+def test_file_input_all_defaults():
+    prefix = 'pre'
+    actual = sagemaker.local.local_session.file_input(fileUri=prefix)
+    expected = \
+        {
+            'DataSource': {
+                'FileDataSource': {
+                    'FileDataDistributionType': 'FullyReplicated',
+                    'FileUri': prefix
+                }
+            }
+        }
+    assert actual.config == expected
+
+
+def test_file_input_content_type():
+    prefix = 'pre'
+    actual = sagemaker.local.local_session.file_input(fileUri=prefix, content_type='text/csv')
+    expected = \
+        {
+            'DataSource': {
+                'FileDataSource': {
+                    'FileDataDistributionType': 'FullyReplicated',
+                    'FileUri': prefix
+                }
+            },
+            'ContentType': 'text/csv'
+        }
+    assert actual.config == expected