Add support for additional files

CHANGELOG.rst

@@ -2,6 +2,11 @@
 CHANGELOG
 =========
 
+1.15.1dev
+=========
+
+* feature: Estimators: source_dir accepts a list of directories
+
 1.15.0
 ======
 
@@ -16,6 +21,7 @@ CHANGELOG
 * feature: Estimator: add support for MetricDefinitions
 * feature: Estimators: add support for Amazon IP Insights algorithm
 
+
 1.14.2
 ======
 

src/sagemaker/chainer/README.rst

@@ -145,10 +145,10 @@ Optional arguments
 
 The following are optional arguments. When you create a ``Chainer`` object, you can specify these as keyword arguments.
 
--  ``source_dir`` Path (absolute or relative) to a directory with any
-   other training source code dependencies including the entry point
-   file. Structure within this directory will be preserved when training
-   on SageMaker.
+-  ``source_dir`` Single path (absolute or relative) or a list of paths
+    to directories with any other training source code dependencies
+    aside from the entry point file (default: None). The structures
+    within this directories are preserved when training on Amazon SageMaker.
 -  ``hyperparameters`` Hyperparameters that will be used for training.
    Will be made accessible as a dict[str, str] to the training code on
    SageMaker. For convenience, accepts other types besides str, but

src/sagemaker/chainer/estimator.py

