feature: include sm-data-distributed and upgrade dependencies (#410)

Author: metrizable

CODE_OF_CONDUCT.md

@@ -0,0 +1,4 @@
+## Code of Conduct
+This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
+For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
+opensource-codeofconduct@amazon.com with any additional questions or comments.

NOTICE

@@ -0,0 +1 @@
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

buildspec-release.yml

@@ -12,14 +12,14 @@ phases:
       # run unit tests
       - AWS_ACCESS_KEY_ID= AWS_SECRET_ACCESS_KEY= AWS_SESSION_TOKEN=
         AWS_CONTAINER_CREDENTIALS_RELATIVE_URI= AWS_DEFAULT_REGION=
-        tox -e py27,py36,py37 --parallel all -- test/unit
+        tox -e py36,py37 --parallel all -- test/unit
 
       # run local integ tests
       #- $(aws ecr get-login --no-include-email --region us-west-2)
-      #- IGNORE_COVERAGE=- tox -e py27,py36 -- test/integ/local
+      #- IGNORE_COVERAGE=- tox -e py37 -- test/integ/local
 
       # run sagemaker integ tests
-      #- IGNORE_COVERAGE=- tox -e py27,py36 -- test/integ/sagemaker
+      #- IGNORE_COVERAGE=- tox -e py37 -- test/integ/sagemaker
 
       # generate the distribution package
       - python3 setup.py sdist

buildspec.yml

@@ -9,7 +9,7 @@ env:
     GITHUB_REPO: 'sagemaker-tensorflow-container'
     DLC_ACCOUNT: '763104351884'
     SETUP_FILE: 'setup_cmds.sh'
-    SETUP_CMDS: '#!/bin/bash\npython3.6 -m pip install --upgrade pip\npython3.6 -m pip install -U -e .\npython3.6 -m pip install -U -e .[test]'
+    SETUP_CMDS: '#!/bin/bash\npython3.6 -m pip install --upgrade pip==20.3.1\npython3.6 -m pip install -U .\npython3.6 -m pip install -U .[test]'
 
 phases:
   pre_build:
@@ -30,7 +30,7 @@ phases:
       - tox -e flake8,twine
 
       # run unit tests
-      - tox -e py27,py36,py37 --parallel all test/unit
+      - tox -e py36,py37 --parallel all test/unit
 
       # define tags
       - GENERIC_TAG="$FRAMEWORK_VERSION-tensorflow-$BUILD_ID"

setup.py

@@ -36,15 +36,14 @@ def read_version():
     "pytest-cov",
     "pytest-xdist",
     "mock",
-    "sagemaker==1.51.3",
+    "sagemaker[local]>=2",
     "tensorflow<2.4",
     "docker-compose",
-    "boto3==1.10.50",
-    "six==1.13.0",
+    "boto3==1.16.34",
     "python-dateutil>=2.1,<2.8.1",
-    "botocore==1.13.50",
+    "botocore==1.19.34",
     "requests-mock",
-    "awscli==1.16.314",
+    "awscli==1.18.194",
 ]
 
 if sys.version_info.major > 2:
@@ -68,12 +67,11 @@ def read_version():
         "Natural Language :: English",
         "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python",
