def GetStateStream(self, request, context=None):
job = self._jobs[request.job_id]
state_queue = queue.Queue()
job.add_state_change_callback(lambda state: state_queue.put(state))
try:
current_state = state_queue.get()
except queue.Empty:
current_state = job.state
yield beam_job_api_pb2.GetJobStateResponse(state=current_state)
while current_state not in TERMINAL_STATES:
current_state = state_queue.get(block=True)
yield beam_job_api_pb2.GetJobStateResponse(state=current_state)
def GetMessageStream(self, request, context=None):
job = self._jobs[request.job_id]
log_queue = queue.Queue()
if job._last_log_message:
# This is likely to contain important information, like errors for
# an already failed job.
# TODO: Decide on proper semantics for the message stream of a
# long-running or completed job.
yield job._last_log_message
job.add_log_callback(log_queue.put)
job.add_state_change_callback(lambda state: log_queue.put(
beam_job_api_pb2.JobMessagesResponse(
state_response=beam_job_api_pb2.GetJobStateResponse(state=state
))))
current_state = job.state
while current_state not in TERMINAL_STATES:
msg = log_queue.get(block=True)
yield msg
if msg.HasField('state_response'):
current_state = msg.state_response.state
try:
while True:
yield log_queue.get(block=False)
except queue.Empty:
pass
def GetStateStream(self, request, context=None):
"""Yields state transitions since the stream started.
"""
if request.job_id not in self._jobs:
raise LookupError("Job {} does not exist".format(request.job_id))
job = self._jobs[request.job_id]
for state in job.get_state_stream():
yield beam_job_api_pb2.GetJobStateResponse(state=state)
def GetMessageStream(self, request, context=None):
"""Yields messages since the stream started.
"""
if request.job_id not in self._jobs:
raise LookupError("Job {} does not exist".format(request.job_id))
job = self._jobs[request.job_id]
for msg in job.get_message_stream():
if isinstance(msg, int):
resp = beam_job_api_pb2.JobMessagesResponse(
state_response=beam_job_api_pb2.GetJobStateResponse(state=msg))
else:
resp = beam_job_api_pb2.JobMessagesResponse(message_response=msg)
yield resp
def __init__(self, job_id, pipeline_options, pipeline_proto):
super(BeamJob, self).__init__()
self._job_id = job_id
self._pipeline_options = pipeline_options
self._pipeline_proto = pipeline_proto
self._log_queue = queue.Queue()
self._state_change_callbacks = [
lambda new_state: self._log_queue.put(
beam_job_api_pb2.JobMessagesResponse(
state_response=beam_job_api_pb2.GetJobStateResponse(
state=new_state)))
]
self._state = None
self.state = beam_job_api_pb2.JobState.STARTING
self.daemon = True
def GetState(self, request, context=None):
return beam_job_api_pb2.GetJobStateResponse(
state=self._jobs[request.job_id].state)
def test_end_to_end(self, http_mock):
with temp_name(suffix='fake.jar') as fake_jar:
# Create the jar file with some trivial contents.
with zipfile.ZipFile(fake_jar, 'w') as zip:
with zip.open('FakeClass.class', 'w') as fout:
fout.write(b'[original_contents]')
job_server = flink_uber_jar_job_server.FlinkUberJarJobServer(
'http://flink', fake_jar)
# Prepare the job.
prepare_response = job_server.Prepare(
beam_job_api_pb2.PrepareJobRequest(
job_name='job', pipeline=beam_runner_api_pb2.Pipeline()))
channel = grpc.insecure_channel(
prepare_response.artifact_staging_endpoint.url)
retrieval_token = beam_artifact_api_pb2_grpc.ArtifactStagingServiceStub(
channel).CommitManifest(
beam_artifact_api_pb2.CommitManifestRequest(
staging_session_token=prepare_response.
staging_session_token,
manifest=beam_artifact_api_pb2.Manifest())
).retrieval_token
channel.close()
# Now actually run the job.
http_mock.post('http://flink/v1/jars/upload',
json={'filename': '/path/to/jar/nonce'})
http_mock.post('http://flink/v1/jars/nonce/run',
json={'jobid': 'some_job_id'})
job_server.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id,
retrieval_token=retrieval_token))
# Check the status until the job is "done" and get all error messages.
http_mock.get('http://flink/v1/jobs/some_job_id/execution-result',
[{
'json': {
'status': {
'id': 'IN_PROGRESS'
}
}
}, {
'json': {
'status': {
'id': 'IN_PROGRESS'
}
}
}, {
'json': {
'status': {
'id': 'COMPLETED'
}
}
}])
http_mock.get('http://flink/v1/jobs/some_job_id',
json={'state': 'FINISHED'})
http_mock.delete('http://flink/v1/jars/nonce')
state_stream = job_server.GetStateStream(
beam_job_api_pb2.GetJobStateRequest(
job_id=prepare_response.preparation_id))
self.assertEqual([s.state for s in state_stream], [
beam_job_api_pb2.JobState.RUNNING,
beam_job_api_pb2.JobState.DONE
])
http_mock.get('http://flink/v1/jobs/some_job_id/exceptions',
json={
'all-exceptions': [{
'exception': 'exc_text',
'timestamp': 0
}]
})
message_stream = job_server.GetMessageStream(
beam_job_api_pb2.JobMessagesRequest(
job_id=prepare_response.preparation_id))
self.assertEqual([m for m in message_stream], [
beam_job_api_pb2.JobMessagesResponse(
message_response=beam_job_api_pb2.JobMessage(
message_id='message0',
time='0',
importance=beam_job_api_pb2.JobMessage.
MessageImportance.JOB_MESSAGE_ERROR,
message_text='exc_text')),
beam_job_api_pb2.JobMessagesResponse(
state_response=beam_job_api_pb2.GetJobStateResponse(
state=beam_job_api_pb2.JobState.DONE)),
])
def getMessageStream(self, request, context):
pipeline_result = self.jobs[request.jobId]
pipeline_result.wait_until_finish()
yield beam_job_api_pb2.JobMessagesResponse(
stateResponse=beam_job_api_pb2.GetJobStateResponse(
state=self._map_state_to_jobState(pipeline_result.state)))
def getState(self, request, context):
pipeline_result = self.jobs[request.jobId]
return beam_job_api_pb2.GetJobStateResponse(
state=self._map_state_to_jobState(pipeline_result.state))