def execute(self, context) -> None:
"""
Executes the Armada Operator.
Runs an Armada job and calls the job_service_client for polling.
:param context: The airflow context.
:return: None
"""
job = self.armada_client.submit_jobs(
queue=self.queue,
job_set_id=self.job_set_id,
job_request_items=self.job_request_items,
)
try:
job_id = job.job_response_items[0].job_id
armada_logger.info("Running Armada job %s with id %s", self.name, job_id)
except Exception:
raise AirflowException("Armada has issues submitting job")
job_state, job_message = search_for_job_complete(
job_service_client=self.job_service,
queue=self.queue,
job_set_id=self.job_set_id,
airflow_task_name=self.name,
job_id=job_id,
)
armada_logger.info(
"Armada Job finished with %s and message: %s", job_state, job_message
)
airflow_error(job_state, self.name, job_id)
def test_two_jobs():
job_set_name = f"test-{uuid.uuid1()}"
first_job = no_auth_client.submit_jobs(
queue="test",
job_set_id=job_set_name,
job_request_items=sleep_pod(image="busybox"),
)
first_job_id = first_job.job_response_items[0].job_id
job_state, job_message = search_for_job_complete(
job_service_client=job_service_client,
queue="test",
job_set_id=job_set_name,
airflow_task_name="test",
job_id=first_job_id,
)
assert job_state == JobState.SUCCEEDED
assert job_message == f"Armada test:{first_job_id} succeeded"
second_job = no_auth_client.submit_jobs(
queue="test",
job_set_id=job_set_name,
job_request_items=sleep_pod(image="busybox"),
)
second_job_id = second_job.job_response_items[0].job_id
job_state, job_message = search_for_job_complete(
job_service_client=job_service_client,
queue="test",
job_set_id=job_set_name,
airflow_task_name="test",
job_id=second_job_id,
)
assert job_state == JobState.SUCCEEDED
assert job_message == f"Armada test:{second_job_id} succeeded"
def test_mock_failed_job():
tester_client.submit_jobs(
queue="test",
job_set_id="test",
job_request_items=sleep_job(),
)
job_state, job_message = search_for_job_complete(
job_service_client=tester_jobservice,
queue="test",
job_set_id="test",
airflow_task_name="test-mock",
job_id="test_failed",
)
assert job_state == JobState.FAILED
assert job_message.startswith("Armada test-mock:test_failed failed")
def test_mock_success_job():
tester_client.submit_jobs(
queue="test",
job_set_id="test",
job_request_items=sleep_job(),
)
job_state, job_message = search_for_job_complete(
job_service_client=tester_jobservice,
queue="test",
job_set_id="test",
airflow_task_name="test-mock",
job_id="test_succeeded",
)
assert job_state == JobState.SUCCEEDED
assert job_message == "Armada test-mock:test_succeeded succeeded"
def test_mock_cancelled_job():
tester_client.submit_jobs(
queue="test",
job_set_id="test",
job_request_items=sleep_job(),
)
job_state, job_message = search_for_job_complete(
job_service_client=tester_jobservice,
queue="test",
job_set_id="test",
airflow_task_name="test-mock",
job_id="test_cancelled",
)
assert job_state == JobState.CANCELLED
assert job_message == "Armada test-mock:test_cancelled cancelled"
def test_bad_job():
job_set_name = f"test-{uuid.uuid1()}"
job = no_auth_client.submit_jobs(
queue="test",
job_set_id=job_set_name,
job_request_items=sleep_pod(image="NOTACONTAINER"),
)
job_id = job.job_response_items[0].job_id
job_state, job_message = search_for_job_complete(
job_service_client=job_service_client,
queue="test",
job_set_id=job_set_name,
airflow_task_name="test",
job_id=job_id,
)
assert job_state == JobState.FAILED
assert job_message.startswith(f"Armada test:{job_id} failed")