Fix/SageMaker client construction

[lucasjia-aws]

Description

This PR simplifies SageMaker client construction and includes a couple of small robustness improvements.

Simplify control-plane client construction. SageMakerClient previously loaded a custom botocore service model (from a bundled sample/ directory) as a temporary measure to expose the generic Job APIs before they were available upstream. Those Job APIs (CreateJob / DescribeJob / ListJobs) now ship in the public botocore, so this temporary loader is no longer needed. The client now uses the standard session.client("sagemaker", ...) path, which is simpler and behaves consistently with the other clients in the constructor.

Bump the boto3 floor to >=1.43.20. Since the client now relies on the public botocore for the Job APIs, the minimum boto3 version is raised to the first release where those operations are available (botocore>=1.43.20), so the Job.create() / Job.get() / Job.get_all() paths resolve correctly on a clean install. This is compatible with the package's existing Python >=3.10 requirement.

Testing

• Added unit tests for the env-logging masking behavior.
• Verified the standard session.client("sagemaker") exposes create_job / describe_job / list_jobs at the pinned botocore floor.
• Existing sagemaker-core unit tests pass.

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.

[mujtaba1747]

Env-log masking gap — the container-driver startup logger masks by key name only, so a credential-shaped value under an innocuous key name would be logged in clear text.

The PR is doing regex against some known patterns, but credentials may/may not match that regex.

I wonder what's the best practice to filter out credentials from logs. I would suggest to explore if there any official AWS guidance on this.

[mujtaba1747]

https://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/protect-sensitive-log-data-types.html

[github-actions]

Claude finished @lucasjia-aws's task in 7m 54s —— View job

---

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

• Read the updated PR diff
• Reviewed utils.py singleton revert + callers
• Reviewed value-masking change
• Posted consolidated review

---

The PR changed shape since my last pass. The per-(session, region) singleton keying I reviewed before has been reverted (e9e835e) — SingletonMeta/SageMakerClient are back to the original class-keyed singleton, which the maintainer confirmed is by design. So utils.py now nets to a clean simplification. Two independent concerns remain:

🟠 Removing the sample-model loader may break the create_job path — please confirm botocore coverage

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

