<s> """<STR_LIT>""" <EOL> from __future__ import absolute_import <EOL> import logging <EOL> import random <EOL> import sys <EOL> import time <EOL> import traceback <EOL> from grpc . beta import implementations <EOL> from google . cloud . dataflow . internal import windmill_pb2 <EOL> from google . cloud . dataflow . internal import windmill_service_pb2 <EOL> from google . cloud . dataflow . utils import retry <EOL> from google . cloud . dataflow . worker import executor <EOL> from google . cloud . dataflow . worker import maptask <EOL> from google . cloud . dataflow . worker import windmillio <EOL> from google . cloud . dataflow . worker import windmillstate <EOL> import apitools . base . py as apitools_base <EOL> import google . cloud . dataflow . internal . clients . dataflow as dataflow <EOL> class WindmillClient ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , host , port , request_timeout = <NUM_LIT:10> ) : <EOL> self . host = host <EOL> self . port = port <EOL> self . request_timeout = request_timeout <EOL> channel = implementations . insecure_channel ( host , port ) <EOL> self . stub = ( <EOL> windmill_service_pb2 . beta_create_CloudWindmillServiceV1Alpha1_stub ( <EOL> channel ) ) <EOL> @ retry . with_exponential_backoff ( ) <EOL> def GetWork ( self , request ) : <EOL> return self . stub . GetWork ( request , self . request_timeout ) <EOL> @ retry . with_exponential_backoff ( ) <EOL> def GetData ( self , request ) : <EOL> return self . stub . GetData ( request , self . request_timeout ) <EOL> @ retry . with_exponential_backoff ( ) <EOL> def CommitWork ( self , request ) : <EOL> return self . stub . CommitWork ( request , self . request_timeout ) <EOL> @ retry . with_exponential_backoff ( ) <EOL> def GetConfig ( self , request ) : <EOL> return self . stub . GetConfig ( request , self . request_timeout ) <EOL> @ retry . with_exponential_backoff ( ) <EOL> def ReportStats ( self , request ) : <EOL> return self . stub . ReportStats ( request , self . request_timeout ) <EOL> class StreamingWorker ( object ) : <EOL> """<STR_LIT>""" <EOL> MAX_GET_WORK_FETCH_BYTES = <NUM_LIT:64> << <NUM_LIT:20> <EOL> MAX_GET_WORK_ITEMS = <NUM_LIT:100> <EOL> RETRY_LOCALLY_DELAY = <NUM_LIT> <EOL> def __init__ ( self , properties ) : <EOL> self . project_id = properties [ '<STR_LIT>' ] <EOL> self . job_id = properties [ '<STR_LIT>' ] <EOL> self . worker_id = properties [ '<STR_LIT>' ] <EOL> self . client_id = random . getrandbits ( <NUM_LIT> ) <EOL> windmill_host = properties [ '<STR_LIT>' ] <EOL> windmill_port = int ( properties [ '<STR_LIT>' ] ) <EOL> logging . info ( '<STR_LIT>' , windmill_host , <EOL> windmill_port ) <EOL> self . windmill = WindmillClient ( windmill_host , windmill_port ) <EOL> self . instruction_map = { } <EOL> self . system_name_to_computation_id_map = { } <EOL> def run ( self ) : <EOL> self . running = True <EOL> self . dispatch_loop ( ) <EOL> def get_work ( self ) : <EOL> request = windmill_pb2 . GetWorkRequest ( <EOL> client_id = self . client_id , <EOL> max_items = StreamingWorker . MAX_GET_WORK_ITEMS , <EOL> max_bytes = StreamingWorker . MAX_GET_WORK_FETCH_BYTES ) <EOL> return self . windmill . GetWork ( request ) <EOL> def add_computation ( self , map_task ) : <EOL> computation_id = self . system_name_to_computation_id_map . get ( <EOL> map_task . systemName , map_task . systemName ) <EOL> if computation_id not in self . instruction_map : <EOL> self . instruction_map [ computation_id ] = map_task <EOL> def parse_map_task ( self , serialized_map_task ) : <EOL> return apitools_base . JsonToMessage ( dataflow . MapTask , serialized_map_task ) <EOL> def get_config ( self , computation_id ) : <EOL> """<STR_LIT>""" <EOL> request = windmill_pb2 . GetConfigRequest ( computations = [ computation_id ] ) <EOL> response = self . windmill . GetConfig ( request ) <EOL> for map_entry in response . system_name_to_computation_id_map : <EOL> self . system_name_to_computation_id_map [ <EOL> map_entry . system_name ] = map_entry . computation_id <EOL> for serialized_map_task in response . cloud_works : <EOL> logging . info ( '<STR_LIT>' , computation_id , <EOL> serialized_map_task ) <EOL> self . add_computation ( self . parse_map_task ( serialized_map_task ) ) <EOL> return response <EOL> def dispatch_loop ( self ) : <EOL> while self . running : <EOL> backoff_seconds = <NUM_LIT> <EOL> while self . running : <EOL> work_response = self . get_work ( ) <EOL> if work_response . work : <EOL> break <EOL> time . sleep ( backoff_seconds ) <EOL> backoff_seconds = min ( <NUM_LIT:1.0> , backoff_seconds * <NUM_LIT:2> ) <EOL> for computation_work in work_response . work : <EOL> self . process_computation ( computation_work ) <EOL> def process_computation ( self , computation_work ) : <EOL> computation_id = computation_work . computation_id <EOL> input_data_watermark = windmillio . windmill_to_harness_timestamp ( <EOL> computation_work . input_data_watermark ) <EOL> if computation_id not in self . instruction_map : <EOL> self . get_config ( computation_id ) <EOL> map_task_proto = self . instruction_map [ computation_id ] <EOL> for work_item in computation_work . work : <EOL> retry_locally = True <EOL> while retry_locally : <EOL> try : <EOL> self . process_work_item ( computation_id , map_task_proto , <EOL> input_data_watermark , work_item ) <EOL> break <EOL> except : <EOL> logging . error ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' , computation_id , work_item , traceback . format_exc ( ) ) <EOL> retry_locally = self . report_failure ( computation_id , work_item , <EOL> sys . exc_info ( ) ) <EOL> if retry_locally : <EOL> logging . error ( '<STR_LIT>' <EOL> '<STR_LIT>' , computation_id , <EOL> work_item . key ) <EOL> time . sleep ( StreamingWorker . RETRY_LOCALLY_DELAY ) <EOL> else : <EOL> logging . error ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' , computation_id , work_item . key ) <EOL> def report_failure ( self , computation_id , work_item , exc_info ) : <EOL> """<STR_LIT>""" <EOL> exc_type , exc_value , exc_traceback = exc_info <EOL> messages = list ( line . strip ( ) for line in <EOL> ( traceback . format_exception_only ( exc_type , <EOL> exc_value ) + <EOL> traceback . format_tb ( exc_traceback ) ) ) <EOL> wm_exception = windmill_pb2 . Exception ( stack_frames = messages ) <EOL> report_stats_request = windmill_pb2 . ReportStatsRequest ( <EOL> computation_id = computation_id , <EOL> key = work_item . key , <EOL> sharding_key = work_item . sharding_key , <EOL> work_token = work_item . work_token , <EOL> exceptions = [ wm_exception ] ) <EOL> response = self . windmill . ReportStats ( report_stats_request ) <EOL> return not response . failed <EOL> def process_work_item ( self , computation_id , map_task_proto , <EOL> input_data_watermark , work_item ) : <EOL> """<STR_LIT>""" <EOL> workitem_commit_request = windmill_pb2 . WorkItemCommitRequest ( <EOL> key = work_item . key , <EOL> work_token = work_item . work_token ) <EOL> env = maptask . WorkerEnvironment ( ) <EOL> context = maptask . StreamingExecutionContext ( ) <EOL> reader = windmillstate . WindmillStateReader ( <EOL> computation_id , <EOL> work_item . key , <EOL> work_item . work_token , <EOL> self . windmill ) <EOL> state_internals = windmillstate . WindmillStateInternals ( reader ) <EOL> state = windmillstate . WindmillUnmergedState ( state_internals ) <EOL> output_data_watermark = windmillio . windmill_to_harness_timestamp ( <EOL> work_item . output_data_watermark ) <EOL> context . start ( computation_id , work_item , input_data_watermark , <EOL> output_data_watermark , workitem_commit_request , <EOL> self . windmill , state ) <EOL> map_task_executor = executor . MapTaskExecutor ( ) <EOL> map_task = maptask . decode_map_task ( map_task_proto , env , context ) <EOL> map_task_executor . execute ( map_task ) <EOL> state_internals . persist_to ( workitem_commit_request ) <EOL> commit_request = windmill_pb2 . CommitWorkRequest ( ) <EOL> computation_commit_request = windmill_pb2 . ComputationCommitWorkRequest ( <EOL> computation_id = computation_id , <EOL> requests = [ workitem_commit_request ] ) <EOL> commit_request . requests . extend ( [ computation_commit_request ] ) <EOL> self . windmill . CommitWork ( commit_request ) </s>
