Add docker compose V2 compatibility, see aws#3659 (comment)

Author: Chris Culhane

src/sagemaker/local/image.py

@@ -29,6 +29,7 @@
 import sys
 import tarfile
 import tempfile
+from functools import cache
 
 from distutils.spawn import find_executable
 from threading import Thread
@@ -58,6 +59,32 @@
 
 logger = logging.getLogger(__name__)
 
+@cache
+def get_docker_compose_cmd(*args):
+    cmd = []
+    # check docker compose v2 first
+    try:
+        output = subprocess.check_output(
+            ['docker', 'compose', 'version'],
+            universal_newlines=True,
+        )
+        if 'v2' in output.strip():
+            cmd.extend(['docker', 'compose'])
+
+    except subprocess.CalledProcessError:
+        # check if docker-compose is installed
+        if find_executable("docker-compose") is None:
+            raise ImportError(
+                "'docker-compose' is not installed. "
+                "Local Mode features will not work without docker-compose. "
+                "For more information on how to install 'docker-compose', please, see "
+                "https://docs.docker.com/compose/install/"
+            )
+        cmd.append('docker-compose')
+
+    cmd.extend(args)
+    return cmd
+
 
 class _SageMakerContainer(object):
     """Handle the lifecycle and configuration of a local container execution.
@@ -93,15 +120,6 @@ def __init__(
         """
         from sagemaker.local.local_session import LocalSession
 
-        # check if docker-compose is installed
-        if find_executable("docker-compose") is None:
-            raise ImportError(
-                "'docker-compose' is not installed. "
-                "Local Mode features will not work without docker-compose. "
-                "For more information on how to install 'docker-compose', please, see "
-                "https://docs.docker.com/compose/install/"
-            )
-
         self.sagemaker_session = sagemaker_session or LocalSession()
         self.instance_type = instance_type
         self.instance_count = instance_count
@@ -715,17 +733,15 @@ def _compose(self, detached=False):
         Args:
             detached:
         """
-        compose_cmd = "docker-compose"
-
-        command = [
+        compose_args = [
             compose_cmd,
             "-f",
             os.path.join(self.container_root, DOCKER_COMPOSE_FILENAME),
             "up",
             "--build",
             "--abort-on-container-exit" if not detached else "--detach",  # mutually exclusive
         ]
-
+        command = get_docker_compose_cmd(*compose_args)
         logger.info("docker command: %s", " ".join(command))
         return command
 
@@ -768,7 +784,7 @@ def _create_docker_host(self, host, environment, optml_subdirs, command, volumes
         # to setting --runtime=nvidia in the docker commandline.
         if self.instance_type == "local_gpu":
             host_config["deploy"] = {
-                "resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}
+                "resources": {"reservations": {"devices": [{"driver": "nvidia", "count": "all", "capabilities": ["gpu"]}]}}
             }
 
         if command == "serve":