Skip to content

Commit 2810b03

Browse files
committed
fix: local mode printing of credentials during docker login closes #2180
1 parent 4325fcd commit 2810b03

File tree

2 files changed

+48
-12
lines changed

2 files changed

+48
-12
lines changed

src/sagemaker/local/image.py

+31-5
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,11 @@ def __init__(
112112
self.container = None
113113

114114
def process(
115-
self, processing_inputs, processing_output_config, environment, processing_job_name
115+
self,
116+
processing_inputs,
117+
processing_output_config,
118+
environment,
119+
processing_job_name,
116120
):
117121
"""Run a processing job locally using docker-compose.
118122
@@ -139,7 +143,11 @@ def process(
139143
for host in self.hosts:
140144
_create_processing_config_file_directories(self.container_root, host)
141145
self.write_processing_config_files(
142-
host, environment, processing_inputs, processing_output_config, processing_job_name
146+
host,
147+
environment,
148+
processing_inputs,
149+
processing_output_config,
150+
processing_job_name,
143151
)
144152

145153
self._generate_compose_file(
@@ -379,7 +387,12 @@ def retrieve_artifacts(self, compose_data, output_data_config, job_name):
379387
return os.path.join(output_data, "model.tar.gz")
380388

381389
def write_processing_config_files(
382-
self, host, environment, processing_inputs, processing_output_config, processing_job_name
390+
self,
391+
host,
392+
environment,
393+
processing_inputs,
394+
processing_output_config,
395+
processing_job_name,
383396
):
384397
"""Write the config files for the processing containers.
385398
@@ -1078,8 +1091,21 @@ def _ecr_login_if_needed(boto_session, image):
10781091
token = raw_token.decode("utf-8").strip("AWS:")
10791092
ecr_url = auth["authorizationData"][0]["proxyEndpoint"]
10801093

1081-
cmd = "docker login -u AWS -p %s %s" % (token, ecr_url)
1082-
subprocess.check_output(cmd.split())
1094+
# Log in to ecr, but use communicate to not print creds to the console
1095+
cmd = [
1096+
"docker",
1097+
"login",
1098+
ecr_url,
1099+
"-u",
1100+
"AWS",
1101+
"--password-stdin",
1102+
]
1103+
proc = subprocess.Popen(
1104+
cmd,
1105+
stdin=subprocess.PIPE,
1106+
)
1107+
1108+
proc.communicate(input=token.encode())
10831109

10841110
return True
10851111

tests/unit/test_image.py

+17-7
Original file line numberDiff line numberDiff line change
@@ -743,8 +743,9 @@ def test_ecr_login_image_exists(_check_output, image):
743743
assert result is False
744744

745745

746-
@patch("subprocess.check_output", return_value="".encode("utf-8"))
747-
def test_ecr_login_needed(check_output):
746+
@patch("subprocess.Popen", return_value=Mock(autospec=subprocess.Popen))
747+
@patch("sagemaker.local.image._check_output", return_value="")
748+
def test_ecr_login_needed(mock_check_output, popen):
748749
session_mock = Mock()
749750

750751
token = "very-secure-token"
@@ -765,13 +766,22 @@ def test_ecr_login_needed(check_output):
765766
}
766767
session_mock.client("ecr").get_authorization_token.return_value = response
767768
image = "520713654638.dkr.ecr.us-east-1.amazonaws.com/image-i-need:1.1"
768-
result = sagemaker.local.image._ecr_login_if_needed(session_mock, image)
769+
# What a sucessful login would look like
770+
popen.return_value.communicate.return_value = (None, None)
769771

770-
expected_command = (
771-
"docker login -u AWS -p %s https://520713654638.dkr.ecr.us-east-1.amazonaws.com" % token
772-
)
772+
result = sagemaker.local.image._ecr_login_if_needed(session_mock, image)
773773

774-
check_output.assert_called_with(expected_command.split())
774+
mock_check_output.assert_called_with(f"docker images -q {image}")
775+
expected_command = [
776+
"docker",
777+
"login",
778+
"https://520713654638.dkr.ecr.us-east-1.amazonaws.com",
779+
"-u",
780+
"AWS",
781+
"--password-stdin",
782+
]
783+
popen.assert_called_with(expected_command, stdin=subprocess.PIPE)
784+
popen.return_value.communicate.assert_called_with(input=token.encode())
775785
session_mock.client("ecr").get_authorization_token.assert_called_with(
776786
registryIds=["520713654638"]
777787
)

0 commit comments

Comments
 (0)