Skip to content

Commit 04d6d3b

Browse files
committed
Update docs and tests
1 parent 3934b7b commit 04d6d3b

File tree

2 files changed

+46
-22
lines changed

2 files changed

+46
-22
lines changed

docs/gallery/autogen/pyfunction.py

Lines changed: 46 additions & 21 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,64 @@ 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)
115115

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

121135

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

126-
# monitor until the specified time
140+
start_time = datetime.datetime.now()
127141
while datetime.datetime.now() < time:
128142
print("Waiting...")
129-
await asyncio.sleep(0.5)
143+
await asyncio.sleep(interval)
144+
if (datetime.datetime.now() - start_time).total_seconds() > timeout:
145+
raise TimeoutError("Monitoring timed out.")
130146

131147

132148
inputs = prepare_pyfunction_inputs(
133149
monitor_time,
134-
function_inputs={"time": datetime.datetime.now() + datetime.timedelta(seconds=5)},
150+
function_inputs={"time": datetime.datetime.now() + datetime.timedelta(seconds=5), "interval": 1.0},
135151
)
136152

137-
node = submit(monitor_time, **inputs)
153+
node = submit(PyFunction, **inputs)
138154

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-
#
155+
# %%
156+
# For user's convenience, we provide a dedicated ``MonitorFunction`` class that inherits from ``PyFunction``.
157+
# User only need to write normal function, which returns True when the monitoring condition is met.
158+
159+
from aiida_pythonjob import MonitorPyFunction
160+
161+
162+
def monitor_time(time: datetime.datetime):
163+
# return True when the current time is greater than the target time
164+
return datetime.datetime.now() > time
165+
166+
167+
inputs = prepare_pyfunction_inputs(
168+
monitor_time,
169+
function_inputs={"time": datetime.datetime.now() + datetime.timedelta(seconds=5)},
170+
interval=1.0,
171+
timeout=20.0,
172+
)
173+
174+
node = submit(MonitorPyFunction, **inputs)

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)