diff --git a/adaptive/runner.py b/adaptive/runner.py index 79cc54aa6..0007dd97d 100644 --- a/adaptive/runner.py +++ b/adaptive/runner.py @@ -259,6 +259,11 @@ def _process_futures(self, done_futs): try: y = fut.result() t = time.time() - fut.start_time # total execution time + except asyncio.CancelledError: + # Cleanup + self._to_retry.pop(pid, None) + self._tracebacks.pop(pid, None) + self._id_to_point.pop(pid, None) except Exception as e: self._tracebacks[pid] = traceback.format_exc() self._to_retry[pid] = self._to_retry.get(pid, 0) + 1 @@ -754,6 +759,22 @@ def elapsed_time(self): end_time = time.time() return end_time - self.start_time + def cancel_point( + self, point: Any | None = None, future: asyncio.Future | None = None + ): + """Cancel a point that is currently being evaluated. + + Parameters + ---------- + point + The point that should be cancelled. + """ + if point is None and future is None: + raise ValueError("Either point or future must be given") + if future is None: + future = next(fut for fut, p in self.pending_points if p == point) + future.cancel() + def add_periodic_callback( self, method: Callable[[AsyncRunner]],