fix(sagemaker-core): use caller session for sagemaker control-plane c…

[lucasjia-aws]

…lient

SageMakerClient.init built the "sagemaker" control-plane client from a fresh default botocore session (botocore.session.get_session()) instead of the caller-provided session. The custom data loader for the pre-GA Job APIs needs a botocore session to attach to, and the new session it created silently dropped the caller's credentials/profile.

Since SageMakerClient is a singleton, that default-credential client was reused process-wide, so control-plane calls went out under the ambient default AWS profile. When the execution role lived in a different account than the default profile, CreateTrainingJob failed with "RoleArn: Cross-account pass role is not allowed". This regressed V2 behavior and was introduced in the MTRL launch (#5919, first shipped in v3.13.0).

Register the custom data loader on the caller session's underlying botocore session (session._session) and build the sagemaker client from that session, so it keeps the caller's credentials while still loading the custom service model. The other clients already used the passed session and are unchanged.

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

Claude finished @lucasjia-aws's task in 1m 32s —— View job

---

Code Review — PR #6107

• Read the PR diff (/tmp/pr.diff)
• Read surrounding code in sagemaker-core/src/sagemaker/core/utils/utils.py
• Check callers / existing test coverage
• Post review

---

Summary

The change is correct and well-targeted. Building the sagemaker control-plane client from a fresh botocore.session.get_session() did drop the caller's credentials/profile, and since SageMakerClient is a singleton that stale client was reused process-wide — matching the cross-account RoleArn regression described in #6069. Registering the custom data_loader on the caller session's underlying botocore session (session._session) and building the client from session preserves credentials while still loading the pre-GA service model. boto3.Session.client() delegates to self._session.create_client(), which resolves the data_loader component, so the custom sample/ model still loads. 👍

The comment block added (lines 386–391, 404) clearly documents the rationale, which is valuable given the non-obvious _session access.

Observations

1. Side effect on the caller's session (minor, worth a note)
session._session.register_component('data_loader', loader) mutates the caller-provided boto3 session in place, replacing its data loader for all clients created from that session thereafter — not just the sagemaker client. The previous code isolated the custom loader inside a throwaway custom_session, so this is a behavioral change. In practice it's benign because include_default_search_paths=True keeps the standard models resolvable, and the region/service coverage is a superset. Still, if a caller reuses the same session elsewhere and inspects/relies on its default loader, this now affects them. Given the # TODO: Remove post-launch intent, this is an acceptable trade-off, but calling it out.

2. Accessing the private _session attribute
session._session is technically private, but it is the conventional way to reach a boto3 Session's underlying botocore session and the added comment documents it. Fine for a temporary shim.

3. No test covers the changed behavior
There's no unit test asserting that the control-plane client is built from the caller's session (e.g. that credentials/profile are preserved and register_component is invoked on session._session). The existing test (tests/unit/generated/test_resources.py:83) passes a Mock session, so session._session.register_component(...) silently succeeds and wouldn't catch a regression back to a fresh session. A small test — passing a real/stubbed boto3.Session and asserting sagemaker_client derives from it (or that register_component was called on session._session) — would lock in the fix. Optional but recommended for a regression of this kind.

Nothing blocking. LGTM aside from the optional test suggestion.

[lucasjia-aws]

related to issue #6069