address PR comments

Author: Jonathan Esterhazy

src/sagemaker/cli/__init__.py

@@ -1,4 +1,4 @@
-from sagemaker.cli.host import host
-from sagemaker.cli.train import train
+import sagemaker.cli.mxnet
+import sagemaker.cli.tensorflow
 
-__all__ = [host, train]
+__all__ = [mxnet, tensorflow]

src/sagemaker/cli/host.py renamed to src/sagemaker/cli/common.py

@@ -1,28 +1,27 @@
 from __future__ import absolute_import
 
+import json
+import logging
 import os
 import shutil
 import tarfile
 import tempfile
 
 import sagemaker
 
+logger = logging.getLogger(__name__)
 
-def host(args):
-    return HostingCommand(args).start()
 
-
-class HostingCommand(object):
+class HostCommand(object):
     def __init__(self, args):
         self.endpoint_name = args.job_name
         self.bucket = args.bucket_name  # may be None
         self.role_name = args.role_name
+        self.python = args.python
         self.data = args.data
         self.script = args.script
-        self.python = args.python
         self.instance_type = args.instance_type
         self.instance_count = args.instance_count
-        self.framework = 'tensorflow' if args.tf else 'mxnet' if args.mx else 'undefined'
         self.environment = {k: v for k, v in (kv.split('=') for kv in args.env)}
 
         self.session = sagemaker.Session()
@@ -51,16 +50,7 @@ def create_model_archive(src):
         return archive
 
     def create_model(self, model_url):
-        if self.framework == 'tensorflow':
-            from sagemaker.tensorflow.model import TensorFlowModel
-            return TensorFlowModel(model_data=model_url, role=self.role_name, entry_point=self.script,
-                                   name=self.endpoint_name, env=self.environment)
-        elif self.framework == 'mxnet':
-            from sagemaker.mxnet.model import MXNetModel
-            return MXNetModel(model_data=model_url, role=self.role_name, entry_point=self.script,
-                              py_version=self.python, name=self.endpoint_name, env=self.environment)
-        else:
-            raise ValueError('unsupported framework value: {}'.format(self.framework))
+        raise NotImplementedError  # subclasses must override
 
     def start(self):
         model_url = self.upload_model()
@@ -69,3 +59,42 @@ def start(self):
                                  instance_type=self.instance_type)
 
         return predictor
+
+
+class TrainCommand(object):
+    def __init__(self, args):
+        self.job_name = args.job_name
+        self.bucket = args.bucket_name  # may be None
+        self.role_name = args.role_name
+        self.python = args.python
+        self.data = args.data
+        self.script = args.script
+        self.instance_type = args.instance_type
+        self.instance_count = args.instance_count
+        self.hyperparameters = self.load_hyperparameters(args.hyperparameters)
+
+        self.session = sagemaker.Session()
+
+    @staticmethod
+    def load_hyperparameters(src):
+        hp = {}
+        if src and os.path.exists(src):
+            with open(src, 'r') as f:
+                hp = json.load(f)
+        return hp
+
+    def upload_training_data(self):
+        prefix = '{}/data'.format(self.job_name)
+        data_url = self.session.upload_data(path=self.data, bucket=self.bucket, key_prefix=prefix)
+        return data_url
+
+    def create_estimator(self):
+        raise NotImplementedError  # subclasses must override
+
+    def start(self):
+        data_url = self.upload_training_data()
+        estimator = self.create_estimator()
+        estimator.fit(data_url)
+        logger.debug('code location: {}'.format(estimator.uploaded_code.s3_prefix))
+        logger.debug('model location: {}{}/output/model.tar.gz'.format(estimator.output_path,
+                                                                      estimator._current_job_name))

src/sagemaker/cli/main.py

@@ -10,53 +10,66 @@
 
 
 def parse_arguments(args):
