|
37 | 37 | import struct |
38 | 38 | import threading |
39 | 39 | import time |
40 | | -from collections import defaultdict |
| 40 | +from collections import defaultdict, deque |
41 | 41 | from dataclasses import dataclass, field |
42 | 42 | from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer |
43 | 43 | from typing import Any, Callable, Dict, List, Optional, Tuple |
@@ -619,6 +619,16 @@ def __init__( |
619 | 619 | # persisted, so a restart still over budget re-sends the stop. |
620 | 620 | self.token_budget = max(0, int(token_budget)) |
621 | 621 | self._budget_stop_sent = False |
| 622 | + # Rolling-rate samples for the token-budget progress / ETA — (monotonic |
| 623 | + # ts, total_tokens) appended whenever a heartbeat advances the aggregate. |
| 624 | + # In-memory only (a live gauge; resets on restart). The window is wide on |
| 625 | + # purpose: heartbeats are ~30 s apart, so a short window would see only a |
| 626 | + # sample or two and give a jittery rate. 300 s spans 6+ heartbeats per |
| 627 | + # worker for a stable estimate while still reflecting a sustained |
| 628 | + # slowdown within a few minutes. Not a CLI knob — tune here if ever |
| 629 | + # needed. maxlen bounds memory if heartbeats are unusually dense. |
| 630 | + self._token_samples: deque = deque(maxlen=512) |
| 631 | + self._token_rate_window_s: float = 300.0 |
622 | 632 | # Verbose per-round sync logging (off by default). Server-authoritative: |
623 | 633 | # advertised in /info so each worker gates its own per-round sync log to |
624 | 634 | # match, and the server gates its per-round outer-step log. A targeted |
@@ -2837,6 +2847,7 @@ def _handle_heartbeat(self, handler: BaseHTTPRequestHandler): |
2837 | 2847 | # the EMA math). |
2838 | 2848 | if worker_stats: |
2839 | 2849 | self._stats.update(worker_id, worker_stats) |
| 2850 | + self._record_token_sample() |
2840 | 2851 | self._maybe_trigger_token_budget() |
2841 | 2852 |
|
2842 | 2853 | # Compute DyLU recommendation if enabled |
@@ -2869,6 +2880,76 @@ def _handle_heartbeat(self, handler: BaseHTTPRequestHandler): |
2869 | 2880 | if worker_stats: |
2870 | 2881 | self._maybe_log_stats() |
2871 | 2882 |
|
| 2883 | + def _record_token_sample(self): |
| 2884 | + """Append a (monotonic ts, total_tokens) sample for the rolling rate. |
| 2885 | +
|
| 2886 | + No-op without a token budget (the rate/ETA is only surfaced for budgeted |
| 2887 | + runs). Lock-free on purpose: ``deque.append`` and the ``total_tokens`` |
| 2888 | + int read are each atomic under the GIL, and different workers' heartbeats |
| 2889 | + race here (this is NOT serial-per-worker like ``pending_command``). A |
| 2890 | + cross-worker read could append two samples a hair out of order; the |
| 2891 | + ``dt<=0`` / ``dtok<0`` guards in ``_rolling_token_rate`` absorb that, and |
| 2892 | + a momentarily stale gauge is fine. |
| 2893 | + """ |
| 2894 | + if self.token_budget <= 0: |
| 2895 | + return |
| 2896 | + self._token_samples.append((time.monotonic(), self._stats.total_tokens)) |
| 2897 | + |
| 2898 | + def _rolling_token_rate(self) -> Optional[float]: |
| 2899 | + """Tokens/second over the recent window, or None if not yet estimable. |
| 2900 | +
|
| 2901 | + Uses the oldest sample within ``_token_rate_window_s`` as the baseline |
| 2902 | + (or the oldest available if the buffer is younger than the window), so a |
| 2903 | + rate appears after the second heartbeat rather than waiting a full |
| 2904 | + window. Returns None until there are two spanning samples. |
| 2905 | + """ |
| 2906 | + samples = list(self._token_samples) |
| 2907 | + if len(samples) < 2: |
| 2908 | + return None |
| 2909 | + now_t, now_tok = samples[-1] |
| 2910 | + cutoff = now_t - self._token_rate_window_s |
| 2911 | + base_t, base_tok = samples[0] |
| 2912 | + for t, tok in samples: |
| 2913 | + if t >= cutoff: |
| 2914 | + base_t, base_tok = t, tok |
| 2915 | + break |
| 2916 | + # After a long stall the only in-window sample can be the tail itself |
| 2917 | + # (base collapses onto ``now``). Fall back to the previous sample so we |
| 2918 | + # still produce a rate — a stall then reads as ~0 tok/s (the slowdown we |
| 2919 | + # most want visible), not a blank gauge. len>=2 guaranteed above. |
| 2920 | + if base_t >= now_t: |
| 2921 | + base_t, base_tok = samples[-2] |
| 2922 | + dt = now_t - base_t |
| 2923 | + dtok = now_tok - base_tok |
| 2924 | + if dt <= 0 or dtok < 0: |
| 2925 | + return None |
| 2926 | + return dtok / dt |
| 2927 | + |
| 2928 | + def _token_progress(self) -> Optional[Dict[str, Any]]: |
| 2929 | + """Token-budget progress + rolling-rate ETA, or None when no budget. |
| 2930 | +
|
| 2931 | + Shape (all keys present when a budget is set): |
| 2932 | + tokens_completed, token_budget, fraction (0..1, clamped), |
| 2933 | + tokens_per_second (None until estimable), eta_seconds (None when the |
| 2934 | + rate is unknown or zero), rate_window_seconds, budget_stop_sent. |
| 2935 | + """ |
| 2936 | + if self.token_budget <= 0: |
| 2937 | + return None |
| 2938 | + done = int(self._stats.total_tokens) |
| 2939 | + budget = int(self.token_budget) |
| 2940 | + rate = self._rolling_token_rate() |
| 2941 | + remaining = max(0, budget - done) |
| 2942 | + eta = remaining / rate if (rate and rate > 0) else None |
| 2943 | + return { |
| 2944 | + "tokens_completed": done, |
| 2945 | + "token_budget": budget, |
| 2946 | + "fraction": min(1.0, done / budget) if budget > 0 else 0.0, |
| 2947 | + "tokens_per_second": rate, |
| 2948 | + "eta_seconds": eta, |
| 2949 | + "rate_window_seconds": self._token_rate_window_s, |
| 2950 | + "budget_stop_sent": self._budget_stop_sent, |
| 2951 | + } |
| 2952 | + |
2872 | 2953 | def _maybe_trigger_token_budget(self): |
2873 | 2954 | """When the aggregated global token count reaches ``token_budget``, |
2874 | 2955 | relay ``save_and_stop`` to every worker (once). The command rides each |
@@ -3176,6 +3257,10 @@ def _handle_status(self, handler: BaseHTTPRequestHandler): |
3176 | 3257 | if self.token_budget > 0: |
3177 | 3258 | response["token_budget"] = self.token_budget |
3178 | 3259 | response["budget_stop_sent"] = self._budget_stop_sent |
| 3260 | + # Progress + rolling-rate ETA (tokens_completed, fraction, |
| 3261 | + # tokens_per_second, eta_seconds, rate_window_seconds) for the |
| 3262 | + # webui progress bar and the CLI status line. |
| 3263 | + response["token_progress"] = self._token_progress() |
3179 | 3264 |
|
3180 | 3265 | if self.async_mode: |
3181 | 3266 | response["total_submissions"] = self._total_submissions |
|
0 commit comments