Skip to content

Commit a90223b

Browse files
fluffy314cursoragent
authored andcommitted
fix(autoresearch): show strategy inference progress
Expose Strategy Prefill and decode activity so the AutoResearch planning phase remains observable before the live GAN experiment begins. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent e4c67bd commit a90223b

2 files changed

Lines changed: 109 additions & 1 deletion

File tree

autoresearch/prefill/supervisor.py

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,76 @@ def parse_research_verdict(output: str, candidate_id: str) -> dict:
139139
}
140140

141141

142+
class StrategyPrefillHeartbeat:
143+
def __init__(
144+
self,
145+
dashboard: str = "http://127.0.0.1:8090",
146+
interval_s: float = 10.0,
147+
) -> None:
148+
self.dashboard = dashboard.rstrip("/")
149+
self.interval_s = interval_s
150+
self._stop = threading.Event()
151+
self._thread: threading.Thread | None = None
152+
self._baseline: dict = {}
153+
self._last: tuple | None = None
154+
155+
def __enter__(self):
156+
try:
157+
self._baseline = _json_request(
158+
f"{self.dashboard}/v1/network/summary",
159+
).get("prefill", {})
160+
except Exception as exc:
161+
print(
162+
"[autoresearch] Strategy Prefill telemetry warning: "
163+
f"{type(exc).__name__}: {exc}",
164+
flush=True,
165+
)
166+
self._thread = threading.Thread(target=self._run, daemon=True)
167+
self._thread.start()
168+
return self
169+
170+
def __exit__(self, *_exc) -> None:
171+
self._stop.set()
172+
if self._thread is not None:
173+
self._thread.join(timeout=self.interval_s + 2)
174+
self._emit()
175+
176+
def _run(self) -> None:
177+
while not self._stop.wait(self.interval_s):
178+
self._emit()
179+
180+
def _delta(self, current: dict, name: str) -> int:
181+
return max(
182+
0,
183+
int(current.get(name, 0)) - int(self._baseline.get(name, 0)),
184+
)
185+
186+
def _emit(self) -> None:
187+
try:
188+
current = _json_request(
189+
f"{self.dashboard}/v1/network/summary",
190+
).get("prefill", {})
191+
except Exception:
192+
return
193+
total = self._delta(current, "remote_job_tokens_total")
194+
computed = self._delta(current, "remote_job_tokens_computed")
195+
state = (
196+
computed,
197+
total,
198+
self._delta(current, "remote_hits"),
199+
self._delta(current, "tokens_reused"),
200+
)
201+
if not total or state == self._last:
202+
return
203+
self._last = state
204+
percent = min(100.0, 100.0 * computed / total)
205+
print(
206+
f"[autoresearch] Strategy Prefill: {computed}/{total} tokens "
207+
f"({percent:.1f}%) · remote_hits={state[2]} reused={state[3]}",
208+
flush=True,
209+
)
210+
211+
142212
def propose_candidate(
143213
*,
144214
address: str,
@@ -175,15 +245,29 @@ def propose_candidate(
175245
enable_thinking=False,
176246
)
177247
generated: list[int] = []
248+
print(
249+
f"[autoresearch] Strategy Prefill: 0/{len(ids)} tokens (0.0%)",
250+
flush=True,
251+
)
178252
with Client(address) as client:
179253
with client.create_session(
180254
eos_token_ids=_resolve_eos_token_ids(tokenizer),
181255
client_label="autoresearch-strategy",
182256
) as session:
183-
session.append(ids)
257+
with StrategyPrefillHeartbeat():
258+
session.append(ids)
259+
print(
260+
f"[autoresearch] Strategy Prefill complete: {len(ids)} tokens",
261+
flush=True,
262+
)
184263
while len(generated) < 2048:
185264
before = len(generated)
186265
generated.extend(int(token) for token in session.generate(max_tokens=64))
266+
print(
267+
f"[autoresearch] Strategy Decode: {len(generated)} tokens "
268+
f"stop_reason={session.last_stop_reason}",
269+
flush=True,
270+
)
187271
if session.last_stop_reason != 1:
188272
break
189273
if len(generated) == before:

tests/inference_engine/bench/test_autoresearch_supervisor.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
read_results,
77
render_candidate,
88
should_keep,
9+
StrategyPrefillHeartbeat,
910
validate_candidate,
1011
)
1112
from pathlib import Path
@@ -164,6 +165,29 @@ def fake_run(command, **kwargs):
164165
assert "a[i+1]='128'" in remote
165166

166167

168+
def test_strategy_prefill_heartbeat_reports_delta(monkeypatch, capsys):
169+
heartbeat = StrategyPrefillHeartbeat(interval_s=0.01)
170+
heartbeat._baseline = {
171+
"remote_job_tokens_total": 100,
172+
"remote_job_tokens_computed": 100,
173+
"remote_hits": 2,
174+
"tokens_reused": 20,
175+
}
176+
monkeypatch.setattr(
177+
"autoresearch.prefill.supervisor._json_request",
178+
lambda _url: {"prefill": {
179+
"remote_job_tokens_total": 300,
180+
"remote_job_tokens_computed": 228,
181+
"remote_hits": 3,
182+
"tokens_reused": 84,
183+
}},
184+
)
185+
heartbeat._emit()
186+
output = capsys.readouterr().out
187+
assert "128/200 tokens (64.0%)" in output
188+
assert "remote_hits=1 reused=64" in output
189+
190+
167191
def test_supervisor_predeploys_before_real_strategy_proposal():
168192
source = (
169193
Path(__file__).resolve().parents[3]

0 commit comments

Comments
 (0)