The revert removes not just the dev-only endpoint/region override (good — that's the point of the PR) but also the custom botocore data-loader that registered sagemaker-core/sample/sagemaker/2017-07-24/service-2.json. The removed code's own comment says that sample model supplies pre-GA Job APIs "(CreateJob, DescribeJob, ListJobs, etc.) not yet in the public botocore release."

There is a live generated caller that depends on those:

# resources.py:18263  (Job.create)
response = client.create_job(**operation_input_args)

With the loader gone, client is a stock session.client("sagemaker"). If CreateJob/DescribeJob/ListJobs are not in the pinned botocore (boto3>=1.42.2), Job.create() / Job.get() / Job.get_all() will raise AttributeError/UnknownOperationException at runtime. There is no unit test exercising create_job, so CI won't catch it.

This is fine if those APIs have since shipped in the public botocore release (which is the stated exit condition in the old TODO). Please confirm that's the case — a quick check that boto3.Session().client("sagemaker") exposes create_job at the pinned version — before merging. If they haven't shipped yet, removing the loader is premature and the sample model still needs to be registered (just without the env-var endpoint override).

🟡 Value-masking: _value_is_sensitive only runs on the string branch (unchanged, awareness only)

.../container_drivers/scripts/environment.py:277

As @mujtaba1747 noted, a credential-shaped string nested inside a dict value still only gets masked by key name (via mask_sensitive_info), not by value shape. That matches prior behavior and is a reasonable scope boundary. The regex approach is also inherently best-effort — the CloudWatch sensitive-data guidance linked in review is a good longer-term direction (data-protection policies / managed data identifiers) but is out of scope here. Fail-safe direction is correct: a miss logs a value that would have been logged anyway; a match only masks. No change required for this PR.

✅ Value-masking change looks good

• AKIA/ASIA-style access-key-ID regex is correct (4-char prefix + 16 = 20 chars, \b-anchored).
• Ordering is right: key-name check → value shape → dict handling.
• Good new tests: credential-shaped value under an innocuous key is masked, and a plain endpoint URL stays visible.
• The two driver-script copies (sagemaker-core / sagemaker-train) remain byte-identical.

🟢 CodeQL alert on test_environment.py — dismissible false positive

The "Incomplete URL substring sanitization" alert on https://example.com:9999 is a false positive: the test does a strict tuple-equality assertion, no URL parsing/substring sanitization. Safe to dismiss as "used in tests".

Verdict: the masking change is solid and well-tested. The one thing I'd confirm before merge is the 🟠 item — that the public botocore at the pinned version actually provides the Job APIs the removed sample-model loader used to supply, so the create_job path doesn't regress.

[github-actions]

Claude finished @lucasjia-aws's task in 2m 16s —— View job

---

Claude review — Fix/aml 509849 (re-review after boto3 floor bump)

• Read the updated PR diff
• Verify the boto3 floor bump addresses the prior 🟠 Job-API concern
• Re-check the value-masking change
• Post consolidated review

---

The new commit (5225d88) bumps the boto3 floor to >=1.43.20 in sagemaker-core/pyproject.toml, which directly targets the one 🟠 item from my last review — the removed sample-model loader that used to supply the pre-GA Job APIs. This is the right way to close that gap: pin the floor to the public-botocore version where CreateJob/DescribeJob/ListJobs shipped, rather than keep the custom data-loader.

🟢 Job-API concern — addressed, with one thing to confirm at runtime

sagemaker-core/pyproject.toml:9 → boto3>=1.43.20,<2.0.0

What I was able to verify from the checked-out repo:

• The env-var endpoint override and the custom botocore data-loader are fully removed from utils.py — no residual SAGEMAKER_ENDPOINT/register_component/extra_search_paths references remain in src/.
• The live caller path is intact: Job.create → resources.py:18263 client.create_job(...), Job.get → describe_job, Job.get_all → list_jobs. All now resolve against a stock session.client("sagemaker").
• The bundled sample/sagemaker/2017-07-24/service-2.json still defines CreateJob/DescribeJob/ListJobs, but it's no longer registered — so from this version on, those operations must come from public botocore.

What I couldn't verify in this environment: botocore isn't installed here and web fetch is disabled, so I can't independently confirm 1.43.20's bundled SageMaker model actually exposes create_job. The commit message asserts it does, which I'm taking as the maintainer having checked. One quick sanity check before merge, if not already done:

import boto3; boto3.Session().client("sagemaker", region_name="us-west-2").create_job

should not raise. If that holds at boto3==1.43.20, the Job path no longer regresses and this concern is fully closed.

🟡 Minor — the now-unused sample/ models are still in the tree

With the loader gone, sample/sagemaker/.../service-2.json (and the sibling sample models) are no longer wired into anything at runtime. Not blocking — leaving them is harmless and they may still be useful as reference — but worth a follow-up to delete the dead sagemaker sample model so it doesn't drift or mislead. Out of scope for this PR.

✅ Value-masking change — still good

• AKIA/ASIA-style access-key-ID regex is correct (4-char prefix + 16 = 20 chars, \b-anchored), plus the aws_session_token / x-amz-security-token patterns.
• Ordering is right: key-name check → value shape → dict handling.
• Fail-safe direction correct (a false positive only masks; it never leaks).
• Good tests: credential-shaped value under an innocuous key is masked; a plain endpoint URL stays visible. The two driver-script copies (sagemaker-core / sagemaker-train) remain byte-identical.
• Known scope boundary (unchanged): _value_is_sensitive runs only on the string branch, so a credential-shaped string nested in a dict is still masked by key name only. Reasonable to leave as-is.

🟢 CodeQL alert on test_environment.py — dismissible false positive

"Incomplete URL substring sanitization" on https://example.com:9999 — the test does a strict tuple-equality assertion, no URL parsing/substring sanitization. Safe to dismiss as "used in tests".

Verdict: the prior blocking concern is addressed by the floor bump, and the masking change remains solid and well-tested. The only pre-merge step I'd suggest is the one-line runtime confirmation above that create_job exists at boto3==1.43.20 (I couldn't run it here). The two 🟡/🟢 notes are non-blocking.

