diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index 6a3d733..ace3c82 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -20,6 +20,7 @@ _ensure_services, _json_request, ) +from inference_engine.bench.prefill_fleet_report import summarize_stages def install_signal_protection() -> None: @@ -33,6 +34,18 @@ def ignore_sigterm(signum, _frame): signal.signal(signal.SIGTERM, ignore_sigterm) +def _telemetry_request(url: str, **kwargs): + try: + return _json_request(url, timeout=2, **kwargs) + except Exception as exc: + print( + f"[telemetry-warning] {type(exc).__name__}: {exc}; " + "inference will continue", + flush=True, + ) + return None + + class TokenPrinter: def __init__(self, tokenizer, label: str) -> None: self.tokenizer = tokenizer @@ -134,8 +147,15 @@ def main() -> int: eos_ids = _resolve_eos_token_ids(tokenizer) api_key = Path(args.api_key_file).expanduser().read_text().strip() + telemetry_state = {"degraded": False, "last_stats": {}} + def get_stats(): - return _json_request(f"{args.dashboard}/v1/network/prefill") + stats = _telemetry_request(f"{args.dashboard}/v1/network/prefill") + if stats is None: + telemetry_state["degraded"] = True + return dict(telemetry_state["last_stats"]) + telemetry_state["last_stats"] = stats + return stats print( "Kakeya Agent GAN REPL ready. Type a prompt; /quit exits.\n" @@ -156,7 +176,8 @@ def get_stats(): print("[bye]") break run_nonce = uuid.uuid4().hex - run = _json_request( + telemetry_state["degraded"] = False + run = _telemetry_request( f"{args.dashboard}/v1/network/benchmarks", api_key=api_key, method="POST", @@ -171,7 +192,8 @@ def get_stats(): }, }, ) - run_id = run["id"] + remote_run = run is not None + run_id = run["id"] if remote_run else f"local_{run_nonce[:16]}" try: generator_messages = [ { @@ -223,14 +245,15 @@ def get_stats(): generator_actual, generator_text, ) - if not generator_stage["ok"]: + if not generator_stage["ok"] and not telemetry_state["degraded"]: raise RuntimeError("Generator KV gate failed") - _json_request( - f"{args.dashboard}/v1/network/benchmarks/{run_id}", - api_key=api_key, - method="PATCH", - body={"stages": [generator_stage]}, - ) + if remote_run: + _telemetry_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={"stages": [generator_stage]}, + ) evidence, evidence_metrics = build_critic_evidence( tokenizer, @@ -301,19 +324,25 @@ def get_stats(): critic_text, extra_metrics=evidence_metrics, ) - if not critic_stage["ok"]: + if not critic_stage["ok"] and not telemetry_state["degraded"]: raise RuntimeError("Critic KV gate failed") - completed = _json_request( - f"{args.dashboard}/v1/network/benchmarks/{run_id}", - api_key=api_key, - method="PATCH", - body={ - "stages": [critic_stage], - "status": "completed", - "finished_at": time.time(), - }, + completed = None + if remote_run: + completed = _telemetry_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={ + "stages": [critic_stage], + "status": "completed", + "finished_at": time.time(), + }, + ) + summary = ( + completed["summary"] + if completed is not None + else summarize_stages([generator_stage, critic_stage]) ) - summary = completed["summary"] print( "[metrics] " f"KV hit={summary['workload_kv_token_hit_rate']:.1%} " @@ -324,12 +353,13 @@ def get_stats(): flush=True, ) except Exception as exc: - _json_request( - f"{args.dashboard}/v1/network/benchmarks/{run_id}", - api_key=api_key, - method="PATCH", - body={"status": "failed", "finished_at": time.time()}, - ) + if remote_run: + _telemetry_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={"status": "failed", "finished_at": time.time()}, + ) print(f"[error] {type(exc).__name__}: {exc}", flush=True) return 0 diff --git a/scripts/benchmark_prefill_architecture.py b/scripts/benchmark_prefill_architecture.py index 55942d5..c68248f 100644 --- a/scripts/benchmark_prefill_architecture.py +++ b/scripts/benchmark_prefill_architecture.py @@ -20,13 +20,20 @@ ) -def _json_request(url: str, *, api_key: str = "", method: str = "GET", body=None): +def _json_request( + url: str, + *, + api_key: str = "", + method: str = "GET", + body=None, + timeout: float = 10, +): data = None if body is None else json.dumps(body).encode() headers = {"Content-Type": "application/json"} if api_key: headers["X-API-Key"] = api_key request = urllib.request.Request(url, data=data, headers=headers, method=method) - with urllib.request.urlopen(request, timeout=10) as response: + with urllib.request.urlopen(request, timeout=timeout) as response: return json.load(response) diff --git a/tests/inference_engine/bridge/test_agent_gan_repl.py b/tests/inference_engine/bridge/test_agent_gan_repl.py index e208bb9..4847e8e 100644 --- a/tests/inference_engine/bridge/test_agent_gan_repl.py +++ b/tests/inference_engine/bridge/test_agent_gan_repl.py @@ -6,6 +6,7 @@ PrefillHeartbeat, TokenPrinter, _stage, + _telemetry_request, install_signal_protection, ) @@ -124,3 +125,17 @@ def test_stage_includes_evidence_window_metrics(): extra_metrics={"critic_omitted_tokens": 100}, ) assert stage["critic_omitted_tokens"] == 100 + + +def test_telemetry_timeout_warns_without_stopping_inference( + monkeypatch, + capsys, +): + def timeout(*_args, **_kwargs): + raise TimeoutError("timed out") + + monkeypatch.setattr("scripts.agent_gan_repl._json_request", timeout) + assert _telemetry_request("http://dashboard/metrics") is None + output = capsys.readouterr().out + assert "telemetry-warning" in output + assert "inference will continue" in output