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_async_job(self) -> None:
mock_run_result = mock.Mock()
mock_run_result.has_job = True
mock_run_result.job_id.return_value = '123'
mock_run_result.state = 'PENDING'
pipeline = beam.Pipeline(
runner=runners.DirectRunner(),
options=job_options.JobOptions(namespace=self.namespace))
with self.swap_to_always_return(pipeline, 'run',
value=mock_run_result):
run = jobs_manager.run_job(WorkingJob, False, pipeline=pipeline)
self.assertEqual(run.dataflow_job_id, '123')
self.assertEqual(run.latest_job_state, 'PENDING')
def test_async_job_that_does_not_start(self) -> None:
mock_run_result = mock.Mock()
mock_run_result.has_job = False
mock_run_result.job_id.return_value = None
mock_run_result.state = 'UNKNOWN'
pipeline = beam.Pipeline(
runner=runners.DirectRunner(),
options=job_options.JobOptions(namespace=self.namespace))
with self.swap_to_always_return(pipeline, 'run',
value=mock_run_result):
run = jobs_manager.run_job(WorkingJob, False, pipeline=pipeline)
self.assertIsNone(run.dataflow_job_id)
self.assertEqual(run.latest_job_state, 'FAILED')
result = beam_job_services.get_beam_job_run_result(run.id)
self.assertIn('Failed to deploy WorkingJob', result.stderr)
def __init__(self, *args: Any, **kwargs: Any) -> None:
super(PipelinedTestBase, self).__init__(*args, **kwargs)
self.pipeline = test_pipeline.TestPipeline(
runner=runners.DirectRunner(),
options=job_options.JobOptions(namespace=self.namespace))
self._pipeline_context_stack: Optional[contextlib.ExitStack] = None
def test_unsupported_values(self) -> None:
with self.assertRaisesRegex(
ValueError,
r'Unsupported option\(s\)'): # type: ignore[no-untyped-call]
job_options.JobOptions(a='a', b='b')
def test_overwritten_values(self) -> None:
options = job_options.JobOptions(namespace='abc')
self.assertEqual(options.namespace, 'abc')
def test_default_values(self) -> None:
options = job_options.JobOptions()
self.assertIsNone(options.namespace)