Poll transient GDI sub-DAG states

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:

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()

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()