Setting the file. One moment. 10
Setup DevOps Agent
24.7
Code Output Guide · references
RDS Oracle
references/dataset-transformation/scripts/transformation_tools.py
references/dataset-transformation/scripts/transformation_tools.py
Python·152 lines·5 KB
11
ProcessingInput,
12 ProcessingOutput,
13 ProcessingS3Input,
14 ProcessingS3Output,
15)
16
17set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)
18
19
20def _get_session(region=None):
21 """Create a SageMaker Session, optionally pinned to a region."""
22 return Session(boto_session=boto3.Session(region_name=region) if region else None)
23
24
25def execute_transformation_job(
26 transform_script_path,
27 dataset_source_s3,
28 dataset_output_s3,
29 instance_type="ml.m5.xlarge",
30 region=None,
31 execution_role=None,
32 base_job_name="dataset-transformation",
33 image_uri=None,
34):
35 """
36 Execute a dataset transformation script as a SageMaker Processing Job
37 using the V3 SDK FrameworkProcessor.
38
39 The entire directory containing the script is uploaded as source_dir,
40 so transform_fn.py (and any other dependencies) are included automatically.
41
42 Args:
43 transform_script_path: Local path to the transformation script (e.g., "<project-dir>/scripts/transform.py")
44 dataset_source_s3: S3 URI of the input dataset
45 dataset_output_s3: S3 URI for the transformed output dataset
46 instance_type: ML instance type (default: ml.m5.xlarge)
47 region: AWS region (auto-detected if None)
48 execution_role: IAM role ARN (auto-detected if None)
49 base_job_name: Prefix for the Processing Job name
50 image_uri: Docker image URI for the processing container.
51 If None, uses the SKLearn processing image.
52 """
53 if not execution_role:
54 execution_role = get_execution_role()
55
56 sagemaker_session = _get_session(region)
57
58 if not region:
59 region = sagemaker_session.boto_region_name
60
61 # Use SKLearn processing image as default (includes pandas)
62 if not image_uri:
63 image_uri = image_uris.retrieve(
64 framework="sklearn",
65 region=region,
66 version="1.2-1",
67 instance_type=instance_type,
68 )
69
70 source_dir = os.path.dirname(os.path.abspath(transform_script_path))
71 script_name = os.path.basename(transform_script_path)
72
73 processor = FrameworkProcessor(
74 role=execution_role,
75 image_uri=image_uri,
76 command=["python3"],
77 instance_count=1,
78 instance_type=instance_type,
79 base_job_name=base_job_name,
80 sagemaker_session=sagemaker_session,
81 )
82
83 input_local_path = "/opt/ml/processing/input"
84 output_local_path = "/opt/ml/processing/output"
85 input_filename = os.path.basename(dataset_source_s3.rstrip("/"))
86
87 processor.run(
88 code=script_name,
89 source_dir=source_dir,
90 arguments=[
91 "--input",
92 os.path.join(input_local_path, input_filename),
93 "--output",
94 os.path.join(output_local_path, input_filename),
95 ],
96 inputs=[
97 ProcessingInput(
98 input_name="dataset",
99 s3_input=ProcessingS3Input(
100 s3_uri=dataset_source_s3,
101 local_path=input_local_path,
102 s3_data_type="S3Prefix",
103 s3_input_mode="File",
104 ),
105 )
106 ],
107 outputs=[
108 ProcessingOutput(
109 output_name="transformed",
110 s3_output=ProcessingS3Output(
111 s3_uri=dataset_output_s3,
112 local_path=output_local_path,
113 s3_upload_mode="EndOfJob",
114 ),
115 )
116 ],
117 wait=False,
118 )
119
120 print(
121 f"Processing job '{processor.latest_job.processing_job_name}' submitted. Output will be at: {dataset_output_s3}"
122 )
123
124
125def describe_transformation_job(job_name, region=None):
126 """
127 Describe a SageMaker Processing Job by name.
128
129 Args:
130 job_name: The name of the processing job to describe.
131 region: AWS region (auto-detected if None).
132
133 Returns:
134 dict: Job details including status, inputs, outputs, and timing info.
135 """
136 sagemaker_session = _get_session(region)
137
138 job = ProcessingJob.get(
139 processing_job_name=job_name,
140 session=sagemaker_session.boto_session,
141 )
142
143 details = job.refresh().__dict__
144 return {
145 "job_name": job_name,
146 "status": details.get("processing_job_status"),
147 "failure_reason": details.get("failure_reason"),
148 "creation_time": str(details.get("creation_time", "")),
149 "processing_end_time": str(details.get("processing_end_time", "")),
150 "inputs": details.get("processing_inputs", []),
151 "outputs": getattr(details.get("processing_output_config"), "outputs", []),
152 }