From 7bd109165dcba3e8f7cc5e0b4c188ac638fa20be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Wed, 4 Nov 2026 07:39:41 -0500 Subject: [PATCH] fix(server): extend the http keep-alive window to 360 seconds --- docs/cli.md | 1 + python/freetoken/server/api_server.py | 10 ++++--- python/freetoken/server/args.py | 10 +++++++ tests/server/test_api_server.py | 42 +++++++++++++++++++++++++++ 4 files changed, 59 insertions(+), 4 deletions(-) create mode 100644 tests/server/test_api_server.py diff --git a/docs/cli.md b/docs/cli.md index 2738dc4ff9..1dee25ff7f 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -41,6 +41,7 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--http-keep-alive-s` | 360 | HTTP idle keep-alive timeout in seconds. Set above the longest client pool idle timeout (e.g. Bun/opencode: 300 s) so the client closes first; uvicorn does not advertise a `Keep-Alive` timeout or buffer the close | | `--api-key` | `$FREETOKEN_API_KEY`, else off | Require the key on every route except `/health`, as `Authorization: Bearer ` or `x-api-key: `; other requests get a 401 | | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | | `--max-running-requests` | 4 | Max concurrently running requests | diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index ce04b34102..acf437a0ff 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -967,7 +967,7 @@ def _flag_shutdown(signum, frame) -> None: signal.signal(sig, _flag_shutdown) -def _serve_and_run_shell(host: str, port: int, api_key: str | None) -> None: +def _serve_and_run_shell(host: str, port: int, api_key: str | None, http_keep_alive_s: int) -> None: """Shell mode: serve the API here, and attach the terminal client to it over the loopback. The shell is an ordinary API client (see ``freetoken.shell``), so shell mode is just @@ -991,7 +991,9 @@ def _serve_and_run_shell(host: str, port: int, api_key: str | None) -> None: netloc = f"[{host}]:{port}" if ":" in host else f"{host}:{port}" origin = resolve_server_url(f"http://{netloc}").origin - server = uvicorn.Server(uvicorn.Config(app, host=host, port=port, access_log=False)) + server = uvicorn.Server(uvicorn.Config( + app, host=host, port=port, access_log=False, timeout_keep_alive=http_keep_alive_s, + )) thread = threading.Thread(target=server.run, name="freetoken-uvicorn", daemon=True) thread.start() _install_shell_stop_handlers() @@ -1129,7 +1131,7 @@ def _on_meta(meta: dict) -> None: ).start() if run_shell: - _serve_and_run_shell(host, port, config.api_key) + _serve_and_run_shell(host, port, config.api_key, config.http_keep_alive_s) return # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group. - uvicorn.run(app, host=host, port=port) + uvicorn.run(app, host=host, port=port, timeout_keep_alive=config.http_keep_alive_s) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 3d75143e2e..4d72cba41c 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -38,6 +38,7 @@ def _nvfp4_entry(value: str) -> str: class ServerArgs(SchedulerConfig): server_host: str = "127.0.0.1" server_port: int = 1919 + http_keep_alive_s: int = 360 # kept out of repr: parse_args logs the whole ServerArgs api_key: str | None = field(default=None, repr=False) num_tokenizer: int = 0 @@ -400,6 +401,15 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The port number for the server to listen on.", ) + parser.add_argument( + "--http-keep-alive-s", + type=int, + default=ServerArgs.http_keep_alive_s, + help="HTTP idle keep-alive timeout in seconds (default: 360). Must exceed the " + "longest client pool idle timeout (e.g. Bun/opencode: 300 s), since uvicorn " + "does not advertise a Keep-Alive timeout or buffer the close.", + ) + parser.add_argument( "--api-key", type=str, diff --git a/tests/server/test_api_server.py b/tests/server/test_api_server.py new file mode 100644 index 0000000000..cc2ce7bc4e --- /dev/null +++ b/tests/server/test_api_server.py @@ -0,0 +1,42 @@ +"""The HTTP idle timeout must reach uvicorn in both server entry paths.""" + +from __future__ import annotations + +import queue +from unittest.mock import patch + +import pytest +import uvicorn +from fastapi import FastAPI + +from freetoken.server import api_server +from freetoken.server.args import parse_args +from freetoken.server.supervisor import BackendHandle + + +@pytest.mark.parametrize("run_shell", [False, True]) +@pytest.mark.parametrize("argv, expected", [([], 360), (["--http-keep-alive-s", "420"], 420)]) +def test_http_keep_alive_reaches_uvicorn(argv, expected, run_shell, monkeypatch): + config, _ = parse_args([ + "--model", "/models/anon", "--dtype", "bfloat16", + "--tool-call-parser", "llama3", "--reasoning-parser", "off", + "--sampling-defaults", "none", *argv, + ], run_shell=run_shell) + + def check_config(self, config): + assert config.timeout_keep_alive == expected + assert config.limit_concurrency is None + raise RuntimeError("stop before binding HTTP") + + monkeypatch.setattr(api_server, "app", FastAPI()) + monkeypatch.setattr(api_server, "_GLOBAL_STATE", None) + monkeypatch.setattr(api_server, "_MODEL_SAMPLING", {}) + monkeypatch.setattr(api_server, "init_request_logging", lambda: None) + monkeypatch.setattr(api_server, "install_polling_access_log_filter", lambda: None) + monkeypatch.setattr(api_server, "ZmqAsyncPullQueue", lambda *a, **kw: None) + monkeypatch.setattr(api_server, "ZmqAsyncPushQueue", lambda *a, **kw: None) + backend = BackendHandle(ack_queue=queue.Queue(), processes=[], expected_acks=0) + # Stop at the real HTTP config, before any socket, worker or terminal starts. + with patch.object(api_server.threading.Thread, "start"), patch.object(uvicorn.Server, "__init__", check_config): + with pytest.raises(RuntimeError, match="stop before binding HTTP"): + api_server.run_api_server(config, lambda: backend, run_shell=run_shell)