feat: Local Mode - Add Support for Docker Compose V2 (#4111)

Author: mufaddal-rohawala

doc/overview.rst

@@ -1550,14 +1550,15 @@ them to your local environment. This is a great way to test your deep learning s
 managed training or hosting environments. Local Mode is supported for frameworks images (TensorFlow, MXNet, Chainer, PyTorch,
 and Scikit-Learn) and images you supply yourself.
 
-You can install necessary dependencies for this feature using pip; local mode also requires docker-compose which you can
-install using the following steps (More info - https://github.com/docker/compose#where-to-get-docker-compose ):
+You can install necessary dependencies for this feature using pip.
 
 ::
 
     pip install 'sagemaker[local]' --upgrade
-    curl -L "https://github.com/docker/compose/releases/download/v2.7.0/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose
-    chmod +x /usr/local/bin/docker-compose
+
+
+Additionally, Local Mode also requires Docker Compose V2. Follow the guidelines in https://docs.docker.com/compose/install/ to install.
+Make sure to have a Compose Version compatible with your Docker Engine installation. Check Docker Engine release notes https://docs.docker.com/engine/release-notes to find a compatible version.
 
 If you want to keep everything local, and not use Amazon S3 either, you can enable "local code" in one of two ways:
 

src/sagemaker/local/image.py

@@ -94,14 +94,8 @@ 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.compose_cmd_prefix = _SageMakerContainer._get_compose_cmd_prefix()
         self.sagemaker_session = sagemaker_session or LocalSession()
         self.instance_type = instance_type
         self.instance_count = instance_count
@@ -118,6 +112,51 @@ def __init__(
         self.container_root = None
         self.container = None
 
+    @staticmethod
+    def _get_compose_cmd_prefix():
+        """Gets the Docker Compose command.
+
+        The method initially looks for 'docker compose' v2
+        executable, if not found looks for 'docker-compose' executable.
+
+        Returns:
+            Docker Compose executable split into list.
+
+        Raises:
+            ImportError: If Docker Compose executable was not found.
+        """
+        compose_cmd_prefix = []
+
+        output = None
+        try:
+            output = subprocess.check_output(
+                ["docker", "compose", "version"],
+                stderr=subprocess.DEVNULL,
+                encoding="UTF-8",
+            )
+        except subprocess.CalledProcessError:
+            logger.info(
+                "'Docker Compose' is not installed. "
+                "Proceeding to check for 'docker-compose' CLI."
+            )
+
+        if output and "v2" in output.strip():
+            logger.info("'Docker Compose' found using Docker CLI.")
+            compose_cmd_prefix.extend(["docker", "compose"])
+            return compose_cmd_prefix
+
+        if find_executable("docker-compose") is not None:
+            logger.info("'Docker Compose' found using Docker Compose CLI.")
+            compose_cmd_prefix.extend(["docker-compose"])
+            return compose_cmd_prefix
+
+        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/"
+        )
+
     def process(
         self,
         processing_inputs,
@@ -715,19 +754,20 @@ def _compose(self, detached=False):
         Args:
             detached:
         """
-        compose_cmd = "docker-compose"
+        compose_cmd = self.compose_cmd_prefix
 
         command = [
-            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
         ]
 
-        logger.info("docker command: %s", " ".join(command))
-        return command
+        compose_cmd.extend(command)
+
+        logger.info("docker command: %s", " ".join(compose_cmd))
+        return compose_cmd
 
     def _create_docker_host(self, host, environment, optml_subdirs, command, volumes):
         """Creates the docker host configuration.

tests/scripts/run-notebook-test.sh

@@ -141,9 +141,6 @@ echo "set SAGEMAKER_ROLE_ARN=$SAGEMAKER_ROLE_ARN"
 ./amazon-sagemaker-examples/sagemaker-python-sdk/scikit_learn_randomforest/Sklearn_on_SageMaker_end2end.ipynb \
 ./amazon-sagemaker-examples/sagemaker-pipelines/tabular/abalone_build_train_deploy/sagemaker-pipelines-preprocess-train-evaluate-batch-transform.ipynb \
 
-# Skipping test until fix in example notebook to move to new conda environment
-#./amazon-sagemaker-examples/advanced_functionality/kmeans_bring_your_own_model/kmeans_bring_your_own_model.ipynb \
-
 # Skipping test until fix in example notebook to install docker-compose is complete
 #./amazon-sagemaker-examples/sagemaker-python-sdk/tensorflow_moving_from_framework_mode_to_script_mode/tensorflow_moving_from_framework_mode_to_script_mode.ipynb \
 

tests/unit/sagemaker/local/test_local_image.py

@@ -28,7 +28,6 @@
 import pytest
 import yaml
 from mock import patch, Mock, MagicMock
-
 import sagemaker
 from sagemaker.local.image import _SageMakerContainer, _Volume, _aws_credentials
 
@@ -91,6 +90,38 @@ def sagemaker_session():
     return sms
 
 
+@patch("subprocess.check_output", Mock(return_value="Docker Compose version v2.0.0-rc.3"))
+def test_get_compose_cmd_prefix_with_docker_cli():
+    compose_cmd_prefix = _SageMakerContainer._get_compose_cmd_prefix()
+    assert compose_cmd_prefix == ["docker", "compose"]
+
+
+@patch(
+    "subprocess.check_output",
+    side_effect=subprocess.CalledProcessError(returncode=1, cmd="docker compose version"),
+)
+@patch("sagemaker.local.image.find_executable", Mock(return_value="/usr/bin/docker-compose"))
+def test_get_compose_cmd_prefix_with_docker_compose_cli(check_output):
+    compose_cmd_prefix = _SageMakerContainer._get_compose_cmd_prefix()
+    assert compose_cmd_prefix == ["docker-compose"]
+
+
+@patch(
+    "subprocess.check_output",
+    side_effect=subprocess.CalledProcessError(returncode=1, cmd="docker compose version"),
+)
+@patch("sagemaker.local.image.find_executable", Mock(return_value=None))
+def test_get_compose_cmd_prefix_raises_import_error(check_output):
+    with pytest.raises(ImportError) as e:
+        _SageMakerContainer._get_compose_cmd_prefix()
+    assert (
+        "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/" in str(e)
+    )
+
+
 def test_sagemaker_container_hosts_should_have_lowercase_names():
     random.seed(a=42)
 
@@ -333,6 +364,10 @@ def test_check_output():
 @patch("sagemaker.local.image._stream_output", Mock())
 @patch("sagemaker.local.image._SageMakerContainer._cleanup")
 @patch("sagemaker.local.image._SageMakerContainer.retrieve_artifacts")
+@patch(
+    "sagemaker.local.image._SageMakerContainer._get_compose_cmd_prefix",
+    Mock(return_value=["docker-compose"]),
+)
 @patch("sagemaker.local.data.get_data_source_instance")
 @patch("subprocess.Popen")
 def test_train(
@@ -438,6 +473,10 @@ def test_train_with_hyperparameters_without_job_name(
 @patch("sagemaker.local.image._stream_output", side_effect=RuntimeError("this is expected"))
 @patch("sagemaker.local.image._SageMakerContainer._cleanup")
 @patch("sagemaker.local.image._SageMakerContainer.retrieve_artifacts")
+@patch(
+    "sagemaker.local.image._SageMakerContainer._get_compose_cmd_prefix",
+    Mock(return_value=["docker-compose"]),
+)
 @patch("sagemaker.local.data.get_data_source_instance")
 @patch("subprocess.Popen", Mock())
 def test_train_error(
@@ -475,6 +514,10 @@ def test_train_error(
 @patch("sagemaker.local.local_session.LocalSession", Mock())
 @patch("sagemaker.local.image._stream_output", Mock())
 @patch("sagemaker.local.image._SageMakerContainer._cleanup", Mock())
+@patch(
+    "sagemaker.local.image._SageMakerContainer._get_compose_cmd_prefix",
+    Mock(return_value=["docker-compose"]),
+)
 @patch("sagemaker.local.data.get_data_source_instance")
 @patch("subprocess.Popen", Mock())
 def test_train_local_code(get_data_source_instance, tmpdir, sagemaker_session):
@@ -528,6 +571,10 @@ def test_train_local_code(get_data_source_instance, tmpdir, sagemaker_session):
 @patch("sagemaker.local.local_session.LocalSession", Mock())
 @patch("sagemaker.local.image._stream_output", Mock())
 @patch("sagemaker.local.image._SageMakerContainer._cleanup", Mock())
+@patch(
+    "sagemaker.local.image._SageMakerContainer._get_compose_cmd_prefix",
+    Mock(return_value=["docker-compose"]),
+)
 @patch("sagemaker.local.data.get_data_source_instance")
 @patch("subprocess.Popen", Mock())
 def test_train_local_intermediate_output(get_data_source_instance, tmpdir, sagemaker_session):