Support for docker compose > v2

[sateeshmannar]

PySDK Version

• PySDK V2 (2.x)
• PySDK V3 (3.x)

Describe the bug
Unable to run sagemaker training in "local" mode with docker version v5.

Component: sagemaker.train.local.local_container._LocalContainer._get_compose_cmd_prefix

_get_compose_cmd_prefix() only recognizes Docker Compose v2 by
checking "v2" in output. Docker Compose v5.x (and v3, v4, or any
future major version) is silently rejected, causing::

ImportError: Docker Compose is not installed.
Local Mode features will not work without docker compose.

even though docker compose is fully installed and functional.

The substring check "v2" in output is too narrow. Every Compose
version with a major number other than 2 is treated as "not installed".

To reproduce
A clear, step-by-step set of instructions to reproduce the bug.

Required for Test:
docker pull docker:latest
docker tag image to docker:latest

# Make sure the latest version is > 2
docker run --rm docker:latest docker compose version
# Docker Compose version v5.1.1

The provided code need to be complete and runnable, if additional data is needed, please include them in the issue.

import re
import subprocess

import mock



def test_local_train_triggers_compose_v5_bug(tmp_path):
 """Invoke sagemaker local-mode training to prove Docker Compose v5 triggers the bug.

 This test exercises the real call chain:
 ModelTrainer.train()
 → _LocalContainer.train()
 → _LocalContainer._generate_compose_command()
 → _LocalContainer._get_compose_cmd_prefix() ← bug lives here

 The bug: ``_get_compose_cmd_prefix`` only accepts output containing "v2".
 Any Compose v3/v4/v5 output causes an ``ImportError`` even though Docker
 Compose is fully installed and functional.

 Assumes ``docker:latest`` is already pulled on the local workstation.
 """
 import os
 import boto3
 from moto import mock_aws
 from sagemaker.core.local import LocalSession
 from sagemaker.core.training.configs import Compute, InputData, OutputDataConfig, SourceCode
 from sagemaker.train.model_trainer import ModelTrainer, Mode

 #Assumes a local docker:latest image is already pulled, which contains a real compose version string in its output. If this test starts failing due to a compose version update, update this image to one that contains a compose version string that triggers the bug (v3+).
 PUBLIC_DOCKER_IMAGE = "docker:latest"

 # Step 1: get the real compose version string from the local docker:latest image
 result = subprocess.run( # nosec B603 B607
 ["docker", "run", "--rm", PUBLIC_DOCKER_IMAGE, "docker", "compose", "version"],
 capture_output=True,
 text=True,
 timeout=30, # image assumed already pulled — no pull delay
 )
 assert result.returncode == 0, (
 f"Could not get compose version from {PUBLIC_DOCKER_IMAGE}:\n{result.stderr}"
 )
 real_version_output = result.stdout # e.g. "Docker Compose version v5.1.1\n"

 match = re.search(r"v(\d+)", real_version_output.strip())
 assert match is not None, f"Could not parse version from: {real_version_output!r}"
 major_version = int(match.group(1))

 if major_version == 2:
 print(
 f"docker:latest ships Compose v{major_version} — bug only triggers on v3+. "
 "Update PUBLIC_DOCKER_IMAGE to an image that ships Compose v3+."
 )

 # Step 2: build a minimal LocalSession backed by moto
 with mock_aws():
 boto_session = boto3.Session(region_name="us-east-1")
 s3 = boto_session.client("s3")
 s3.create_bucket(Bucket="sagemaker-bug-repro")

 sm_session = LocalSession(
 boto_session=boto_session,
 default_bucket="sagemaker-bug-repro",
 sagemaker_config={
 "SchemaVersion": "1.0",
 "SageMaker": {"PythonSDK": {"Modules": {"TelemetryOptOut": True}}},
 },
 )
 sm_session.config = {
 "local": {"local_code": True, "container_root": str(tmp_path)}
 }
 sm_session._default_bucket = "sagemaker-bug-repro"

 # Step 3: build a ModelTrainer pointing at the public docker:latest image.
 # We mock subprocess.check_output so the SDK receives the *real* version
 # string from the public image instead of whatever is installed on this host.
 trainer = ModelTrainer(
 training_image=PUBLIC_DOCKER_IMAGE,
 role="arn:aws:iam::123456789012:role/fake-role",
 compute=Compute(instance_type="local_cpu", instance_count=1),
 sagemaker_session=sm_session,
 output_data_config=OutputDataConfig(s3_output_path="s3://sagemaker-bug-repro/output"),
 base_job_name="compose-v5-bug",
 training_mode=Mode.LOCAL_CONTAINER,
 local_container_root=str(tmp_path),
 source_code=SourceCode(command="true"), # no-op command
 )

 # Step 4: patch subprocess so _get_compose_cmd_prefix sees the real
 # version string from docker:latest (v5.x), then assert the bug fires.
 with (
 mock.patch("subprocess.check_output", return_value=real_version_output),
 mock.patch("shutil.which", return_value=None),
 ):
 trainer.train(wait=False, logs=False)

