feature: tensorflow 2.3 support (#397)

* feature: tensorflow 2.3 support

* update test container framework versions

* use tf training 2.2 py3 image

* update dlc image tag to py37

Co-authored-by: Chuyang Deng <chuyangd@amazon.com>

Author: chuyang-deng

buildspec.yml

@@ -2,7 +2,7 @@ version: 0.2
 
 env:
   variables:
-    FRAMEWORK_VERSION: '2.1.0'
+    FRAMEWORK_VERSION: '2.2.0'
     CPU_INSTANCE_TYPE: 'ml.c4.xlarge'
     GPU_INSTANCE_TYPE: 'ml.p2.xlarge'
     ECR_REPO: 'sagemaker-test'

setup.py

@@ -30,7 +30,7 @@ def read_version():
 
 
 test_dependencies = ['tox', 'flake8', 'pytest', 'pytest-cov', 'pytest-xdist', 'mock',
-                     'sagemaker==1.51.3', 'tensorflow<2.0', 'docker-compose', 'boto3==1.10.50',
+                     'sagemaker==1.51.3', 'tensorflow<2.4', 'docker-compose', 'boto3==1.10.50',
                      'six==1.13.0', 'python-dateutil>=2.1,<2.8.1', 'botocore==1.13.50',
                      'requests-mock', 'awscli==1.16.314']
 

test/conftest.py

@@ -68,7 +68,7 @@ def pytest_addoption(parser):
     parser.addoption('--docker-base-name', default='sagemaker-tensorflow-training')
     parser.addoption('--tag', default=None)
     parser.addoption('--region', default='us-west-2')
-    parser.addoption('--framework-version', default='2.1.0')
+    parser.addoption('--framework-version', default='2.2.0')
     parser.addoption('--processor', default='cpu', choices=['cpu', 'gpu', 'cpu,gpu'])
     parser.addoption('--py-version', default='3', choices=['2', '3', '2,3'])
     parser.addoption('--account-id', default='142577830533')

test/container/2.1.0/Dockerfile.dlc.cpu renamed to test/container/2.2.0/Dockerfile.dlc.cpu

@@ -1,5 +1,5 @@
 ARG region
-FROM 763104351884.dkr.ecr.$region.amazonaws.com/tensorflow-training:2.1.0-cpu-py2
+FROM 763104351884.dkr.ecr.$region.amazonaws.com/tensorflow-training:2.2.0-cpu-py37
 
 COPY dist/sagemaker_tensorflow_training-*.tar.gz /sagemaker_tensorflow_training.tar.gz
 RUN pip install --upgrade --no-cache-dir /sagemaker_tensorflow_training.tar.gz && \

test/container/2.1.0/Dockerfile.dlc.gpu renamed to test/container/2.2.0/Dockerfile.dlc.gpu

@@ -1,5 +1,5 @@
 ARG region
-FROM 763104351884.dkr.ecr.$region.amazonaws.com/tensorflow-training:2.1.0-gpu-py3
+FROM 763104351884.dkr.ecr.$region.amazonaws.com/tensorflow-training:2.2.0-gpu-py37
 
 COPY dist/sagemaker_tensorflow_training-*.tar.gz /sagemaker_tensorflow_training.tar.gz
 RUN pip install --upgrade --no-cache-dir /sagemaker_tensorflow_training.tar.gz && \

test/container/2.1.0/Dockerfile.tf renamed to test/container/2.2.0/Dockerfile.tf

@@ -1,4 +1,4 @@
-FROM tensorflow/tensorflow:2.1.0-gpu-py3
+FROM tensorflow/tensorflow:2.3.0-gpu
 
 ENV SAGEMAKER_TRAINING_MODULE sagemaker_tensorflow_container.training:main