Skip to content

[Pylib] add at_tsim case for simulation callbacks #1894

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
tdavidcl wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

[Pylib] add at_tsim case for simulation callbacks #1894

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ebc8cd1
[Pylib] add at_tsim case for simulation callbacks
tdavidcl Jun 25, 2026
b37aad3
Merge branch 'main' into at_tile_callback
tdavidcl Jun 26, 2026
4563eb6
Merge branch 'main' into at_tile_callback
tdavidcl Jul 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 54 additions & 9 deletions src/pylib/shamrock/utils/SimulationRunner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,13 @@
# ----------------------------


def callback(*, tsim_interval=None, iter_count_interval=None, walltime_interval=None):
def callback(
*,
tsim_interval=None,
iter_count_interval=None,
walltime_interval=None,
at_tsim=None,
):
"""
Decorator to mark a function as a simulation callback.

Expand All @@ -20,24 +26,40 @@ def callback(*, tsim_interval=None, iter_count_interval=None, walltime_interval=
def analysis(self, icallback):
print("analysis", icallback)

@callback(at_tsim=[1.0, 5.0, 10.0])
def snapshot(self, icallback):
print("snapshot", icallback)

Args:
tsim_interval: The time step of the callback.
iter_count_interval: The iteration count interval of the callback.
walltime_interval: The walltime interval of the callback.
at_tsim: Exact simulation time(s) at which to trigger the callback (float or list of floats).
Returns:
The decorated function.
"""

# at least one of the intervals must be provided
if tsim_interval is None and iter_count_interval is None and walltime_interval is None:
if (
tsim_interval is None
and iter_count_interval is None
and walltime_interval is None
and at_tsim is None
):
raise ValueError("At least one of the intervals must be provided")

if at_tsim is not None:
if isinstance(at_tsim, (int, float)):
at_tsim = [float(at_tsim)]
else:
at_tsim = sorted(float(t) for t in at_tsim)
Comment on lines +50 to +54

@gemini-code-assist gemini-code-assist Bot Jun 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If at_tsim is passed as a string (e.g., "10"), it is not an instance of (int, float), so the code falls back to the else block and attempts to iterate over it. Since strings are iterable in Python, this will silently parse "10" as ['1', '0'] and convert them to [1.0, 0.0], leading to unexpected behavior.

To prevent this, we should include str in the isinstance check so that single string values are correctly converted to a single float (or raise a ValueError if the string is not a valid number).

Suggested change
if at_tsim is not None:
if isinstance(at_tsim, (int, float)):
at_tsim = [float(at_tsim)]
else:
at_tsim = sorted(float(t) for t in at_tsim)
if at_tsim is not None:
if isinstance(at_tsim, (int, float, str)):
at_tsim = [float(at_tsim)]
else:
at_tsim = sorted(float(t) for t in at_tsim)


def deco(func):
func.__simulation_callback__ = {
"func_name": func.__name__,
"tsim_interval": tsim_interval,
"iter_count_interval": iter_count_interval,
"walltime_interval": walltime_interval,
"at_tsim": at_tsim,
}

return func
Expand Down Expand Up @@ -150,28 +172,45 @@ class CallbackInfo:
tsim_interval: float | None = None
iter_count_interval: int | None = None
walltime_interval: float | None = None
at_tsim: list[float] | None = None


class CallbackState:
def __init__(self, info: CallbackInfo, tsim_start: float):
self.info = info
self.counter = 0
self.next_tsim = tsim_start if info.tsim_interval is not None else None

candidates = []
if info.tsim_interval is not None:
candidates.append(tsim_start + info.tsim_interval)
if info.at_tsim is not None:
future = [t for t in info.at_tsim if t > tsim_start]
if future:
candidates.append(min(future))
self.next_tsim = min(candidates) if candidates else None
Comment on lines +183 to +190

