Skip to content

fix: keep the Transformer's configured session when waiting on or stopping a transform job - #6509

Merged
mohamedzeidan2021 merged 1 commit into
aws:masterfrom
mohamedzeidan2021:fix/issue-6503-transformer-session
Oct 21, 2026
Merged

mohamedzeidan2021 merged 1 commit into
aws:masterfrom
mohamedzeidan2021:fix/issue-6503-transformer-session

Conversation

@mohamedzeidan2021

Copy link
Copy Markdown
Collaborator

Issue

Fixes #6503

Transformer.transform() submits the transform job through the configured sagemaker_session, then builds the resource it will poll with no session at all:

self.latest_transform_job = TransformJob(**transformed)

if wait:
    self.latest_transform_job.wait(logs=logs)

A session-less TransformJob builds its SageMaker client from the process default boto session (Base.get_sagemaker_client(session=None, region_name=None)), so when the configured session and the default session differ by account or region, the job is created correctly and the wait then fails — ResourceNotFound, NoRegionError, or an authorization error. Transformer.wait() and stop_transform_job() inherit the same fault through the same latest_transform_job.

Local mode is worse: LocalSagemakerClient.create_transform_job runs the job to completion in-process and stores it in LocalSagemakerClient._transform_jobs. The job never exists in AWS, so the wait called a real AWS endpoint looking for it.

The correct pattern already exists in the same class — delete_model() and _retrieve_image_uri() both pass self.sagemaker_session.boto_session to Model.get(). transform() does not.

This is the third recurrence of the #6069 session-propagation defect (#6069 → #6459 ModelBuilder reload → #6503 Transformer wait). The #6069 fix added Base._set_client_context() precisely so a resource keeps talking to the session it was created with; this call site never adopted it.

Fix

Three changes in sagemaker-core/src/sagemaker/core/transformer.py:

  1. Bind the configured session to the resource. transform() now calls _set_client_context(session=..., region=...) on the TransformJob it builds, using the same boto_session / boto_region_name submission used. refresh(), wait() and stop() go through Base._get_client(), so all three now sign against the configured account and region. This reuses the mechanism the SageMakerClient CustomSession uses default values instead of passed user session #6069 fix introduced rather than adding a second one, and — unlike the TransformJob.get(...) the issue suggests — it costs no extra DescribeTransformJob call.

  2. Never poll AWS for a local job. transform() records whether the submission ran in local mode; wait() and stop_transform_job() return early for a local job, since it already ran to completion inside create_transform_job and does not exist in AWS. transform(wait=True) now routes through self.wait() so the local guard applies on both paths.

  3. attach() passed the wrong session type. TransformJob.get() wants a boto3 session, but attach() handed it the sagemaker Session, so the reattached resource was bound to an object that cannot produce a client. It now passes boto_session and boto_region_name, and keeps the resource describe already returned instead of issuing a second, identical DescribeTransformJob. Not reported in Transformer.transform(wait=True) ignores its configured session while waiting #6503, but it is the same defect in the same class, and the ask on the parent ticket is the pattern rather than the single call site.

Before (cross-account, configured session in account A, default session in account B):

CreateTransformJob  -> account A   (configured session)
DescribeTransformJob -> account B  (default session)  => ResourceNotFound

After: both go to account A.

Testing

Reproduced first with an offline script derived from the repro in #6503 — Base.get_sagemaker_client patched to record the (session, region) of every client request, so nothing calls AWS. Pre-fix the wait requests (None, None), the process default; post-fix it requests (<configured boto session>, 'us-east-2').

  • New sagemaker-core/tests/unit/test_transformer_session_propagation.py (13 tests), modelled on the existing test_resource_client_context.py from the SageMakerClient CustomSession uses default values instead of passed user session #6069 fix:

    • the resource is bound to the configured session and region after transform();
    • the wait inside transform(wait=True), a later wait(), and stop_transform_job() each use only the configured session/region;
    • end to end through SageMakerClient with a default client built first, asserting the bound client is signed with the configured credentials (AKIDTARGET) and is not the default client — this is what the cross-account failure actually comes down to;
    • local mode makes no client request at all on either wait() or stop_transform_job(), and a remote transform is not mistaken for a local one;
    • attach() passes boto_session/boto_region_name and reuses the described resource;
    • the _ensure_last_transform_job guards still raise.

    Against pre-fix source 10 of these 13 fail; all pass with the fix.

  • Existing sagemaker-core/tests/unit/test_transformer.py: 30 passed, unchanged.

  • Full sagemaker-core/tests/unit: 3988 passed, 21 skipped.

  • black (100) clean, flake8 (120) clean, pylint 9.81/10 — identical to the unmodified baseline on the same checkout, with no message in the changed file.

No integration test: the failure only surfaces when the configured and default sessions differ by account, which single-account CI credentials cannot express. The unit tests assert the session and credentials the client is actually built with, which is the part that regressed.

Backwards compatibility

No public signature or default changes. wait() gained a real docstring in place of "Placeholder docstring".

Two behaviour changes, both correcting failures:

  • Remote waits/stops now sign against the configured session instead of the default one. Identical where the two sessions share an account and region — which is every case that worked before.
  • A local transform's wait() and stop_transform_job() are now no-ops that log, rather than AWS calls that fail. The local job is already terminal when create_transform_job returns, so there was never anything to wait for.

attach() now issues one DescribeTransformJob instead of two, and the attached latest_transform_job is the described resource rather than a second fetch of the same job.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI code review

This PR fixes session/region propagation in Transformer: transform() now binds the configured boto_session/boto_region_name to the TransformJob resource via _set_client_context so wait()/stop() poll the correct account and region; it adds a local-mode guard so local jobs are never polled over AWS; and it fixes attach() to pass boto_session (not the SageMaker Session) to TransformJob.get() while reusing the already-described resource. The changes match the generated-resource client-context mechanism (_get_client/_set_client_context) and the established pattern from prior work, and come with a thorough new unit test module. I found no correctness defects; the implementation is consistent with the repo's conventions. Note this fix substantially overlaps the already-documented PR #6411 (same transform() binding and same attach() boto-session fix); the only net-new behavior here is the local-mode short-circuit — worth coordinating to avoid duplicate/conflicting changes.

Reviewed commit 4771c41. Automated review; verify before acting.

# the resource that describe returned instead of describing it a second time.
transform_job = TransformJob.get(
transform_job_name=transform_job_name, session=sagemaker_session
transform_job_name=transform_job_name,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Low: The transform() session-binding and the attach() boto-session fix here appear to duplicate work already described in PR #6411, which made the same two changes (_set_client_context in transform() and session=boto_session, region=boto_region_name in attach()). Only the local-mode guard (_latest_transform_job_is_local) is net-new. Consider rebasing on / coordinating with that change to avoid conflicting or redundant commits. Source: #6411

@mohamedzeidan2021
mohamedzeidan2021 merged commit c9ec795 into aws:master Oct 21, 2026
28 of 30 checks passed

This branch was previously deployed

1 inactive deployment
auto-approve — 4771c41d Deployed Oct 20, 2026 by mohamedzeidan2021 via wait-for-approval #704
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Transformer.transform(wait=True) ignores its configured session while waiting

2 participants