Task failure only fails dependent tasks

arhall0 · arhall0 · commit 801fc56bf983 · 2025-02-20T12:00:36.000-07:00
Currently no changes to build failure
diff --git a/beeflow/common/gdb/gdb_driver.py b/beeflow/common/gdb/gdb_driver.py
@@ -279,7 +279,7 @@ def set_task_output_glob(self, task, output_id, glob):
     def workflow_completed(self):
         """Determine if a workflow has completed.
 
-        A workflow has completed if each of its final tasks has state 'COMPLETED'.
+        A workflow has completed if each of its final tasks has finished or failed.
 
         :rtype: bool
         """
diff --git a/beeflow/common/gdb/neo4j_cypher.py b/beeflow/common/gdb/neo4j_cypher.py
@@ -716,7 +716,9 @@ def reset_workflow_id(tx, old_id, new_id):
 
 
 def final_tasks_completed(tx, wf_id):
-    """Return true if each of a workflow's final Task nodes has state 'COMPLETED'.
+    """Return true if each of a workflow's final Task nodes is in a finished state.
+
+    Finished states: 'COMPLETED', 'FAILED', 'SUBMIT_FAIL', 'BUILD_FAIL', or 'DEP_FAIL'.
 
     :param wf_id: the workflow's id
     :type wf_id: str
@@ -725,7 +727,8 @@ def final_tasks_completed(tx, wf_id):
     restart = "|RESTARTED_FROM" if get_workflow_by_id(tx, wf_id)['restart'] else ""
     not_completed_query = ("MATCH (m:Metadata)-[:DESCRIBES]->(t:Task {workflow_id: $wf_id}) "
                            f"WHERE NOT (t)<-[:DEPENDS_ON{restart}]-(:Task) "
-                           "AND m.state <> 'COMPLETED' "
+                           "AND NOT m.state IN "
+                           "['COMPLETED', 'FAILED', 'SUBMIT_FAIL', 'BUILD_FAIL', 'DEP_FAIL'] "
                            "RETURN t IS NOT NULL LIMIT 1")
 
     # False if at least one task with state not 'COMPLETED'
diff --git a/beeflow/common/gdb/neo4j_driver.py b/beeflow/common/gdb/neo4j_driver.py
@@ -1,3 +1,5 @@
+# pylint: disable=W0221
+
 """Neo4j interface module.
 
 Connection requires a valid URI, Username, and Password.
@@ -415,7 +417,7 @@ def set_task_output_glob(self, task, output_id, glob):
     def workflow_completed(self, workflow_id):
         """Determine if a workflow in the Neo4j database has completed.
 
-        A workflow has completed if each of its final task nodes have state 'COMPLETED'.
+        A workflow has completed if each of its final tasks has finished or failed.
         :param workflow_id: the workflow id
         :type workflow_id: str
         :rtype: bool
diff --git a/beeflow/tests/test_neo4j_cypher.py b/beeflow/tests/test_neo4j_cypher.py
@@ -1,5 +1,6 @@
 """Tests for neo4j_cypher module."""
 
+from unittest.mock import call
 import pytest
 from beeflow.common.gdb import neo4j_cypher
 
@@ -54,3 +55,26 @@ def test_get_task_hints(mocker, hints, expected):
     )
     result = neo4j_cypher.get_task_hints(tx, "")
     assert result == expected
+
+
+@pytest.mark.parametrize(
+    "restart, single, expected",
+    [
+        (True, None, True),
+        (False, None, True),
+        (True, True, False),
+        (False, True, False),
+    ],
+)
+def test_final_tasks_completed(mocker, restart, single, expected):
+    """Regression test final_tasks_completed."""
+    tx = mocker.MagicMock()
+    tx.run().single.return_value = single
+    mocker.patch(
+        "beeflow.common.gdb.neo4j_cypher.get_workflow_by_id",
+        return_value={"restart": restart},
+    )
+    result = neo4j_cypher.final_tasks_completed(tx, "WFID")
+    assert ("|RESTARTED_FROM" in tx.run.mock_calls[1].args[0]) == restart
+    assert tx.run.mock_calls[2] == call().single()
+    assert result == expected
diff --git a/beeflow/tests/test_wf_update.py b/beeflow/tests/test_wf_update.py
@@ -70,3 +70,71 @@ def test_archive_archived_wf(mocker, wf_state):
             f"in state {wf_state}."
         )
     )