@gemini-code-assist gemini-code-assist Bot Jun 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This change introduces a regression where tsim_interval callbacks are no longer triggered at the start of the simulation (tsim_start). In the original code, self.next_tsim was initialized to tsim_start, allowing the callback to run at t = 0.0 before any simulation steps. By changing it to tsim_start + info.tsim_interval, the initial callback is skipped.

Additionally, if at_tsim contains tsim_start (e.g., 0.0), it will also be skipped because of the strict inequality t > tsim_start.

To fix both issues and preserve the original behavior, we should initialize candidates with tsim_start for tsim_interval, and use t >= tsim_start for at_tsim.

Suggested change
candidates = []
if info.tsim_interval is not None:
candidates.append(tsim_start + info.tsim_interval)
if info.at_tsim is not None:
future = [t for t in info.at_tsim if t > tsim_start]
if future:
candidates.append(min(future))
self.next_tsim = min(candidates) if candidates else None
candidates = []
if info.tsim_interval is not None:
candidates.append(tsim_start)
if info.at_tsim is not None:
future = [t for t in info.at_tsim if t >= tsim_start]
if future:
candidates.append(min(future))
self.next_tsim = min(candidates) if candidates else None


self.next_iter_count = 0 if info.iter_count_interval is not None else None
self.next_walltime = 0.0 if info.walltime_interval is not None else None

def advance(self, t_model: float, iter_count: int, walltime: float):
self.counter += 1

candidates = []
if self.info.tsim_interval is not None:
self.next_tsim = t_model + self.info.tsim_interval
candidates.append(t_model + self.info.tsim_interval)
if self.info.at_tsim is not None:
future = [t for t in self.info.at_tsim if t > t_model]
if future:
candidates.append(min(future))
self.next_tsim = min(candidates) if candidates else None

if self.info.iter_count_interval is not None:
self.next_iter_count = iter_count + self.info.iter_count_interval
if self.info.walltime_interval is not None:
self.next_walltime = walltime + self.info.walltime_interval

rank_0_print(f'[Simulation] Advancing callback "{self.info.name}"')
if self.info.tsim_interval is not None:
if self.next_tsim is not None:
rank_0_print(f" -> t = {t_model} -> {self.next_tsim}")
if self.info.iter_count_interval is not None:
rank_0_print(f" -> iter = {iter_count} -> {self.next_iter_count}")
Expand All @@ -183,7 +222,7 @@ def should_trigger(self, t_model: float, iter_count: int, walltime: float) -> bo

log = []

if self.info.tsim_interval is not None:
if self.next_tsim is not None:
if t_model >= self.next_tsim: # should i add a tolerance here ?
trig = True
log.append(f" -> t = {t_model} >= {self.next_tsim}")
Expand Down Expand Up @@ -230,12 +269,17 @@ class SimulationRunner(metaclass=SimulationMeta):
And can define callbacks (any function decorated with @callback):

< call every tsim = i * time_step >
- @callback(time_step=1.0)
- @callback(tsim_interval=1.0)
def analysis(self, icallback):
rank_0_print("analysis")

< call at exact simulation times >
- @callback(at_tsim=[1.0, 5.0, 10.0])
def snapshot(self, icallback):
rank_0_print("snapshot")

< call when tsim = dt_stop, niter_max is reached or walltime_step is reached >
- @callback(time_step=dt_stop, niter_max=1000, walltime_step=30*60)
- @callback(tsim_interval=dt_stop, iter_count_interval=1000, walltime_interval=30*60)
def do_checkpoint(self, icheckpoint):
self.dump_helper.dump(icheckpoint)

Expand Down Expand Up @@ -264,6 +308,7 @@ def __init__(self, model):
tsim_interval=info["tsim_interval"],
iter_count_interval=info["iter_count_interval"],
walltime_interval=info["walltime_interval"],
at_tsim=info["at_tsim"],
)

self._callbacks.append(copied)
Expand Down
Loading