From d6b7dad795875158f50456593048321196f1fece Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Sun, 8 Nov 2026 21:41:08 -0600 Subject: [PATCH] fix(core): describe transform jobs through the session client in logs_for_transform_job logs_for_transform_job described the job with TransformJob.get(transform_job_name=..., session=sagemaker_session).__dict__. TransformJob.get expects a boto3 Session, so passing the SageMaker Session failed pydantic validation ('Input should be an instance of Session'), and the resource's snake_case __dict__ does not have the PascalCase keys (TransformResources, TransformJobStatus) that _logs_init and _check_job_status read. Describe the job through sagemaker_session.sagemaker_client, as logs_for_processing_job does (#6369). #6509 already fixed the session used by transform()/wait()/stop()/attach(); this covers the log tailing path. --- .../src/sagemaker/core/transformer.py | 11 ++++++--- sagemaker-core/tests/unit/test_transformer.py | 24 +++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/transformer.py b/sagemaker-core/src/sagemaker/core/transformer.py index 787978c8d6..d151689c61 100644 --- a/sagemaker-core/src/sagemaker/core/transformer.py +++ b/sagemaker-core/src/sagemaker/core/transformer.py @@ -779,8 +779,12 @@ def logs_for_transform_job(sagemaker_session, job_name, wait=False, poll=10): ValueError: If the transform job fails. """ + # Describe through the session's (region-aware) client: TransformJob.get() expects a + # boto3 Session, and _logs_init/_check_job_status read the API's PascalCase keys. description = _wait_until( - lambda: TransformJob.get(transform_job_name=job_name, session=sagemaker_session).__dict__, + lambda: sagemaker_session.sagemaker_client.describe_transform_job( + TransformJobName=job_name + ), poll, ) @@ -830,8 +834,9 @@ def logs_for_transform_job(sagemaker_session, job_name, wait=False, poll=10): if state == LogState.JOB_COMPLETE: state = LogState.COMPLETE elif time.time() - last_describe_job_call >= 30: - transform_job = TransformJob.get(transform_job_name=job_name, session=sagemaker_session) - description = transform_job.__dict__ if transform_job else None + description = sagemaker_session.sagemaker_client.describe_transform_job( + TransformJobName=job_name + ) last_describe_job_call = time.time() status = description["TransformJobStatus"] diff --git a/sagemaker-core/tests/unit/test_transformer.py b/sagemaker-core/tests/unit/test_transformer.py index e8c954ec71..c610b0547b 100644 --- a/sagemaker-core/tests/unit/test_transformer.py +++ b/sagemaker-core/tests/unit/test_transformer.py @@ -685,3 +685,27 @@ def test_transform_removes_tags_before_building_transform_job(self, mock_session call_kwargs = mock_transform_job_class.call_args[1] assert "tags" not in call_kwargs assert call_kwargs["transform_job_name"] == "test-job" + + +@patch("sagemaker.core.transformer._flush_log_streams") +@patch("sagemaker.core.transformer._logs_init") +def test_logs_for_transform_job_describes_with_session_client( + mock_logs_init, mock_flush, mock_session +): + """logs_for_transform_job describes the job through the session's region-aware client. + + TransformJob.get() expects a boto3 Session (passing the SageMaker Session fails + validation), and _logs_init/_check_job_status read PascalCase describe keys. + """ + from sagemaker.core.transformer import logs_for_transform_job + + description = {"TransformJobName": "test-job", "TransformJobStatus": "Completed"} + mock_session.sagemaker_client.describe_transform_job.return_value = description + mock_logs_init.return_value = (1, [], {}, Mock(), "/aws/sagemaker/TransformJobs", False, Mock()) + + logs_for_transform_job(mock_session, "test-job", wait=True, poll=0) + + mock_session.sagemaker_client.describe_transform_job.assert_called_with( + TransformJobName="test-job" + ) + assert mock_logs_init.call_args[0][1] == description