Raise an exception explicitly when the timeout is exceeded (#390)

pyproject.toml

@@ -91,7 +91,6 @@ workgraph = "aiida_workgraph.cli.cmd_workgraph:workgraph"
 "workgraph.while" = "aiida_workgraph.tasks.builtins:While"
 "workgraph.if" = "aiida_workgraph.tasks.builtins:If"
 "workgraph.select" = "aiida_workgraph.tasks.builtins:Select"
-"workgraph.gather" = "aiida_workgraph.tasks.builtins:Gather"
 "workgraph.set_context" = "aiida_workgraph.tasks.builtins:SetContext"
 "workgraph.get_context" = "aiida_workgraph.tasks.builtins:GetContext"
 "workgraph.time_monitor" = "aiida_workgraph.tasks.monitors:TimeMonitor"

src/aiida_workgraph/engine/workgraph.py

@@ -301,7 +301,6 @@ def setup_ctx_workgraph(self, wgdata: t.Dict[str, t.Any]) -> None:
         self.ctx._tasks = wgdata["tasks"]
         self.ctx._links = wgdata["links"]
         self.ctx._connectivity = wgdata["connectivity"]
-        self.ctx._ctrl_links = wgdata["ctrl_links"]
         self.ctx._workgraph = wgdata
         self.ctx._error_handlers = wgdata["error_handlers"]
 

src/aiida_workgraph/utils/analysis.py

@@ -22,8 +22,13 @@ def __init__(
         """
         self.process = process
         self.restart_process = restart_process
+        wgdata.setdefault("uuid", "")
+        wgdata.setdefault("tasks", {})
+        wgdata.setdefault("links", [])
+        wgdata.setdefault("ctrl_links", [])
+        wgdata.setdefault("error_handlers", {})
+        wgdata.setdefault("context", {})
         self.wgdata = wgdata
-        self.uuid = wgdata["uuid"]
         self.name = wgdata["name"]
         self.wait_to_link()
         self.clean_hanging_links()

src/aiida_workgraph/workgraph.py

@@ -110,14 +110,14 @@ def submit(
         self,
         inputs: Optional[Dict[str, Any]] = None,
         wait: bool = False,
-        timeout: int = 60,
+        timeout: int = 600,
         interval: int = 5,
         metadata: Optional[Dict[str, Any]] = None,
     ) -> aiida.orm.ProcessNode:
         """Submit the AiiDA workgraph process and optionally wait for it to finish.
         Args:
             wait (bool): Wait for the process to finish.
-            timeout (int): The maximum time in seconds to wait for the process to finish. Defaults to 60.
+            timeout (int): The maximum time in seconds to wait for the process to finish. Defaults to 600.
             restart (bool): Restart the process, and reset the modified tasks, then only re-run the modified tasks.
             new (bool): Submit a new process.
         """
@@ -228,11 +228,17 @@ def get_error_handlers(self) -> Dict[str, Any]:
                 task["exit_codes"] = exit_codes
         return error_handlers
 
-    def wait(self, timeout: int = 50, tasks: dict = None, interval: int = 5) -> None:
+    def wait(self, timeout: int = 600, tasks: dict = None, interval: int = 5) -> None:
         """
         Periodically checks and waits for the AiiDA workgraph process to finish until a given timeout.
+
         Args:
-            timeout (int): The maximum time in seconds to wait for the process to finish. Defaults to 50.
+            timeout (int): The maximum time in seconds to wait for the process to finish. Defaults to 600.
+            tasks (dict): Optional; specifies task states to wait for in the format {task_name: [acceptable_states]}.
+            interval (int): The time interval in seconds between checks. Defaults to 5.
+
+        Raises:
+            TimeoutError: If the process does not finish within the given timeout.
         """
         terminating_states = (
             "KILLED",
@@ -245,8 +251,10 @@ def wait(self, timeout: int = 50, tasks: dict = None, interval: int = 5) -> None
         start = time.time()
         self.update()
         finished = False
+
         while not finished:
             self.update()
+
             if tasks is not None:
                 states = []
                 for name, value in tasks.items():
@@ -255,9 +263,17 @@ def wait(self, timeout: int = 50, tasks: dict = None, interval: int = 5) -> None
                 finished = all(states)
             else:
                 finished = self.state in terminating_states
+
+            if finished:
+                print(f"Process {self.process.pk} finished with state: {self.state}")
+                return
+
             time.sleep(interval)
+
             if time.time() - start > timeout:
-                break
+                raise TimeoutError(
+                    f"Timeout reached after {timeout} seconds while waiting for the WorkGraph: {self.process.pk}. "
+                )
 
     def update(self) -> None:
         """

tests/conftest.py

@@ -242,12 +242,43 @@ def wg_engine(decorated_add, add_code) -> WorkGraph:
 
 
 @pytest.fixture
-def finished_process_node():
-    """Return a finished process node."""
-
-    node = WorkflowNode()
-    node.set_process_state("finished")
-    node.set_exit_status(0)
-    node.seal()
-    node.store()
-    return node
+def create_process_node():
+    """Return a process node."""
+
+    def process_node(state="finished", exit_status=0):
+        """Return a finished process node."""
+
+        node = WorkflowNode()
+        node.set_process_state(state)
+        node.set_exit_status(exit_status)
+        node.seal()
+        node.store()
+        return node
+
+    return process_node
+
+
+@pytest.fixture
+def create_workgraph_process_node():
+    """Return a process node."""
+
+    def process_node(state="finished", exit_status=0):
+        """Return a finished process node."""
+        from aiida_workgraph.engine.workgraph import WorkGraphEngine
+
+        process = WorkGraphEngine(
+            inputs={
+                "wg": {
+                    "name": "test",
+                    "state": "",
+                }
+            }
+        )
+        node = process.node
+        node.set_process_state(state)
+        node.set_exit_status(exit_status)
+        node.seal()
+        node.store()
+        return node
+
+    return process_node

tests/test_action.py

@@ -39,10 +39,10 @@ def test_pause_play_task(wg_calcjob):
     assert wg.tasks.add2.outputs.sum.value == 9
 
 
-def test_pause_play_error_handler(wg_calcjob, finished_process_node):
+def test_pause_play_error_handler(wg_calcjob, create_process_node):
     wg = wg_calcjob
     wg.name = "test_pause_play_error_handler"
-    wg.process = finished_process_node
+    wg.process = create_process_node(state="finished", exit_status=0)
     try:
         wg.pause_tasks(["add1"])
     except Exception as e:

tests/test_workgraph.py

@@ -181,3 +181,14 @@ def test_workgraph_group_outputs(decorated_add):
     wg.run()
     assert wg.process.outputs.sum.value == 5
     # assert wg.process.outputs.add1.result.value == 5
+
+
+@pytest.mark.usefixtures("started_daemon_client")
+def test_wait_timeout(create_workgraph_process_node):
+    wg = WorkGraph()
+    wg.process = create_workgraph_process_node(state="running")
+    with pytest.raises(
+        TimeoutError,
+        match="Timeout reached after 1 seconds while waiting for the WorkGraph:",
+    ):
+        wg.wait(timeout=1, interval=1)