Skip to content

Commit 69e264f

Browse files
authored
fix: add describe_transform_job in session class (#1507)
* fix: add describe_transform_job in session class
1 parent bc78615 commit 69e264f

File tree

2 files changed

+17
-7
lines changed

2 files changed

+17
-7
lines changed

src/sagemaker/session.py

+12
Original file line numberDiff line numberDiff line change
@@ -2573,6 +2573,18 @@ def wait_for_tuning_job(self, job, poll=5):
25732573
self._check_job_status(job, desc, "HyperParameterTuningJobStatus")
25742574
return desc
25752575

2576+
def describe_transform_job(self, job_name):
2577+
"""Calls the DescribeTransformJob API for the given job name
2578+
and returns the response.
2579+
2580+
Args:
2581+
job_name (str): The name of the transform job to describe.
2582+
2583+
Returns:
2584+
dict: A dictionary response with the transform job description.
2585+
"""
2586+
return self.sagemaker_client.describe_transform_job(TransformJobName=job_name)
2587+
25762588
def wait_for_transform_job(self, job, poll=5):
25772589
"""Wait for an Amazon SageMaker transform job to complete.
25782590

tests/integ/test_transformer.py

+5-7
Original file line numberDiff line numberDiff line change
@@ -97,8 +97,8 @@ def test_transform_mxnet(
9797
):
9898
transformer.wait()
9999

100-
job_desc = transformer.sagemaker_session.sagemaker_client.describe_transform_job(
101-
TransformJobName=transformer.latest_transform_job.name
100+
job_desc = transformer.sagemaker_session.describe_transform_job(
101+
job_name=transformer.latest_transform_job.name
102102
)
103103
assert kms_key_arn == job_desc["TransformResources"]["VolumeKmsKeyId"]
104104
assert output_filter == job_desc["DataProcessing"]["OutputFilter"]
@@ -323,8 +323,8 @@ def test_stop_transform_job(mxnet_estimator, mxnet_transform_input, cpu_instance
323323

324324
transformer.stop_transform_job()
325325

326-
desc = transformer.latest_transform_job.sagemaker_session.sagemaker_client.describe_transform_job(
327-
TransformJobName=latest_transform_job_name
326+
desc = transformer.latest_transform_job.sagemaker_session.describe_transform_job(
327+
job_name=latest_transform_job_name
328328
)
329329
assert desc["TransformJobStatus"] == "Stopped"
330330

@@ -393,9 +393,7 @@ def test_transform_tf_kms_network_isolation(sagemaker_session, cpu_instance_type
393393
)
394394
assert model_desc["EnableNetworkIsolation"]
395395

396-
job_desc = sagemaker_session.sagemaker_client.describe_transform_job(
397-
TransformJobName=job_name
398-
)
396+
job_desc = sagemaker_session.describe_transform_job(job_name=job_name)
399397
assert job_desc["TransformOutput"]["S3OutputPath"] == output_path
400398
assert job_desc["TransformOutput"]["KmsKeyId"] == kms_key
401399
assert job_desc["TransformResources"]["VolumeKmsKeyId"] == kms_key

0 commit comments

Comments
 (0)