@@ -64,9 +64,9 @@ def __init__(self, entry_point, use_mpi=None, num_processes=None, process_slots_
                 set to the number of GPUs on the instance (on GPU instances), or one (on CPU instances).
             additional_mpi_options (str): String of options to the 'mpirun' command used to run the entry point.
                 For example, '-X NCCL_DEBUG=WARN' will pass that option string to the mpirun command.
-            source_dir (str): Path (absolute or relative) to a directory with any other training
-                source code dependencies aside from tne entry point file (default: None). Structure within this
-                directory are preserved when training on Amazon SageMaker.
+            source_dir (str or [str]): Single path (absolute or relative) or a list of paths to directories with
+                any other training source code dependencies aside from the entry point file (default: None).
+                The structures within this directories are preserved when training on Amazon SageMaker.
             hyperparameters (dict): Hyperparameters that will be used for training (default: None).
                 The hyperparameters are made accessible as a dict[str, str] to the training code on SageMaker.
                 For convenience, this accepts other types for keys and values, but ``str()`` will be called

src/sagemaker/estimator.py

@@ -638,9 +638,9 @@ def __init__(self, entry_point, source_dir=None, hyperparameters=None, enable_cl
         Args:
             entry_point (str): Path (absolute or relative) to the Python source file which should be executed
                 as the entry point to training. This should be compatible with either Python 2.7 or Python 3.5.
-            source_dir (str): Path (absolute or relative) to a directory with any other training
-                source code dependencies aside from tne entry point file (default: None). Structure within this
-                directory are preserved when training on Amazon SageMaker.
+            source_dir (str or [str]): Single path (absolute or relative) or a list of paths to directories with
+                any other training source code dependencies aside from the entry point file (default: None).
+                The structures within this directories are preserved when training on Amazon SageMaker.
             hyperparameters (dict): Hyperparameters that will be used for training (default: None).
                 The hyperparameters are made accessible as a dict[str, str] to the training code on SageMaker.
                 For convenience, this accepts other types for keys and values, but ``str()`` will be called
@@ -657,7 +657,14 @@ def __init__(self, entry_point, source_dir=None, hyperparameters=None, enable_cl
             **kwargs: Additional kwargs passed to the ``EstimatorBase`` constructor.
         """
         super(Framework, self).__init__(**kwargs)
-        self.source_dir = source_dir
+
+        if isinstance(source_dir, list):
+            self.source_dir = source_dir[0]
+            self._additional_files = source_dir[1:]
+        else:
+            self.source_dir = source_dir
+            self._additional_files = []
+
         self.entry_point = entry_point
         if enable_cloudwatch_metrics:
             warnings.warn('enable_cloudwatch_metrics is now deprecated and will be removed in the future.',
@@ -724,15 +731,17 @@ def _stage_user_code_in_s3(self):
                                   bucket=code_bucket,
                                   s3_key_prefix=code_s3_prefix,
                                   script=self.entry_point,
-                                  directory=self.source_dir)
+                                  directory=self.source_dir,
+                                  additional_files=self._additional_files)
 
     def _model_source_dir(self):
         """Get the appropriate value to pass as source_dir to model constructor on deploying
 
         Returns:
             str: Either a local or an S3 path pointing to the source_dir to be used for code by the model to be deployed
         """
-        return self.source_dir if self.sagemaker_session.local_mode else self.uploaded_code.s3_prefix
+        source_dir = [self.source_dir] + self._additional_files
+        return source_dir if self.sagemaker_session.local_mode else self.uploaded_code.s3_prefix
 
     def hyperparameters(self):
         """Return the hyperparameters as a dictionary to use for training.

src/sagemaker/fw_utils.py

@@ -107,7 +107,7 @@ def validate_source_dir(script, directory):
     return True
 
 
-def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory):
+def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory, additional_files=None):
     """Pack and upload source files to S3 only if directory is empty or local.
 
     Note:
@@ -118,31 +118,43 @@ def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory):
         bucket (str): S3 bucket to which the compressed file is uploaded.
         s3_key_prefix (str): Prefix for the S3 key.
         script (str): Script filename.
-        directory (str): Directory containing the source file. If it starts with "s3://", no action is taken.
+        directory (str or None): Directory containing the source file. If it starts with "s3://", no action is taken.
 
     Returns:
         sagemaker.fw_utils.UserCode: An object with the S3 bucket and key (S3 prefix) and script name.
     """
-    if directory:
-        if directory.lower().startswith("s3://"):
-            return UploadedCode(s3_prefix=directory, script_name=os.path.basename(script))
-        else:
-            script_name = script
-            source_files = [os.path.join(directory, name) for name in os.listdir(directory)]
+    key = '%s/sourcedir.tar.gz' % s3_key_prefix
+
+    if directory and directory.lower().startswith("s3://"):
+        return UploadedCode(s3_prefix=directory, script_name=os.path.basename(script))
     else:
-        # If no directory is specified, the script parameter needs to be a valid relative path.
-        os.path.exists(script)
-        script_name = os.path.basename(script)
-        source_files = [script]
+        source_files = _list_root_files(script, directory, additional_files)
+        _upload_code(session, bucket, key, source_files)
+
+        script_name = script if directory else os.path.basename(script)
+        return UploadedCode(s3_prefix='s3://%s/%s' % (bucket, key), script_name=script_name)
 
-    s3 = session.resource('s3')
-    key = '{}/{}'.format(s3_key_prefix, 'sourcedir.tar.gz')
 
+def _upload_code(session, bucket, key, source_files):
     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)
+    try:
+        session.resource('s3').Object(bucket, key).upload_file(tar_file)
+    finally:
+        os.remove(tar_file)
+
+
+def _list_root_files(script, directory, additional_files):
+    additional_files = additional_files or []
+    basedir = directory if directory else os.path.dirname(script)
+    files = [basedir] + additional_files
+
+    for file in files:
+        if os.path.isfile(file):
+            yield file
+        else:
+            for name in os.listdir(file):
+                yield os.path.join(file, name)
 
 
 def framework_name_from_image(image_name):

src/sagemaker/model.py

