From 231bae614082df594e4eb5d4888513bb02d4a169 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 28 Jul 2026 18:29:05 -0700 Subject: [PATCH 1/6] fix(sagemaker-core): remove dev-only endpoint override and fix client singleton pinning Remove the temporary SAGEMAKER_ENDPOINT / SAGEMAKER_RUNTIME_ENDPOINT / SAGEMAKER_STAGE overrides and the custom botocore service-model loader from SageMakerClient. These were internal beta/gamma testing shims (marked "# TODO: Remove post-launch") added while the generic Job APIs were not yet in public botocore. Those APIs (CreateJob/DescribeJob/ListJobs) now ship in public botocore, so the standard session.client("sagemaker", ...) path is used and the env-driven endpoint redirection is no longer read by the SDK. Also fix SingletonMeta so SageMakerClient is keyed by (session, region) instead of by class alone. Previously the first instance was pinned process-wide, so a later call with a different session/region silently reused the original client. Parameterless SageMakerClient() calls still share a single instance. Harden the container-driver env logger to mask values that look like AWS credential material (access key IDs, session tokens) even when the key name contains none of SENSITIVE_KEYWORDS. --- .../container_drivers/scripts/environment.py | 16 ++++ .../src/sagemaker/core/utils/utils.py | 74 +++++++------------ .../unit/modules/train/test_environment.py | 20 +++++ 3 files changed, 62 insertions(+), 48 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py b/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py index 897b1f8af4..33f7318573 100644 --- a/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py +++ b/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py @@ -18,6 +18,7 @@ import subprocess import json import os +import re import sys from pathlib import Path import logging @@ -62,6 +63,19 @@ SENSITIVE_KEYWORDS = ["SECRET", "PASSWORD", "KEY", "TOKEN", "PRIVATE", "CREDS", "CREDENTIALS"] HIDDEN_VALUE = "******" +SENSITIVE_VALUE_PATTERNS = [ + re.compile(r"\b(?:AKIA|ASIA|AROA|AIDA|AGPA|ANPA|ANVA|ASCA)[A-Z0-9]{16}\b"), + re.compile(r"(?i)\baws_session_token\b"), + re.compile(r"(?i)x-amz-security-token"), +] + + +def _value_is_sensitive(value) -> bool: + """Return True if a value looks like credential material, regardless of key name.""" + if not isinstance(value, str): + return False + return any(pattern.search(value) for pattern in SENSITIVE_VALUE_PATTERNS) + def num_cpus() -> int: """Return the number of CPUs available in the current container. @@ -260,6 +274,8 @@ def log_key_value(key: str, value: str): """Log a key-value pair, masking sensitive values if necessary.""" if any(keyword.lower() in key.lower() for keyword in SENSITIVE_KEYWORDS): logger.info("%s=%s", key, HIDDEN_VALUE) + elif _value_is_sensitive(value): + logger.info("%s=%s", key, HIDDEN_VALUE) elif isinstance(value, dict): masked_value = mask_sensitive_info(value) logger.info("%s=%s", key, json.dumps(masked_value)) diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index 5ac288eab7..eaca183ed1 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -318,6 +318,9 @@ def __bool__(self): class SingletonMeta(type): """ Singleton metaclass. Ensures that a single instance of a class using this metaclass is created. + + A class may define a ``_singleton_key(*args, **kwargs)`` classmethod returning + a hashable key to cache instances per ``(class, key)`` instead of per class. """ _instances = {} @@ -327,10 +330,14 @@ def __call__(cls, *args, **kwargs): Overrides the call method to return an existing instance of the class if it exists, or create a new one if it doesn't. """ - if cls not in cls._instances: + key = cls + key_fn = getattr(cls, "_singleton_key", None) + if key_fn is not None: + key = (cls, key_fn(*args, **kwargs)) + if key not in cls._instances: instance = super().__call__(*args, **kwargs) - cls._instances[cls] = instance - return cls._instances[cls] + cls._instances[key] = instance + return cls._instances[key] class SageMakerClient(metaclass=SingletonMeta): @@ -338,10 +345,22 @@ class SageMakerClient(metaclass=SingletonMeta): A singleton class for creating a SageMaker client. """ + @staticmethod + def _singleton_key(session: Session = None, region_name: str = None, config: Config = None): + """Cache instances per (session, region).""" + session_key = id(session) if session is not None else None + return (session_key, region_name) + @classmethod def reset(cls): - """Reset the singleton instance so the next call re-reads env vars.""" - SingletonMeta._instances.pop(cls, None) + """Reset cached singleton instances for this class.""" + keys = [ + k + for k in SingletonMeta._instances + if k == cls or (isinstance(k, tuple) and k[0] is cls) + ] + for key in keys: + SingletonMeta._instances.pop(key, None) def __init__( self, @@ -368,53 +387,12 @@ def __init__( self.config = Config(user_agent_extra=get_user_agent_extra_suffix()) self.session = session self.region_name = region_name - # Read region from environment variable, default to us-west-2 - import os - env_region = os.environ.get('SAGEMAKER_REGION', region_name) - env_stage = os.environ.get('SAGEMAKER_STAGE', 'prod') # default to gamma - logger.info(f"Runs on sagemaker {env_stage}, region:{env_region}") - - endpoint_url = os.environ.get('SAGEMAKER_ENDPOINT') - runtime_endpoint_url = os.environ.get('SAGEMAKER_RUNTIME_ENDPOINT') - - # TODO: Remove post-launch. This loads a custom botocore service model - # (from the 'sample/' directory) that includes pre-GA Job APIs - # (CreateJob, DescribeJob, ListJobs, etc.) not yet in the public - # botocore release. Once these APIs ship in the official botocore - # service-2.json, this custom loader is unnecessary and the standard - # session.client("sagemaker", endpoint_url=...) path will work. - # - # NOTE: The custom data loader is registered on the caller-provided - # session's underlying botocore session so that the "sagemaker" - # control-plane client keeps the caller's credentials/profile. Using a - # fresh botocore.session.get_session() here would silently fall back to - # the default AWS profile and break cross-account calls (issue #6069). - import botocore.loaders - import pathlib - - # Look for bundled service model inside the package first, - # fall back to source-tree layout (sample/ as sibling of src/) - _bundled = pathlib.Path(__file__).resolve().parent.parent / "data" / "sample" - _source_tree = pathlib.Path(__file__).resolve().parent.parent.parent.parent.parent / "sample" - sample_model_dir = str(_bundled if _bundled.exists() else _source_tree) - loader = botocore.loaders.Loader( - extra_search_paths=[sample_model_dir], - include_default_search_paths=True, - ) - # boto3.Session exposes its underlying botocore session as `_session`. - session._session.register_component('data_loader', loader) self.sagemaker_client = session.client( - "sagemaker", - region_name=env_region, - endpoint_url=endpoint_url, - config=self.config, + "sagemaker", region_name, config=self.config ) - self.sagemaker_runtime_client = session.client( - "sagemaker-runtime", region_name, - endpoint_url=runtime_endpoint_url, - config=self.config, + "sagemaker-runtime", region_name, config=self.config ) self.sagemaker_featurestore_runtime_client = session.client( "sagemaker-featurestore-runtime", region_name, config=self.config diff --git a/sagemaker-core/tests/unit/modules/train/test_environment.py b/sagemaker-core/tests/unit/modules/train/test_environment.py index e80eeef005..6b02e4e759 100644 --- a/sagemaker-core/tests/unit/modules/train/test_environment.py +++ b/sagemaker-core/tests/unit/modules/train/test_environment.py @@ -164,6 +164,26 @@ def test_log_key_value_sensitive(self, mock_logger): call_args = mock_logger.info.call_args[0] assert "******" in str(call_args) + @patch("sagemaker.core.modules.train.container_drivers.scripts.environment.logger") + def test_log_key_value_credential_shaped_value_under_innocuous_key(self, mock_logger): + """Values that look like AWS credentials are masked even under a non-sensitive key name.""" + log_key_value("MY_INNOCENT_VAR", "ASIAABCDEFGHIJKLMNOP") + + mock_logger.info.assert_called_once() + call_args = mock_logger.info.call_args[0] + assert "******" in str(call_args) + assert "ASIAABCDEFGHIJKLMNOP" not in str(call_args) + + @patch("sagemaker.core.modules.train.container_drivers.scripts.environment.logger") + def test_log_key_value_non_secret_url_visible(self, mock_logger): + """A plain endpoint URL is not a secret and remains visible.""" + log_key_value("SAGEMAKER_ENDPOINT", "https://example.com:9999") + + mock_logger.info.assert_called_once() + call_args = mock_logger.info.call_args[0] + assert "https://example.com:9999" in str(call_args) + assert "******" not in str(call_args) + @patch("sagemaker.core.modules.train.container_drivers.scripts.environment.logger") def test_log_key_value_dict(self, mock_logger): """Test log_key_value with dictionary value""" From ad7a9c43c7c333cfc0aa84058e98d982fe273710 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 28 Jul 2026 18:29:14 -0700 Subject: [PATCH 2/6] fix(sagemaker-train): mask credential-shaped values in container env logging The startup env logger masked only by key name, so a sensitive value under a key name without a SENSITIVE_KEYWORDS substring was logged verbatim. Add value-based detection for AWS credential material (access key IDs, session tokens) so such values are masked regardless of key name. --- .../container_drivers/scripts/environment.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py b/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py index 897b1f8af4..33f7318573 100644 --- a/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py +++ b/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py @@ -18,6 +18,7 @@ import subprocess import json import os +import re import sys from pathlib import Path import logging @@ -62,6 +63,19 @@ SENSITIVE_KEYWORDS = ["SECRET", "PASSWORD", "KEY", "TOKEN", "PRIVATE", "CREDS", "CREDENTIALS"] HIDDEN_VALUE = "******" +SENSITIVE_VALUE_PATTERNS = [ + re.compile(r"\b(?:AKIA|ASIA|AROA|AIDA|AGPA|ANPA|ANVA|ASCA)[A-Z0-9]{16}\b"), + re.compile(r"(?i)\baws_session_token\b"), + re.compile(r"(?i)x-amz-security-token"), +] + + +def _value_is_sensitive(value) -> bool: + """Return True if a value looks like credential material, regardless of key name.""" + if not isinstance(value, str): + return False + return any(pattern.search(value) for pattern in SENSITIVE_VALUE_PATTERNS) + def num_cpus() -> int: """Return the number of CPUs available in the current container. @@ -260,6 +274,8 @@ def log_key_value(key: str, value: str): """Log a key-value pair, masking sensitive values if necessary.""" if any(keyword.lower() in key.lower() for keyword in SENSITIVE_KEYWORDS): logger.info("%s=%s", key, HIDDEN_VALUE) + elif _value_is_sensitive(value): + logger.info("%s=%s", key, HIDDEN_VALUE) elif isinstance(value, dict): masked_value = mask_sensitive_info(value) logger.info("%s=%s", key, json.dumps(masked_value)) From 18d1e7ae786cc7809fbcea45c80e059e242455f1 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Wed, 29 Jul 2026 08:48:08 -0700 Subject: [PATCH 3/6] fix(sagemaker-core): address review feedback on client singleton Make SageMakerClient._singleton_key tolerant of extra constructor arguments by accepting *args/**kwargs, so callers passing additional keywords (such as the legacy service_name=) no longer raise a TypeError during key computation on cache-hit paths. Document why id(session) is a stable cache key: the cached instance retains the session reference, preventing garbage collection and id reuse while the entry lives. Add unit tests for the singleton keying and reset(): distinct instance per (session, region), reuse across repeated no-arg calls, reuse for an identical session/region, and reset() clearing all keyed entries for the class. Rework the env-logging test assertion to compare the logged arguments exactly instead of a URL substring check, removing the CodeQL incomplete-URL-sanitization finding on the test file. --- .../src/sagemaker/core/utils/utils.py | 9 +++- .../tests/unit/generated/test_utils.py | 54 +++++++++++++++++++ .../unit/modules/train/test_environment.py | 6 +-- 3 files changed, 64 insertions(+), 5 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index eaca183ed1..957a93a8ec 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -346,8 +346,13 @@ class SageMakerClient(metaclass=SingletonMeta): """ @staticmethod - def _singleton_key(session: Session = None, region_name: str = None, config: Config = None): - """Cache instances per (session, region).""" + def _singleton_key(session: Session = None, region_name: str = None, *args, **kwargs): + """Cache instances per (session, region). + + id(session) is stable because the cached instance holds a reference to + the session, so the object cannot be garbage-collected while the entry + lives. + """ session_key = id(session) if session is not None else None return (session_key, region_name) diff --git a/sagemaker-core/tests/unit/generated/test_utils.py b/sagemaker-core/tests/unit/generated/test_utils.py index 0928f09b68..fe704ef007 100644 --- a/sagemaker-core/tests/unit/generated/test_utils.py +++ b/sagemaker-core/tests/unit/generated/test_utils.py @@ -432,3 +432,57 @@ def test_unassigned_in_conditional(self): # Should work with not assert not u + + +class TestSageMakerClientSingleton: + """Tests for SageMakerClient singleton keying and reset.""" + + def setup_method(self): + SageMakerClient.reset() + + def teardown_method(self): + SageMakerClient.reset() + + def test_parameterless_calls_reuse_same_instance(self): + assert SageMakerClient() is SageMakerClient() + + def test_distinct_session_produces_distinct_instance(self): + import boto3 + + c1 = SageMakerClient(session=boto3.Session(region_name="us-west-2"), region_name="us-west-2") + c2 = SageMakerClient(session=boto3.Session(region_name="us-west-2"), region_name="us-west-2") + assert c1 is not c2 + + def test_distinct_region_produces_distinct_instance(self): + import boto3 + + session = boto3.Session(region_name="us-west-2") + c1 = SageMakerClient(session=session, region_name="us-west-2") + c2 = SageMakerClient(session=session, region_name="eu-west-1") + assert c1 is not c2 + + def test_same_session_and_region_reuses_instance(self): + import boto3 + + session = boto3.Session(region_name="us-west-2") + c1 = SageMakerClient(session=session, region_name="us-west-2") + c2 = SageMakerClient(session=session, region_name="us-west-2") + assert c1 is c2 + + def test_reset_clears_all_keyed_instances(self): + import boto3 + + SageMakerClient() + SageMakerClient(session=boto3.Session(region_name="us-west-2"), region_name="us-west-2") + SageMakerClient.reset() + remaining = [ + k + for k in SingletonMeta._instances + if k is SageMakerClient or (isinstance(k, tuple) and k[0] is SageMakerClient) + ] + assert remaining == [] + + def test_singleton_key_tolerates_extra_kwargs(self): + # Extra kwargs (e.g. service_name) must not break key computation. + key = SageMakerClient._singleton_key(session=None, region_name="us-west-2", service_name="sagemaker") + assert key == (None, "us-west-2") diff --git a/sagemaker-core/tests/unit/modules/train/test_environment.py b/sagemaker-core/tests/unit/modules/train/test_environment.py index 6b02e4e759..18ca3ba92a 100644 --- a/sagemaker-core/tests/unit/modules/train/test_environment.py +++ b/sagemaker-core/tests/unit/modules/train/test_environment.py @@ -177,12 +177,12 @@ def test_log_key_value_credential_shaped_value_under_innocuous_key(self, mock_lo @patch("sagemaker.core.modules.train.container_drivers.scripts.environment.logger") def test_log_key_value_non_secret_url_visible(self, mock_logger): """A plain endpoint URL is not a secret and remains visible.""" - log_key_value("SAGEMAKER_ENDPOINT", "https://example.com:9999") + endpoint = "https://example.com:9999" + log_key_value("SAGEMAKER_ENDPOINT", endpoint) mock_logger.info.assert_called_once() call_args = mock_logger.info.call_args[0] - assert "https://example.com:9999" in str(call_args) - assert "******" not in str(call_args) + assert call_args == ("%s=%s", "SAGEMAKER_ENDPOINT", endpoint) @patch("sagemaker.core.modules.train.container_drivers.scripts.environment.logger") def test_log_key_value_dict(self, mock_logger): From e9e835ee349ea2aa239d219762e95b449c87a858 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Wed, 29 Jul 2026 13:45:55 -0700 Subject: [PATCH 4/6] revert(sagemaker-core): keep SageMakerClient singleton keyed by class Revert the (session, region) singleton keying introduced earlier in this PR. The security concern behind Finding 3 (a poisoned endpoint being pinned) is already resolved by removing the SAGEMAKER_ENDPOINT override in Finding 1, so there is no attacker-controlled endpoint left to pin. Per team design, the SDK operates with a single session/region per process, so a class-keyed singleton is the intended behavior. Restores the original SingletonMeta, SageMakerClient.reset(), and drops the singleton keying unit tests. Finding 1 (endpoint override removal) and Finding 5 (env-log value masking) are unchanged. --- .../src/sagemaker/core/utils/utils.py | 34 ++---------- .../tests/unit/generated/test_utils.py | 54 ------------------- 2 files changed, 5 insertions(+), 83 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index 957a93a8ec..9f916902f4 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -318,9 +318,6 @@ def __bool__(self): class SingletonMeta(type): """ Singleton metaclass. Ensures that a single instance of a class using this metaclass is created. - - A class may define a ``_singleton_key(*args, **kwargs)`` classmethod returning - a hashable key to cache instances per ``(class, key)`` instead of per class. """ _instances = {} @@ -330,14 +327,10 @@ def __call__(cls, *args, **kwargs): Overrides the call method to return an existing instance of the class if it exists, or create a new one if it doesn't. """ - key = cls - key_fn = getattr(cls, "_singleton_key", None) - if key_fn is not None: - key = (cls, key_fn(*args, **kwargs)) - if key not in cls._instances: + if cls not in cls._instances: instance = super().__call__(*args, **kwargs) - cls._instances[key] = instance - return cls._instances[key] + cls._instances[cls] = instance + return cls._instances[cls] class SageMakerClient(metaclass=SingletonMeta): @@ -345,27 +338,10 @@ class SageMakerClient(metaclass=SingletonMeta): A singleton class for creating a SageMaker client. """ - @staticmethod - def _singleton_key(session: Session = None, region_name: str = None, *args, **kwargs): - """Cache instances per (session, region). - - id(session) is stable because the cached instance holds a reference to - the session, so the object cannot be garbage-collected while the entry - lives. - """ - session_key = id(session) if session is not None else None - return (session_key, region_name) - @classmethod def reset(cls): - """Reset cached singleton instances for this class.""" - keys = [ - k - for k in SingletonMeta._instances - if k == cls or (isinstance(k, tuple) and k[0] is cls) - ] - for key in keys: - SingletonMeta._instances.pop(key, None) + """Reset the singleton instance.""" + SingletonMeta._instances.pop(cls, None) def __init__( self, diff --git a/sagemaker-core/tests/unit/generated/test_utils.py b/sagemaker-core/tests/unit/generated/test_utils.py index fe704ef007..0928f09b68 100644 --- a/sagemaker-core/tests/unit/generated/test_utils.py +++ b/sagemaker-core/tests/unit/generated/test_utils.py @@ -432,57 +432,3 @@ def test_unassigned_in_conditional(self): # Should work with not assert not u - - -class TestSageMakerClientSingleton: - """Tests for SageMakerClient singleton keying and reset.""" - - def setup_method(self): - SageMakerClient.reset() - - def teardown_method(self): - SageMakerClient.reset() - - def test_parameterless_calls_reuse_same_instance(self): - assert SageMakerClient() is SageMakerClient() - - def test_distinct_session_produces_distinct_instance(self): - import boto3 - - c1 = SageMakerClient(session=boto3.Session(region_name="us-west-2"), region_name="us-west-2") - c2 = SageMakerClient(session=boto3.Session(region_name="us-west-2"), region_name="us-west-2") - assert c1 is not c2 - - def test_distinct_region_produces_distinct_instance(self): - import boto3 - - session = boto3.Session(region_name="us-west-2") - c1 = SageMakerClient(session=session, region_name="us-west-2") - c2 = SageMakerClient(session=session, region_name="eu-west-1") - assert c1 is not c2 - - def test_same_session_and_region_reuses_instance(self): - import boto3 - - session = boto3.Session(region_name="us-west-2") - c1 = SageMakerClient(session=session, region_name="us-west-2") - c2 = SageMakerClient(session=session, region_name="us-west-2") - assert c1 is c2 - - def test_reset_clears_all_keyed_instances(self): - import boto3 - - SageMakerClient() - SageMakerClient(session=boto3.Session(region_name="us-west-2"), region_name="us-west-2") - SageMakerClient.reset() - remaining = [ - k - for k in SingletonMeta._instances - if k is SageMakerClient or (isinstance(k, tuple) and k[0] is SageMakerClient) - ] - assert remaining == [] - - def test_singleton_key_tolerates_extra_kwargs(self): - # Extra kwargs (e.g. service_name) must not break key computation. - key = SageMakerClient._singleton_key(session=None, region_name="us-west-2", service_name="sagemaker") - assert key == (None, "us-west-2") From 5225d88324154898b94d50233f179b0232d7fc74 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Wed, 29 Jul 2026 14:55:33 -0700 Subject: [PATCH 5/6] fix(sagemaker-core): bump boto3 floor to >=1.43.20 for Job APIs Removing the custom sample-model loader makes the control-plane client a stock session.client("sagemaker"), so the generic Job APIs must be present in the resolved botocore. CreateJob/DescribeJob/ListJobs first ship in botocore 1.43.20; the previous floor (boto3>=1.42.2) allowed botocore versions without them, which would regress Job.create()/get()/get_all() at runtime. Raise the pin to boto3>=1.43.20 (which requires botocore>=1.43.20) to guarantee the Job APIs are available. --- sagemaker-core/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sagemaker-core/pyproject.toml b/sagemaker-core/pyproject.toml index 7129e9fe98..accbfc9816 100644 --- a/sagemaker-core/pyproject.toml +++ b/sagemaker-core/pyproject.toml @@ -12,7 +12,7 @@ authors = [ readme = "README.rst" dependencies = [ # Add your dependencies here (Include lower and upper bounds as applicable) - "boto3>=1.42.2,<2.0.0", + "boto3>=1.43.20,<2.0.0", "pydantic>=2.0.0,<3.0.0", "PyYAML>=6.0, <7.0", "jsonschema<5.0.0", From b5885f21b0368acd3dfa636c1abbd9c5551b87b2 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Wed, 29 Jul 2026 16:02:25 -0700 Subject: [PATCH 6/6] revert(sagemaker-train): drop container env-log value masking Remove the value-based masking added earlier in this PR for the container -driver env logger. The behavior it targeted was raised as a forensic signal rather than a vulnerability, and comprehensive credential masking in logs is handled by the platform-side CloudWatch Logs data protection policies. Restores environment.py (sagemaker-train and the sagemaker-core mirror) and its tests to their original state. --- .../container_drivers/scripts/environment.py | 16 --------------- .../unit/modules/train/test_environment.py | 20 ------------------- .../container_drivers/scripts/environment.py | 16 --------------- 3 files changed, 52 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py b/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py index 33f7318573..897b1f8af4 100644 --- a/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py +++ b/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py @@ -18,7 +18,6 @@ import subprocess import json import os -import re import sys from pathlib import Path import logging @@ -63,19 +62,6 @@ SENSITIVE_KEYWORDS = ["SECRET", "PASSWORD", "KEY", "TOKEN", "PRIVATE", "CREDS", "CREDENTIALS"] HIDDEN_VALUE = "******" -SENSITIVE_VALUE_PATTERNS = [ - re.compile(r"\b(?:AKIA|ASIA|AROA|AIDA|AGPA|ANPA|ANVA|ASCA)[A-Z0-9]{16}\b"), - re.compile(r"(?i)\baws_session_token\b"), - re.compile(r"(?i)x-amz-security-token"), -] - - -def _value_is_sensitive(value) -> bool: - """Return True if a value looks like credential material, regardless of key name.""" - if not isinstance(value, str): - return False - return any(pattern.search(value) for pattern in SENSITIVE_VALUE_PATTERNS) - def num_cpus() -> int: """Return the number of CPUs available in the current container. @@ -274,8 +260,6 @@ def log_key_value(key: str, value: str): """Log a key-value pair, masking sensitive values if necessary.""" if any(keyword.lower() in key.lower() for keyword in SENSITIVE_KEYWORDS): logger.info("%s=%s", key, HIDDEN_VALUE) - elif _value_is_sensitive(value): - logger.info("%s=%s", key, HIDDEN_VALUE) elif isinstance(value, dict): masked_value = mask_sensitive_info(value) logger.info("%s=%s", key, json.dumps(masked_value)) diff --git a/sagemaker-core/tests/unit/modules/train/test_environment.py b/sagemaker-core/tests/unit/modules/train/test_environment.py index 18ca3ba92a..e80eeef005 100644 --- a/sagemaker-core/tests/unit/modules/train/test_environment.py +++ b/sagemaker-core/tests/unit/modules/train/test_environment.py @@ -164,26 +164,6 @@ def test_log_key_value_sensitive(self, mock_logger): call_args = mock_logger.info.call_args[0] assert "******" in str(call_args) - @patch("sagemaker.core.modules.train.container_drivers.scripts.environment.logger") - def test_log_key_value_credential_shaped_value_under_innocuous_key(self, mock_logger): - """Values that look like AWS credentials are masked even under a non-sensitive key name.""" - log_key_value("MY_INNOCENT_VAR", "ASIAABCDEFGHIJKLMNOP") - - mock_logger.info.assert_called_once() - call_args = mock_logger.info.call_args[0] - assert "******" in str(call_args) - assert "ASIAABCDEFGHIJKLMNOP" not in str(call_args) - - @patch("sagemaker.core.modules.train.container_drivers.scripts.environment.logger") - def test_log_key_value_non_secret_url_visible(self, mock_logger): - """A plain endpoint URL is not a secret and remains visible.""" - endpoint = "https://example.com:9999" - log_key_value("SAGEMAKER_ENDPOINT", endpoint) - - mock_logger.info.assert_called_once() - call_args = mock_logger.info.call_args[0] - assert call_args == ("%s=%s", "SAGEMAKER_ENDPOINT", endpoint) - @patch("sagemaker.core.modules.train.container_drivers.scripts.environment.logger") def test_log_key_value_dict(self, mock_logger): """Test log_key_value with dictionary value""" diff --git a/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py b/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py index 33f7318573..897b1f8af4 100644 --- a/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py +++ b/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py @@ -18,7 +18,6 @@ import subprocess import json import os -import re import sys from pathlib import Path import logging @@ -63,19 +62,6 @@ SENSITIVE_KEYWORDS = ["SECRET", "PASSWORD", "KEY", "TOKEN", "PRIVATE", "CREDS", "CREDENTIALS"] HIDDEN_VALUE = "******" -SENSITIVE_VALUE_PATTERNS = [ - re.compile(r"\b(?:AKIA|ASIA|AROA|AIDA|AGPA|ANPA|ANVA|ASCA)[A-Z0-9]{16}\b"), - re.compile(r"(?i)\baws_session_token\b"), - re.compile(r"(?i)x-amz-security-token"), -] - - -def _value_is_sensitive(value) -> bool: - """Return True if a value looks like credential material, regardless of key name.""" - if not isinstance(value, str): - return False - return any(pattern.search(value) for pattern in SENSITIVE_VALUE_PATTERNS) - def num_cpus() -> int: """Return the number of CPUs available in the current container. @@ -274,8 +260,6 @@ def log_key_value(key: str, value: str): """Log a key-value pair, masking sensitive values if necessary.""" if any(keyword.lower() in key.lower() for keyword in SENSITIVE_KEYWORDS): logger.info("%s=%s", key, HIDDEN_VALUE) - elif _value_is_sensitive(value): - logger.info("%s=%s", key, HIDDEN_VALUE) elif isinstance(value, dict): masked_value = mask_sensitive_info(value) logger.info("%s=%s", key, json.dumps(masked_value))