Skip to content

Commit

Permalink
Add Runner.add_periodic_callback
Browse files Browse the repository at this point in the history
  • Loading branch information
basnijholt committed Nov 23, 2022
1 parent 28d4c35 commit c9ce912
Showing 1 changed file with 29 additions and 7 deletions.
36 changes: 29 additions & 7 deletions adaptive/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -649,6 +649,7 @@ def __init__(

self.task = self.ioloop.create_task(self._run())
self.saving_task = None
self.callbacks = []
if in_ipynb() and not self.ioloop.is_running():
warnings.warn(
"The runner has been scheduled, but the asyncio "
Expand Down Expand Up @@ -753,6 +754,31 @@ def elapsed_time(self):
end_time = time.time()
return end_time - self.start_time

def add_periodic_callback(
self,
method: Callable[[AsyncRunner]],
interval: int = 30,
):
"""Start a periodic callback that calls the given method on the runner.
Parameters
----------
method : callable
The method to call periodically.
interval : int
The interval in seconds between the calls.
"""

async def _callback():
while self.status() == "running":
method(self)
await asyncio.sleep(interval)
method(self) # one last time

task = self.ioloop.create_task(_callback())
self.callbacks.append(task)
return task

def start_periodic_saving(
self,
save_kwargs: dict[str, Any] | None = None,
Expand Down Expand Up @@ -781,6 +807,8 @@ def start_periodic_saving(
... save_kwargs=dict(fname='data/test.pickle'),
... interval=600)
"""
if self.saving_task is not None:
raise RuntimeError("Already saving.")

def default_save(learner):
learner.save(**save_kwargs)
Expand All @@ -790,13 +818,7 @@ def default_save(learner):
if save_kwargs is None:
raise ValueError("Must provide `save_kwargs` if method=None.")

async def _saver():
while self.status() == "running":
method(self.learner)
await asyncio.sleep(interval)
method(self.learner) # one last time

self.saving_task = self.ioloop.create_task(_saver())
self.saving_task = self.add_periodic_callback(method, interval=interval)
return self.saving_task


Expand Down

0 comments on commit c9ce912

Please sign in to comment.