fix: remove only the 'AWS:' prefix from the ECR authorization token

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

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

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"""

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

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

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")