<s> """<STR_LIT>""" <EOL> import os <EOL> import time <EOL> import logging <EOL> import luigi <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> _dataproc_client = None <EOL> try : <EOL> import httplib2 <EOL> import oauth2client . client <EOL> from googleapiclient import discovery <EOL> from googleapiclient . errors import HttpError <EOL> DEFAULT_CREDENTIALS = oauth2client . client . GoogleCredentials . get_application_default ( ) <EOL> _dataproc_client = discovery . build ( '<STR_LIT>' , '<STR_LIT>' , credentials = DEFAULT_CREDENTIALS , http = httplib2 . Http ( ) ) <EOL> except ImportError : <EOL> logger . warning ( "<STR_LIT>" ) <EOL> def get_dataproc_client ( ) : <EOL> return _dataproc_client <EOL> def set_dataproc_client ( client ) : <EOL> global _dataproc_client <EOL> _dataproc_client = client <EOL> class _DataprocBaseTask ( luigi . Task ) : <EOL> gcloud_project_id = luigi . Parameter ( significant = False , positional = False ) <EOL> dataproc_cluster_name = luigi . Parameter ( significant = False , positional = False ) <EOL> dataproc_region = luigi . Parameter ( default = "<STR_LIT>" , significant = False , positional = False ) <EOL> dataproc_client = get_dataproc_client ( ) <EOL> class DataprocBaseTask ( _DataprocBaseTask ) : <EOL> """<STR_LIT>""" <EOL> _job = None <EOL> _job_name = None <EOL> _job_id = None <EOL> def submit_job ( self , job_config ) : <EOL> self . _job = self . dataproc_client . projects ( ) . regions ( ) . jobs ( ) . submit ( projectId = self . gcloud_project_id , region = self . dataproc_region , body = job_config ) . execute ( ) <EOL> self . _job_id = self . _job [ '<STR_LIT>' ] [ '<STR_LIT>' ] <EOL> return self . _job <EOL> def submit_spark_job ( self , jars , main_class , job_args = [ ] ) : <EOL> job_config = { "<STR_LIT>" : { <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . dataproc_cluster_name <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT:args>" : job_args , <EOL> "<STR_LIT>" : main_class , <EOL> "<STR_LIT>" : jars <EOL> } <EOL> } } <EOL> self . submit_job ( job_config ) <EOL> self . _job_name = os . path . basename ( self . _job [ '<STR_LIT>' ] [ '<STR_LIT>' ] ) <EOL> logger . info ( "<STR_LIT>" . format ( self . _job_name , self . _job_id ) ) <EOL> return self . _job <EOL> def submit_pyspark_job ( self , job_file , extra_files = [ ] , job_args = [ ] ) : <EOL> job_config = { "<STR_LIT>" : { <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . dataproc_cluster_name <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : job_file , <EOL> "<STR_LIT>" : extra_files , <EOL> "<STR_LIT:args>" : job_args <EOL> } <EOL> } } <EOL> self . submit_job ( job_config ) <EOL> self . _job_name = os . path . basename ( self . _job [ '<STR_LIT>' ] [ '<STR_LIT>' ] ) <EOL> logger . info ( "<STR_LIT>" . format ( self . _job_name , self . _job_id ) ) <EOL> return self . _job <EOL> def wait_for_job ( self ) : <EOL> if self . _job is None : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> while True : <EOL> job_result = self . dataproc_client . projects ( ) . regions ( ) . jobs ( ) . get ( projectId = self . gcloud_project_id , region = self . dataproc_region , jobId = self . _job_id ) . execute ( ) <EOL> status = job_result [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] <EOL> logger . info ( "<STR_LIT>" . format ( status , self . _job_name , self . _job_id ) ) <EOL> if status == '<STR_LIT>' : <EOL> break <EOL> if status == '<STR_LIT>' : <EOL> raise Exception ( job_result [ '<STR_LIT:status>' ] [ '<STR_LIT>' ] ) <EOL> time . sleep ( <NUM_LIT:5> ) <EOL> class DataprocSparkTask ( DataprocBaseTask ) : <EOL> """<STR_LIT>""" <EOL> main_class = luigi . Parameter ( ) <EOL> jars = luigi . Parameter ( default = "<STR_LIT>" ) <EOL> job_args = luigi . Parameter ( default = "<STR_LIT>" ) <EOL> def run ( self ) : <EOL> self . submit_spark_job ( main_class = self . main_class , <EOL> jars = self . jars . split ( "<STR_LIT:U+002C>" ) if self . jars else [ ] , <EOL> job_args = self . job_args . split ( "<STR_LIT:U+002C>" ) if self . job_args else [ ] ) <EOL> self . wait_for_job ( ) <EOL> class DataprocPysparkTask ( DataprocBaseTask ) : <EOL> """<STR_LIT>""" <EOL> job_file = luigi . Parameter ( ) <EOL> extra_files = luigi . Parameter ( default = "<STR_LIT>" ) <EOL> job_args = luigi . Parameter ( default = "<STR_LIT>" ) <EOL> def run ( self ) : <EOL> self . submit_pyspark_job ( job_file = "<STR_LIT>" , <EOL> extra_files = self . extra_files . split ( "<STR_LIT:U+002C>" ) if self . extra_files else [ ] , <EOL> job_args = self . job_args . split ( "<STR_LIT:U+002C>" ) if self . job_args else [ ] ) <EOL> self . wait_for_job ( ) <EOL> class CreateDataprocClusterTask ( _DataprocBaseTask ) : <EOL> """<STR_LIT>""" <EOL> gcloud_zone = luigi . Parameter ( default = "<STR_LIT>" ) <EOL> gcloud_network = luigi . Parameter ( default = "<STR_LIT:default>" ) <EOL> master_node_type = luigi . Parameter ( default = "<STR_LIT>" ) <EOL> master_disk_size = luigi . Parameter ( default = "<STR_LIT:100>" ) <EOL> worker_node_type = luigi . Parameter ( default = "<STR_LIT>" ) <EOL> worker_disk_size = luigi . Parameter ( default = "<STR_LIT:100>" ) <EOL> worker_normal_count = luigi . Parameter ( default = "<STR_LIT:2>" ) <EOL> worker_preemptible_count = luigi . Parameter ( default = "<STR_LIT:0>" ) <EOL> def _get_cluster_status ( self ) : <EOL> return self . dataproc_client . projects ( ) . regions ( ) . clusters ( ) . get ( projectId = self . gcloud_project_id , region = self . dataproc_region , clusterName = self . dataproc_cluster_name ) . execute ( ) <EOL> def complete ( self ) : <EOL> try : <EOL> self . _get_cluster_status ( ) <EOL> return True <EOL> except HttpError as e : <EOL> if e . resp . status == <NUM_LIT> : <EOL> return False <EOL> else : <EOL> raise e <EOL> def run ( self ) : <EOL> base_uri = "<STR_LIT>" . format ( self . gcloud_project_id ) <EOL> cluster_conf = { <EOL> "<STR_LIT>" : self . dataproc_cluster_name , <EOL> "<STR_LIT>" : self . gcloud_project_id , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : base_uri + "<STR_LIT>" + self . gcloud_network , <EOL> "<STR_LIT>" : base_uri + "<STR_LIT>" + self . gcloud_zone , <EOL> "<STR_LIT>" : [ <EOL> "<STR_LIT>" <EOL> ] <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : <NUM_LIT:1> , <EOL> "<STR_LIT>" : base_uri + "<STR_LIT>" + self . gcloud_zone + "<STR_LIT>" + self . master_node_type , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . master_disk_size , <EOL> "<STR_LIT>" : <NUM_LIT:0> <EOL> } <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . worker_normal_count , <EOL> "<STR_LIT>" : base_uri + "<STR_LIT>" + self . gcloud_zone + "<STR_LIT>" + self . worker_node_type , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . worker_disk_size , <EOL> "<STR_LIT>" : <NUM_LIT:0> <EOL> } <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . worker_preemptible_count , <EOL> "<STR_LIT>" : True <EOL> } <EOL> } <EOL> } <EOL> self . dataproc_client . projects ( ) . regions ( ) . clusters ( ) . create ( projectId = self . gcloud_project_id , region = self . dataproc_region , body = cluster_conf ) . execute ( ) <EOL> while True : <EOL> time . sleep ( <NUM_LIT:10> ) <EOL> cluster_status = self . _get_cluster_status ( ) <EOL> status = cluster_status [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] <EOL> logger . info ( "<STR_LIT>" . format ( self . dataproc_cluster_name , status ) ) <EOL> if status == '<STR_LIT>' : <EOL> break <EOL> if status == '<STR_LIT>' : <EOL> raise Exception ( cluster_status [ '<STR_LIT:status>' ] [ '<STR_LIT>' ] ) <EOL> class DeleteDataprocClusterTask ( _DataprocBaseTask ) : <EOL> """<STR_LIT>""" <EOL> def _get_cluster_status ( self ) : <EOL> try : <EOL> return self . dataproc_client . projects ( ) . regions ( ) . clusters ( ) . get ( projectId = self . gcloud_project_id , region = self . dataproc_region , <EOL> clusterName = self . dataproc_cluster_name , fields = "<STR_LIT:status>" ) . execute ( ) <EOL> except HttpError as e : <EOL> if e . resp . status == <NUM_LIT> : <EOL> return None <EOL> else : <EOL> raise e <EOL> def complete ( self ) : return self . _get_cluster_status ( ) is None <EOL> def run ( self ) : <EOL> self . dataproc_client . projects ( ) . regions ( ) . clusters ( ) . delete ( projectId = self . gcloud_project_id , region = self . dataproc_region , clusterName = self . dataproc_cluster_name ) . execute ( ) <EOL> while True : <EOL> time . sleep ( <NUM_LIT:10> ) <EOL> status = self . _get_cluster_status ( ) <EOL> if status is None : <EOL> logger . info ( "<STR_LIT>" . format ( self . dataproc_cluster_name ) ) <EOL> break <EOL> logger . info ( "<STR_LIT>" . format ( self . dataproc_cluster_name , status [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] ) ) </s>
