diff --git a/src/sagemaker/local/image.py b/src/sagemaker/local/image.py index 07b4c015db..51b20a3bd8 100644 --- a/src/sagemaker/local/image.py +++ b/src/sagemaker/local/image.py @@ -14,6 +14,7 @@ from __future__ import absolute_import import base64 +import copy import errno import json import logging @@ -670,7 +671,13 @@ def _generate_compose_file(self, command, additional_volumes=None, additional_en raise e yaml_content = yaml.dump(content, default_flow_style=False) - logger.info("docker compose file: \n%s", yaml_content) + # Mask all environment vars for logging, could contain secrects. + masked_content = copy.deepcopy(content) + for _, service_data in masked_content["services"].items(): + service_data["environment"] = ["[Masked]" for _ in service_data["environment"]] + + masked_content_for_logging = yaml.dump(masked_content, default_flow_style=False) + logger.info("docker compose file: \n%s", masked_content_for_logging) with open(docker_compose_path, "w") as f: f.write(yaml_content) diff --git a/tests/unit/test_image.py b/tests/unit/test_image.py index 4d1a65309f..4568d0742c 100644 --- a/tests/unit/test_image.py +++ b/tests/unit/test_image.py @@ -19,6 +19,7 @@ from botocore.credentials import Credentials import base64 +import logging import json import os import subprocess @@ -332,12 +333,14 @@ def test_check_output(): @patch("sagemaker.local.data.get_data_source_instance") @patch("subprocess.Popen") def test_train( - popen, get_data_source_instance, retrieve_artifacts, cleanup, tmpdir, sagemaker_session + popen, get_data_source_instance, retrieve_artifacts, cleanup, tmpdir, sagemaker_session, caplog ): data_source = Mock() data_source.get_root_dir.return_value = "foo" get_data_source_instance.return_value = data_source + caplog.set_level(logging.INFO) + directories = [str(tmpdir.mkdir("container-root")), str(tmpdir.mkdir("data"))] with patch( "sagemaker.local.image._SageMakerContainer._create_tmp_folder", side_effect=directories @@ -388,6 +391,7 @@ def test_train( retrieve_artifacts.assert_called_once() cleanup.assert_called_once() + assert "[Masked]" in caplog.text @patch("sagemaker.local.local_session.LocalSession", Mock()) @@ -579,7 +583,8 @@ def test_container_does_not_enable_nvidia_docker_for_cpu_containers(sagemaker_se @patch("sagemaker.local.image._SageMakerContainer._prepare_serving_volumes", Mock(return_value=[])) @patch("shutil.copy", Mock()) @patch("shutil.copytree", Mock()) -def test_serve(tmpdir, sagemaker_session): +def test_serve(tmpdir, sagemaker_session, caplog): + caplog.set_level(logging.INFO) with patch( "sagemaker.local.image._SageMakerContainer._create_tmp_folder", return_value=str(tmpdir.mkdir("container-root")), @@ -601,6 +606,7 @@ def test_serve(tmpdir, sagemaker_session): for h in sagemaker_container.hosts: assert config["services"][h]["image"] == image assert config["services"][h]["command"] == "serve" + assert "[Masked]" in caplog.text @patch("sagemaker.local.image._HostingContainer.run", Mock())