support nvidia-docker2 natively in local mode.

Local Mode wit GPU relied on the environment having
the nvidia runtime as default. This changes that so that
runtime: nvidia is passed in docker-compose.yaml allowing
local mode to work with a vanilla nvidia-docker2 installation.

Author: Ignacio Quintero

CHANGELOG.rst

@@ -6,6 +6,7 @@ CHANGELOG
 =========
 
 * feature: Local Mode: Add support for Batch Inference
+* enhancement: Local Mode: support nvidia-docker2 natively
 
 1.11.2
 ======

src/sagemaker/local/image.py

@@ -362,8 +362,8 @@ def _generate_compose_file(self, command, additional_volumes=None, additional_en
         }
 
         content = {
-            # Some legacy hosts only support the 2.1 format.
-            'version': '2.1',
+            # Use version 2.3 as a minimum so that we can specify the runtime
+            'version': '2.3',
             'services': services,
             'networks': {
                 'sagemaker-local': {'name': 'sagemaker-local'}
@@ -415,6 +415,11 @@ def _create_docker_host(self, host, environment, optml_subdirs, command, volumes
             }
         }
 
+        # for GPU support pass in nvidia as the runtime, this is equivalent
+        # to setting --runtime=nvidia in the docker commandline.
+        if self.instance_type == 'local_gpu':
+            host_config['runtime'] = 'nvidia'
+
         if command == 'serve':
             serving_port = sagemaker.utils.get_config_value('local.serving_port',
                                                             self.sagemaker_session.config) or 8080