-
Notifications
You must be signed in to change notification settings - Fork 6.6k
feat(dataproc): create pyspark nodegroup cluster sample #13513
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,81 @@ | ||||||||||||
#!/usr/bin/env python | ||||||||||||
|
||||||||||||
# Copyright 2025 Google LLC | ||||||||||||
# | ||||||||||||
# Licensed under the Apache License, Version 2.0 (the "License"); | ||||||||||||
# you may not use this file except in compliance with the License. | ||||||||||||
# You may obtain a copy of the License at | ||||||||||||
# | ||||||||||||
# http://www.apache.org/licenses/LICENSE-2.0 | ||||||||||||
# | ||||||||||||
# Unless required by applicable law or agreed to in writing, software | ||||||||||||
# distributed under the License is distributed on an "AS IS" BASIS, | ||||||||||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||||||||||||
# See the License for the specific language governing permissions and | ||||||||||||
# limitations under the License. | ||||||||||||
|
||||||||||||
# This sample walks a user through submitting a Spark job to a | ||||||||||||
# Dataproc driver node group cluster using the Dataproc | ||||||||||||
# client library. | ||||||||||||
|
||||||||||||
# Usage: | ||||||||||||
# python submit_pyspark_job_to_driver_node_group_cluster.py \ | ||||||||||||
# --project_id <PROJECT_ID> --region <REGION> \ | ||||||||||||
# --cluster_name <CLUSTER_NAME> | ||||||||||||
Comment on lines
+22
to
+24
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The
Suggested change
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||||||||||||
|
||||||||||||
# [START dataproc_submit_pyspark_job_to_driver_node_group_cluster] | ||||||||||||
|
||||||||||||
import re | ||||||||||||
|
||||||||||||
|
||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||||||||||||
from google.cloud import dataproc_v1 as dataproc | ||||||||||||
from google.cloud import storage | ||||||||||||
|
||||||||||||
|
||||||||||||
def submit_job(project_id, region, cluster_name): | ||||||||||||
glasnt marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||
# Create the job client. | ||||||||||||
job_client = dataproc.JobControllerClient( | ||||||||||||
client_options={"api_endpoint": f"{region}-dataproc.googleapis.com:443"} | ||||||||||||
) | ||||||||||||
|
||||||||||||
driver_scheduling_config = dataproc.DriverSchedulingConfig( | ||||||||||||
memory_mb=2048, # Example memory in MB | ||||||||||||
vcores=2, # Example number of vcores | ||||||||||||
) | ||||||||||||
|
||||||||||||
# Create the job config. 'main_jar_file_uri' can also be a | ||||||||||||
# Google Cloud Storage URL. | ||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||||||||||||
job = { | ||||||||||||
"placement": {"cluster_name": cluster_name}, | ||||||||||||
"pyspark_job": { | ||||||||||||
"main_python_file_uri": "gs://dataproc-examples/pyspark/hello-world/hello-world.py" | ||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Where does this object come from? Is there open source code for this object? |
||||||||||||
}, | ||||||||||||
"driver_scheduling_config": driver_scheduling_config | ||||||||||||
} | ||||||||||||
|
||||||||||||
operation = job_client.submit_job_as_operation( | ||||||||||||
request={"project_id": project_id, "region": region, "job": job} | ||||||||||||
) | ||||||||||||
response = operation.result() | ||||||||||||
|
||||||||||||
# Dataproc job output gets saved to the Google Cloud Storage bucket | ||||||||||||
# allocated to the job. Use a regex to obtain the bucket and blob info. | ||||||||||||
matches = re.match("gs://(.*?)/(.*)", response.driver_output_resource_uri) | ||||||||||||
glasnt marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||
|
||||||||||||
output = ( | ||||||||||||
storage.Client() | ||||||||||||
.get_bucket(matches.group(1)) | ||||||||||||
.blob(f"{matches.group(2)}.000000000") | ||||||||||||
.download_as_bytes() | ||||||||||||
.decode("utf-8") | ||||||||||||
) | ||||||||||||
|
||||||||||||
print(f"Job finished successfully: {output}") | ||||||||||||
|
||||||||||||
if __name__ == "__main__": | ||||||||||||
|
||||||||||||
my_project_id = "your_cluster" # <-- REPLACE THIS | ||||||||||||
my_region = "us-central1" # <-- REPLACE THIS | ||||||||||||
my_cluster_name = "your-node-group-cluster" # <-- REPLACE THIS | ||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The script uses hardcoded placeholder values, which requires users to manually edit the file before running it. To make the sample easier to use and consistent with other samples in the repository, it should parse command-line arguments. Additionally, the placeholder value for Please consider replacing the current implementation with one that uses if __name__ == "__main__":
import sys
if len(sys.argv) != 4:
sys.exit(
f"Usage: python {sys.argv[0]} <project_id> <region> <cluster_name>"
)
project_id = sys.argv[1]
region = sys.argv[2]
cluster_name = sys.argv[3]
submit_job(project_id, region, cluster_name) There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. corrected the placeholder value for my_project_id. I intend the main function to be as is. |
||||||||||||
|
||||||||||||
submit_job(my_project_id, my_region, my_cluster_name) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This file needs to be moved to the
dataproc/snippets
folder.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Also, please clarify that the existing "submit_spark_job_to_driver_node_group_cluster.py" sample is different to this "submit_pyspark_job_to_driver_node_group_cluster.py` sample.