Setting the file. One moment. Script Processing · AWS AI ML · aws/agent-toolkit-for-aws · Skills Docs93
Creating Amazon Aurora Db Cluster With Instances
Shieldadvanced
Resilience Program Design
Connecting Lambda To DynamoDB
(opens in a new tab)
references/processing-jobs/code_templates/script_processing.py
Python·99 lines·3 KB
import
ScriptProcessor, logs_for_processing_job
11from sagemaker.core.shapes import (
12 ProcessingInput,
13 ProcessingOutput,
14 ProcessingS3Input,
15 ProcessingS3Output,
16)
17
18set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)
19
20REGION = ""
21ROLE_ARN = "[ROLE_ARN]"
22SCRIPT_PATH = "[SCRIPT_PATH]"
23CUSTOM_IMAGE_URI = ""
24IMAGE_FRAMEWORK = "[IMAGE_FRAMEWORK]"
25IMAGE_VERSION = ""
26IMAGE_PYTHON_VERSION = ""
27INPUT_S3_URI = "[INPUT_S3_URI]"
28OUTPUT_S3_URI = "[OUTPUT_S3_URI]"
29INSTANCE_TYPE = "[INSTANCE_TYPE]"
30INSTANCE_COUNT = 1
31VOLUME_SIZE_GB = 30
32MAX_RUNTIME_SECONDS = 86400
33BASE_JOB_NAME = "[BASE_JOB_NAME]"
34ARGUMENTS = []
35WAIT_FOR_COMPLETION = False
36
37boto_session = boto3.Session(region_name=REGION or None)
38region = boto_session.region_name
39sm_client = boto_session.client("sagemaker")
40sm_session = Session(boto_session=boto_session, sagemaker_client=sm_client)
41role = ROLE_ARN or get_execution_role()
42image_uri = CUSTOM_IMAGE_URI or image_uris.retrieve(
43 framework=IMAGE_FRAMEWORK,
44 region=region,
45 version=IMAGE_VERSION or None,
46 py_version=IMAGE_PYTHON_VERSION or None,
47 instance_type=INSTANCE_TYPE,
48)
49print(f"Resolved processing image URI: {image_uri}")
50
51# Cell 3: Build and submit the processor
52processor = ScriptProcessor(
53 image_uri=image_uri,
54 role=role,
55 command=["python3"],
56 instance_type=INSTANCE_TYPE,
57 instance_count=INSTANCE_COUNT,
58 volume_size_in_gb=VOLUME_SIZE_GB,
59 max_runtime_in_seconds=MAX_RUNTIME_SECONDS,
60 base_job_name=BASE_JOB_NAME,
61 sagemaker_session=sm_session,
62)
63processor.run(
64 code=SCRIPT_PATH,
65 inputs=[
66 ProcessingInput(
67 input_name="input",
68 s3_input=ProcessingS3Input(
69 s3_uri=INPUT_S3_URI,
70 local_path="/opt/ml/processing/input",
71 s3_data_type="S3Prefix",
72 s3_input_mode="File",
73 ),
74 )
75 ],
76 outputs=[
77 ProcessingOutput(
78 output_name="output",
79 s3_output=ProcessingS3Output(
80 s3_uri=OUTPUT_S3_URI,
81 local_path="/opt/ml/processing/output",
82 s3_upload_mode="EndOfJob",
83 ),
84 )
85 ],
86 arguments=ARGUMENTS,
87 wait=False,
88 logs=False,
89)
90
91# Cell 4: Report status and optionally wait
92job_name = processor.latest_job.processing_job_name
93job_arn = sm_client.describe_processing_job(ProcessingJobName=job_name)["ProcessingJobArn"]
94print(f"Processing Job Name: {job_name}")
95print(f"Processing Job ARN: {job_arn}")
96
97if WAIT_FOR_COMPLETION:
98 logs_for_processing_job(sm_session, job_name, wait=True)
99 print(f"Output: {OUTPUT_S3_URI}")