def run_job(
job_class: Type[base_jobs.JobBase],
sync: bool,
namespace: Optional[str] = None,
pipeline: Optional[beam.Pipeline] = None
) -> beam_job_models.BeamJobRunModel:
"""Runs the specified job synchronously.
In other words, the function will wait for the job to finish running before
returning a value.
Args:
job_class: type(base_jobs.JobBase). The type of job to run.
sync: bool. Whether to run the job synchronously.
namespace: str. The namespace in which models should be created.
pipeline: Pipeline. The pipeline to run the job upon. If omitted, then a
new pipeline will be used instead.
Returns:
BeamJobRun. Contains metadata related to the execution status of the
job.
Raises:
RuntimeError. Failed to deploy given job to the Dataflow service.
"""
if pipeline is None:
pipeline = beam.Pipeline(
runner=runners.DirectRunner()
if sync else runners.DataflowRunner(),
options=job_options.JobOptions(namespace=namespace))
job = job_class(pipeline)
job_name = job_class.__name__
# NOTE: Exceptions raised within this context are logged and suppressed.
with _job_bookkeeping_context(job_name) as run_model:
_ = job.run() | job_io.PutResults(run_model.id)
run_result = pipeline.run()
if sync:
run_result.wait_until_finish()
run_model.latest_job_state = beam_job_models.BeamJobState.DONE.value
elif run_result.has_job:
run_model.dataflow_job_id = run_result.job_id()
run_model.latest_job_state = run_result.state
else:
raise RuntimeError(
'Failed to deploy %s to the Dataflow service. Please try again '
'after a few minutes.' % job_name)
# NDB operations in Beam do not properly update the context cache
# (this cache is separate for every application thread), thus we clear
# it ourselves.
with datastore_services.get_ndb_context() as ndb_context:
ndb_context.clear_cache()
return run_model
def test_single_output(self) -> None:
messages = [
job_run_result.JobRunResult(stdout='Hello, World!',
stderr='Uh-oh, World!'),
]
self.assert_pcoll_empty(self.pipeline
| beam.Create(messages)
| job_io.PutResults(self.JOB_ID))
result = beam_job_services.get_beam_job_run_result(self.JOB_ID)
self.assertEqual(result.stdout, 'Hello, World!')
self.assertEqual(result.stderr, 'Uh-oh, World!')
def test_sharded_output(self) -> None:
messages = [
job_run_result.JobRunResult(stdout='abc', stderr='123'),
job_run_result.JobRunResult(stdout='def', stderr='456'),
job_run_result.JobRunResult(stdout='ghi', stderr='789'),
]
with self.swap(job_run_result, 'MAX_OUTPUT_BYTES', 11):
self.assert_pcoll_empty(self.pipeline
| beam.Create(messages)
| job_io.PutResults(self.JOB_ID))
result = beam_job_services.get_beam_job_run_result(self.JOB_ID)
self.assertItemsEqual(
result.stdout.split('\n'),
['abc', 'def', 'ghi']) # type: ignore[no-untyped-call]
self.assertItemsEqual(
result.stderr.split('\n'),
['123', '456', '789']) # type: ignore[no-untyped-call]