Setting the file. One moment. Dataset Definition Processing
93
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/dataset_definition_processing.py
Python·127 lines·4 KB
import
ScriptProcessor, logs_for_processing_job
11from sagemaker.core.shapes import (
12 AthenaDatasetDefinition,
13 DatasetDefinition,
14 ProcessingInput,
15 ProcessingOutput,
16 ProcessingS3Output,
17 RedshiftDatasetDefinition,
18)
19
20set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)
21
22REGION = ""
23ROLE_ARN = "[ROLE_ARN]"
24SOURCE_KIND = "athena"
25PROCESSING_SCRIPT = "[QUERY_PROCESSING_SCRIPT]"
26CUSTOM_IMAGE_URI = ""
27IMAGE_FRAMEWORK = "[IMAGE_FRAMEWORK]"
28IMAGE_VERSION = ""
29IMAGE_PYTHON_VERSION = ""
30QUERY_RESULT_S3_URI = "[QUERY_RESULT_S3_URI]"
31OUTPUT_S3_URI = "[OUTPUT_S3_URI]"
32OUTPUT_FORMAT = "PARQUET"
33INSTANCE_TYPE = "[INSTANCE_TYPE]"
34MAX_RUNTIME_SECONDS = 86400
35BASE_JOB_NAME = "[BASE_JOB_NAME]"
36WAIT_FOR_COMPLETION = False
37
38ATHENA_CATALOG = "[ATHENA_CATALOG]"
39ATHENA_DATABASE = "[ATHENA_DATABASE]"
40ATHENA_WORK_GROUP = "primary"
41ATHENA_QUERY = "[ATHENA_QUERY]"
42
43REDSHIFT_CLUSTER = "[REDSHIFT_CLUSTER]"
44REDSHIFT_DATABASE = "[REDSHIFT_DATABASE]"
45REDSHIFT_DB_USER = "[REDSHIFT_DB_USER]"
46REDSHIFT_CLUSTER_ROLE_ARN = "[REDSHIFT_CLUSTER_ROLE_ARN]"
47REDSHIFT_QUERY = "[REDSHIFT_QUERY]"
48
49boto_session = boto3.Session(region_name=REGION or None)
50region = boto_session.region_name
51sm_client = boto_session.client("sagemaker")
52sm_session = Session(boto_session=boto_session, sagemaker_client=sm_client)
53role = ROLE_ARN or get_execution_role()
54image_uri = CUSTOM_IMAGE_URI or image_uris.retrieve(
55 framework=IMAGE_FRAMEWORK,
56 region=region,
57 version=IMAGE_VERSION or None,
58 py_version=IMAGE_PYTHON_VERSION or None,
59 instance_type=INSTANCE_TYPE,
60)
61print(f"Resolved processing image URI: {image_uri}")
62
63# Cell 3: Build the dataset definition and submit
64if SOURCE_KIND == "athena":
65 source = DatasetDefinition(
66 athena_dataset_definition=AthenaDatasetDefinition(
67 catalog=ATHENA_CATALOG,
68 database=ATHENA_DATABASE,
69 work_group=ATHENA_WORK_GROUP,
70 query_string=ATHENA_QUERY,
71 output_s3_uri=QUERY_RESULT_S3_URI,
72 output_format=OUTPUT_FORMAT,
73 ),
74 local_path="/opt/ml/processing/input/query",
75 )
76elif SOURCE_KIND == "redshift":
77 source = DatasetDefinition(
78 redshift_dataset_definition=RedshiftDatasetDefinition(
79 cluster_id=REDSHIFT_CLUSTER,
80 database=REDSHIFT_DATABASE,
81 db_user=REDSHIFT_DB_USER,
82 cluster_role_arn=REDSHIFT_CLUSTER_ROLE_ARN,
83 query_string=REDSHIFT_QUERY,
84 output_s3_uri=QUERY_RESULT_S3_URI,
85 output_format=OUTPUT_FORMAT,
86 ),
87 local_path="/opt/ml/processing/input/query",
88 )
89else:
90 raise ValueError("SOURCE_KIND must be 'athena' or 'redshift'.")
91
92processor = ScriptProcessor(
93 image_uri=image_uri,
94 role=role,
95 command=["python3"],
96 instance_type=INSTANCE_TYPE,
97 instance_count=1,
98 max_runtime_in_seconds=MAX_RUNTIME_SECONDS,
99 base_job_name=BASE_JOB_NAME,
100 sagemaker_session=sm_session,
101)
102processor.run(
103 code=PROCESSING_SCRIPT,
104 inputs=[ProcessingInput(input_name="query", dataset_definition=source)],
105 outputs=[
106 ProcessingOutput(
107 output_name="output",
108 s3_output=ProcessingS3Output(
109 s3_uri=OUTPUT_S3_URI,
110 local_path="/opt/ml/processing/output",
111 s3_upload_mode="EndOfJob",
112 ),
113 )
114 ],
115 wait=False,
116 logs=False,
117)
118
119# Cell 4: Report status and optionally wait
120job_name = processor.latest_job.processing_job_name
121job_arn = sm_client.describe_processing_job(ProcessingJobName=job_name)["ProcessingJobArn"]
122print(f"Processing Job Name: {job_name}")
123print(f"Processing Job ARN: {job_arn}")
124
125if WAIT_FOR_COMPLETION:
126 logs_for_processing_job(sm_session, job_name, wait=True)
127 print(f"Output: {OUTPUT_S3_URI}")