Fix/aml 509849

[lucasjia-aws]

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 2m 51s —— View job

---

Claude review — Fix/aml 509849

• Read the PR diff (/tmp/pr.diff)
• Reviewed changed files and their callers for context
• Posted review feedback

---

I reviewed the four changed files. The change splits into two independent concerns: (1) masking credential-shaped values in container env logging, and (2) reworking SingletonMeta/SageMakerClient to key instances per (session, region) while removing the dev-only endpoint override. The masking change is clean and well-tested. I have a few notes on the singleton rework.

Note: the inline-comment tool isn't available in this run, so findings are consolidated here with file/line references.

🟠 _singleton_key narrows the accepted arguments and can raise on cache-hit paths

sagemaker-core/src/sagemaker/core/utils/utils.py:348-352

The metaclass now evaluates key_fn(*args, **kwargs) on every construction call (before super().__call__). SageMakerClient._singleton_key only accepts (session, region_name, config). There are live (generated) callers that pass service_name=:

# sagemaker-core/src/sagemaker/core/resources.py:14021
client = SageMakerClient(session=session, region_name=region, service_name="sagemaker").client

Previously, once the singleton existed, extra kwargs were silently ignored (a cache hit never touched __init__). Now the same call routes through _singleton_key(... service_name="sagemaker") and raises TypeError: unexpected keyword argument 'service_name' even on what used to be a successful cache hit. Consider making the key function tolerant of the same arguments the constructor accepts, e.g.:

@staticmethod
def _singleton_key(session: Session = None, region_name: str = None, *args, **kwargs):
    session_key = id(session) if session is not None else None
    return (session_key, region_name)

(These service_name=/.client call sites appear to be pre-existing dead/broken paths, but the new code turns a formerly-silent no-op into a hard error, so hardening the signature is cheap insurance.)

🟡 id(session) as a cache key — document the lifetime assumption

sagemaker-core/src/sagemaker/core/utils/utils.py:351

Using id(session) is safe here only because the cached SageMakerClient stores self.session = session, keeping the session object alive so its id() can't be recycled while the entry exists. That's a subtle invariant worth a one-line comment, since id() reuse after GC would otherwise be a latent aliasing bug. Also note two functionally-identical but distinct Session objects now produce two separate cached clients — a behavior change from the old global singleton. If that's intended, fine; worth confirming it doesn't defeat the caching for callers that build a fresh Session each call.

🟡 No test coverage for the riskier change

The masking change got two good new tests. The SingletonMeta keying + SageMakerClient.reset() rework — the higher-risk part of this PR — has none. Suggest adding unit tests for: distinct instances per (session, region), same instance on repeated no-arg calls, and reset() clearing all keyed entries for the class (the new tuple-key filter in reset()).

✅ Value-masking change looks good

.../container_drivers/scripts/environment.py

• The AWS access-key-ID regex ((?:AKIA|ASIA|...)[A-Z0-9]{16}, 20 chars total) is correct, and applying it only after the existing key-name check is the right order.
• Fail-safe direction is good: a false positive just masks a value; it never leaks.
• The two copies (sagemaker-core and sagemaker-train) are byte-identical, consistent with how these driver scripts are already duplicated in the repo.

Nothing blocking on the masking side. The _singleton_key signature is the one item I'd recommend addressing before merge.

[github-actions]

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

---

Claude review — Fix/aml 509849 (re-review after sync)

• Read the updated PR diff (/tmp/pr.diff)
• Verified prior feedback was addressed
• Reviewed the CodeQL alert and remaining edge cases

---

Thanks for the updates — the two items I flagged as most important last round are resolved:

• ✅ _singleton_key signature hardened. It now takes (session=None, region_name=None, *args, **kwargs), so cache-hit calls that pass service_name= (e.g. resources.py:14021) no longer raise TypeError in the key function. A new test (test_singleton_key_tolerates_extra_kwargs) locks this in.
• ✅ id(session) lifetime documented. The docstring now states the invariant (the cached instance holds a reference to the session, so its id() can't be recycled while the entry lives). Good.
• ✅ Singleton keying + reset() now have tests — distinct instance per (session, region), reuse on repeated calls, and reset() clearing all keyed entries.

🟢 CodeQL alert (test_environment.py) — safe to dismiss

The "Incomplete URL substring sanitization" alert on https://example.com:9999 is a false positive. The test does a strict tuple equality assertion (assert call_args == ("%s=%s", "SAGEMAKER_ENDPOINT", endpoint)) — there's no URL parsing or substring-based sanitization happening. No product code is involved. You can dismiss it as "won't fix / used in tests" in the code-scanning UI.

🟡 Minor (non-blocking) — pre-existing broken caller unchanged

resources.py:14021 still calls SageMakerClient(session=..., region_name=..., service_name="sagemaker").client. With the tolerant key function the keying now succeeds, but SageMakerClient has no .client attribute (only sagemaker_client, sagemaker_runtime_client, etc.), so that line would still AttributeError if ever reached. This is pre-existing generated code and out of scope for this PR — just noting it's still latent so it isn't mistaken for fixed.

🟡 Minor — value-masking doesn't recurse into dicts

In log_key_value, _value_is_sensitive(value) only runs on the string branch. A credential-shaped string nested inside a dict value goes through mask_sensitive_info(value), which (as before) masks by key name only — so a credential-shaped value under an innocuous nested key wouldn't be caught. This matches prior behavior and is a reasonable scope boundary; flagging only for awareness.

✅ Value-masking change still looks good

• AKIA/ASIA-style access-key-ID regex is correct (4-char prefix + 16 → 20 chars total).
• Ordering is right: key-name check first, then value shape, then dict handling.
• Fail-safe direction is correct (a false positive only masks; it never leaks), and the two driver-script copies remain byte-identical.

Verdict: the blocking item from the prior review is resolved and the risky change now has test coverage. Nothing blocking remains — the CodeQL alert is a dismissible false positive, and the two 🟡 notes are pre-existing / awareness-only.