Spaces:
Running on Zero
Running on Zero
File size: 3,278 Bytes
c731c95 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """Cooperative cancellation and deadlines for one process's active graph work."""
from __future__ import annotations
import threading
import time
from contextvars import ContextVar
from functools import wraps
from uuid import UUID
class ExecutionStopped(RuntimeError):
def __init__(self, reason: str):
self.reason = reason
self.code = "research_timeout" if reason == "deadline" else "research_cancelled"
super().__init__(
"Research took too long. Try a more focused question."
if reason == "deadline"
else "Research was stopped."
)
class Execution:
def __init__(self, seconds: float):
self.deadline = time.monotonic() + seconds
self.reason = ""
self.stopped = threading.Event()
self.idle = threading.Event()
self.idle.set()
self._lock = threading.RLock()
self._workers = 0
def stop(self, reason: str = "cancelled") -> None:
with self._lock:
if not self.reason:
self.reason = reason
self.stopped.set()
def check(self) -> None:
if time.monotonic() >= self.deadline:
self.stop("deadline")
if self.stopped.is_set():
raise ExecutionStopped(self.reason)
def enter(self) -> None:
with self._lock:
self.check()
self._workers += 1
self.idle.clear()
def leave(self) -> None:
with self._lock:
self._workers -= 1
if not self._workers:
self.idle.set()
execution: ContextVar[Execution | None] = ContextVar("precedent_execution", default=None)
_active: dict[tuple[str, UUID], Execution] = {}
_registry_lock = threading.Lock()
def register(owner: str, ident: UUID, control: Execution) -> bool:
with _registry_lock:
key = (owner, ident)
if key in _active:
return False
_active[key] = control
return True
def unregister(owner: str, ident: UUID, control: Execution) -> None:
with _registry_lock:
if _active.get((owner, ident)) is control:
del _active[(owner, ident)]
def stop_request(owner: str, ident: UUID) -> bool:
with _registry_lock:
control = _active.get((owner, ident))
if control is None:
return False
control.stop()
return True
def check_execution() -> None:
control = execution.get()
if control is not None:
control.check()
def remaining_seconds(maximum: float) -> float:
check_execution()
control = execution.get()
return min(maximum, max(0.01, control.deadline - time.monotonic())) if control else maximum
def pause(seconds: float) -> None:
control = execution.get()
if control:
control.stopped.wait(remaining_seconds(seconds))
control.check()
else:
time.sleep(seconds)
def guarded_node(node):
@wraps(node)
def guarded(*args, **kwargs):
control = execution.get()
if control is None:
return node(*args, **kwargs)
control.enter()
try:
result = node(*args, **kwargs)
control.check()
return result
finally:
control.leave()
return guarded
|