@@ -136,9 +136,8 @@ def __init__(self, model_data, image, role, entry_point, source_dir=None, predic
             role (str): An IAM role name or ARN for SageMaker to access AWS resources on your behalf.
             entry_point (str): Path (absolute or relative) to the Python source file which should be executed
                 as the entry point to model hosting. This should be compatible with either Python 2.7 or Python 3.5.
-            source_dir (str): Path (absolute or relative) to a directory with any other training
-                source code dependencies aside from tne entry point file (default: None). Structure within this
-                directory will be preserved when training on SageMaker.
+            source_dir (str or [str]): Single path (absolute or relative) or a list of paths to directories with
+                any other training source code dependencies aside from the entry point file (default: None).
                 If the directory points to S3, no code will be uploaded and the S3 location will be used instead.
             predictor_cls (callable[string, sagemaker.session.Session]): A function to call to create
                a predictor (default: None). If not None, ``deploy`` will return the result of invoking
@@ -158,8 +157,14 @@ def __init__(self, model_data, image, role, entry_point, source_dir=None, predic
         """
         super(FrameworkModel, self).__init__(model_data, image, role, predictor_cls=predictor_cls, env=env, name=name,
                                              sagemaker_session=sagemaker_session, **kwargs)
+        if isinstance(source_dir, list):
+            self.source_dir = source_dir[0]
+            self._additional_files = source_dir[1:]
+        else:
+            self.source_dir = source_dir
+            self._additional_files = []
+
         self.entry_point = entry_point
-        self.source_dir = source_dir
         self.enable_cloudwatch_metrics = enable_cloudwatch_metrics
         self.container_log_level = container_log_level
         if code_location:
@@ -194,7 +199,8 @@ def _upload_code(self, key_prefix):
                                                     bucket=self.bucket or self.sagemaker_session.default_bucket(),
                                                     s3_key_prefix=key_prefix,
                                                     script=self.entry_point,
-                                                    directory=self.source_dir)
+                                                    directory=self.source_dir,
+                                                    additional_files=self._additional_files)
 
     def _framework_env_vars(self):
         if self.uploaded_code:

src/sagemaker/mxnet/README.rst

@@ -267,10 +267,10 @@ Optional arguments
 
 The following are optional arguments. When you create an ``MXNet`` object, you can specify these as keyword arguments.
 
--  ``source_dir`` Path (absolute or relative) to a directory with any
-   other training source code dependencies including the entry point
-   file. Structure within this directory will be preserved when training
-   on SageMaker.
+-  ``source_dir`` Single path (absolute or relative) or a list of paths
+    to directories with any other training source code dependencies
+    aside from the entry point file (default: None). The structures
+    within this directories are preserved when training on Amazon SageMaker.
 -  ``hyperparameters`` Hyperparameters that will be used for training.
    Will be made accessible as a dict[str, str] to the training code on
    SageMaker. For convenience, accepts other types besides str, but

src/sagemaker/mxnet/estimator.py

@@ -50,9 +50,9 @@ def __init__(self, entry_point, source_dir=None, hyperparameters=None, py_versio
         Args:
             entry_point (str): Path (absolute or relative) to the Python source file which should be executed
                 as the entry point to training. This should be compatible with either Python 2.7 or Python 3.5.
-            source_dir (str): Path (absolute or relative) to a directory with any other training
-                source code dependencies aside from tne entry point file (default: None). Structure within this
-                directory are preserved when training on Amazon SageMaker.
+            source_dir (str or [str]): Single path (absolute or relative) or a list of paths to directories with
+                any other training source code dependencies aside from the entry point file (default: None).
+                The structures within this directories are preserved when training on Amazon SageMaker.
             hyperparameters (dict): Hyperparameters that will be used for training (default: None).
                 The hyperparameters are made accessible as a dict[str, str] to the training code on SageMaker.
                 For convenience, this accepts other types for keys and values, but ``str()`` will be called

src/sagemaker/pytorch/README.rst

@@ -171,10 +171,10 @@ Optional arguments
 
 The following are optional arguments. When you create a ``PyTorch`` object, you can specify these as keyword arguments.
 
--  ``source_dir`` Path (absolute or relative) to a directory with any
-   other training source code dependencies including the entry point
-   file. Structure within this directory will be preserved when training
-   on SageMaker.
+-  ``source_dir`` Single path (absolute or relative) or a list of paths
+    to directories with any other training source code dependencies
+    aside from the entry point file (default: None). The structures
+    within this directories are preserved when training on Amazon SageMaker.
 -  ``hyperparameters`` Hyperparameters that will be used for training.
    Will be made accessible as a dict[str, str] to the training code on
    SageMaker. For convenience, accepts other types besides strings, but

src/sagemaker/pytorch/estimator.py

@@ -47,9 +47,9 @@ def __init__(self, entry_point, source_dir=None, hyperparameters=None, py_versio
         Args:
             entry_point (str): Path (absolute or relative) to the Python source file which should be executed
                 as the entry point to training. This should be compatible with either Python 2.7 or Python 3.5.
-            source_dir (str): Path (absolute or relative) to a directory with any other training
-                source code dependencies aside from tne entry point file (default: None). Structure within this
-                directory are preserved when training on Amazon SageMaker.
+            source_dir (str or [str]): Single path (absolute or relative) or a list of paths to directories with
+                any other training source code dependencies aside from the entry point file (default: None).
+                The structures within this directories are preserved when training on Amazon SageMaker.
             hyperparameters (dict): Hyperparameters that will be used for training (default: None).
                 The hyperparameters are made accessible as a dict[str, str] to the training code on SageMaker.
                 For convenience, this accepts other types for keys and values, but ``str()`` will be called

src/sagemaker/tensorflow/README.rst

@@ -405,10 +405,10 @@ Optional Arguments
 The following are optional arguments. When you create a ``TensorFlow`` object,
 you can specify these as keyword arguments.
 
--  ``source_dir (str)`` Path (absolute or relative) to a directory with any
-   other training source code dependencies including the entry point
-   file. Structure within this directory will be preserved when training
-   on SageMaker.
+-  ``source_dir (str)`` Single path (absolute or relative) or a list of paths
+    to directories with any other training source code dependencies
+    aside from the entry point file (default: None). The structures
+    within this directories are preserved when training on Amazon SageMaker.
 -  ``requirements_file (str)`` Path to a ``requirements.txt`` file. The path should
    be within and relative to ``source_dir``. This is a file containing a list of items to be
    installed using pip install. Details on the format can be found in the

tests/data/pytorch_source_dirs/train.py

@@ -0,0 +1,26 @@
+# Copyright 2018 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 alexa
+
+
+def model_fn(anything):
+    return alexa
+
+
+def predict_fn(input_object, model):
+    return input_object
+
+
+if __name__ == '__main__':
+    with open('/opt/ml/model/answer', 'w') as model:
+        model.write(str(alexa.question('How many roads must a man walk down?')))

tests/integ/test_source_dirs.py

@@ -0,0 +1,40 @@
+# Copyright 2018 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 __future__ import absolute_import
+
+import os
+
+from sagemaker.pytorch.estimator import PyTorch
+from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ.timeout import timeout
+
+
+def test_source_dirs(sagemaker_session, tmpdir):
+    source_dir = os.path.join(DATA_DIR, 'pytorch_source_dirs')
+    lib = os.path.join(str(tmpdir), 'alexa.py')
+
+    with open(lib, 'w') as f:
+        f.write('def question(to_anything): return 42')
+
+    instance_type = 'local'
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        estimator = PyTorch(entry_point='train.py', role='SageMakerRole', source_dir=[source_dir, lib],
+                            py_version=PYTHON_VERSION, train_instance_count=1, train_instance_type=instance_type)
+
+        estimator.fit()
+
+        predictor = estimator.deploy(initial_instance_count=1, instance_type=instance_type)
+
+        predict_response = predictor.predict([24])
+
+        assert predict_response == [24]