+
+
+@pytest.mark.parametrize("job_state", ["FAILED", "SUBMIT_FAIL"])
+def test_handle_state_change_failed_task(mocker, job_state):
+    """Regression test task failure."""
+    state_update = mocker.MagicMock()
+    state_update.job_state = job_state
+    task = mocker.MagicMock()
+    task.name = "TestTask"
+    wfi = mocker.MagicMock()
+    wfi.workflow_completed.return_value = False
+    wfi.cancelled_workflow_completed.return_value = False
+    mock_archive_workflow = mocker.patch(
+        "beeflow.wf_manager.resources.wf_update.archive_workflow"
+    )
+    db = mocker.MagicMock()
+    mock_log_info = mocker.patch("logging.Logger.info")
+    mock_set_dependent_tasks_dep_fail = mocker.patch(
+        "beeflow.wf_manager.resources.wf_update.set_dependent_tasks_dep_fail"
+    )
+    workflow_update = wf_update.WFUpdate()
+    workflow_update.handle_state_change(state_update, task, wfi, db)
+    mock_log_info.assert_any_call("Task TestTask failed")
+    mock_set_dependent_tasks_dep_fail.assert_called_once()
+    mock_archive_workflow.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "completed, cancelled_completed, wf_state",
+    [
+        (True, False, ""),
+        (False, True, "Cancelled"),
+        (False, False, "Cancelled"),
+        (False, True, ""),
+    ],
+)
+def test_handle_state_change_completed_wf(
+    mocker, completed, cancelled_completed, wf_state
+):
+    """Regression test when workflow is complete."""
+    state_update = mocker.MagicMock()
+    task = mocker.MagicMock()
+    wfi = mocker.MagicMock()
+    wfi.workflow_completed.return_value = completed
+    wfi.cancelled_workflow_completed.return_value = cancelled_completed
+    wfi.get_workflow_state.return_value = wf_state
+    wfi.workflow_id = "TESTID"
+    mock_archive_workflow = mocker.patch(
+        "beeflow.wf_manager.resources.wf_update.archive_workflow"
+    )
+    db = mocker.MagicMock()
+    mock_log_info = mocker.patch("logging.Logger.info")
+    workflow_update = wf_update.WFUpdate()
+    workflow_update.handle_state_change(state_update, task, wfi, db)
+    print(mock_log_info.mock_calls)
+    if completed:
+        mock_log_info.assert_any_call("Workflow TESTID Completed")
+        mock_log_info.assert_any_call("Workflow Archived")
+        mock_archive_workflow.assert_called_once()
+    elif cancelled_completed and wf_state == "Cancelled":
+        mock_log_info.assert_any_call(
+            "Scheduled tasks for cancelled workflow TESTID completed"
+        )
+        mock_log_info.assert_any_call("Workflow Archived")
+        mock_archive_workflow.assert_called_once()
+    else:
+        mock_log_info.assert_not_called()
+        mock_archive_workflow.assert_not_called()
diff --git a/beeflow/wf_manager/resources/wf_update.py b/beeflow/wf_manager/resources/wf_update.py
@@ -134,40 +134,37 @@ def handle_checkpoint_restart(self, state_update, task, wfi, db):
 
     def handle_state_change(self, state_update, task, wfi, db):
         """Handle a normal state change for a task."""
+        wf_state = wfi.get_workflow_state()
         if state_update.job_state == 'COMPLETED':
             for output in task.outputs:
                 if output.glob is not None:
                     wfi.set_task_output(task, output.id, output.glob)
                 else:
                     wfi.set_task_output(task, output.id, "temp")
             tasks = wfi.finalize_task(task)
-            wf_state = wfi.get_workflow_state()
             if tasks and wf_state not in ('PAUSED', 'Cancelled'):
                 wf_utils.schedule_submit_tasks(state_update.wf_id, tasks)
 
-            if wfi.workflow_completed():
-                wf_id = wfi.workflow_id
-                log.info(f"Workflow {wf_id} Completed")
-                archive_workflow(db, state_update.wf_id)
-                log.info('Workflow Completed')
-            elif wf_state == 'Cancelled' and wfi.cancelled_workflow_completed():
-                wf_id = wfi.workflow_id
-                log.info(f"Scheduled tasks for cancelled workflow {wf_id} completed")
-                archive_workflow(db, wf_id, final_state=wf_state)
-                log.info('Workflow Archived')
-
-        # If the job failed and it doesn't include a checkpoint-restart hint,
-        # then fail the entire workflow
+        # If the job failed, fail the dependent tasks
         if state_update.job_state in ['FAILED', 'SUBMIT_FAIL']:
             set_dependent_tasks_dep_fail(db, wfi, state_update.wf_id, task)
-            log.info("Workflow failed")
-            wf_id = wfi.workflow_id
-            archive_fail_workflow(db, wf_id)
+            log.info(f"Task {task.name} failed")
 
         if state_update.job_state == 'BUILD_FAIL':
             log.error(f'Workflow failed due to failed container build for task {task.name}')
             archive_fail_workflow(db, state_update.wf_id)
 
+        if wfi.workflow_completed():
+            wf_id = wfi.workflow_id
+            log.info(f"Workflow {wf_id} Completed")
+            archive_workflow(db, wf_id)
+            log.info('Workflow Archived')
+        elif wf_state == 'Cancelled' and wfi.cancelled_workflow_completed():
+            wf_id = wfi.workflow_id
+            log.info(f"Scheduled tasks for cancelled workflow {wf_id} completed")
+            archive_workflow(db, wf_id, final_state=wf_state)
+            log.info('Workflow Archived')
+
     def update_task_state(self, state_update, db):
         """Update the state of a single task from the task manager."""
         wfi = wf_utils.get_workflow_interface(state_update.wf_id)
