From 521ba97236925349e4dd81b46d03d4714503f9d2 Mon Sep 17 00:00:00 2001 From: superstar54 Date: Sat, 30 Nov 2024 23:30:22 +0100 Subject: [PATCH] fix max number of running process --- aiida_workgraph/engine/awaitable_manager.py | 4 ++-- tests/test_engine.py | 4 +++- tests/test_workgraph.py | 2 +- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/aiida_workgraph/engine/awaitable_manager.py b/aiida_workgraph/engine/awaitable_manager.py index 414933a5..20124c5a 100644 --- a/aiida_workgraph/engine/awaitable_manager.py +++ b/aiida_workgraph/engine/awaitable_manager.py @@ -85,8 +85,8 @@ def resolve_awaitable(self, awaitable: Awaitable, value: Any) -> None: raise AssertionError(f"Unsupported awaitable action: {awaitable.action}") awaitable.resolved = True - # remove awaitabble from the list - self._awaitables = [a for a in self._awaitables if a.pk != awaitable.pk] + # remove awaitabble from the list, and use the same list reference + self._awaitables[:] = [a for a in self._awaitables if a.pk != awaitable.pk] if not self.process.has_terminated(): # the process may be terminated, for example, if the process was killed or excepted diff --git a/tests/test_engine.py b/tests/test_engine.py index 2b322d7b..04c4fea5 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -32,6 +32,7 @@ def test_reset_node(wg_engine: WorkGraph) -> None: assert len(wg.process.base.extras.get("_workgraph_queue")) == 1 +@pytest.mark.usefixtures("started_daemon_client") def test_max_number_jobs(add_code) -> None: from aiida_workgraph import WorkGraph from aiida.orm import Int @@ -46,6 +47,7 @@ def test_max_number_jobs(add_code) -> None: ) # Set the maximum number of running jobs inside the WorkGraph wg.max_number_jobs = 2 - wg.submit(wait=True, timeout=100) + wg.submit(wait=True, timeout=40) report = get_workchain_report(wg.process, "REPORT") assert "tasks ready to run: add2" in report + wg.tasks["add2"].outputs["sum"].value == 2 diff --git a/tests/test_workgraph.py b/tests/test_workgraph.py index 1f00221d..419fa11f 100644 --- a/tests/test_workgraph.py +++ b/tests/test_workgraph.py @@ -130,7 +130,7 @@ def test_pause_play_task(wg_calcjob): wg.pause_tasks(["add2"]) wg.play_tasks(["add1"]) # wait for the workgraph to launch add2 - wg.wait(tasks={"add2": ["CREATED"]}, timeout=20) + wg.wait(tasks={"add2": ["CREATED"]}, timeout=40) assert wg.tasks["add2"].node.process_state.value.upper() == "CREATED" assert wg.tasks["add2"].node.process_status == "Paused through WorkGraph" # I disabled the following lines because the test is not stable