def test_run_beam_job_using_job_class(self) -> None:
model = beam_job_services.create_beam_job_run_model('NoOpJob')
with self.swap_to_always_return(jobs_manager, 'run_job', value=model):
run = beam_job_services.run_beam_job(job_class=NoOpJob)
self.assertEqual(
beam_job_services.get_beam_job_run_from_model(model).to_dict(),
run.to_dict())
def test_cancel_beam_job(self) -> None:
model = beam_job_services.create_beam_job_run_model(
'NoOpJob', dataflow_job_id='123')
model.put()
with self.swap_to_always_return(jobs_manager, 'cancel_job'):
run = beam_job_services.cancel_beam_job(model.id)
self.assertEqual(
run.to_dict(),
beam_job_services.get_beam_job_run_from_model(model).to_dict())
def test_put_starts_new_job(self) -> None:
model = beam_job_services.create_beam_job_run_model('FooJob')
with self.swap_to_always_return(jobs_manager, 'run_job', value=model):
response = self.put_json( # type: ignore[no-untyped-call]
'/beam_job_run', {'job_name': 'FooJob'},
csrf_token=self.get_new_csrf_token()) # type: ignore[no-untyped-call]
self.assertEqual(
response,
beam_job_services.get_beam_job_run_from_model(model).to_dict())
def test_run_beam_job(self) -> None:
run_model = beam_job_services.create_beam_job_run_model('WorkingJob')
get_job_class_by_name_swap = self.swap_to_always_return(
jobs_registry, 'get_job_class_by_name')
run_job_swap = self.swap_to_always_return(
jobs_manager, 'run_job', value=run_model)
with get_job_class_by_name_swap, run_job_swap:
run = beam_job_services.run_beam_job('WorkingJob')
self.assertEqual(
beam_job_services.get_beam_job_run_from_model(run_model).to_dict(),
run.to_dict())
def test_failing_job(self) -> None:
run = jobs_manager.run_job_sync('FailingJob', [],
namespace=self.namespace)
self.assertEqual(run.job_state, 'FAILED')
run_model = beam_job_models.BeamJobRunModel.get(
run.job_id) # type: ignore[attr-defined]
self.assertEqual(
run.to_dict(),
beam_job_services.get_beam_job_run_from_model(run_model).to_dict())
self.assertIn(
'uh-oh',
beam_job_services.get_beam_job_run_result(run.job_id).stderr)
def test_working_job(self) -> None:
run = jobs_manager.run_job_sync('WorkingJob', [],
namespace=self.namespace)
self.assertEqual(run.job_state, 'DONE')
run_model = beam_job_models.BeamJobRunModel.get(
run.job_id) # type: ignore[attr-defined]
self.assertEqual(
run.to_dict(),
beam_job_services.get_beam_job_run_from_model(run_model).to_dict())
self.assertEqual(
beam_job_services.get_beam_job_run_result(run.job_id).to_dict(), {
'stdout': 'o',
'stderr': 'e'
})
def run_job_sync(
job_name: str,
job_args: List[str],
namespace: Optional[str] = None) -> beam_job_domain.BeamJobRun:
"""Runs the specified job synchronously.
In other words, the function will wait for the job to finish running before
returning a value.
Args:
job_name: str. The name of the job to run.
job_args: list(str). The arguments to the job's run() method.
namespace: str. The namespace in which models should be created.
Returns:
BeamJobRun. Contains metadata related to the execution status of the
job.
"""
job_pipeline = beam.Pipeline(
runner=runners.DirectRunner(),
options=job_options.JobOptions(namespace=namespace))
job_class = registry.get_job_class_by_name(job_name)
job = job_class(job_pipeline)
run_model = beam_job_services.create_beam_job_run_model(job_name, job_args)
try:
with job_pipeline:
unused_pdone = job.run(*job_args) | job_io.PutResults(run_model.id)
except Exception as exception:
run_model.latest_job_state = beam_job_models.BeamJobState.FAILED.value
# If the pipeline fails to put the results into storage, then we'll
# explicitly write them to storage by using the caught exception.
result_model = beam_job_services.create_beam_job_run_result_model(
run_model.id, '', python_utils.UNICODE(exception))
result_model.put()
else:
run_model.latest_job_state = beam_job_models.BeamJobState.DONE.value
finally:
run_model.put()
return beam_job_services.get_beam_job_run_from_model(run_model)
