diff --git a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist index 0bc8eb92..17c530e8 100644 --- a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist +++ b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist @@ -38,7 +38,7 @@ --cache-link-mbps10000 --cache-default-rtt-ms0.55 --remote-prefill-min-tokens0 - --prefill-worker-timeout-s300 + --prefill-worker-timeout-s900 --prefill-policyremote-required --network-http-host127.0.0.1 --network-http-port8090 diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 7b55fae9..5f0623bf 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -335,10 +335,11 @@ streams `generator>` and `critic>` tokens as they arrive, followed by KV hit rate, decode tok/s, generation latency and E2E tok/s. `/quit` exits; services remain running. Reports remain redacted and appear in the Benchmarks tab. Generation uses 64-token streaming chunks and automatically continues to the -model's EOS. `--max-response-tokens` defaults to 512 as an explicit safety cap; -reaching it marks the stage incomplete instead of presenting a truncated answer -as successful. The Critic receives the Generator completion status and must not -penalize an honest statement that an open problem has no accepted proof. +model's EOS. `--max-response-tokens` defaults to `0` (no client-side cap); +operators may explicitly set a positive cap, which marks the stage incomplete +instead of presenting a truncated answer as successful. The Critic receives the +Generator completion status and must not penalize an honest statement that an +open problem has no accepted proof. The REPL ignores external `SIGTERM`; its shell supervisor restarts signal-based exits. Only `/quit`, `/exit`, or EOF is treated as approval to stop. diff --git a/scripts/agent_gan_inference_demo.py b/scripts/agent_gan_inference_demo.py index 8d894a3d..91720868 100644 --- a/scripts/agent_gan_inference_demo.py +++ b/scripts/agent_gan_inference_demo.py @@ -51,11 +51,19 @@ def _infer( append_done = time.perf_counter() first_at = None generated = [] - response_limit = int(max_response_tokens or output_tokens) + response_limit = ( + int(output_tokens) + if max_response_tokens is None + else int(max_response_tokens) or None + ) stop_reason = "unknown" - while len(generated) < response_limit: + while response_limit is None or len(generated) < response_limit: before_count = len(generated) - chunk = min(output_tokens, response_limit - len(generated)) + chunk = ( + output_tokens + if response_limit is None + else min(output_tokens, response_limit - len(generated)) + ) for token in s.generate(max_tokens=chunk): generated.append(int(token)) if on_token is not None: @@ -73,7 +81,11 @@ def _infer( if len(generated) == before_count: stop_reason = "no_progress" break - if len(generated) >= response_limit and stop_reason == "max_tokens": + if ( + response_limit is not None + and len(generated) >= response_limit + and stop_reason == "max_tokens" + ): stop_reason = "client_safety_limit" done = time.perf_counter() after = get_stats() @@ -106,7 +118,12 @@ def main() -> int: "reliably; larger values require more worker memory or timeout.", ) parser.add_argument("--output-tokens", type=int, default=64) - parser.add_argument("--max-response-tokens", type=int, default=512) + parser.add_argument( + "--max-response-tokens", + type=int, + default=0, + help="Optional client response cap; 0 means generate until model EOS.", + ) parser.add_argument("--report", default="/tmp/kakeya-agent-gan-demo.json") parser.add_argument("--skip-ensure", action="store_true") args = parser.parse_args() diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index e9e8adcf..1c603e38 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -76,7 +76,12 @@ def main() -> int: parser.add_argument("--api-key-file", default="~/.kakeya/network_api_key") parser.add_argument("--tokenizer-id", required=True) parser.add_argument("--output-tokens", type=int, default=64) - parser.add_argument("--max-response-tokens", type=int, default=512) + parser.add_argument( + "--max-response-tokens", + type=int, + default=0, + help="Optional client response cap; 0 means generate until model EOS.", + ) parser.add_argument("--skip-ensure", action="store_true") args = parser.parse_args() diff --git a/tests/inference_engine/bridge/test_agent_gan_demo.py b/tests/inference_engine/bridge/test_agent_gan_demo.py index 82edc4db..d339b295 100644 --- a/tests/inference_engine/bridge/test_agent_gan_demo.py +++ b/tests/inference_engine/bridge/test_agent_gan_demo.py @@ -67,7 +67,7 @@ def test_infer_continues_chunks_until_eos(): 2, lambda: {}, on_token=lambda values: streamed.append(list(values)), - max_response_tokens=10, + max_response_tokens=0, ) assert tokens == [1, 2, 3] assert streamed == [[1], [1, 2], [1, 2, 3]] diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py index 3044b677..f305f806 100644 --- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py +++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py @@ -48,6 +48,10 @@ def test_two_mac_deployment_uses_allens_as_prefill_only(): "--prefill-policyremote-required" in plist ) + assert ( + "--prefill-worker-timeout-s900" + in plist + ) assert ( "--cache-tenant-idprivate-fleet" in plist