Setting the file. One moment. Spark 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/spark_processing.py
Python·118 lines·4 KB
logs_for_processing_job
11from sagemaker.core.shapes import (
12 ProcessingInput,
13 ProcessingOutput,
14 ProcessingS3Input,
15 ProcessingS3Output,
16)
17from sagemaker.core.spark.processing import PySparkProcessor, SparkJarProcessor
18
19set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)
20
21REGION = ""
22ROLE_ARN = "[ROLE_ARN]"
23APPLICATION_KIND = "pyspark"
24SPARK_APPLICATION = "[SPARK_APPLICATION]"
25SPARK_MAIN_CLASS = "[SPARK_MAIN_CLASS]"
26SPARK_VERSION = ""
27SPARK_PYTHON_VERSION = ""
28SPARK_CONTAINER_VERSION = ""
29SPARK_EVENT_LOGS_S3_URI = "[SPARK_EVENT_LOGS_S3_URI]"
30INPUT_S3_URI = "[INPUT_S3_URI]"
31OUTPUT_S3_URI = "[OUTPUT_S3_URI]"
32INSTANCE_TYPE = "[INSTANCE_TYPE]"
33INSTANCE_COUNT = 2
34VOLUME_SIZE_GB = 100
35MAX_RUNTIME_SECONDS = 86400
36BASE_JOB_NAME = "[BASE_JOB_NAME]"
37ARGUMENTS = []
38SUBMIT_PY_FILES = []
39SUBMIT_JARS = []
40WAIT_FOR_COMPLETION = False
41
42boto_session = boto3.Session(region_name=REGION or None)
43sm_client = boto_session.client("sagemaker")
44sm_session = Session(boto_session=boto_session, sagemaker_client=sm_client)
45role = ROLE_ARN or get_execution_role()
46
47# Cell 3: Build and submit the Spark processor
48processor_kwargs = {
49 "role": role,
50 "framework_version": SPARK_VERSION or None,
51 "py_version": SPARK_PYTHON_VERSION or None,
52 "container_version": SPARK_CONTAINER_VERSION or None,
53 "instance_type": INSTANCE_TYPE,
54 "instance_count": INSTANCE_COUNT,
55 "volume_size_in_gb": VOLUME_SIZE_GB,
56 "max_runtime_in_seconds": MAX_RUNTIME_SECONDS,
57 "base_job_name": BASE_JOB_NAME,
58 "sagemaker_session": sm_session,
59}
60if APPLICATION_KIND == "pyspark":
61 processor = PySparkProcessor(**processor_kwargs)
62 dependency_kwargs = {
63 "submit_py_files": SUBMIT_PY_FILES,
64 "submit_jars": SUBMIT_JARS,
65 }
66elif APPLICATION_KIND == "jar":
67 if not SPARK_MAIN_CLASS:
68 raise ValueError("SPARK_MAIN_CLASS is required for JAR applications.")
69 processor = SparkJarProcessor(**processor_kwargs)
70 dependency_kwargs = {
71 "submit_class": SPARK_MAIN_CLASS,
72 "submit_jars": SUBMIT_JARS,
73 }
74else:
75 raise ValueError("APPLICATION_KIND must be 'pyspark' or 'jar'.")
76
77print(f"Resolved Spark image URI: {processor.image_uri}")
78
79processor.run(
80 submit_app=SPARK_APPLICATION,
81 inputs=[
82 ProcessingInput(
83 input_name="input",
84 s3_input=ProcessingS3Input(
85 s3_uri=INPUT_S3_URI,
86 local_path="/opt/ml/processing/input",
87 s3_data_type="S3Prefix",
88 s3_input_mode="File",
89 s3_data_distribution_type="ShardedByS3Key",
90 ),
91 )
92 ],
93 outputs=[
94 ProcessingOutput(
95 output_name="output",
96 s3_output=ProcessingS3Output(
97 s3_uri=OUTPUT_S3_URI,
98 local_path="/opt/ml/processing/output",
99 s3_upload_mode="EndOfJob",
100 ),
101 )
102 ],
103 arguments=ARGUMENTS,
104 spark_event_logs_s3_uri=SPARK_EVENT_LOGS_S3_URI,
105 wait=False,
106 logs=False,
107 **dependency_kwargs,
108)
109
110# Cell 4: Report status and optionally wait
111job_name = processor.latest_job.processing_job_name
112job_arn = sm_client.describe_processing_job(ProcessingJobName=job_name)["ProcessingJobArn"]
113print(f"Processing Job Name: {job_name}")
114print(f"Processing Job ARN: {job_arn}")
115
116if WAIT_FOR_COMPLETION:
117 logs_for_processing_job(sm_session, job_name, wait=True)
118 print(f"Output: {OUTPUT_S3_URI}")