from google.api_core import retry
from google.cloud import dataproc_v1
from google.api_core.client_options import ClientOptions

import time

project = '. . .'
region = 'us-central1'

# Create Dataproc batch
client_options = ClientOptions(api_endpoint="{}-dataproc.googleapis.com:443".format(region))
client = dataproc_v1.BatchControllerClient(client_options=client_options)
batch = dataproc_v1.Batch()
batch.spark_batch.main_class = 'org.apache.spark.examples.SparkPi'
batch.spark_batch.jar_file_uris = [
  'file:///usr/lib/spark/examples/jars/spark-examples.jar',
]
batch.spark_batch.args = ['1']
batch.runtime_config.properties = {
  'spark.executor.instances': '2',
}

parent = f"projects/{project}/locations/{region}"
request = dataproc_v1.CreateBatchRequest(
  parent=parent,  # type: ignore
  batch=batch,  # type: ignore
)

# Make the request
print("Creating batch")
operation = client.create_batch(request=request)  # type: ignore
print("Batch created")

# This takes quite a while, waiting on GCP response to resolve
retry = retry.Retry(initial=10, maximum=10, multiplier=1.0, deadline=600)
start = time.time()
response = operation.result(retry=retry)
print("Batch finished")
print(f"Took {time.time() - start} seconds to finish batch job")
print(response)

while not operation.done(retry=None):
  time.sleep(10)
response = operation.metadata