-        "Programming Language :: Python :: 2.7",
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
     ],
     install_requires=[
-        "sagemaker-training>=3.6.4",
+        "sagemaker-training>=3.7.0",
         "numpy",
         "scipy",
         "sklearn",

src/sagemaker_tensorflow_container/training.py

@@ -27,6 +27,7 @@
 logger = logging.getLogger(__name__)
 
 SAGEMAKER_PARAMETER_SERVER_ENABLED = "sagemaker_parameter_server_enabled"
+SAGEMAKER_DISTRIBUTED_DATAPARALLEL_ENABLED = "sagemaker_distributed_dataparallel_enabled"
 MODEL_DIR = "/opt/ml/model"
 
 
@@ -136,6 +137,9 @@ def train(env, cmd_args):
     parameter_server_enabled = env.additional_framework_parameters.get(
         SAGEMAKER_PARAMETER_SERVER_ENABLED, False
     )
+    sagemaker_distributed_dataparallel_enabled = env.additional_framework_parameters.get(
+        SAGEMAKER_DISTRIBUTED_DATAPARALLEL_ENABLED, False
+    )
     if len(env.hosts) > 1 and parameter_server_enabled:
 
         tf_config = _build_tf_config(hosts=env.hosts, current_host=env.current_host)
@@ -155,6 +159,8 @@ def train(env, cmd_args):
 
         if mpi_enabled:
             runner_type = runner.MPIRunnerType
+        elif sagemaker_distributed_dataparallel_enabled:
+            runner_type = runner.SMDataParallelRunnerType
         else:
             runner_type = runner.ProcessRunnerType
 

test/integration/local/test_horovod.py

@@ -51,12 +51,11 @@ def _test_distributed_training_horovod(
     estimator = TensorFlow(
         entry_point=os.path.join(RESOURCE_PATH, "hvdbasic", "train_hvd_basic.py"),
         role="SageMakerRole",
-        train_instance_type=instance_type,
+        instance_type=instance_type,
         sagemaker_session=session,
-        train_instance_count=instances,
-        image_name=image_uri,
+        instance_count=instances,
+        image_uri=image_uri,
         output_path=output_path,
-        framework_version=framework_version,
         hyperparameters={
             "sagemaker_mpi_enabled": True,
             "sagemaker_network_interface_name": "eth0",

test/integration/local/test_training.py

@@ -102,10 +102,10 @@ def run_tf_training(
     estimator = TensorFlow(
         entry_point=script,
         role="SageMakerRole",
-        train_instance_count=instance_count,
-        train_instance_type=instance_type,
+        instance_count=instance_count,
+        instance_type=instance_type,
         sagemaker_session=sagemaker_local_session,
-        image_name=image_uri,
+        image_uri=image_uri,
         model_dir="/opt/ml/model",
         output_path=output_path,
         hyperparameters=hyperparameters,

test/integration/sagemaker/test_horovod.py

@@ -31,12 +31,11 @@ def test_distributed_training_horovod(
     estimator = TensorFlow(
         entry_point=os.path.join(RESOURCE_PATH, "mnist", "horovod_mnist.py"),
         role="SageMakerRole",
-        train_instance_type=instance_type,
-        train_instance_count=2,
-        image_name=image_uri,
+        instance_type=instance_type,
+        instance_count=2,
+        image_uri=image_uri,
         framework_version=framework_version,
         py_version="py3",
-        script_mode=True,
         hyperparameters={
             "sagemaker_mpi_enabled": True,
             "sagemaker_mpi_custom_mpi_options": mpi_options,
@@ -64,12 +63,11 @@ def test_distributed_training_horovod_with_env_vars(
     estimator = TensorFlow(
         entry_point=os.path.join(RESOURCE_PATH, "hvdbasic", "train_hvd_env_vars.py"),
         role="SageMakerRole",
-        train_instance_type=instance_type,
-        train_instance_count=2,
-        image_name=image_uri,
+        instance_type=instance_type,
+        instance_count=2,
+        image_uri=image_uri,
         framework_version=framework_version,
         py_version="py3",
-        script_mode=True,
         hyperparameters={
             "sagemaker_mpi_enabled": True,
             "sagemaker_mpi_custom_mpi_options": mpi_options,

test/integration/sagemaker/test_mnist.py

@@ -31,12 +31,11 @@ def test_mnist(sagemaker_session, image_uri, instance_type, framework_version):
     estimator = TensorFlow(
         entry_point=script,
         role="SageMakerRole",
-        train_instance_type=instance_type,
-        train_instance_count=1,
+        instance_type=instance_type,
+        instance_count=1,
         sagemaker_session=sagemaker_session,
-        image_name=image_uri,
+        image_uri=image_uri,
         framework_version=framework_version,
-        script_mode=True,
     )
     inputs = estimator.sagemaker_session.upload_data(
         path=os.path.join(resource_path, "mnist", "data"), key_prefix="scriptmode/mnist"
@@ -51,12 +50,11 @@ def test_distributed_mnist_no_ps(sagemaker_session, image_uri, instance_type, fr
     estimator = TensorFlow(
         entry_point=script,
         role="SageMakerRole",
-        train_instance_count=2,
-        train_instance_type=instance_type,
+        instance_count=2,
+        instance_type=instance_type,
         sagemaker_session=sagemaker_session,
-        image_name=image_uri,
+        image_uri=image_uri,
         framework_version=framework_version,
-        script_mode=True,
     )
     inputs = estimator.sagemaker_session.upload_data(
         path=os.path.join(resource_path, "mnist", "data"), key_prefix="scriptmode/mnist"
@@ -72,12 +70,11 @@ def test_distributed_mnist_ps(sagemaker_session, image_uri, instance_type, frame
         entry_point=script,
         role="SageMakerRole",
         hyperparameters={"sagemaker_parameter_server_enabled": True},
-        train_instance_count=2,
-        train_instance_type=instance_type,
+        instance_count=2,
+        instance_type=instance_type,
         sagemaker_session=sagemaker_session,
-        image_name=image_uri,
+        image_uri=image_uri,
         framework_version=framework_version,
-        script_mode=True,
     )
     inputs = estimator.sagemaker_session.upload_data(
         path=os.path.join(resource_path, "mnist", "data-distributed"),
@@ -95,10 +92,10 @@ def test_tuning(sagemaker_session, image_uri, instance_type, framework_version):
     estimator = TensorFlow(
         entry_point=script,
         role="SageMakerRole",
-        train_instance_type=instance_type,
-        train_instance_count=1,
+        instance_type=instance_type,
+        instance_count=1,
         sagemaker_session=sagemaker_session,
-        image_name=image_uri,
+        image_uri=image_uri,
         framework_version=framework_version,
         script_mode=True,
     )

test/integration/sagemaker/test_tuning_model_dir.py

@@ -28,9 +28,9 @@ def test_model_dir_with_training_job_name(
     estimator = TensorFlow(
         entry_point=script,
         role="SageMakerRole",
-        train_instance_type=instance_type,
-        train_instance_count=1,
-        image_name=image_uri,
+        instance_type=instance_type,
+        instance_count=1,
+        image_uri=image_uri,
         framework_version=framework_version,
         py_version="py3",
         sagemaker_session=sagemaker_session,

tox.ini

@@ -4,7 +4,7 @@
 # and then run "tox" from this directory.
 
 [tox]
-envlist = py27,py36,py37,flake8
+envlist = py36,py37,flake8
 skip_missing_interpreters = False
 
 [travis]