diff --git a/sagemaker-core/src/sagemaker/core/local/image.py b/sagemaker-core/src/sagemaker/core/local/image.py index 9e8e732302..739ffc567a 100644 --- a/sagemaker-core/src/sagemaker/core/local/image.py +++ b/sagemaker-core/src/sagemaker/core/local/image.py @@ -860,7 +860,11 @@ def _create_docker_host( if self.container_entrypoint: host_config["entrypoint"] = self.container_entrypoint if self.container_arguments: - host_config["entrypoint"] = host_config["entrypoint"] + self.container_arguments + if "entrypoint" in host_config: + host_config["entrypoint"] = host_config["entrypoint"] + self.container_arguments + else: + # No custom entrypoint: pass the arguments to the image's own entrypoint. + host_config["command"] = self.container_arguments if self.is_studio: host_config["network_mode"] = "sagemaker" diff --git a/sagemaker-core/tests/unit/local/test_image.py b/sagemaker-core/tests/unit/local/test_image.py index 79940405b4..80d60b1f81 100644 --- a/sagemaker-core/tests/unit/local/test_image.py +++ b/sagemaker-core/tests/unit/local/test_image.py @@ -1252,6 +1252,39 @@ def test_container_with_entrypoint(self, mock_session_class, mock_get_compose): assert container.container_entrypoint == ["/bin/bash"] assert container.container_arguments == ["script.sh"] + @pytest.mark.parametrize( + "entrypoint, arguments, expected", + [ + (None, ["--input", "data"], {"command": ["--input", "data"]}), + (["python3", "run.py"], None, {"entrypoint": ["python3", "run.py"]}), + ( + ["python3", "run.py"], + ["--input", "data"], + {"entrypoint": ["python3", "run.py", "--input", "data"]}, + ), + ], + ) + @patch("sagemaker.core.local.utils.check_for_studio", return_value=False) + @patch("sagemaker.core.local.image._SageMakerContainer._get_compose_cmd_prefix") + def test_create_docker_host_processing_entrypoint_and_arguments( + self, mock_get_compose, mock_studio, entrypoint, arguments, expected + ): + """Processing arguments are passed with or without a custom entrypoint""" + mock_get_compose.return_value = ["docker", "compose"] + container = _SageMakerContainer( + "local", + 1, + "test-image:latest", + Mock(config={}), + container_entrypoint=entrypoint, + container_arguments=arguments, + ) + + host_config = container._create_docker_host("algo-1", [], set(), "process", []) + + for key in ("command", "entrypoint"): + assert host_config.get(key) == expected.get(key) + @patch("subprocess.check_output") def test_get_compose_cmd_prefix_v2(self, mock_check_output): """Test getting docker compose v2 command"""