Local mode support for file:// URI as the input for training data, bypassing uploading to/downloading from S3.

Author: van Roekel, Gertjan

src/sagemaker/estimator.py

@@ -12,6 +12,7 @@
 # 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
@@ -21,13 +22,18 @@
 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
+
+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
 
 
@@ -321,6 +327,13 @@ def start_new(cls, estimator, inputs):
             sagemaker.estimator.Framework: Constructed object that captures all information about the started job.
         """
 
+        local_mode = estimator.local_mode
+
+        # Allow file:// input only in local mode
+        if isinstance(inputs, str) and inputs.startswith('file://'):
+            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)
@@ -343,12 +356,14 @@ def start_new(cls, estimator, inputs):
     def _format_inputs_to_input_config(inputs):
         input_dict = {}
         if isinstance(inputs, string_types):
-            input_dict['training'] = _TrainingJob._format_s3_uri_input(inputs)
+            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_s3_uri_input(v)
+                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))
 
@@ -360,15 +375,20 @@ 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 and 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):

src/sagemaker/local/image.py

@@ -93,7 +93,14 @@ 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:
-            uri = channel['DataSource']['S3DataSource']['S3Uri']
+
+            if channel['DataSource'] and 'S3DataSource' in channel['DataSource']:
+                uri = channel['DataSource']['S3DataSource']['S3Uri']
+            elif channel['DataSource'] and 'FileDataSource' in channel['DataSource']:
+                uri = channel['DataSource']['FileDataSource']['FileUri']
+            else:
+                raise ValueError('Need channel[\'DataSource\'] to have [\'S3DataSource\'] or [\'FileDataSource\']')
+
             parsed_uri = urlparse(uri)
             key = parsed_uri.path.lstrip('/')
 
@@ -104,8 +111,13 @@ def train(self, input_data_config, hyperparameters):
             if parsed_uri.scheme == 's3':
                 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:')
+                volumes.append(_Volume(path, channel=channel_name))
             else:
-                volumes.append(_Volume(uri, channel=channel_name))
+                raise ValueError('Unknown URI scheme {}'.format(parsed_uri.scheme))
 
         # Create the configuration files for each container that we will create
         # Each container will map the additional local volumes (if any).

src/sagemaker/local/local_session.py

@@ -56,15 +56,22 @@ def create_training_job(self, TrainingJobName, AlgorithmSpecification, RoleArn,
                                                    AlgorithmSpecification['TrainingImage'], self.sagemaker_session)
 
         for channel in InputDataConfig:
-            data_distribution = channel['DataSource']['S3DataSource']['S3DataDistributionType']
+
+            if channel['DataSource'] and 'S3DataSource' in channel['DataSource']:
+                data_distribution = channel['DataSource']['S3DataSource']['S3DataDistributionType']
+            elif channel['DataSource'] and 'FileDataSource' in channel['DataSource']:
+                data_distribution = channel['DataSource']['FileDataSource']['FileDataDistributionType']
+            else:
+                raise ValueError('Need channel[\'DataSource\'] to have [\'S3DataSource\'] or [\'FileDataSource\']')
+
             if data_distribution != 'FullyReplicated':
                 raise RuntimeError("DataDistribution: %s is not currently supported in Local Mode" %
                                    data_distribution)
 
         self.s3_model_artifacts = self.train_container.train(InputDataConfig, HyperParameters)
 
     def describe_training_job(self, TrainingJobName):
-        """Describe a local traininig job.
+        """Describe a local training job.
 
         Args:
             TrainingJobName (str): Not used in this implmentation.
@@ -171,3 +178,43 @@ def logs_for_job(self, job_name, wait=False, poll=5):
         # override logs_for_job() as it doesn't need to perform any action
         # 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``.
+    """
+
+    def __init__(self, fileUri):
+        """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
+