diff --git a/sagemaker-core/src/sagemaker/core/local/image.py b/sagemaker-core/src/sagemaker/core/local/image.py index 9e8e732302..6137f30586 100644 --- a/sagemaker-core/src/sagemaker/core/local/image.py +++ b/sagemaker-core/src/sagemaker/core/local/image.py @@ -1214,7 +1214,7 @@ def _ecr_login_if_needed(boto_session, image): authorization_data = auth["authorizationData"][0] raw_token = base64.b64decode(authorization_data["authorizationToken"]) - token = raw_token.decode("utf-8").strip("AWS:") + token = raw_token.decode("utf-8").split(":", 1)[1] ecr_url = auth["authorizationData"][0]["proxyEndpoint"] # Log in to ecr, but use communicate to not print creds to the console diff --git a/sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py b/sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py index 7df9409b86..fcdec0a896 100644 --- a/sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py +++ b/sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py @@ -507,7 +507,7 @@ def _ecr_login_if_needed(self): authorization_data = auth["authorizationData"][0] raw_token = base64.b64decode(authorization_data["authorizationToken"]) - token = raw_token.decode("utf-8").strip("AWS:") + token = raw_token.decode("utf-8").split(":", 1)[1] ecr_url = auth["authorizationData"][0]["proxyEndpoint"] # Log in to ecr, but use communicate to not print creds to the console diff --git a/sagemaker-core/tests/unit/local/test_image.py b/sagemaker-core/tests/unit/local/test_image.py index 79940405b4..fafcfa0dd2 100644 --- a/sagemaker-core/tests/unit/local/test_image.py +++ b/sagemaker-core/tests/unit/local/test_image.py @@ -317,6 +317,29 @@ def test_ecr_login_needed(self, mock_check_output, mock_popen): assert result is True mock_popen.assert_called_once() + @patch("subprocess.Popen") + @patch("sagemaker.core.local.image._check_output") + def test_ecr_login_removes_only_aws_prefix_from_token(self, mock_check_output, mock_popen): + """Test only the 'AWS:' prefix is removed, not trailing A/W/S/: characters""" + mock_check_output.return_value = "" + mock_session = Mock() + mock_ecr = Mock() + mock_ecr.get_authorization_token.return_value = { + "authorizationData": [ + { + # base64 encoded "AWS:abcWAS": the password ends in characters of "AWS:" + "authorizationToken": "QVdTOmFiY1dBUw==", + "proxyEndpoint": "https://123456789012.dkr.ecr.us-west-2.amazonaws.com", + } + ] + } + mock_session.client.return_value = mock_ecr + ecr_image = "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest" + + _ecr_login_if_needed(mock_session, ecr_image) + + mock_popen.return_value.communicate.assert_called_once_with(input=b"abcWAS") + class TestPullImage: """Test cases for _pull_image function""" diff --git a/sagemaker-core/tests/unit/modules/local_core/test_local_container.py b/sagemaker-core/tests/unit/modules/local_core/test_local_container.py index cf5f4c847e..fb9c6437ca 100644 --- a/sagemaker-core/tests/unit/modules/local_core/test_local_container.py +++ b/sagemaker-core/tests/unit/modules/local_core/test_local_container.py @@ -450,6 +450,42 @@ def test_ecr_login_if_needed_with_ecr_image( assert result is True + @patch("sagemaker.core.modules.local_core.local_container.subprocess.Popen") + @patch("sagemaker.core.modules.local_core.local_container._check_output") + def test_ecr_login_removes_only_aws_prefix_from_token( + self, mock_check_output, mock_popen, mock_session, basic_channel + ): + """Test only the 'AWS:' prefix is removed, not trailing A/W/S/: characters""" + container = _LocalContainer( + training_job_name="test-job", + instance_type="local", + instance_count=1, + image="123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest", + container_root="/tmp/test", + input_data_config=[basic_channel], + environment={}, + hyper_parameters={}, + container_entrypoint=[], + container_arguments=[], + sagemaker_session=mock_session, + ) + mock_check_output.return_value = "" + ecr_client = Mock() + ecr_client.get_authorization_token.return_value = { + "authorizationData": [ + { + # base64 encoded "AWS:abcWAS": the password ends in characters of "AWS:" + "authorizationToken": "QVdTOmFiY1dBUw==", + "proxyEndpoint": "https://123456789012.dkr.ecr.us-west-2.amazonaws.com", + } + ] + } + mock_session.boto_session.client.return_value = ecr_client + + container._ecr_login_if_needed() + + mock_popen.return_value.communicate.assert_called_once_with(input=b"abcWAS") + @patch("sagemaker.core.modules.local_core.local_container._check_output") def test_ecr_login_if_needed_with_local_image( self, mock_check_output, mock_session, basic_channel diff --git a/sagemaker-train/src/sagemaker/train/local/local_container.py b/sagemaker-train/src/sagemaker/train/local/local_container.py index ddc21a09d8..1e55f96d01 100644 --- a/sagemaker-train/src/sagemaker/train/local/local_container.py +++ b/sagemaker-train/src/sagemaker/train/local/local_container.py @@ -515,7 +515,7 @@ def _ecr_login_if_needed(self): authorization_data = auth["authorizationData"][0] raw_token = base64.b64decode(authorization_data["authorizationToken"]) - token = raw_token.decode("utf-8").strip("AWS:") + token = raw_token.decode("utf-8").split(":", 1)[1] ecr_url = auth["authorizationData"][0]["proxyEndpoint"] # Log in to ecr, but use communicate to not print creds to the console diff --git a/sagemaker-train/tests/unit/train/local/test_local_container.py b/sagemaker-train/tests/unit/train/local/test_local_container.py index be3ed26cfa..73f486f6d1 100644 --- a/sagemaker-train/tests/unit/train/local/test_local_container.py +++ b/sagemaker-train/tests/unit/train/local/test_local_container.py @@ -10,13 +10,14 @@ # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import subprocess from sagemaker.train.local.local_container import _rmtree, _LocalContainer from sagemaker.core.shapes import DataSource, S3DataSource from sagemaker.core.shapes import Channel +from sagemaker.core.helper.session_helper import Session IMAGE = "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.1-cpu-py310" @@ -223,3 +224,35 @@ def test_get_compose_cmd_prefix_standalone_fallback( mock_which.return_value = "/usr/local/bin/docker-compose" result = container._get_compose_cmd_prefix() assert result == ["docker-compose"] + + +class TestEcrLoginIfNeeded: + """Test cases for _LocalContainer._ecr_login_if_needed.""" + + @patch("sagemaker.train.local.local_container.subprocess.Popen") + @patch("sagemaker.train.local.local_container._check_output") + def test_ecr_login_removes_only_aws_prefix_from_token( + self, mock_check_output, mock_popen, _basic_channel + ): + """Only the 'AWS:' prefix is removed, not trailing A/W/S/: characters.""" + mock_session = Mock(spec=Session) + mock_session.boto_session = Mock() + container = _make_container(_basic_channel) + container.image = IMAGE + container.sagemaker_session = mock_session + mock_check_output.return_value = "" + ecr_client = Mock() + ecr_client.get_authorization_token.return_value = { + "authorizationData": [ + { + # base64 encoded "AWS:abcWAS": the password ends in characters of "AWS:" + "authorizationToken": "QVdTOmFiY1dBUw==", + "proxyEndpoint": "https://123456789012.dkr.ecr.us-west-2.amazonaws.com", + } + ] + } + mock_session.boto_session.client.return_value = ecr_client + + container._ecr_login_if_needed() + + mock_popen.return_value.communicate.assert_called_once_with(input=b"abcWAS")