[lucasjia-aws]

Re: Claude review

Fixed. create_job/describe_job/list_jobs first ship in botocore 1.43.20 (verified: 1.43.17/1.43.19 lack them, 1.43.20 has them; 1.43.0 and all of 1.42.x lack them). The old boto3>=1.42.2 pin allowed botocore below that, so the create_job path could break at runtime. Bumped the pin to boto3>=1.43.20 (requires botocore>=1.43.20). Python is fine — boto3 1.43.x needs >=3.10 and sagemaker-core already requires >=3.10.

Other notes acknowledged: dict-nested value masking is an intentional scope boundary; the CodeQL alert is a false positive (strict tuple-equality assertion, no URL sanitization).

[mujtaba1747]

Is this intentional?

[github-actions]

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

---

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

• Read the updated PR diff
• Verify utils.py simplification + callers
• Confirm boto3 floor bump
• Post consolidated review

---

The PR narrowed since the last pass. The env-log value-masking change was reverted (b5885f2), so this PR now nets to two files — the utils.py simplification and the boto3 floor bump — and is a clean, focused change. Re-verified from the checked-out branch:

✅ utils.py simplification is correct and complete

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

• The dev-only env overrides (SAGEMAKER_REGION / SAGEMAKER_STAGE / SAGEMAKER_ENDPOINT / SAGEMAKER_RUNTIME_ENDPOINT) and the custom botocore data-loader (register_component('data_loader', ...) on session._session) are fully removed. Grep over src/ finds no residual references — the only SAGEMAKER_REGION* hit is the unrelated SAGEMAKER_REGION_PARAM_NAME constant.
• The client now uses the standard path: session.client("sagemaker", region_name, config=self.config). The positional region_name maps correctly to boto3's Session.client(service_name, region_name=None, ...) signature, and this matches the sibling clients (sagemaker-runtime, sagemaker-featurestore-runtime, sagemaker-metrics) already in the constructor.
• The reset() docstring was trimmed to match (it no longer re-reads env vars) — consistent with the removal.

✅ boto3 floor bump closes the Job-API gap

sagemaker-core/pyproject.toml:15 → boto3>=1.43.20,<2.0.0

The removed data-loader used to supply the pre-GA Job APIs. The live callers still depend on them — Job.create → resources.py:18263 client.create_job(...), Job.get → describe_job (18313/18354), Job.get_all → list_jobs (18620) — all now resolving against stock session.client("sagemaker"). Pinning the floor to the first public-botocore release that ships those operations (1.43.20) is the right way to close the gap rather than keeping the custom loader. The author confirmed the version boundary (1.43.17/1.43.19 lack the ops, 1.43.20 has them), and this is compatible with the existing Python >=3.10 requirement. This also answers @mujtaba1747's "Is this intentional?" on pyproject.toml — yes, and it's the load-bearing part of the change.

🟡 Minor (non-blocking, pre-existing) — broken service_name=/.client caller unchanged

resources.py:14021 still calls SageMakerClient(session=..., region_name=..., service_name="sagemaker").client. __init__ accepts no service_name kwarg and there is no .client attribute (only sagemaker_client, etc.), so this generated line would TypeError/AttributeError if ever reached. It's pre-existing generated code, untouched by this PR and out of scope — noting only so it isn't mistaken for introduced-here or fixed-here. The template source (tools/templates.py:109) has the same shape.

🟡 Minor — now-unused sample/ models still in the tree

With the loader gone, the bundled sample/sagemaker/.../service-2.json is no longer wired into anything at runtime. Harmless to leave, but a good follow-up cleanup so it doesn't drift. Out of scope here.

Verdict: clean, well-scoped simplification. The removal is complete (no residual env/loader references), and the floor bump correctly guarantees the Job APIs come from public botocore. Nothing blocking — the two 🟡 notes are pre-existing / follow-up only.