Add docker-compose and unit tests.

docker-compose was not a dependency of the SDk but it is required for
local mode.

Author: Ignacio Quintero

setup.py

@@ -44,8 +44,8 @@ def read(fname):
       ],
 
       # Declare minimal set for installation
-      install_requires=['boto3>=1.4.8', 'numpy>=1.9.0', 'protobuf>=3.1', 'scipy>=0.19.0', 'urllib3>=1.2',
-                        'PyYAML>=3.2', 'protobuf3-to-dict>=0.1.5'],
+      install_requires=['boto3>=1.4.8', 'numpy>=1.9.0', 'protobuf>=3.1', 'scipy>=0.19.0', 'urllib3 >=1.21, <1.23',
+                        'PyYAML>=3.2', 'protobuf3-to-dict>=0.1.5', 'docker-compose>=1.21.0'],
 
       extras_require={
           'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'pytest-xdist',

tests/unit/test_image.py

@@ -334,6 +334,27 @@ def test_train_local_code(download_folder, _cleanup, popen, _stream_output,
                 assert '%s:/opt/ml/shared' % shared_folder_path in volumes
 
 
+def test_container_has_gpu_support(tmpdir, sagemaker_session):
+    instance_count = 1
+    image = 'my-image'
+    sagemaker_container = _SageMakerContainer('local_gpu', instance_count, image,
+                                              sagemaker_session=sagemaker_session)
+
+    docker_host = sagemaker_container._create_docker_host('host-1', {}, set(), 'train', [])
+    assert 'runtime' in docker_host
+    assert docker_host['runtime'] == 'nvidia'
+
+
+def test_container_does_not_enable_nvidia_docker_for_cpu_containers(tmpdir, sagemaker_session):
+    instance_count = 1
+    image = 'my-image'
+    sagemaker_container = _SageMakerContainer('local', instance_count, image,
+                                              sagemaker_session=sagemaker_session)
+
+    docker_host = sagemaker_container._create_docker_host('host-1', {}, set(), 'train', [])
+    assert 'runtime' not in docker_host
+
+
 @patch('sagemaker.local.image._HostingContainer.run')
 @patch('shutil.copy')
 @patch('shutil.copytree')