Expected behavior
sagemaker train in local mode should not error

Screenshots or logs --- STACK Trace

.tox/py313/lib/python3.13/site-packages/sagemaker/core/telemetry/telemetry_logging.py:187: in wrapper 
 raise caught_ex 
.tox/py313/lib/python3.13/site-packages/sagemaker/core/telemetry/telemetry_logging.py:153: in wrapper 
 response = func(*args, **kwargs) 
 ^^^^^^^^^^^^^^^^^^^^^ 
.tox/py313/lib/python3.13/site-packages/sagemaker/core/workflow/pipeline_context.py:346: in wrapper 
 return run_func(*args, **kwargs) 
 ^^^^^^^^^^^^^^^^^^^^^^^^^ 
.tox/py313/lib/python3.13/site-packages/pydantic/_internal/_validate_call.py:39: in wrapper_function 
 return wrapper(*args, **kwargs) 
 ^^^^^^^^^^^^^^^^^^^^^^^^ 
.tox/py313/lib/python3.13/site-packages/pydantic/_internal/_validate_call.py:136: in __call__ 
 res = self.__pydantic_validator__.validate_python(pydantic_core.ArgsKwargs(args, kwargs)) 
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 
.tox/py313/lib/python3.13/site-packages/sagemaker/train/model_trainer.py:813: in train 
 local_container.train(wait) 
.tox/py313/lib/python3.13/site-packages/sagemaker/train/local/local_container.py:237: in train 
 compose_command = self._generate_compose_command(wait) 
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 
.tox/py313/lib/python3.13/site-packages/sagemaker/train/local/local_container.py:479: in _generate_compose_command 
 _compose_cmd_prefix = self._get_compose_cmd_prefix() 
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
self = _LocalContainer(training_job_name='a-test-20260408132342', instance_type='local_cpu', instance_count=1, image='glm-sta...er_arguments=['-c', 'chmod +x /opt/ml/input/data/sm_drivers/sm_train.sh && /opt/ml/input/data/sm_drivers/sm_train.sh']) 
 def _get_compose_cmd_prefix(self) -> List[str]: 
 """Gets the Docker Compose command. 
 
 The method initially looks for 'docker compose' v2 
 executable, if not found looks for 'docker-compose' executable. 
 
 Returns: 
 List[str]: 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 shutil.which("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/" 
 ) 
E 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/ 
.tox/py313/lib/python3.13/site-packages/sagemaker/train/local/local_container.py:638: ImportError

# ═══════════════════════════════════════════════════════════════════════════
# PROPOSED FIX – standalone validation of the corrected logic
# ═══════════════════════════════════════════════════════════════════════════


def _fixed_compose_check(output: str | None) -> bool:
 """Proposed replacement for the ``"v2" in output`` check.

 Accepts any Docker Compose plugin version >= 2.0.0.

 Examples
 --------
 >>> _fixed_compose_check("Docker Compose version v5.1.1")
 True
 >>> _fixed_compose_check("Docker Compose version v2.27.0")
 True
 >>> _fixed_compose_check("Docker Compose version v1.29.2")
 False
 >>> _fixed_compose_check("")
 False
 >>> _fixed_compose_check(None)
 False
 """
 if not output:
 return False
 match = re.search(r"v(\d+)", output.strip())
 return match is not None and int(match.group(1)) >= 2

System information
A description of your system. Please provide:

• SageMaker Python SDK version: sagemaker-train (>=1.6.0,<2.0.0) / sagemaker-core (>=2.7.1,<3.0.0)
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): NA
• Framework version: sagemaker-train (>=1.6.0,<2.0.0) / sagemaker-core (>=2.7.1,<3.0.0)
• Python version: 3.13
• CPU or GPU: CPU
• Custom Docker image (Y/N): N

Additional context
Add any other context about the problem here.