From a30996865a8b27dda21cb1ec0d000c1f7c9da201 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Wed, 26 Aug 2026 21:11:57 -0400 Subject: [PATCH] feat(server): add TLS certificate support --- docs/cli.md | 11 +++- python/freetoken/server/api_server.py | 13 ++++- python/freetoken/server/args.py | 23 ++++++++ tests/server/test_tls_args.py | 75 +++++++++++++++++++++++++++ 4 files changed, 120 insertions(+), 2 deletions(-) create mode 100644 tests/server/test_tls_args.py diff --git a/docs/cli.md b/docs/cli.md index ff4af382d..89e4ffad1 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -40,6 +40,8 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--ssl-certfile` | disabled | PEM certificate chain for HTTPS; requires `--ssl-keyfile` | +| `--ssl-keyfile` | disabled | PEM private key for HTTPS; requires `--ssl-certfile` | | `--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 | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | @@ -48,6 +50,14 @@ parsers all resolve automatically from the checkpoint and the GPU. | `--cuda-graph-max-bs`, `--graph` | = max running requests | Max batch size captured as CUDA graphs | | `--decode-log-interval` | 40 | Scheduler status line every N decode steps | +To serve HTTPS directly, provide the certificate and private key together: + +```bash +ft serve --model ... --host 0.0.0.0 \ + --ssl-certfile /etc/ssl/example/fullchain.pem \ + --ssl-keyfile /etc/ssl/example/privkey.pem +``` + ### Choosing a GPU For example, a machine with an RTX 5090 and an RTX 3060 Ti: @@ -172,4 +182,3 @@ profile that `ft serve --moe-backend auto` and `--moe-hybrid-max-fetch -1` then - What to measure: `--dtype`, `--model`, `--formats`, `--isa`. - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe by that factor. - diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc854..04cbda2ac 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -917,6 +917,17 @@ def _serve_and_run_shell(host: str, port: int) -> None: _reap_backend_workers(_GLOBAL_STATE.backend_processes) +def _uvicorn_tls_kwargs(config: ServerArgs) -> dict[str, str]: + """Return uvicorn's HTTPS settings after the CLI has validated the pair.""" + if not config.ssl_certfile: + return {} + assert config.ssl_keyfile is not None + return { + "ssl_certfile": config.ssl_certfile, + "ssl_keyfile": config.ssl_keyfile, + } + + def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_shell: bool) -> None: """ Run the frontend API server (FastAPI + uvicorn) and wire it to the tokenizer process via ZMQ. @@ -1037,4 +1048,4 @@ def _on_meta(meta: dict) -> None: _serve_and_run_shell(host, port) 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, **_uvicorn_tls_kwargs(config)) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b68193..fbc68fa68 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -15,6 +15,8 @@ class ServerArgs(SchedulerConfig): server_host: str = "127.0.0.1" server_port: int = 1919 + ssl_certfile: str | None = None + ssl_keyfile: str | None = None num_tokenizer: int = 0 silent_output: bool = False # The terminal shell is attached to this server (ft shell --model / ft serve --shell-mode). @@ -304,6 +306,20 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The port number for the server to listen on.", ) + parser.add_argument( + "--ssl-certfile", + type=str, + default=ServerArgs.ssl_certfile, + help="PEM certificate chain for HTTPS. Requires --ssl-keyfile.", + ) + + parser.add_argument( + "--ssl-keyfile", + type=str, + default=ServerArgs.ssl_keyfile, + help="PEM private key for HTTPS. Requires --ssl-certfile.", + ) + parser.add_argument( "--cuda-graph-max-bs", "--graph", @@ -642,6 +658,10 @@ def _infer_reasoning_parser(model_path: str) -> str | None: # resolve some arguments run_shell |= kwargs.pop("shell_mode") kwargs["shell_mode"] = run_shell + if bool(kwargs["ssl_certfile"]) != bool(kwargs["ssl_keyfile"]): + parser.error("--ssl-certfile and --ssl-keyfile must be provided together") + if run_shell and kwargs["ssl_certfile"]: + parser.error("TLS is not supported with --shell-mode") if run_shell: kwargs["cuda_graph_max_bs"] = 1 kwargs["max_running_req"] = 1 @@ -649,6 +669,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if kwargs["model_path"].startswith("~"): kwargs["model_path"] = os.path.expanduser(kwargs["model_path"]) + for tls_path in ("ssl_certfile", "ssl_keyfile"): + if kwargs[tls_path] and kwargs[tls_path].startswith("~"): + kwargs[tls_path] = os.path.expanduser(kwargs[tls_path]) if kwargs["served_model_name"] is None: kwargs["served_model_name"] = ( diff --git a/tests/server/test_tls_args.py b/tests/server/test_tls_args.py new file mode 100644 index 000000000..44be28eba --- /dev/null +++ b/tests/server/test_tls_args.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from freetoken.server.args import parse_args + + +class _Config: + def to_dict(self) -> dict: + return {"architectures": ["DeepseekV4ForCausalLM"], "torch_dtype": "bfloat16"} + + +def _parse(extra: list[str]): + with patch("freetoken.utils.cached_load_hf_config", lambda _path: _Config()): + return parse_args(["--model", "/models/anon", *extra]) + + +def test_tls_certificate_and_key_are_parsed_as_a_pair(): + args, run_shell = _parse( + ["--ssl-certfile", "/certs/fullchain.pem", "--ssl-keyfile", "/certs/privkey.pem"] + ) + + assert run_shell is False + assert args.ssl_certfile == "/certs/fullchain.pem" + assert args.ssl_keyfile == "/certs/privkey.pem" + + +@pytest.mark.parametrize( + "single_flag", + [ + ["--ssl-certfile", "/certs/fullchain.pem"], + ["--ssl-keyfile", "/certs/privkey.pem"], + ], +) +def test_tls_rejects_an_incomplete_certificate_pair(single_flag): + with pytest.raises(SystemExit, match="2"): + _parse(single_flag) + + +def test_tls_rejects_shell_mode(): + with pytest.raises(SystemExit, match="2"): + _parse( + [ + "--shell-mode", + "--ssl-certfile", + "/certs/fullchain.pem", + "--ssl-keyfile", + "/certs/privkey.pem", + ] + ) + + +def test_tls_is_forwarded_to_uvicorn(): + from freetoken.server.api_server import _uvicorn_tls_kwargs + + config = SimpleNamespace( + ssl_certfile="/certs/fullchain.pem", + ssl_keyfile="/certs/privkey.pem", + ) + + assert _uvicorn_tls_kwargs(config) == { + "ssl_certfile": "/certs/fullchain.pem", + "ssl_keyfile": "/certs/privkey.pem", + } + + +def test_plain_http_keeps_uvicorn_tls_disabled(): + from freetoken.server.api_server import _uvicorn_tls_kwargs + + config = SimpleNamespace(ssl_certfile=None, ssl_keyfile=None) + + assert _uvicorn_tls_kwargs(config) == {}