Fix/aml 509849

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))

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,21 +330,42 @@ 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):
     """
     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 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 +392,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

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")

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."""
+        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"""

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))