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