fix: pass the configured session to ModelBuilder's Model reloads

[mohamedzeidan2021]

Issue #, if available: Fixes #6459

Description

ModelBuilder.build() created the SageMaker Model through the configured sagemaker_session and then reloaded it with a session-less Model.get():

self.sagemaker_session.create_model(**create_model_args)   # configured session
...
return Model.get(model_name=self.model_name, region=self.region)   # no session

Model.get() resolves its client through Base.get_sagemaker_client(), and with session=None that helper constructs a brand-new boto3.Session from the ambient credential chain (SageMakerClient.__init__: "No boto3 session provided. Creating a new session."). So CreateModel is signed with the configured session's credentials while the follow-up DescribeModel is signed with the process default ones. Cross-account, the model is created in the target account and the reload fails with ValidationException: Could not find model ..., exactly as reported. Single-account setups hide the defect because both sessions resolve to the same account. This is the same class of bug as #6069; the fixes for that issue (#6107 / #6120) covered the Endpoint.get() call sites on the deploy path but not the Model reload on the build path.

The fix passes self.sagemaker_session.boto_session to Model.get(), which is the pattern the already-fixed Endpoint.get() call sites in this same file use. It is also more than cosmetic for the returned object: Model.get() calls _set_client_context(session=..., region=...), so the Model that build() hands back now keeps the configured session for its own refresh(), delete() and wait_* calls instead of silently falling back to the default account.

Widened beyond the one-line patch in the issue. The suggested diff fixes _create_sagemaker_model() only. Grepping every resource lookup on the build path found two more Model.get() calls with the identical shape — the model name is read through self.sagemaker_session.sagemaker_client and then described with no session, so they are guaranteed to be cross-account inconsistent:

Method	Reached from
_create_sagemaker_model()	build() → _create_model() (the reported path)
_find_reusable_model()	build(reuse_resources=True)
_get_model_for_endpoint()	Model-behind-endpoint lookup (currently no in-tree caller)

All three are fixed here. Nothing depends on the old behaviour: no unit test in the repo asserts a session-less Model.get() call, and the PipelineSession branch returns before the reload, so pipeline builds are untouched (covered by a new test).

Deliberately not widened, to keep this change on one path: model_builder.py still has session-less HubContent.get() (×2), TrainingJob.get() (×3), ModelPackage.get() (×1), and, on the Nova/LoRA model-customization deploy path, Endpoint.get() (×2, incl. _does_endpoint_exist) and InferenceComponent.get() (×2). Those are the same bug class on different resources and a different path, and fixing them needs their own tests; happy to follow up in a separate PR if you'd like them bundled.

Testing

New file sagemaker-serve/tests/unit/test_model_builder_session_propagation.py, 5 tests, modelled on the existing issue-scoped file test_model_builder_pipeline_variable_image_uri.py. It uses two real, distinguishable boto3.Session objects (a "configured" one and a "default" one) so the assertion is about which session object the reload actually receives:

1. Model.get() receives session=<configured boto session> after CreateModel (_create_sagemaker_model).
2. Deeper check: with Base.get_sagemaker_client intercepted, the client that issues DescribeModel is built from the configured session and the right region — this is the assertion that fails with session=None on master.
3. A PipelineSession build still returns before the reload and never calls Model.get() (guards the path that already worked).
4. _find_reusable_model() passes the configured session.
5. _get_model_for_endpoint() passes the configured session.

pytest tests/unit/test_model_builder_session_propagation.py
5 passed

Negative check (the installed sagemaker 3.25.0 model_builder.py restored to its released content, new tests unchanged):

4 failed, 1 passed
  FAILED ...::test_reload_passes_configured_session_to_model_get
  FAILED ...::test_reload_describe_client_is_built_from_configured_session
  FAILED ...::test_find_reusable_model_passes_configured_session
  FAILED ...::test_get_model_for_endpoint_passes_configured_session

(The one that passes either way is the PipelineSession regression guard, by design.)

Full package regression, sagemaker-serve/tests/unit:

with fix:     70 failed, 2075 passed, 25 skipped
fix reverted: 70 failed, 2070 passed, 25 skipped   (new test file excluded)

The two FAILED lists are identical; the +5 is the new tests. All 70 are pre-existing and environmental in this sandbox: 65 fail with botocore ClientError: ExpiredToken (they call live AWS) and 5 with ModuleNotFoundError: No module named 'tensorflow'.

Lint on both changed files: black --line-length 100 --check clean, flake8 --max-line-length 120 clean, pydocstyle clean on the src file. pylint --rcfile=.pylintrc: model_builder.py 9.84/10 before and after with an identical message set (one pre-existing E0203 line number shifts); package-wide src/sagemaker 9.89/10 before and after (+0.00), 83 messages in both runs — no new findings.

Not verified: there are no usable AWS credentials in my environment, so I did not run a live cross-account round trip of the reproduction script from the issue. The evidence above is offline and is about which boto3 session object the reload's client is constructed from; the account-level failure mode is inferred from that plus the CloudTrail evidence in the issue.

Merge Checklist

Put an x in the boxes that apply.

General

• I have read the CONTRIBUTING doc
• I certify that the changes I'm introducing will be backward compatible, and I have discussed concerns about this, if any, with the Python SDK team
• I used the commit message format described in CONTRIBUTING
• I have passed the region in to all S3 and STS clients that I've initialized as part of this change
• I have updated any necessary documentation, including READMEs and API docs (if appropriate)

Tests

• I have added tests that prove my fix is effective or that my feature works (if appropriate)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes
• I have checked that my tests are not configured for a specific region or account (if appropriate)
• I have used unique_name_from_base to create resource names in integ tests (if appropriate)

On the region item: this change initializes no S3 or STS clients. The only client it affects is the SageMaker client created inside Model.get(), which now receives both the configured session and region=self.region. The documentation item is unticked because the change is a bug fix with no user-facing API or doc surface, and no integ tests were added (hence the last item).

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]

🤖 AI code review

This PR fixes a cross-account bug in ModelBuilder by passing the configured sagemaker_session.boto_session to three Model.get() reload/lookup call sites (_create_sagemaker_model, _find_reusable_model, _get_model_for_endpoint), matching the already-fixed Endpoint.get() pattern elsewhere in the same file. I confirmed Model.get() accepts a session kwarg and that the existing Endpoint.get() sites use exactly this self.sagemaker_session.boto_session form, so the change is correct and consistent. New unit tests cover all three paths plus the PipelineSession early-return. The change looks clean; no issues found.

_{Reviewed commit 8249ac5. Automated review; verify before acting.}