diff --git a/src/sagemaker/local/image.py b/src/sagemaker/local/image.py index 7b177e1ca8..38d963bfda 100644 --- a/src/sagemaker/local/image.py +++ b/src/sagemaker/local/image.py @@ -112,7 +112,11 @@ def __init__( self.container = None def process( - self, processing_inputs, processing_output_config, environment, processing_job_name + self, + processing_inputs, + processing_output_config, + environment, + processing_job_name, ): """Run a processing job locally using docker-compose. @@ -139,7 +143,11 @@ def process( for host in self.hosts: _create_processing_config_file_directories(self.container_root, host) self.write_processing_config_files( - host, environment, processing_inputs, processing_output_config, processing_job_name + host, + environment, + processing_inputs, + processing_output_config, + processing_job_name, ) self._generate_compose_file( @@ -381,7 +389,12 @@ def retrieve_artifacts(self, compose_data, output_data_config, job_name): return os.path.join(output_data, "model.tar.gz") def write_processing_config_files( - self, host, environment, processing_inputs, processing_output_config, processing_job_name + self, + host, + environment, + processing_inputs, + processing_output_config, + processing_job_name, ): """Write the config files for the processing containers. @@ -1080,8 +1093,14 @@ def _ecr_login_if_needed(boto_session, image): token = raw_token.decode("utf-8").strip("AWS:") ecr_url = auth["authorizationData"][0]["proxyEndpoint"] - cmd = "docker login -u AWS -p %s %s" % (token, ecr_url) - subprocess.check_output(cmd.split()) + # Log in to ecr, but use communicate to not print creds to the console + cmd = f"docker login {ecr_url} -u AWS --password-stdin".split() + proc = subprocess.Popen( + cmd, + stdin=subprocess.PIPE, + ) + + proc.communicate(input=token.encode()) return True diff --git a/tests/unit/test_image.py b/tests/unit/test_image.py index 27fe99202f..020f648834 100644 --- a/tests/unit/test_image.py +++ b/tests/unit/test_image.py @@ -753,8 +753,9 @@ def test_ecr_login_image_exists(_check_output, image): assert result is False -@patch("subprocess.check_output", return_value="".encode("utf-8")) -def test_ecr_login_needed(check_output): +@patch("subprocess.Popen", return_value=Mock(autospec=subprocess.Popen)) +@patch("sagemaker.local.image._check_output", return_value="") +def test_ecr_login_needed(mock_check_output, popen): session_mock = Mock() token = "very-secure-token" @@ -775,13 +776,22 @@ def test_ecr_login_needed(check_output): } session_mock.client("ecr").get_authorization_token.return_value = response image = "520713654638.dkr.ecr.us-east-1.amazonaws.com/image-i-need:1.1" - result = sagemaker.local.image._ecr_login_if_needed(session_mock, image) + # What a sucessful login would look like + popen.return_value.communicate.return_value = (None, None) - expected_command = ( - "docker login -u AWS -p %s https://520713654638.dkr.ecr.us-east-1.amazonaws.com" % token - ) + result = sagemaker.local.image._ecr_login_if_needed(session_mock, image) - check_output.assert_called_with(expected_command.split()) + mock_check_output.assert_called_with(f"docker images -q {image}") + expected_command = [ + "docker", + "login", + "https://520713654638.dkr.ecr.us-east-1.amazonaws.com", + "-u", + "AWS", + "--password-stdin", + ] + popen.assert_called_with(expected_command, stdin=subprocess.PIPE) + popen.return_value.communicate.assert_called_with(input=token.encode()) session_mock.client("ecr").get_authorization_token.assert_called_with( registryIds=["520713654638"] )