From e758f29ea6640bdc1c7f2f9c838f90b7724ae1b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nguy=E1=BB=85n=20Ng=E1=BB=8Dc=20B=C3=ADch=20Ch=C4=83m?= <248288664+Bichcham2003@users.noreply.github.com> Date: Mon, 11 May 2026 19:40:26 +0700 Subject: [PATCH] Redact CodeArtifact index from pip command logs --- .../core/utils/install_requirements.py | 35 +++++++++++++++++-- .../tests/unit/test_install_requirements.py | 28 ++++++++++++++- 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/utils/install_requirements.py b/sagemaker-core/src/sagemaker/core/utils/install_requirements.py index 9849ac3593..5df6e9ede9 100644 --- a/sagemaker-core/src/sagemaker/core/utils/install_requirements.py +++ b/sagemaker-core/src/sagemaker/core/utils/install_requirements.py @@ -40,6 +40,7 @@ import re import subprocess import sys +import tempfile logger = logging.getLogger(__name__) @@ -111,6 +112,20 @@ def _login_awscli(region, account, domain, repo): ) +def _write_pip_config(index): + """Write the authenticated package index to a temporary pip config file.""" + config = tempfile.NamedTemporaryFile( + mode="w", prefix="sagemaker-pip-", suffix=".conf", delete=False + ) + try: + config.write("[global]\n") + config.write(f"index-url = {index}\n") + finally: + config.close() + os.chmod(config.name, 0o600) + return config.name + + def configure_pip(auth_method=CodeArtifactAuthMethod.AUTO): """Configure pip for CodeArtifact if ``CA_REPOSITORY_ARN`` is set. @@ -183,10 +198,26 @@ def install_requirements( python_executable = python_executable or sys.executable pip_cmd = [python_executable, "-m", "pip", "install", "-r", requirements_file] index = configure_pip(auth_method=auth_method) + env = None + pip_config = None if index: - pip_cmd.extend(["-i", index]) + pip_config = _write_pip_config(index) + env = os.environ.copy() + env["PIP_CONFIG_FILE"] = pip_config logger.info("Running: %s", " ".join(pip_cmd)) - subprocess.check_call(pip_cmd) + try: + if env is not None: + subprocess.check_call(pip_cmd, env=env) + else: + subprocess.check_call(pip_cmd) + finally: + if pip_config: + try: + os.remove(pip_config) + except FileNotFoundError: + pass + except OSError as e: + logger.warning("Failed to remove temporary pip config file %s: %s", pip_config, e) def main(): diff --git a/sagemaker-core/tests/unit/test_install_requirements.py b/sagemaker-core/tests/unit/test_install_requirements.py index ac026d6e8c..8281057085 100644 --- a/sagemaker-core/tests/unit/test_install_requirements.py +++ b/sagemaker-core/tests/unit/test_install_requirements.py @@ -12,6 +12,8 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import +import logging +import os import subprocess import sys from unittest import mock @@ -187,10 +189,34 @@ def test_without_codeartifact(self): mock_call.assert_called_once_with(_pip_cmd()) def test_with_codeartifact_index(self): + captured = {} + + def fake_check_call(cmd, env=None): + captured["cmd"] = cmd + captured["env"] = env + with open(env["PIP_CONFIG_FILE"], "r") as config: + captured["config"] = config.read() + + with mock.patch(f"{_MODULE}.configure_pip", return_value=EXPECTED_INDEX): + with mock.patch("subprocess.check_call", side_effect=fake_check_call): + install_requirements("reqs.txt") + + assert captured["cmd"] == _pip_cmd() + assert captured["env"]["PIP_CONFIG_FILE"] + assert EXPECTED_INDEX in captured["config"] + assert not os.path.exists(captured["env"]["PIP_CONFIG_FILE"]) + + def test_codeartifact_index_not_logged_or_passed_in_argv(self, caplog): + caplog.set_level(logging.INFO, logger=_MODULE) with mock.patch(f"{_MODULE}.configure_pip", return_value=EXPECTED_INDEX): with mock.patch("subprocess.check_call") as mock_call: install_requirements("reqs.txt") - mock_call.assert_called_once_with(_pip_cmd("-i", EXPECTED_INDEX)) + + assert FAKE_TOKEN not in caplog.text + assert EXPECTED_INDEX not in caplog.text + argv = mock_call.call_args[0][0] + assert FAKE_TOKEN not in " ".join(argv) + assert EXPECTED_INDEX not in " ".join(argv) def test_with_cli_fallback_no_index_flag(self): with mock.patch(f"{_MODULE}.configure_pip", return_value=None):