Skip to content

Commit c4d71f5

Browse files
authored
fix: mask creds from docker commands in local mode. Closes #2118 (#2146)
1 parent e08c04e commit c4d71f5

File tree

2 files changed

+16
-3
lines changed

2 files changed

+16
-3
lines changed

src/sagemaker/local/image.py

+8-1
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from __future__ import absolute_import
1515

1616
import base64
17+
import copy
1718
import errno
1819
import json
1920
import logging
@@ -670,7 +671,13 @@ def _generate_compose_file(self, command, additional_volumes=None, additional_en
670671
raise e
671672

672673
yaml_content = yaml.dump(content, default_flow_style=False)
673-
logger.info("docker compose file: \n%s", yaml_content)
674+
# Mask all environment vars for logging, could contain secrects.
675+
masked_content = copy.deepcopy(content)
676+
for _, service_data in masked_content["services"].items():
677+
service_data["environment"] = ["[Masked]" for _ in service_data["environment"]]
678+
679+
masked_content_for_logging = yaml.dump(masked_content, default_flow_style=False)
680+
logger.info("docker compose file: \n%s", masked_content_for_logging)
674681
with open(docker_compose_path, "w") as f:
675682
f.write(yaml_content)
676683

tests/unit/test_image.py

+8-2
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from botocore.credentials import Credentials
2020

2121
import base64
22+
import logging
2223
import json
2324
import os
2425
import subprocess
@@ -332,12 +333,14 @@ def test_check_output():
332333
@patch("sagemaker.local.data.get_data_source_instance")
333334
@patch("subprocess.Popen")
334335
def test_train(
335-
popen, get_data_source_instance, retrieve_artifacts, cleanup, tmpdir, sagemaker_session
336+
popen, get_data_source_instance, retrieve_artifacts, cleanup, tmpdir, sagemaker_session, caplog
336337
):
337338
data_source = Mock()
338339
data_source.get_root_dir.return_value = "foo"
339340
get_data_source_instance.return_value = data_source
340341

342+
caplog.set_level(logging.INFO)
343+
341344
directories = [str(tmpdir.mkdir("container-root")), str(tmpdir.mkdir("data"))]
342345
with patch(
343346
"sagemaker.local.image._SageMakerContainer._create_tmp_folder", side_effect=directories
@@ -388,6 +391,7 @@ def test_train(
388391

389392
retrieve_artifacts.assert_called_once()
390393
cleanup.assert_called_once()
394+
assert "[Masked]" in caplog.text
391395

392396

393397
@patch("sagemaker.local.local_session.LocalSession", Mock())
@@ -579,7 +583,8 @@ def test_container_does_not_enable_nvidia_docker_for_cpu_containers(sagemaker_se
579583
@patch("sagemaker.local.image._SageMakerContainer._prepare_serving_volumes", Mock(return_value=[]))
580584
@patch("shutil.copy", Mock())
581585
@patch("shutil.copytree", Mock())
582-
def test_serve(tmpdir, sagemaker_session):
586+
def test_serve(tmpdir, sagemaker_session, caplog):
587+
caplog.set_level(logging.INFO)
583588
with patch(
584589
"sagemaker.local.image._SageMakerContainer._create_tmp_folder",
585590
return_value=str(tmpdir.mkdir("container-root")),
@@ -601,6 +606,7 @@ def test_serve(tmpdir, sagemaker_session):
601606
for h in sagemaker_container.hosts:
602607
assert config["services"][h]["image"] == image
603608
assert config["services"][h]["command"] == "serve"
609+
assert "[Masked]" in caplog.text
604610

605611

606612
@patch("sagemaker.local.image._HostingContainer.run", Mock())

0 commit comments

Comments
 (0)