From 8010a45414752795ece67a59b669f5a64ebd915b Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Mon, 14 Sep 2026 14:35:10 -0700 Subject: [PATCH] fix: SageMakerClient honors the passed boto3 session and config SageMakerClient was a process-wide singleton: whichever instance was created first was returned to every later caller, and the session passed on those later calls was silently ignored. Any code path that created a client with default-chain credentials before the user's own session was seen (bare SageMakerClient() calls, a second profile, an assumed role) then signed every control-plane call with the wrong identity. This is the residual half of #5986 and the cause of the cross-account "RoleArn: Cross-account pass role is not allowed" failure in #6069. The config argument was also computed and then discarded, so caller timeouts and retry settings never reached the boto3 clients. Replace the singleton with a cache keyed on the constructor arguments. An explicit session (or region/config) always yields a client built from exactly those arguments; SageMakerClient() with no arguments keeps returning the process default, and the first instance created still becomes that default, so the configure-once pattern used by generated stop() methods and existing callers is unchanged. Keyed entries are a small LRU (32) so a service minting a session per request does not retain every session and its four boto3 clients; the default entry is never evicted. Construction is serialized with a lock so concurrent first callers share one instance. The caller's config is merged into the client config with the SDK user-agent suffix appended rather than replacing it. SingletonMeta is left in place for CloudWatchLogsClient. Fixes #5986 Fixes #6069 --- X-AI-Prompt: Fix the SageMakerClient singleton so an explicitly passed boto3 session is honored (#5986/#6069) X-AI-Tool: Kiro --- .../src/sagemaker/core/utils/utils.py | 98 +++++++- .../tests/unit/test_sagemaker_client.py | 230 ++++++++++++++++++ 2 files changed, 320 insertions(+), 8 deletions(-) create mode 100644 sagemaker-core/tests/unit/test_sagemaker_client.py diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index 9f916902f4..d662408821 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -16,6 +16,8 @@ import os import re import subprocess +import threading +from collections import OrderedDict from boto3.session import Session from botocore.config import Config @@ -333,15 +335,92 @@ def __call__(cls, *args, **kwargs): return cls._instances[cls] -class SageMakerClient(metaclass=SingletonMeta): +class _ClientCacheMeta(type): """ - A singleton class for creating a SageMaker client. + Metaclass that caches instances per constructor arguments. + + A call with no arguments returns the process-wide default instance, creating it + from the default credential chain if none exists yet. A call with an explicit + ``session``, ``region_name`` or ``config`` returns an instance built from exactly + those arguments, cached so repeated calls with the same objects are cheap. The + first instance ever created also becomes the default, so code that configures + ``SageMakerClient(session=...)`` once up front and then relies on bare + ``SageMakerClient()`` calls keeps working. + + The keyed entries form a small LRU: once ``_MAX_KEYED_ENTRIES`` distinct + argument sets have been seen, the least recently used one is dropped, so a + long-lived process that mints a fresh session per unit of work does not retain + every session and its four boto3 clients forever. The default entry is never + evicted. Construction is serialized by a lock so concurrent first callers for + the same key share one instance. + + This replaces a plain singleton, which returned whichever instance was created + first and silently ignored the session passed on every later call. + """ + + _DEFAULT_KEY = "default" + _MAX_KEYED_ENTRIES = 32 + _instances: Dict[type, "OrderedDict[Any, Any]"] = {} + _lock = threading.RLock() + + def __call__(cls, *args, **kwargs): + """Return a cached instance for these arguments, creating it if needed.""" + key = cls._cache_key(*args, **kwargs) + with cls._lock: + cache = cls._instances.setdefault(cls, OrderedDict()) + instance = cache.get(key) + if instance is not None: + if key != cls._DEFAULT_KEY: + cache.move_to_end(key) + return instance + instance = super().__call__(*args, **kwargs) + cache[key] = instance + # The first client created in the process becomes the default used by + # argument-less calls, matching the previous singleton behaviour. + cache.setdefault(cls._DEFAULT_KEY, instance) + _ClientCacheMeta._evict(cache, cls._DEFAULT_KEY, cls._MAX_KEYED_ENTRIES) + return instance + + @staticmethod + def _evict(cache, default_key, max_keyed_entries): + """Drop least recently used keyed entries beyond the cap; keep the default.""" + keyed = [k for k in cache if k != default_key] + for key in keyed[: max(0, len(keyed) - max_keyed_entries)]: + del cache[key] + + +class SageMakerClient(metaclass=_ClientCacheMeta): + """ + Cached factory for SageMaker boto3 clients. + + Clients are cached per (session, region_name, config). Passing an explicit boto3 + ``session`` always yields clients signed with that session's credentials, even if + a client for a different session was created earlier in the process. Calling + ``SageMakerClient()`` with no arguments returns the process default. """ + @classmethod + def _cache_key(cls, session: Session = None, region_name: str = None, config: Config = None): + """Build the cache key for a constructor call. + + Sessions and configs are keyed by identity: two distinct boto3 sessions must + never share a client even if they look alike, because credentials live on the + session object. The cached instance keeps references to both, so the ids + cannot be recycled while the entry is alive. + """ + if session is None and region_name is None and config is None: + return cls._DEFAULT_KEY + return ( + id(session) if session is not None else None, + region_name, + id(config) if config is not None else None, + ) + @classmethod def reset(cls): - """Reset the singleton instance.""" - SingletonMeta._instances.pop(cls, None) + """Drop every cached instance, including the default.""" + with _ClientCacheMeta._lock: + _ClientCacheMeta._instances.pop(cls, None) def __init__( self, @@ -365,13 +444,16 @@ def __init__( logger.debug("No config provided. Using default config.") config = Config(retries={"max_attempts": 10, "mode": "standard"}) - self.config = Config(user_agent_extra=get_user_agent_extra_suffix()) + # Keep the caller's config object alive: the cache key uses its identity. + self._base_config = config + user_agent_extra = get_user_agent_extra_suffix() + if config.user_agent_extra: + user_agent_extra = f"{config.user_agent_extra} {user_agent_extra}" + self.config = config.merge(Config(user_agent_extra=user_agent_extra)) self.session = session self.region_name = region_name - self.sagemaker_client = session.client( - "sagemaker", region_name, config=self.config - ) + self.sagemaker_client = session.client("sagemaker", region_name, config=self.config) self.sagemaker_runtime_client = session.client( "sagemaker-runtime", region_name, config=self.config ) diff --git a/sagemaker-core/tests/unit/test_sagemaker_client.py b/sagemaker-core/tests/unit/test_sagemaker_client.py new file mode 100644 index 0000000000..8e9f92c66b --- /dev/null +++ b/sagemaker-core/tests/unit/test_sagemaker_client.py @@ -0,0 +1,230 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Tests for SageMakerClient caching and session handling. + +Regression coverage for aws/sagemaker-python-sdk#5986 and #6069: a client built +for one boto3 session must never be handed to a caller that passed a different +session, and the ``config`` argument must actually reach the boto3 clients. +""" + +from __future__ import absolute_import + +import boto3 +import pytest +from botocore.config import Config + +from sagemaker.core.utils.utils import SageMakerClient + +REGION_A = "us-west-2" +REGION_B = "us-east-1" + + +def _session(access_key, region): + """Offline boto3 session carrying distinguishable static credentials.""" + return boto3.Session( + aws_access_key_id=access_key, + aws_secret_access_key="secret-" + access_key, + region_name=region, + ) + + +def _access_key(boto_client): + return boto_client._request_signer._credentials.access_key + + +@pytest.fixture(autouse=True) +def _isolated_cache(): + SageMakerClient.reset() + yield + SageMakerClient.reset() + + +def test_explicit_session_is_honored_after_another_client_exists(): + """The second session must not receive the first session's client (#5986, #6069).""" + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + session_b = _session("AKIASECONDSESSION000", REGION_B) + + client_a = SageMakerClient(session=session_a) + client_b = SageMakerClient(session=session_b) + + assert client_a is not client_b + assert client_b.session is session_b + assert client_b.region_name == REGION_B + for service in ( + "sagemaker", + "sagemaker-runtime", + "sagemaker-featurestore-runtime", + "sagemaker-metrics", + ): + assert _access_key(client_b.get_client(service)) == "AKIASECONDSESSION000" + assert _access_key(client_a.get_client(service)) == "AKIAFIRSTSESSION0000" + + +def test_same_session_returns_cached_instance(): + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + + assert SageMakerClient(session=session_a) is SageMakerClient(session=session_a) + assert SageMakerClient(session=session_a, region_name=REGION_A) is SageMakerClient( + session=session_a, region_name=REGION_A + ) + + +def test_same_session_different_region_gets_distinct_client(): + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + + client_a = SageMakerClient(session=session_a, region_name=REGION_A) + client_b = SageMakerClient(session=session_a, region_name=REGION_B) + + assert client_a is not client_b + assert client_b.sagemaker_client.meta.region_name == REGION_B + + +def test_bare_call_returns_first_created_instance_as_default(): + """Configure-once pattern: SageMakerClient(session=...) then bare SageMakerClient().""" + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + + configured = SageMakerClient(session=session_a) + + assert SageMakerClient() is configured + assert SageMakerClient() is configured + + +def test_bare_call_default_is_not_replaced_by_later_explicit_session(): + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + session_b = _session("AKIASECONDSESSION000", REGION_B) + + default = SageMakerClient(session=session_a) + SageMakerClient(session=session_b) + + assert SageMakerClient() is default + + +def test_bare_call_creates_default_from_default_chain(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAENVIRONMENT00000") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret") + monkeypatch.setenv("AWS_DEFAULT_REGION", REGION_A) + + default = SageMakerClient() + + assert SageMakerClient() is default + assert default.region_name == REGION_A + assert _access_key(default.sagemaker_client) == "AKIAENVIRONMENT00000" + + +def test_config_argument_is_applied_and_user_agent_suffix_appended(): + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + config = Config(connect_timeout=3, read_timeout=7, user_agent_extra="mytool/1.0") + + client = SageMakerClient(session=session_a, config=config) + + meta_config = client.sagemaker_client.meta.config + assert meta_config.connect_timeout == 3 + assert meta_config.read_timeout == 7 + assert meta_config.user_agent_extra.startswith("mytool/1.0 ") + assert "sagemaker" in meta_config.user_agent_extra.lower() + + +def test_default_config_keeps_retry_settings(): + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + + client = SageMakerClient(session=session_a) + + # botocore normalizes {"max_attempts": N} to {"total_max_attempts": N + 1}. + retries = client.sagemaker_client.meta.config.retries + assert retries["total_max_attempts"] == 11 + assert retries["mode"] == "standard" + + +def test_distinct_config_objects_get_distinct_clients(): + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + config_1 = Config(connect_timeout=1) + config_2 = Config(connect_timeout=2) + + client_1 = SageMakerClient(session=session_a, config=config_1) + client_2 = SageMakerClient(session=session_a, config=config_2) + + assert client_1 is not client_2 + assert client_1.sagemaker_client.meta.config.connect_timeout == 1 + assert client_2.sagemaker_client.meta.config.connect_timeout == 2 + + +def test_reset_clears_default_and_keyed_entries(): + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + client_a = SageMakerClient(session=session_a) + default = SageMakerClient() + + SageMakerClient.reset() + + assert SageMakerClient(session=session_a) is not client_a + assert SageMakerClient() is not default + + +def _cached_entries(): + from sagemaker.core.utils.utils import _ClientCacheMeta + + return _ClientCacheMeta._instances.get(SageMakerClient, {}) + + +def test_keyed_cache_is_bounded_and_default_survives_eviction(): + from sagemaker.core.utils.utils import _ClientCacheMeta + + cap = _ClientCacheMeta._MAX_KEYED_ENTRIES + sessions = [_session(f"AKIA{i:016d}", REGION_A) for i in range(cap + 5)] + + default = SageMakerClient(session=sessions[0]) + clients = [SageMakerClient(session=s) for s in sessions] + + entries = _cached_entries() + keyed = [k for k in entries if k != _ClientCacheMeta._DEFAULT_KEY] + assert len(keyed) == cap + # The oldest keyed entries were dropped, the newest kept. + assert SageMakerClient(session=sessions[-1]) is clients[-1] + assert SageMakerClient(session=sessions[0]) is not clients[0] + # The process default is pinned regardless of eviction. + assert SageMakerClient() is default + + +def test_recently_used_entry_is_kept_over_stale_ones(): + from sagemaker.core.utils.utils import _ClientCacheMeta + + cap = _ClientCacheMeta._MAX_KEYED_ENTRIES + sessions = [_session(f"AKIA{i:016d}", REGION_A) for i in range(cap)] + clients = [SageMakerClient(session=s) for s in sessions] + + # Touch the oldest entry so it becomes most recently used, then overflow by one. + assert SageMakerClient(session=sessions[0]) is clients[0] + SageMakerClient(session=_session("AKIAOVERFLOW00000000", REGION_A)) + + assert SageMakerClient(session=sessions[0]) is clients[0] + assert SageMakerClient(session=sessions[1]) is not clients[1] + + +def test_concurrent_first_calls_share_one_instance(): + import threading + + session_a = _session("AKIAFIRSTSESSION0000", REGION_A) + results = [] + start = threading.Barrier(8) + + def build(): + start.wait() + results.append(SageMakerClient(session=session_a)) + + threads = [threading.Thread(target=build) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(results) == 8 + assert all(r is results[0] for r in results)