local mode: improve training input/output (aws#449)

output_path is now honored:
Can be either file:// or s3:// - This also changes the default
behavior of local mode to use the SDK provided default S3 bucket
if nothing is passed. This makes it easier for customers to create
models in SageMaker too since their Model Artifacts will already be
a tarfile in S3.

input_channel content_type is now honored in the same way as SageMaker. If it is not provided it is not passed to the container. Before we were always passing 'application/octet-stream'

Author: iquintero

CHANGELOG.rst

@@ -2,11 +2,13 @@
 CHANGELOG
 =========
 
-1.13.0
-======
+========
+1.13.dev
+========
 
 * feature: Estimator: add input mode to training channels
 * feature: Estimator: add model_uri and model_channel_name parameters
+* enhancement: Local Mode: support output_path. Can be either file:// or s3://
 
 1.12.0
 ======

src/sagemaker/estimator.py

@@ -118,6 +118,9 @@ def __init__(self, role, train_instance_count, train_instance_type,
 
         self.base_job_name = base_job_name
         self._current_job_name = None
+        if (not self.sagemaker_session.local_mode
+                and output_path and output_path.startswith('file://')):
+            raise RuntimeError('file:// output paths are only supported in Local Mode')
         self.output_path = output_path
         self.output_kms_key = output_kms_key
         self.latest_training_job = None

src/sagemaker/fw_utils.py

@@ -14,12 +14,10 @@
 
 import os
 import re
-import tarfile
-import tempfile
 from collections import namedtuple
 from six.moves.urllib.parse import urlparse
 
-from sagemaker.utils import name_from_image
+import sagemaker.utils
 
 """This module contains utility functions shared across ``Framework`` components."""
 
@@ -128,14 +126,9 @@ def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory):
     s3 = session.resource('s3')
     key = '{}/{}'.format(s3_key_prefix, 'sourcedir.tar.gz')
 
-    with tempfile.TemporaryFile() as f:
-        with tarfile.open(mode='w:gz', fileobj=f) as t:
-            for sf in source_files:
-                # Add all files from the directory into the root of the directory structure of the tar
-                t.add(sf, arcname=os.path.basename(sf))
-        # Need to reset the file descriptor position after writing to prepare for read
-        f.seek(0)
-        s3.Object(bucket, key).put(Body=f)
+    tar_file = sagemaker.utils.create_tar_file(source_files)
+    s3.Object(bucket, key).upload_file(tar_file)
+    os.remove(tar_file)
 
     return UploadedCode(s3_prefix='s3://{}/{}'.format(bucket, key), script_name=script_name)
 
@@ -226,7 +219,8 @@ def model_code_key_prefix(code_location_key_prefix, model_name, image):
     Returns:
         str: the key prefix to be used in uploading code
     """
-    return '/'.join(filter(None, [code_location_key_prefix, model_name or name_from_image(image)]))
+    training_job_name = sagemaker.utils.name_from_image(image)
+    return '/'.join(filter(None, [code_location_key_prefix, model_name or training_job_name]))
 
 
 def empty_framework_version_warning(default_version):

src/sagemaker/local/data.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import os
+import platform
 import sys
 import tempfile
 from abc import ABCMeta
@@ -162,6 +163,12 @@ def __init__(self, bucket, prefix, sagemaker_session):
             root_dir = os.path.abspath(root_dir)
 
         working_dir = tempfile.mkdtemp(dir=root_dir)
+        # Docker cannot mount Mac OS /var folder properly see
+        # https://forums.docker.com/t/var-folders-isnt-mounted-properly/9600
+        # Only apply this workaround if the user didn't provide an alternate storage root dir.
+        if root_dir is None and platform.system() == 'Darwin':
+            working_dir = '/private{}'.format(working_dir)
+
         sagemaker.utils.download_folder(bucket, prefix, working_dir, sagemaker_session)
         self.files = LocalFileDataSource(working_dir)
 

src/sagemaker/local/entities.py

@@ -46,23 +46,28 @@ def __init__(self, container):
         self.start_time = None
         self.end_time = None
 
-    def start(self, input_data_config, hyperparameters, job_name):
+    def start(self, input_data_config, output_data_config, hyperparameters, job_name):
         for channel in input_data_config:
             if channel['DataSource'] and 'S3DataSource' in channel['DataSource']:
                 data_distribution = channel['DataSource']['S3DataSource']['S3DataDistributionType']
+                data_uri = channel['DataSource']['S3DataSource']['S3Uri']
             elif channel['DataSource'] and 'FileDataSource' in channel['DataSource']:
                 data_distribution = channel['DataSource']['FileDataSource']['FileDataDistributionType']
+                data_uri = channel['DataSource']['FileDataSource']['FileUri']
             else:
                 raise ValueError('Need channel[\'DataSource\'] to have [\'S3DataSource\'] or [\'FileDataSource\']')
 
+            # use a single Data URI - this makes handling S3 and File Data easier down the stack
+            channel['DataUri'] = data_uri
+
             if data_distribution != 'FullyReplicated':
                 raise RuntimeError('DataDistribution: %s is not currently supported in Local Mode' %
                                    data_distribution)
 
         self.start = datetime.datetime.now()
         self.state = self._TRAINING
 
-        self.model_artifacts = self.container.train(input_data_config, hyperparameters, job_name)
+        self.model_artifacts = self.container.train(input_data_config, output_data_config, hyperparameters, job_name)
         self.end = datetime.datetime.now()
         self.state = self._COMPLETED
 
@@ -298,7 +303,7 @@ def _perform_batch_inference(self, input_data, output_data, **kwargs):
                     if 'AssembleWith' in output_data and output_data['AssembleWith'] == 'Line':
                         f.write(b'\n')
 
-        move_to_destination(working_dir, output_data['S3OutputPath'], self.local_session)
+        move_to_destination(working_dir, output_data['S3OutputPath'], self.name, self.local_session)
         self.container.stop_serving()
 
 

src/sagemaker/local/image.py

@@ -33,6 +33,9 @@
 import yaml
 
 import sagemaker
+import sagemaker.local.data
+import sagemaker.local.utils
+import sagemaker.utils
 
 CONTAINER_PREFIX = 'algo'
 DOCKER_COMPOSE_FILENAME = 'docker-compose.yaml'
@@ -78,7 +81,7 @@ def __init__(self, instance_type, instance_count, image, sagemaker_session=None)
         self.container_root = None
         self.container = None
 
-    def train(self, input_data_config, hyperparameters, job_name):
+    def train(self, input_data_config, output_data_config, hyperparameters, job_name):
         """Run a training job locally using docker-compose.
         Args:
             input_data_config (dict): The Input Data Configuration, this contains data such as the
@@ -126,23 +129,17 @@ def train(self, input_data_config, hyperparameters, job_name):
             msg = "Failed to run: %s, %s" % (compose_command, str(e))
             raise RuntimeError(msg)
 
-        s3_artifacts = self.retrieve_artifacts(compose_data)
+        artifacts = self.retrieve_artifacts(compose_data, output_data_config, job_name)
 
         # free up the training data directory as it may contain
         # lots of data downloaded from S3. This doesn't delete any local
         # data that was just mounted to the container.
-        _delete_tree(data_dir)
-        _delete_tree(shared_dir)
-        # Also free the container config files.
-        for host in self.hosts:
-            container_config_path = os.path.join(self.container_root, host)
-            _delete_tree(container_config_path)
-
-        self._cleanup()
-        # Print our Job Complete line to have a simmilar experience to training on SageMaker where you
+        dirs_to_delete = [data_dir, shared_dir]
+        self._cleanup(dirs_to_delete)
+        # Print our Job Complete line to have a similar experience to training on SageMaker where you
         # see this line at the end.
         print('===== Job Complete =====')
-        return s3_artifacts
+        return artifacts
 
     def serve(self, model_dir, environment):
         """Host a local endpoint using docker-compose.
@@ -188,7 +185,7 @@ def stop_serving(self):
         # for serving we can delete everything in the container root.
         _delete_tree(self.container_root)
 
-    def retrieve_artifacts(self, compose_data):
+    def retrieve_artifacts(self, compose_data, output_data_config, job_name):
         """Get the model artifacts from all the container nodes.
 
         Used after training completes to gather the data from all the individual containers. As the
@@ -201,26 +198,49 @@ def retrieve_artifacts(self, compose_data):
         Returns: Local path to the collected model artifacts.
 
         """
-        # Grab the model artifacts from all the Nodes.
-        s3_artifacts = os.path.join(self.container_root, 's3_artifacts')
-        os.mkdir(s3_artifacts)
+        # We need a directory to store the artfiacts from all the nodes
+        # and another one to contained the compressed final artifacts
+        artifacts = os.path.join(self.container_root, 'artifacts')
+        compressed_artifacts = os.path.join(self.container_root, 'compressed_artifacts')
+        os.mkdir(artifacts)
+
+        model_artifacts = os.path.join(artifacts, 'model')
+        output_artifacts = os.path.join(artifacts, 'output')
 
-        s3_model_artifacts = os.path.join(s3_artifacts, 'model')
-        s3_output_artifacts = os.path.join(s3_artifacts, 'output')
-        os.mkdir(s3_model_artifacts)
-        os.mkdir(s3_output_artifacts)
+        artifact_dirs = [model_artifacts, output_artifacts, compressed_artifacts]
+        for d in artifact_dirs:
+            os.mkdir(d)
 
+        # Gather the artifacts from all nodes into artifacts/model and artifacts/output
         for host in self.hosts:
             volumes = compose_data['services'][str(host)]['volumes']
-
             for volume in volumes:
                 host_dir, container_dir = volume.split(':')
                 if container_dir == '/opt/ml/model':
-                    sagemaker.local.utils.recursive_copy(host_dir, s3_model_artifacts)
+                    sagemaker.local.utils.recursive_copy(host_dir, model_artifacts)
                 elif container_dir == '/opt/ml/output':
-                    sagemaker.local.utils.recursive_copy(host_dir, s3_output_artifacts)
+                    sagemaker.local.utils.recursive_copy(host_dir, output_artifacts)
 
-        return s3_model_artifacts
+        # Tar Artifacts -> model.tar.gz and output.tar.gz
+        model_files = [os.path.join(model_artifacts, name) for name in os.listdir(model_artifacts)]
+        output_files = [os.path.join(output_artifacts, name) for name in os.listdir(output_artifacts)]
+        sagemaker.utils.create_tar_file(model_files, os.path.join(compressed_artifacts, 'model.tar.gz'))
+        sagemaker.utils.create_tar_file(output_files, os.path.join(compressed_artifacts, 'output.tar.gz'))
+
+        if output_data_config['S3OutputPath'] == '':
+            output_data = 'file://%s' % compressed_artifacts
+        else:
+            # Now we just need to move the compressed artifacts to wherever they are required
+            output_data = sagemaker.local.utils.move_to_destination(
+                compressed_artifacts,
+                output_data_config['S3OutputPath'],
+                job_name,
+                self.sagemaker_session)
+
+        _delete_tree(model_artifacts)
+        _delete_tree(output_artifacts)
+
+        return os.path.join(output_data, 'model.tar.gz')
 
     def write_config_files(self, host, hyperparameters, input_data_config):
         """Write the config files for the training containers.
@@ -235,17 +255,22 @@ def write_config_files(self, host, hyperparameters, input_data_config):
         Returns: None
 
         """
-
         config_path = os.path.join(self.container_root, host, 'input', 'config')
 
         resource_config = {
             'current_host': host,
             'hosts': self.hosts
         }
 
-        json_input_data_config = {
-            c['ChannelName']: {'ContentType': 'application/octet-stream'} for c in input_data_config
-        }
+        print(input_data_config)
+        json_input_data_config = {}
+        for c in input_data_config:
+            channel_name = c['ChannelName']
+            json_input_data_config[channel_name] = {
+                'TrainingInputMode': 'File'
+            }
+            if 'ContentType' in c:
+                json_input_data_config[channel_name]['ContentType'] = c['ContentType']
 
         _write_json_file(os.path.join(config_path, 'hyperparameters.json'), hyperparameters)
         _write_json_file(os.path.join(config_path, 'resourceconfig.json'), resource_config)
@@ -261,29 +286,13 @@ def _prepare_training_volumes(self, data_dir, 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']:
-                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('/')
-
+            uri = channel['DataUri']
             channel_name = channel['ChannelName']
             channel_dir = os.path.join(data_dir, channel_name)
             os.mkdir(channel_dir)
 
-            if parsed_uri.scheme == 's3':
-                bucket_name = parsed_uri.netloc
-                sagemaker.utils.download_folder(bucket_name, key, channel_dir, self.sagemaker_session)
-            elif parsed_uri.scheme == 'file':
-                path = parsed_uri.path
-                volumes.append(_Volume(path, channel=channel_name))
-            else:
-                raise ValueError('Unknown URI scheme {}'.format(parsed_uri.scheme))
+            data_source = sagemaker.local.data.get_data_source_instance(uri, self.sagemaker_session)
+            volumes.append(_Volume(data_source.get_root_dir(), channel=channel_name))
 
         # If there is a training script directory and it is a local directory,
         #  mount it to the container.
@@ -301,25 +310,20 @@ def _prepare_serving_volumes(self, model_location):
         volumes = []
         host = self.hosts[0]
         # Make the model available to the container. If this is a local file just mount it to
-        # the container as a volume. If it is an S3 location download it and extract the tar file.
+        # the container as a volume. If it is an S3 location, the DataSource will download it, we
+        # just need to extract the tar file.
         host_dir = os.path.join(self.container_root, host)
         os.makedirs(host_dir)
 
-        if model_location.startswith('s3'):
-            container_model_dir = os.path.join(self.container_root, host, 'model')
-            os.makedirs(container_model_dir)
+        model_data_source = sagemaker.local.data.get_data_source_instance(
+            model_location, self.sagemaker_session)
 
-            parsed_uri = urlparse(model_location)
-            filename = os.path.basename(parsed_uri.path)
-            tar_location = os.path.join(container_model_dir, filename)
-            sagemaker.utils.download_file(parsed_uri.netloc, parsed_uri.path, tar_location, self.sagemaker_session)
+        for filename in model_data_source.get_file_list():
+            if tarfile.is_tarfile(filename):
+                with tarfile.open(filename) as tar:
+                    tar.extractall(path=model_data_source.get_root_dir())
 
-            if tarfile.is_tarfile(tar_location):
-                with tarfile.open(tar_location) as tar:
-                    tar.extractall(path=container_model_dir)
-            volumes.append(_Volume(container_model_dir, '/opt/ml/model'))
-        else:
-            volumes.append(_Volume(model_location, '/opt/ml/model'))
+        volumes.append(_Volume(model_data_source.get_root_dir(), '/opt/ml/model'))
 
         return volumes
 
@@ -368,7 +372,6 @@ def _generate_compose_file(self, command, additional_volumes=None, additional_en
             'networks': {
                 'sagemaker-local': {'name': 'sagemaker-local'}
             }
-
         }
 
         docker_compose_path = os.path.join(self.container_root, DOCKER_COMPOSE_FILENAME)
@@ -469,9 +472,15 @@ def _build_optml_volumes(self, host, subdirs):
 
         return volumes
 
-    def _cleanup(self):
-        # we don't need to cleanup anything at the moment
-        pass
+    def _cleanup(self, dirs_to_delete=None):
+        if dirs_to_delete:
+            for d in dirs_to_delete:
+                _delete_tree(d)
+
+        # Free the container config files.
+        for host in self.hosts:
+            container_config_path = os.path.join(self.container_root, host)
+            _delete_tree(container_config_path)
 
 
 class _HostingContainer(Thread):
@@ -610,7 +619,7 @@ def _aws_credentials(session):
                 'AWS_SECRET_ACCESS_KEY=%s' % (str(secret_key))
             ]
         elif not _aws_credentials_available_in_metadata_service():
-            logger.warn("Using the short-lived AWS credentials found in session. They might expire while running.")
+            logger.warning("Using the short-lived AWS credentials found in session. They might expire while running.")
             return [
                 'AWS_ACCESS_KEY_ID=%s' % (str(access_key)),
                 'AWS_SECRET_ACCESS_KEY=%s' % (str(secret_key)),

src/sagemaker/local/local_session.py

@@ -71,7 +71,7 @@ def create_training_job(self, TrainingJobName, AlgorithmSpecification, InputData
                                         AlgorithmSpecification['TrainingImage'], self.sagemaker_session)
         training_job = _LocalTrainingJob(container)
         hyperparameters = kwargs['HyperParameters'] if 'HyperParameters' in kwargs else {}
-        training_job.start(InputDataConfig, hyperparameters, TrainingJobName)
+        training_job.start(InputDataConfig, OutputDataConfig, hyperparameters, TrainingJobName)
 
         LocalSagemakerClient._training_jobs[TrainingJobName] = training_job