Repository navigation
fix: keep the Transformer's configured session when waiting on or stopping a transform job - #6509
Conversation
…pping a transform job
There was a problem hiding this comment.
🤖 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, |
There was a problem hiding this comment.
🟡 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
Issue
Fixes #6503
Transformer.transform()submits the transform job through the configuredsagemaker_session, then builds the resource it will poll with no session at all:A session-less
TransformJobbuilds 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()andstop_transform_job()inherit the same fault through the samelatest_transform_job.Local mode is worse:
LocalSagemakerClient.create_transform_jobruns the job to completion in-process and stores it inLocalSagemakerClient._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 passself.sagemaker_session.boto_sessiontoModel.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:Bind the configured session to the resource.
transform()now calls_set_client_context(session=..., region=...)on theTransformJobit builds, using the sameboto_session/boto_region_namesubmission used.refresh(),wait()andstop()go throughBase._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 theTransformJob.get(...)the issue suggests — it costs no extraDescribeTransformJobcall.Never poll AWS for a local job.
transform()records whether the submission ran in local mode;wait()andstop_transform_job()return early for a local job, since it already ran to completion insidecreate_transform_joband does not exist in AWS.transform(wait=True)now routes throughself.wait()so the local guard applies on both paths.attach()passed the wrong session type.TransformJob.get()wants a boto3 session, butattach()handed it the sagemakerSession, so the reattached resource was bound to an object that cannot produce a client. It now passesboto_sessionandboto_region_name, and keeps the resourcedescribealready returned instead of issuing a second, identicalDescribeTransformJob. Not reported inTransformer.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):
After: both go to account A.
Testing
Reproduced first with an offline script derived from the repro in #6503 —
Base.get_sagemaker_clientpatched 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 existingtest_resource_client_context.pyfrom the SageMakerClient CustomSession uses default values instead of passed user session #6069 fix:transform();transform(wait=True), a laterwait(), andstop_transform_job()each use only the configured session/region;SageMakerClientwith 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;wait()orstop_transform_job(), and a remote transform is not mistaken for a local one;attach()passesboto_session/boto_region_nameand reuses the described resource;_ensure_last_transform_jobguards 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,pylint9.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:
wait()andstop_transform_job()are now no-ops that log, rather than AWS calls that fail. The local job is already terminal whencreate_transform_jobreturns, so there was never anything to wait for.attach()now issues oneDescribeTransformJobinstead of two, and the attachedlatest_transform_jobis the described resource rather than a second fetch of the same job.