Skip to content

Commit 5ed3fe7

Browse files
committed
Update docs and tests
1 parent 3934b7b commit 5ed3fe7

File tree

2 files changed

+53
-21
lines changed

2 files changed

+53
-21
lines changed

docs/gallery/autogen/pyfunction.py

Lines changed: 53 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ def generate_structures(element: str, factors: list) -> dict:
9393

9494
from aiida.engine import submit
9595
import datetime
96-
from aiida_pythonjob import prepare_pyfunction_inputs
96+
from aiida_pythonjob import prepare_pyfunction_inputs, PyFunction
9797

9898

9999
@pyfunction()
@@ -111,39 +111,72 @@ async def add_async(x, y, time: float):
111111
function_inputs={"x": 2, "y": 3, "time": 2.0},
112112
)
113113

114-
node = submit(add_async, **inputs)
114+
node = submit(PyFunction, **inputs)
115+
# sphinx_gallery_start_ignore
116+
assert node.is_finished_ok
117+
# sphinx_gallery_end_ignore
115118

116-
# %%
117-
# Here is an example to monitor external events or conditions without blocking.
119+
# %%#
120+
# Killing an async process
121+
# ~~~~~~~~~~~~~~~~~~~~~~~~
122+
# Since async functions run as regular AiiDA processes, they can be controlled and killed
123+
# programmatically. This is useful for managing long-running or stuck tasks.
124+
# You can kill a running async function using the AiiDA command line interface.
125+
#
126+
# .. code-block:: bash
127+
#
128+
# $ verdi process kill <pk>
129+
#
130+
# Monitor external events
131+
# ------------------------
132+
#
133+
# Async functions are particularly useful for monitoring external events or conditions without blocking the AiiDA daemon.
118134
# Here is an example that waits until a specified time.
119135
#
120136

121137

122-
@pyfunction()
123-
async def monitor_time(time: datetime.datetime):
138+
async def monitor_time(time: datetime.datetime, interval: float = 0.5, timeout: float = 60.0):
139+
"""Monitor the current time until it reaches the specified target time."""
124140
import asyncio
125141

126-
# monitor until the specified time
142+
start_time = datetime.datetime.now()
127143
while datetime.datetime.now() < time:
128144
print("Waiting...")
129-
await asyncio.sleep(0.5)
145+
await asyncio.sleep(interval)
146+
if (datetime.datetime.now() - start_time).total_seconds() > timeout:
147+
raise TimeoutError("Monitoring timed out.")
148+
149+
150+
inputs = prepare_pyfunction_inputs(
151+
monitor_time,
152+
function_inputs={"time": datetime.datetime.now() + datetime.timedelta(seconds=5), "interval": 1.0},
153+
)
154+
155+
node = submit(PyFunction, **inputs)
156+
# sphinx_gallery_start_ignore
157+
assert node.is_finished_ok
158+
# sphinx_gallery_end_ignore
159+
# %%
160+
# For user's convenience, we provide a dedicated ``MonitorFunction`` class that inherits from ``PyFunction``.
161+
# User only need to write normal function, which returns True when the monitoring condition is met.
162+
163+
from aiida_pythonjob import MonitorPyFunction
164+
165+
166+
def monitor_time(time: datetime.datetime):
167+
# return True when the current time is greater than the target time
168+
return datetime.datetime.now() > time
130169

131170

132171
inputs = prepare_pyfunction_inputs(
133172
monitor_time,
134173
function_inputs={"time": datetime.datetime.now() + datetime.timedelta(seconds=5)},
174+
interval=1.0,
175+
timeout=20.0,
135176
)
136177

137-
node = submit(monitor_time, **inputs)
178+
node = submit(MonitorPyFunction, **inputs)
138179

139-
# %%#
140-
# Killing an async process
141-
# ------------------------
142-
# Since async functions run as regular AiiDA processes, they can be controlled and killed
143-
# programmatically. This is useful for managing long-running or stuck tasks.
144-
# You can kill a running async function using the AiiDA command line interface.
145-
#
146-
# .. code-block:: bash
147-
#
148-
# $ verdi process kill <pk>
149-
#
180+
# sphinx_gallery_start_ignore
181+
assert node.is_finished_ok
182+
# sphinx_gallery_end_ignore

tests/test_monitor.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ def test_async_function_runs_and_returns_result():
1616
)
1717
result, node = run_get_node(MonitorPyFunction, **inputs)
1818
assert node.is_finished_ok
19-
assert "result" in result
2019
# The actual monitor function returns None
2120
assert result["result"].value is None
2221

0 commit comments

Comments
 (0)