def run_pipeline(self, pipeline):
# Java has different expectations about coders
# (windowed in Fn API, but *un*windowed in runner API), whereas the
# FnApiRunner treats them consistently, so we must guard this.
# See also BEAM-2717.
proto_context = pipeline_context.PipelineContext(
default_environment_url=self._docker_image)
proto_pipeline = pipeline.to_runner_api(context=proto_context)
if self._runner_api_address:
for pcoll in proto_pipeline.components.pcollections.values():
if pcoll.coder_id not in proto_context.coders:
coder = coders.registry.get_coder(pickler.loads(pcoll.coder_id))
pcoll.coder_id = proto_context.coders.get_id(coder)
proto_context.coders.populate_map(proto_pipeline.components.coders)
# Some runners won't detect the GroupByKey transform unless it has no
# subtransforms. Remove all sub-transforms until BEAM-4605 is resolved.
for _, transform_proto in list(
proto_pipeline.components.transforms.items()):
if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
for sub_transform in transform_proto.subtransforms:
del proto_pipeline.components.transforms[sub_transform]
del transform_proto.subtransforms[:]
job_service = self._create_job_service()
prepare_response = job_service.Prepare(
beam_job_api_pb2.PrepareJobRequest(
job_name='job', pipeline=proto_pipeline))
run_response = job_service.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id))
return PipelineResult(job_service, run_response.job_id)
def prepare(self, proto_pipeline):
# type: (beam_runner_api_pb2.Pipeline) -> beam_job_api_pb2.PrepareJobResponse
"""Prepare the job on the job service"""
return self.job_service.Prepare(beam_job_api_pb2.PrepareJobRequest(
job_name='job',
pipeline=proto_pipeline,
pipeline_options=self.get_pipeline_options()),
timeout=self.timeout)
def run(self, pipeline):
job_service = self._get_job_service()
prepare_response = job_service.Prepare(
beam_job_api_pb2.PrepareJobRequest(
job_name='job', pipeline=pipeline.to_runner_api()))
run_response = job_service.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id))
return PipelineResult(job_service, run_response.job_id)
def run_pipeline(self, pipeline):
docker_image = (
pipeline.options.view_as(PortableOptions).harness_docker_image
or self.default_docker_image())
job_endpoint = pipeline.options.view_as(PortableOptions).job_endpoint
if not job_endpoint:
raise ValueError(
'job_endpoint should be provided while creating runner.')
proto_context = pipeline_context.PipelineContext(
default_environment_url=docker_image)
proto_pipeline = pipeline.to_runner_api(context=proto_context)
if not self.is_embedded_fnapi_runner:
# Java has different expectations about coders
# (windowed in Fn API, but *un*windowed in runner API), whereas the
# embedded FnApiRunner treats them consistently, so we must guard this
# for now, until FnApiRunner is fixed.
# See also BEAM-2717.
for pcoll in proto_pipeline.components.pcollections.values():
if pcoll.coder_id not in proto_context.coders:
# This is not really a coder id, but a pickled coder.
coder = coders.registry.get_coder(
pickler.loads(pcoll.coder_id))
pcoll.coder_id = proto_context.coders.get_id(coder)
proto_context.coders.populate_map(proto_pipeline.components.coders)
# Some runners won't detect the GroupByKey transform unless it has no
# subtransforms. Remove all sub-transforms until BEAM-4605 is resolved.
for _, transform_proto in list(
proto_pipeline.components.transforms.items()):
if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
for sub_transform in transform_proto.subtransforms:
del proto_pipeline.components.transforms[sub_transform]
del transform_proto.subtransforms[:]
job_service = beam_job_api_pb2_grpc.JobServiceStub(
grpc.insecure_channel(job_endpoint))
prepare_response = job_service.Prepare(
beam_job_api_pb2.PrepareJobRequest(job_name='job',
pipeline=proto_pipeline))
if prepare_response.artifact_staging_endpoint.url:
stager = portable_stager.PortableStager(
grpc.insecure_channel(
prepare_response.artifact_staging_endpoint.url),
prepare_response.staging_session_token)
retrieval_token, _ = stager.stage_job_resources(
pipeline._options, staging_location='')
else:
retrieval_token = None
run_response = job_service.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id,
retrieval_token=retrieval_token))
return PipelineResult(job_service, run_response.job_id)
def test_end_to_end(self):
job_service = local_job_service.LocalJobServicer()
job_service.start_grpc_server()
# this logic is taken roughly from PortableRunner.run_pipeline()
# Prepare the job.
prepare_response = job_service.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()
state_stream = job_service.GetStateStream(
beam_job_api_pb2.GetJobStateRequest(
job_id=prepare_response.preparation_id))
# If there's an error, we don't always get it until we try to read.
# Fortunately, there's always an immediate current state published.
# state_results.append(next(state_stream))
state_stream = increment_iter(state_stream)
message_stream = job_service.GetMessageStream(
beam_job_api_pb2.JobMessagesRequest(
job_id=prepare_response.preparation_id))
job_service.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id,
retrieval_token=retrieval_token))
state_results = list(state_stream)
message_results = list(message_stream)
expected_states = [
beam_job_api_pb2.JobState.STOPPED,
beam_job_api_pb2.JobState.STARTING,
beam_job_api_pb2.JobState.RUNNING,
beam_job_api_pb2.JobState.DONE,
]
self.assertEqual([s.state for s in state_results], expected_states)
self.assertEqual([s.state_response.state for s in message_results],
expected_states)
def run_pipeline(self, pipeline):
# Java has different expectations about coders
# (windowed in Fn API, but *un*windowed in runner API), whereas the
# FnApiRunner treats them consistently, so we must guard this.
# See also BEAM-2717.
proto_context = pipeline_context.PipelineContext(
default_environment_url=self._docker_image)
proto_pipeline = pipeline.to_runner_api(context=proto_context)
if self._runner_api_address:
for pcoll in proto_pipeline.components.pcollections.values():
if pcoll.coder_id not in proto_context.coders:
coder = coders.registry.get_coder(
pickler.loads(pcoll.coder_id))
pcoll.coder_id = proto_context.coders.get_id(coder)
proto_context.coders.populate_map(proto_pipeline.components.coders)
# Some runners won't detect the GroupByKey transform unless it has no
# subtransforms. Remove all sub-transforms until BEAM-4605 is resolved.
for _, transform_proto in list(
proto_pipeline.components.transforms.items()):
if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
for sub_transform in transform_proto.subtransforms:
del proto_pipeline.components.transforms[sub_transform]
del transform_proto.subtransforms[:]
job_service = self._create_job_service()
prepare_response = job_service.Prepare(
beam_job_api_pb2.PrepareJobRequest(job_name='job',
pipeline=proto_pipeline))
if prepare_response.artifact_staging_endpoint.url:
# Must commit something to get a retrieval token,
# committing empty manifest for now.
# TODO(BEAM-3883): Actually stage required files.
artifact_service = beam_artifact_api_pb2_grpc.ArtifactStagingServiceStub(
grpc.insecure_channel(
prepare_response.artifact_staging_endpoint.url))
commit_manifest = artifact_service.CommitManifest(
beam_artifact_api_pb2.CommitManifestRequest(
manifest=beam_artifact_api_pb2.Manifest(),
staging_session_token=prepare_response.
staging_session_token))
retrieval_token = commit_manifest.retrieval_token
else:
retrieval_token = None
run_response = job_service.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id,
retrieval_token=retrieval_token))
return PipelineResult(job_service, run_response.job_id)
def send_prepare_request(max_retries=5):
num_retries = 0
while True:
try:
# This reports channel is READY but connections may fail
# Seems to be only an issue on Mac with port forwardings
grpc.channel_ready_future(channel).result()
return job_service.Prepare(
beam_job_api_pb2.PrepareJobRequest(
job_name='job', pipeline=proto_pipeline,
pipeline_options=job_utils.dict_to_struct(options)))
except grpc._channel._Rendezvous as e:
num_retries += 1
if num_retries > max_retries:
raise e
def run_pipeline(self, pipeline):
docker_image = (
pipeline.options.view_as(PortableOptions).harness_docker_image
or self.default_docker_image())
job_endpoint = pipeline.options.view_as(PortableOptions).job_endpoint
if not job_endpoint:
raise ValueError(
'job_endpoint should be provided while creating runner.')
proto_context = pipeline_context.PipelineContext(
default_environment_url=docker_image)
proto_pipeline = pipeline.to_runner_api(context=proto_context)
# Some runners won't detect the GroupByKey transform unless it has no
# subtransforms. Remove all sub-transforms until BEAM-4605 is resolved.
for _, transform_proto in list(
proto_pipeline.components.transforms.items()):
if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
for sub_transform in transform_proto.subtransforms:
del proto_pipeline.components.transforms[sub_transform]
del transform_proto.subtransforms[:]
job_service = beam_job_api_pb2_grpc.JobServiceStub(
grpc.insecure_channel(job_endpoint))
prepare_response = job_service.Prepare(
beam_job_api_pb2.PrepareJobRequest(job_name='job',
pipeline=proto_pipeline))
if prepare_response.artifact_staging_endpoint.url:
# Must commit something to get a retrieval token,
# committing empty manifest for now.
# TODO(BEAM-3883): Actually stage required files.
artifact_service = beam_artifact_api_pb2_grpc.ArtifactStagingServiceStub(
grpc.insecure_channel(
prepare_response.artifact_staging_endpoint.url))
commit_manifest = artifact_service.CommitManifest(
beam_artifact_api_pb2.CommitManifestRequest(
manifest=beam_artifact_api_pb2.Manifest(),
staging_session_token=prepare_response.
staging_session_token))
retrieval_token = commit_manifest.retrieval_token
else:
retrieval_token = None
run_response = job_service.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id,
retrieval_token=retrieval_token))
return PipelineResult(job_service, run_response.job_id)
def run_pipeline(self, pipeline):
# Java has different expectations about coders
# (windowed in Fn API, but *un*windowed in runner API), whereas the
# FnApiRunner treats them consistently, so we must guard this.
# See also BEAM-2717.
proto_pipeline, proto_context = pipeline.to_runner_api(
return_context=True)
if self._runner_api_address:
for pcoll in proto_pipeline.components.pcollections.values():
if pcoll.coder_id not in proto_context.coders:
coder = coders.registry.get_coder(
pickler.loads(pcoll.coder_id))
pcoll.coder_id = proto_context.coders.get_id(coder)
proto_context.coders.populate_map(proto_pipeline.components.coders)
job_service = self._get_job_service()
prepare_response = job_service.Prepare(
beam_job_api_pb2.PrepareJobRequest(job_name='job',
pipeline=proto_pipeline))
run_response = job_service.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id))
return PipelineResult(job_service, run_response.job_id)
def run_pipeline(self, pipeline, options):
portable_options = options.view_as(PortableOptions)
# TODO: https://issues.apache.org/jira/browse/BEAM-5525
# portable runner specific default
if options.view_as(SetupOptions).sdk_location == 'default':
options.view_as(SetupOptions).sdk_location = 'container'
# This is needed as we start a worker server if one is requested
# but none is provided.
if portable_options.environment_type == 'LOOPBACK':
use_loopback_process_worker = options.view_as(
DebugOptions).lookup_experiment('use_loopback_process_worker',
False)
portable_options.environment_config, server = (
worker_pool_main.BeamFnExternalWorkerPoolServicer.start(
state_cache_size=sdk_worker_main._get_state_cache_size(
options),
use_process=use_loopback_process_worker))
cleanup_callbacks = [functools.partial(server.stop, 1)]
else:
cleanup_callbacks = []
proto_pipeline = pipeline.to_runner_api(
default_environment=PortableRunner._create_environment(
portable_options))
# Some runners won't detect the GroupByKey transform unless it has no
# subtransforms. Remove all sub-transforms until BEAM-4605 is resolved.
for _, transform_proto in list(
proto_pipeline.components.transforms.items()):
if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
for sub_transform in transform_proto.subtransforms:
del proto_pipeline.components.transforms[sub_transform]
del transform_proto.subtransforms[:]
# Preemptively apply combiner lifting, until all runners support it.
# These optimizations commute and are idempotent.
pre_optimize = options.view_as(DebugOptions).lookup_experiment(
'pre_optimize', 'lift_combiners').lower()
if not options.view_as(StandardOptions).streaming:
flink_known_urns = frozenset([
common_urns.composites.RESHUFFLE.urn,
common_urns.primitives.IMPULSE.urn,
common_urns.primitives.FLATTEN.urn,
common_urns.primitives.GROUP_BY_KEY.urn
])
if pre_optimize == 'none':
pass
elif pre_optimize == 'all':
proto_pipeline = fn_api_runner_transforms.optimize_pipeline(
proto_pipeline,
phases=[
fn_api_runner_transforms.
annotate_downstream_side_inputs,
fn_api_runner_transforms.
annotate_stateful_dofns_as_roots,
fn_api_runner_transforms.fix_side_input_pcoll_coders,
fn_api_runner_transforms.lift_combiners,
fn_api_runner_transforms.expand_sdf,
fn_api_runner_transforms.fix_flatten_coders,
# fn_api_runner_transforms.sink_flattens,
fn_api_runner_transforms.greedily_fuse,
fn_api_runner_transforms.read_to_impulse,
fn_api_runner_transforms.extract_impulse_stages,
fn_api_runner_transforms.remove_data_plane_ops,
fn_api_runner_transforms.sort_stages
],
known_runner_urns=flink_known_urns)
else:
phases = []
for phase_name in pre_optimize.split(','):
# For now, these are all we allow.
if phase_name in 'lift_combiners':
phases.append(
getattr(fn_api_runner_transforms, phase_name))
else:
raise ValueError(
'Unknown or inapplicable phase for pre_optimize: %s'
% phase_name)
proto_pipeline = fn_api_runner_transforms.optimize_pipeline(
proto_pipeline,
phases=phases,
known_runner_urns=flink_known_urns,
partial=True)
job_service = self.create_job_service(options)
# fetch runner options from job service
# retries in case the channel is not ready
def send_options_request(max_retries=5):
num_retries = 0
while True:
try:
# This reports channel is READY but connections may fail
# Seems to be only an issue on Mac with port forwardings
return job_service.DescribePipelineOptions(
beam_job_api_pb2.DescribePipelineOptionsRequest(),
timeout=portable_options.job_server_timeout)
except grpc.FutureTimeoutError:
# no retry for timeout errors
raise
except grpc._channel._Rendezvous as e:
num_retries += 1
if num_retries > max_retries:
raise e
time.sleep(1)
options_response = send_options_request()
def add_runner_options(parser):
for option in options_response.options:
try:
# no default values - we don't want runner options
# added unless they were specified by the user
add_arg_args = {
'action': 'store',
'help': option.description
}
if option.type == beam_job_api_pb2.PipelineOptionType.BOOLEAN:
add_arg_args['action'] = 'store_true'\
if option.default_value != 'true' else 'store_false'
elif option.type == beam_job_api_pb2.PipelineOptionType.INTEGER:
add_arg_args['type'] = int
elif option.type == beam_job_api_pb2.PipelineOptionType.ARRAY:
add_arg_args['action'] = 'append'
parser.add_argument("--%s" % option.name, **add_arg_args)
except Exception as e:
# ignore runner options that are already present
# only in this case is duplicate not treated as error
if 'conflicting option string' not in str(e):
raise
_LOGGER.debug("Runner option '%s' was already added" %
option.name)
all_options = options.get_all_options(
add_extra_args_fn=add_runner_options)
# TODO: Define URNs for options.
# convert int values: https://issues.apache.org/jira/browse/BEAM-5509
p_options = {
'beam:option:' + k + ':v1': (str(v) if type(v) == int else v)
for k, v in all_options.items() if v is not None
}
prepare_request = beam_job_api_pb2.PrepareJobRequest(
job_name='job',
pipeline=proto_pipeline,
pipeline_options=job_utils.dict_to_struct(p_options))
_LOGGER.debug('PrepareJobRequest: %s', prepare_request)
prepare_response = job_service.Prepare(
prepare_request, timeout=portable_options.job_server_timeout)
artifact_endpoint = (portable_options.artifact_endpoint
if portable_options.artifact_endpoint else
prepare_response.artifact_staging_endpoint.url)
if artifact_endpoint:
stager = portable_stager.PortableStager(
grpc.insecure_channel(artifact_endpoint),
prepare_response.staging_session_token)
retrieval_token, _ = stager.stage_job_resources(
options, staging_location='')
else:
retrieval_token = None
try:
state_stream = job_service.GetStateStream(
beam_job_api_pb2.GetJobStateRequest(
job_id=prepare_response.preparation_id),
timeout=portable_options.job_server_timeout)
# If there's an error, we don't always get it until we try to read.
# Fortunately, there's always an immediate current state published.
state_stream = itertools.chain([next(state_stream)], state_stream)
message_stream = job_service.GetMessageStream(
beam_job_api_pb2.JobMessagesRequest(
job_id=prepare_response.preparation_id),
timeout=portable_options.job_server_timeout)
except Exception:
# TODO(BEAM-6442): Unify preparation_id and job_id for all runners.
state_stream = message_stream = None
# Run the job and wait for a result, we don't set a timeout here because
# it may take a long time for a job to complete and streaming
# jobs currently never return a response.
run_response = job_service.Run(
beam_job_api_pb2.RunJobRequest(
preparation_id=prepare_response.preparation_id,
retrieval_token=retrieval_token))
if state_stream is None:
state_stream = job_service.GetStateStream(
beam_job_api_pb2.GetJobStateRequest(
job_id=run_response.job_id))
message_stream = job_service.GetMessageStream(
beam_job_api_pb2.JobMessagesRequest(
job_id=run_response.job_id))
return PipelineResult(job_service, run_response.job_id, message_stream,
state_stream, cleanup_callbacks)
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 test_end_to_end(self, http_mock):
submission_id = "submission-id"
worker_host_port = "workerhost:12345"
worker_id = "worker-id"
server_spark_version = "1.2.3"
def spark_submission_status_response(state):
return {
'json': {
"action": "SubmissionStatusResponse",
"driverState": state,
"serverSparkVersion": server_spark_version,
"submissionId": submission_id,
"success": "true",
"workerHostPort": worker_host_port,
"workerId": worker_id
}
}
with temp_name(suffix='fake.jar') as fake_jar:
with zipfile.ZipFile(fake_jar, 'w') as zip:
with zip.open('spark-version-info.properties', 'w') as fout:
fout.write(b'version=4.5.6')
options = pipeline_options.SparkRunnerOptions()
options.spark_job_server_jar = fake_jar
job_server = spark_uber_jar_job_server.SparkUberJarJobServer(
'http://host:6066', options)
# 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.LegacyArtifactStagingServiceStub(
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://host:6066/v1/submissions/create',
json={
"action": "CreateSubmissionResponse",
"message":
"Driver successfully submitted as submission-id",
"serverSparkVersion": "1.2.3",
"submissionId": "submission-id",
"success": "true"
})
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://host:6066/v1/submissions/status/submission-id', [
spark_submission_status_response('RUNNING'),
spark_submission_status_response('RUNNING'), {
'json': {
"action": "SubmissionStatusResponse",
"driverState": "ERROR",
"message": "oops",
"serverSparkVersion": "1.2.3",
"submissionId": submission_id,
"success": "true",
"workerHostPort": worker_host_port,
"workerId": worker_id
}
}
])
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.STOPPED,
beam_job_api_pb2.JobState.RUNNING,
beam_job_api_pb2.JobState.RUNNING,
beam_job_api_pb2.JobState.FAILED
])
message_stream = job_server.GetMessageStream(
beam_job_api_pb2.JobMessagesRequest(
job_id=prepare_response.preparation_id))
def get_item(x):
if x.HasField('message_response'):
return x.message_response
else:
return x.state_response.state
self.assertEqual([get_item(m) for m in message_stream], [
beam_job_api_pb2.JobState.STOPPED,
beam_job_api_pb2.JobState.RUNNING,
beam_job_api_pb2.JobMessage(
message_id='message0',
time='0',
importance=beam_job_api_pb2.JobMessage.MessageImportance.
JOB_MESSAGE_ERROR,
message_text="oops"),
beam_job_api_pb2.JobState.FAILED,
])