-    # common arguments
-    common_parser = argparse.ArgumentParser(add_help=False)
-
-    # image-related settings
-    image_mtx = common_parser.add_mutually_exclusive_group(required=True)
-    image_mtx.add_argument('--tf', help='use a TensorFlow container image', action='store_true')
-    image_mtx.add_argument('--mx', help='use an MXNet container image', action='store_true')
+    parser = argparse.ArgumentParser(description='Launch SageMaker training jobs or hosting endpoints')
+    parser.set_defaults(func=lambda x: parser.print_usage())
 
-    # path to data and script files
+    # common args for training/hosting/all frameworks
+    common_parser = argparse.ArgumentParser(add_help=False)
     common_parser.add_argument('--data', help='path to training data or model files', type=str, default='./data')
     common_parser.add_argument('--script', help='path to script', type=str, default='./script.py')
     common_parser.add_argument('--job-name', help='job or endpoint name', type=str, default=None)
     common_parser.add_argument('--bucket-name', help='S3 bucket', type=str, default=None)
     common_parser.add_argument('--role-name', help='SageMaker execution role name', type=str,
                                default='AmazonSageMakerFullAccess')
+    common_parser.add_argument('--python', help='python version', type=str, default='py2')
 
     instance_group = common_parser.add_argument_group('instance settings')
     instance_group.add_argument('--instance-type', type=str, help='instance type', default='ml.m4.xlarge')
     instance_group.add_argument('--instance-count', type=int, help='instance count', default=1)
 
-    image_group = common_parser.add_argument_group('other container image settings')
-    image_group.add_argument('--python', help='python version (mxnet only)', type=str, default='py2')
-
-    parser = argparse.ArgumentParser(description='Launch SageMaker training jobs or hosting endpoints')
-    parser.set_defaults(func=lambda x: parser.print_usage())
-
-    log_group = parser.add_argument_group('log settings')
+    log_group = common_parser.add_argument_group('optional log settings')
     log_group.add_argument('--log-level', help='log level for this command', type=str, default='info')
     log_group.add_argument('--botocore-log-level', help='log level for botocore', type=str, default='warning')
 
+    # common training args
+    common_train_parser = argparse.ArgumentParser(add_help=False)
+    common_train_parser.add_argument('--hyperparameters', help='path to training hyperparameters file',
+                                     type=str, default='./hyperparameters.json')
+
+    # common hosting args
+    common_host_parser = argparse.ArgumentParser(add_help=False)
+    common_host_parser.add_argument('--env', help='hosting environment variable(s)', type=str, nargs='*', default=[])
+
     subparsers = parser.add_subparsers()
-    train_parser = subparsers.add_parser('train', help='start a training job', parents=[common_parser])
-    train_group = train_parser.add_argument_group('training settings')
-    train_group.add_argument('--hyperparameters', help='path to training hyperparameters file',
-                             type=str, default='./hyperparameters.json')
-    train_group.add_argument('--training-steps',
-                             help='number of training steps (tensorflow only)', type=int, default=None)
-    train_group.add_argument('--evaluation-steps',
-                             help='number of evaluation steps (tensorflow only)', type=int, default=None)
-    train_parser.set_defaults(mode='train')
-    train_parser.set_defaults(func=sagemaker.cli.train)
-
-    host_parser = subparsers.add_parser('host', help='start a hosting endpoint', parents=[common_parser])
-    host_group = host_parser.add_argument_group('hosting settings')
-    host_group.add_argument('--env', help='hosting environment variable(s)', type=str, nargs='*', default=[])
-    train_parser.set_defaults(mode='host')
-    host_parser.set_defaults(func=sagemaker.cli.host)
+
+    # framework/algo subcommands
+    mxnet_parser = subparsers.add_parser('mxnet', help='use MXNet', parents=[])
+    mxnet_subparsers = mxnet_parser.add_subparsers()
+    mxnet_train_parser = mxnet_subparsers.add_parser('train',
+                                                     help='start a training job',
+                                                     parents=[common_parser, common_train_parser])
+    mxnet_train_parser.set_defaults(func=sagemaker.cli.mxnet.train)
+
+    mxnet_host_parser = mxnet_subparsers.add_parser('host',
+                                                    help='start a hosting endpoint',
+                                                    parents=[common_parser, common_host_parser])
+    mxnet_host_parser.set_defaults(func=sagemaker.cli.mxnet.host)
+
+    tensorflow_parser = subparsers.add_parser('tensorflow', help='use TensorFlow', parents=[])
+    tensorflow_subparsers = tensorflow_parser.add_subparsers()
+    tensorflow_train_parser = tensorflow_subparsers.add_parser('train',
+                                                               help='start a training job',
+                                                               parents=[common_parser, common_train_parser])
+    tensorflow_train_parser.add_argument('--training-steps',
+                                         help='number of training steps (tensorflow only)', type=int, default=None)
+    tensorflow_train_parser.add_argument('--evaluation-steps',
+                                         help='number of evaluation steps (tensorflow only)', type=int, default=None)
+    tensorflow_train_parser.set_defaults(func=sagemaker.cli.tensorflow.train)
+
+    tensorflow_host_parser = tensorflow_subparsers.add_parser('host',
+                                                              help='start a hosting endpoint',
+                                                              parents=[common_parser, common_host_parser])
+    tensorflow_host_parser.set_defaults(func=sagemaker.cli.tensorflow.host)
 
     return parser.parse_args(args)
 

