Coverage for scanpath_studio/progress.py: 100%
159 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-07 21:10 +0000
1"""UX-165: a Streamlit-free progress hook for slow work.
3Loaders and figure builders call :func:`report` from inside their loops, and
4the orchestrator calls :func:`step_to` between stages. With no task active both
5return at once, so the headless API and the CLI see no change. With one active
6— ``loading.card`` opens it — they update a thread-safe :class:`Task` that the
7card's timer thread reads.
9The two calls are also the **cancel checkpoint**: after :func:`cancel`, the next
10call in the computing thread raises :class:`Cancelled`, so abandoned work stops
11within one file, frame or chunk instead of running to the end of its step.
13A :func:`report` also says that **real work** is happening (:attr:`Task.worked`),
14which is what a card over work that is cheap on a cache hit waits for before it
15shows (``loading.Card``'s ``reveal_on_work``). So a build such a card covers
16that has no loop to report from calls a bare ``report()`` first thing: it only
17runs on a miss.
19And they are where slow work **lets the server talk**. A CPU-bound build on the
20script thread holds Python's GIL, and the server's event loop needs several
21handoffs of it to send one message, each waiting out the interpreter's switch
22interval — so everything a run queues, a card's own reveal included, used to
23reach the browser only once the build ended. Inside a task, a call therefore
24sleeps for ``YIELD_S`` whenever ``YIELD_EVERY_S`` has passed since the task
25began or last did — so a task quicker than that never sleeps: a sleep releases
26the GIL outright, and the loop drains its queue uncontended.
28Nothing here creates a Streamlit element, deliberately. The replay's frame loop
29runs inside two nested ``st.cache_data`` functions, and Streamlit replays every
30element created inside a cached function on each later hit — a progress bar
31drawn from inside the loop would reappear on every rerun.
32"""
34from __future__ import annotations
36import contextvars
37import threading
38import time
39import weakref
40from collections.abc import Hashable, Iterator, Sequence
41from contextlib import contextmanager
42from dataclasses import dataclass
44#: How often a task's checkpoints hand the GIL over, and for how long — at most
45#: YIELD_S / YIELD_EVERY_S (4%) of the work's time, the first time one interval
46#: into the task. Measured on the cold first replay build: the rail reached the
47#: browser at 0.5 s instead of 3.2 s, and the card showed on time instead of not
48#: at all.
49YIELD_EVERY_S = 0.025
50YIELD_S = 0.001
52# The yield's clock and sleep: module attributes, so a test can replace them
53# without patching the `time` module every thread shares.
54_clock = time.monotonic
55_sleep = time.sleep
58class Cancelled(BaseException):
59 """Raised at a checkpoint once the running task has been cancelled.
61 A ``BaseException``, like Streamlit's own ``StopException``, so the broad
62 ``except Exception`` handlers in the load path (a mapping that fails, a
63 corpus that will not read) cannot mistake a cancel for a data error.
64 """
67@dataclass(frozen=True)
68class Snapshot:
69 """One consistent reading of a :class:`Task`, for a card to draw."""
71 title: str
72 steps: tuple[str, ...]
73 step_seconds: tuple[float | None, ...]
74 current: int
75 done: int | None
76 total: int | None
77 unit: str
78 detail: str | None
79 elapsed: float
80 finished: bool
83class Task:
84 """What one slow job has done so far.
86 Written by the computing thread and read by a card's timer thread, so every
87 field sits behind one lock and readers take a :class:`Snapshot`.
88 """
90 def __init__(self, key: Hashable, *, title: str, steps: Sequence[str] = ()):
91 self.key = key
92 self._lock = threading.Lock()
93 self._title = title
94 self._steps = list(steps)
95 self._seconds: list[float | None] = [None] * len(self._steps)
96 self._current = 0
97 self._done: int | None = None
98 self._total: int | None = None
99 self._unit = ""
100 self._detail: str | None = None
101 self.started = time.monotonic()
102 self._step_started = self.started
103 self._finished = False
104 self._worked = False
105 self._cancelled = threading.Event()
106 # The first yield comes YIELD_EVERY_S into the task, not at its first
107 # checkpoint: a plain rerun's page card makes one, and a quick run must
108 # not pay a sleep for it. Read from `_clock`, the clock `_checkpoint`
109 # compares against (and the one a test replaces).
110 self._last_yield = _clock()
112 @property
113 def cancelled(self) -> bool:
114 return self._cancelled.is_set()
116 @property
117 def finished(self) -> bool:
118 with self._lock:
119 return self._finished
121 @property
122 def worked(self) -> bool:
123 """Has anything reported to this task — real work, i.e. a cache miss?
125 Set by :meth:`report`, never by :meth:`step_to`: the orchestrator moves
126 between steps on a cache hit just the same, while only a build that is
127 actually running reports (a bare ``report()`` is enough). A gated card
128 (``loading.Card``'s ``reveal_on_work``) waits for this.
129 """
130 with self._lock:
131 return self._worked
133 def cancel(self) -> None:
134 self._cancelled.set()
136 def _checkpoint(self) -> None:
137 """Stop a cancelled task here; otherwise let the server send (see the
138 module docstring) at most every ``YIELD_EVERY_S``."""
139 if self._cancelled.is_set():
140 raise Cancelled(self.key)
141 # The slot is claimed under the lock, before the sleep, so two threads
142 # reporting to one joined task take one yield per window between them;
143 # the sleep is outside it, since the card's timer reads snapshots under
144 # this lock.
145 with self._lock:
146 due = _clock() - self._last_yield >= YIELD_EVERY_S
147 if due:
148 self._last_yield = _clock()
149 if due:
150 _sleep(YIELD_S)
152 def report(
153 self,
154 done: int | None = None,
155 total: int | None = None,
156 *,
157 unit: str = "",
158 detail: str | None = None,
159 ) -> None:
160 """Record progress — and that work is happening (:attr:`worked`).
162 A bare ``report()`` (no count, unit or detail) only does the latter,
163 leaving the count as it is: a run that waits on a build another run is
164 counting into the same, joined task must not wipe that count.
165 """
166 self._checkpoint()
167 with self._lock:
168 self._worked = True
169 if done is None and total is None and not unit and detail is None:
170 return
171 self._done = done
172 self._total = total
173 if unit:
174 self._unit = unit
175 if detail is not None:
176 self._detail = detail
178 def step_to(self, index: int, label: str | None = None) -> None:
179 """Start step ``index``, finishing the ones before it.
181 Never moves back: a run that joined a task another run was driving may
182 ask for a step the other already reached.
183 """
184 self._checkpoint()
185 now = time.monotonic()
186 with self._lock:
187 if label and 0 <= index < len(self._steps):
188 self._steps[index] = label
189 if index <= self._current:
190 return
191 for i in range(self._current, min(index, len(self._steps))):
192 if self._seconds[i] is None:
193 self._seconds[i] = (
194 now - self._step_started if i == self._current else 0.0
195 )
196 self._current = index
197 self._step_started = now
198 self._done = self._total = None
199 self._unit = ""
200 self._detail = None
202 def finish(self, *, duration_key: Hashable | None = None) -> float:
203 """Mark every step done; remember the duration under ``duration_key``."""
204 now = time.monotonic()
205 with self._lock:
206 for i in range(len(self._steps)):
207 if self._seconds[i] is None:
208 self._seconds[i] = (
209 now - self._step_started if i == self._current else 0.0
210 )
211 self._current = len(self._steps)
212 self._finished = True
213 seconds = now - self.started
214 if duration_key is not None:
215 with _REGISTRY_LOCK:
216 _DURATIONS[duration_key] = seconds
217 return seconds
219 def snapshot(self) -> Snapshot:
220 with self._lock:
221 return Snapshot(
222 title=self._title,
223 steps=tuple(self._steps),
224 step_seconds=tuple(self._seconds),
225 current=self._current,
226 done=self._done,
227 total=self._total,
228 unit=self._unit,
229 detail=self._detail,
230 elapsed=time.monotonic() - self.started,
231 finished=self._finished,
232 )
235#: Every task key carries the session id, so a plain dict here kept every task of
236#: every session for the life of the server. Weak: a task stays registered
237#: exactly while something holds it — its card during the run, and a superseded
238#: run's thread while that is still computing it, the case joining exists for.
239_REGISTRY: weakref.WeakValueDictionary[Hashable, Task] = weakref.WeakValueDictionary()
240#: One float per duration key (a dataset name), so this one stays a plain dict.
241_DURATIONS: dict[Hashable, float] = {}
242_REGISTRY_LOCK = threading.Lock()
243_ACTIVE: contextvars.ContextVar[Task | None] = contextvars.ContextVar(
244 "scanpath_progress_task", default=None
245)
248def begin(
249 key: Hashable, *, title: str, steps: Sequence[str] = (), fresh: bool = False
250) -> Task:
251 """The task for ``key`` — the one already running, or a new one.
253 Joining is what lets a rerun that interrupted a load keep showing that
254 load's counts instead of starting from zero — while the interrupted run's
255 thread still holds the task (`_REGISTRY` is weak). A finished or cancelled
256 task is never joined: its record is replaced. ``fresh=True`` replaces any
257 existing record unconditionally, even one still mid-flight — for a caller
258 with no stable identity to join across runs in the first place (e.g. a
259 region card opened with no explicit task key).
260 """
261 with _REGISTRY_LOCK:
262 task = None if fresh else _REGISTRY.get(key)
263 if task is None or task.cancelled or task.finished:
264 task = Task(key, title=title, steps=steps)
265 _REGISTRY[key] = task
266 return task
269def activate(task: Task) -> contextvars.Token:
270 return _ACTIVE.set(task)
273def deactivate(token: contextvars.Token) -> None:
274 _ACTIVE.reset(token)
277def active() -> Task | None:
278 return _ACTIVE.get()
281@contextmanager
282def task(key: Hashable, *, title: str, steps: Sequence[str] = ()) -> Iterator[Task]:
283 """:func:`begin` + :func:`activate` for one block."""
284 current = begin(key, title=title, steps=steps)
285 token = activate(current)
286 try:
287 yield current
288 finally:
289 deactivate(token)
292@contextmanager
293def scope() -> Iterator[None]:
294 """No active task inside, whatever an earlier run in this thread left."""
295 token = _ACTIVE.set(None)
296 try:
297 yield
298 finally:
299 _ACTIVE.reset(token)
302def report(
303 done: int | None = None,
304 total: int | None = None,
305 *,
306 unit: str = "",
307 detail: str | None = None,
308) -> None:
309 """Record progress on the active task — a no-op without one.
311 Any call, a bare ``report()`` included, marks the task as having worked;
312 a bare one leaves its count alone (see :meth:`Task.report`).
313 """
314 current = _ACTIVE.get()
315 if current is not None:
316 current.report(done, total, unit=unit, detail=detail)
319def step_to(index: int, label: str | None = None) -> None:
320 """Move the active task to step ``index`` — a no-op without one."""
321 current = _ACTIVE.get()
322 if current is not None:
323 current.step_to(index, label)
326def running(key: Hashable) -> bool:
327 """Is a task running under ``key`` — registered, neither finished nor
328 cancelled?
330 The registry is weak (`_REGISTRY`), so a task is in it only while something
331 holds it: its card during its run, or a superseded run's thread still
332 computing it. A key some run left behind — one that ended by ``st.stop()``,
333 an exception or a rerun in the middle of its work — names no task once
334 nothing holds that task any more, and so no work in flight.
335 """
336 with _REGISTRY_LOCK:
337 current = _REGISTRY.get(key)
338 return current is not None and not current.finished and not current.cancelled
341def cancel(key: Hashable) -> None:
342 """Cancel ``key``'s task: its computing thread stops at its next checkpoint."""
343 with _REGISTRY_LOCK:
344 current = _REGISTRY.get(key)
345 if current is not None:
346 current.cancel()
349def last_duration(key: Hashable) -> float | None:
350 with _REGISTRY_LOCK:
351 return _DURATIONS.get(key)