fix: Support for docker compose > v2 (5739) by aviruthen ยท Pull Request #5740 ยท aws/sagemaker-python-sdk

๐Ÿค– Iteration #1 โ€” Review Comments Addressed

Description

Fix Docker Compose version detection to support v2+ (including v3, v4, v5, etc.)

The _get_compose_cmd_prefix method previously used a narrow substring check "v2" in output to validate the Docker Compose version. This only matched Docker Compose v2.x but silently rejected v3, v4, v5, or any future major version, causing an ImportError even when Docker Compose is fully installed and functional.

The fix replaces the narrow "v2" in output check with a regex-based version extraction (re.search(r"v(\d+)", output)) that accepts any Docker Compose version with a major number >= 2.

The bug existed in three identical copies of the method across three files:

  1. sagemaker-core/src/sagemaker/core/local/image.py โ€” _SageMakerContainer._get_compose_cmd_prefix() (static method)
  2. sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py โ€” _LocalContainer._get_compose_cmd_prefix() (instance method)
  3. sagemaker-train/src/sagemaker/train/local/local_container.py โ€” _LocalContainer._get_compose_cmd_prefix() (instance method)

All three have been fixed with the same approach. The re module was already imported in all three files.

Changes Made

Source files

  • sagemaker-core/src/sagemaker/core/local/image.py: Replaced if output and "v2" in output.strip() with regex-based version check accepting major version >= 2
  • sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py: Same fix
  • sagemaker-train/src/sagemaker/train/local/local_container.py: Same fix

Test files

  • sagemaker-core/tests/unit/modules/local_core/test_local_container.py: Added unit tests for v3, v5, and v1-rejection scenarios
  • sagemaker-train/tests/unit/train/local/test_local_container.py: Added unit tests for v2, v3, v5, v1-fallback, and not-installed scenarios
  • sagemaker-train/tests/integ/train/test_docker_compose_version_detection.py: Integration tests using real Docker Compose installation; replaced hardcoded /tmp/test with tempfile.mkdtemp() per review feedback

Testing

Unit tests verify:

  • Docker Compose v2.x โ†’ accepted (["docker", "compose"])
  • Docker Compose v3.x โ†’ accepted (["docker", "compose"])
  • Docker Compose v5.x โ†’ accepted (["docker", "compose"])
  • Docker Compose v1.x โ†’ rejected, falls through to standalone docker-compose
  • No Docker Compose installed โ†’ raises ImportError

Integration tests verify the fix works against the real Docker Compose installation on the machine.

Comments reviewed: 2
Files modified: sagemaker-core/src/sagemaker/core/local/image.py, sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py, sagemaker-core/tests/unit/modules/local_core/test_local_container.py, sagemaker-train/src/sagemaker/train/local/local_container.py, sagemaker-train/tests/integ/train/test_docker_compose_version_detection.py, sagemaker-train/tests/unit/train/local/test_local_container.py

  • sagemaker-core/src/sagemaker/core/local/image.py: The file already has the fix applied from the previous iteration. No changes needed.
  • sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py: The file already has the fix applied from the previous iteration. No changes needed.
  • sagemaker-core/tests/unit/modules/local_core/test_local_container.py: The file already has the test cases from the previous iteration. No changes needed.
  • sagemaker-train/src/sagemaker/train/local/local_container.py: The file already has the fix applied from the previous iteration. No changes needed.
  • sagemaker-train/tests/integ/train/test_docker_compose_version_detection.py: Replace hardcoded /tmp/test with tempfile.mkdtemp() per reviewer feedback.
  • sagemaker-train/tests/unit/train/local/test_local_container.py: The file already has the test cases from the previous iteration. No changes needed.