diff --git a/dpti/dags/dp_ti_gdi.py b/dpti/dags/dp_ti_gdi.py index fb90aeb6..2105e766 100644 --- a/dpti/dags/dp_ti_gdi.py +++ b/dpti/dags/dp_ti_gdi.py @@ -19,6 +19,7 @@ from airflow.utils.state import State from dpdispatcher import Machine, Resources, Submission, Task +from dpti.dags.utils import is_transient_dag_run_state from dpti.gdi import gdi_main_loop # default_args = {'owner': 'airflow', @@ -187,7 +188,7 @@ def wait_until_end(self): if dag_run_state == State.SUCCESS: print(f"dag_run_state: {dag_run_state}") break - elif dag_run_state == State.RUNNING: + elif is_transient_dag_run_state(dag_run_state): print(f"dag_run_state: {dag_run_state}") time.sleep(30) else: diff --git a/dpti/dags/utils.py b/dpti/dags/utils.py index d9513c45..20b3b843 100644 --- a/dpti/dags/utils.py +++ b/dpti/dags/utils.py @@ -3,6 +3,25 @@ from dpdispatcher import Machine, Resources, Submission +_TRANSIENT_DAG_RUN_STATES = { + "queued", + "scheduled", + "running", + "up_for_retry", + "up_for_reschedule", + "deferred", + "restarting", +} + + +def is_transient_dag_run_state(state): + """Return whether an Airflow DAG run should continue being polled.""" + if state is None: + # A freshly triggered run may not be visible in the metadata DB yet. + return True + state_value = getattr(state, "value", state) + return str(state_value).lower() in _TRANSIENT_DAG_RUN_STATES + def get_empty_submission(job_work_dir, context): # context = get_current_context() diff --git a/tests/test_dag_utils.py b/tests/test_dag_utils.py new file mode 100644 index 00000000..3c0d825e --- /dev/null +++ b/tests/test_dag_utils.py @@ -0,0 +1,23 @@ +import unittest +from enum import Enum + +from dpti.dags.utils import is_transient_dag_run_state + + +class ExampleState(Enum): + QUEUED = "queued" + + +class TestDagRunStates(unittest.TestCase): + def test_healthy_transient_states_are_polled(self): + """Queued, scheduled, running, and not-yet-visible runs are not failures.""" + for state in (None, "queued", "scheduled", "running", ExampleState.QUEUED): + with self.subTest(state=state): + self.assertTrue(is_transient_dag_run_state(state)) + + def test_terminal_failure_is_not_transient(self): + self.assertFalse(is_transient_dag_run_state("failed")) + + +if __name__ == "__main__": + unittest.main()