src/sagemaker/cli/mxnet.py

@@ -0,0 +1,34 @@
+from sagemaker.cli.common import HostCommand, TrainCommand
+
+
+def train(args):
+    MXNetTrainCommand(args).start()
+
+
+def host(args):
+    MXNetHostCommand(args).start()
+
+
+class MXNetTrainCommand(TrainCommand):
+    def __init__(self, args):
+        super(MXNetTrainCommand, self).__init__(args)
+
+    def create_estimator(self):
+        from sagemaker.mxnet.estimator import MXNet
+        return MXNet(self.script,
+                     role=self.role_name,
+                     base_job_name=self.job_name,
+                     train_instance_count=self.instance_count,
+                     train_instance_type=self.instance_type,
+                     hyperparameters=self.hyperparameters,
+                     py_version=self.python)
+
+
+class MXNetHostCommand(HostCommand):
+    def __init__(self, args):
+        super(MXNetHostCommand, self).__init__(args)
+
+    def create_model(self, model_url):
+        from sagemaker.mxnet.model import MXNetModel
+        return MXNetModel(model_data=model_url, role=self.role_name, entry_point=self.script,
+                          py_version=self.python, name=self.endpoint_name, env=self.environment)

src/sagemaker/cli/tensorflow.py

@@ -0,0 +1,38 @@
+from sagemaker.cli.common import HostCommand, TrainCommand
+
+
+def train(args):
+    TensorFlowTrainCommand(args).start()
+
+
+def host(args):
+    TensorFlowHostCommand(args).start()
+
+
+class TensorFlowTrainCommand(TrainCommand):
+    def __init__(self, args):
+        super(TensorFlowTrainCommand, self).__init__(args)
+        self.training_steps = args.training_steps
+        self.evaluation_steps = args.evaluation_steps
+
+    def create_estimator(self):
+        from sagemaker.tensorflow import TensorFlow
+        return TensorFlow(training_steps=self.training_steps,
+                          evaluation_steps=self.evaluation_steps,
+                          py_version=self.python,
+                          entry_point=self.script,
+                          role=self.role_name,
+                          base_job_name=self.job_name,
+                          train_instance_count=self.instance_count,
+                          train_instance_type=self.instance_type,
+                          hyperparameters=self.hyperparameters)
+
+
+class TensorFlowHostCommand(HostCommand):
+    def __init__(self, args):
+        super(TensorFlowHostCommand, self).__init__(args)
+
+    def create_model(self, model_url):
+        from sagemaker.tensorflow.model import TensorFlowModel
+        return TensorFlowModel(model_data=model_url, role=self.role_name, entry_point=self.script,
+                               py_version=self.python, name=self.endpoint_name, env=self.environment)