Docker Compose V2 compatibility for local mode

[dbpprt]

Describe the feature you'd like
Currently local mode relies on docker-compose (v1), v2 is mostly compatible with the commandline and makes installing more straight forward.

How would this feature be used? Please describe.
Allow local mode with Docker Compose V2

Describe alternatives you've considered
Install V1 which will be EoL in June 2023.

[dancsalo]

+1

[pdifranc]

+1

[ozancaglayan]

Hello,

Nowadays all installations of docker are defaulting to compose v2. The fact that this package does not support this and installs the old docker-compose from pip is very confusing. In short, if your system is happily using docker compose v2 already, there is no way to use local mode with sagemaker on that machine.

[ozancaglayan]

This is more or less the change required in backward-compatible way. I tested it and it works. It does not work without setting the count/driver by the way.

diff --git a/src/sagemaker/local/image.py b/src/sagemaker/local/image.py
index adba12d3..1d9a07ce 100644
--- a/src/sagemaker/local/image.py
+++ b/src/sagemaker/local/image.py
@@ -31,6 +31,7 @@ import tarfile
 import tempfile
 
 from distutils.spawn import find_executable
+from functools import cache
 from threading import Thread
 from six.moves.urllib.parse import urlparse
 
@@ -52,6 +53,33 @@ S3_ENDPOINT_URL_ENV_NAME = "S3_ENDPOINT_URL"
 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.
 
@@ -86,15 +114,6 @@ class _SageMakerContainer(object):
         """
         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
@@ -707,17 +726,14 @@ class _SageMakerContainer(object):
         Args:
             detached:
         """
-        compose_cmd = "docker-compose"
-
-        command = [
-            compose_cmd,
+        compose_args = [
             "-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
 
@@ -760,7 +776,7 @@ class _SageMakerContainer(object):
         # 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":

[pdifranc]

@knikure How is #3855 fixing this issue? As a matter of fact, after this merge I cant pip install sagemaker[local] anymore due to a conflict in PyYAML versions

[ozancaglayan]

Yeah no idea, 3855 is totally irrelevant to this issue.