2727
2828import importlib
2929import importlib .util
30+ import json
3031import logging
3132import os
3233import secrets
@@ -262,7 +263,6 @@ def create_app(
262263
263264 # Resolve the signing key once, here, so the worker (which seals catalog
264265 # opaque-data envelopes) and the HTTP state-token machinery share the same
265- # key. When the operator did not configure VGI_SIGNING_KEY, generate an
266266 # key. See resolve_shared_signing_key for why every process serving this
267267 # deployment has to agree on it.
268268 if signing_key is None :
@@ -337,6 +337,32 @@ def serve(
337337 worker_ref : str = typer .Argument (help = "Worker reference: module:Class, module, or ./file.py" ),
338338 # Transport
339339 http : bool = typer .Option (False , "--http" , help = "Serve over HTTP instead of stdin/stdout" ),
340+ server : str = typer .Option (
341+ "waitress" ,
342+ "--server" ,
343+ help = (
344+ "HTTP server: 'waitress' (default, pure Python) or 'granian' "
345+ "(Rust I/O off the GIL; needs the [granian] extra)."
346+ ),
347+ ),
348+ http_threads : int | None = typer .Option (
349+ None ,
350+ "--http-threads" ,
351+ help = (
352+ "Waitress worker threads. The request path is CPU-bound Python, so more "
353+ "threads contend for the GIL rather than adding throughput; raise this only "
354+ "for functions that block on external I/O."
355+ ),
356+ ),
357+ http_workers : int | None = typer .Option (
358+ None ,
359+ "--http-workers" ,
360+ help = (
361+ "Worker processes (granian only). Each is a separate interpreter, so "
362+ "memory scales with it. Requires VGI_SIGNING_KEY when >1 unless vgi-serve "
363+ "mints one for them."
364+ ),
365+ ),
340366 quiet : bool = typer .Option (False , "--quiet" , "-q" , help = "Suppress startup banner (stdio mode)" ),
341367 # Logging
342368 debug : bool = typer .Option (False , "--debug" , help = "Enable DEBUG on all vgi + vgi_rpc loggers" ),
@@ -393,6 +419,10 @@ def serve(
393419 authenticate = _resolve_authenticate ()
394420 oauth_metadata = _resolve_oauth_resource_metadata ()
395421 otel_config = _resolve_otel_config ()
422+ if server not in ("waitress" , "granian" ):
423+ sys .exit (f"--server must be 'waitress' or 'granian', got { server !r} " )
424+ if http_workers is not None and server != "granian" :
425+ sys .exit ("--http-workers only applies to --server granian" )
396426 _serve_http (
397427 worker_cls ,
398428 effective_level = effective_level ,
@@ -406,6 +436,10 @@ def serve(
406436 oauth_resource_metadata = oauth_metadata ,
407437 otel_config = otel_config ,
408438 max_stream_response_bytes = max_stream_response_bytes ,
439+ server = server ,
440+ worker_ref = worker_ref ,
441+ http_workers = http_workers ,
442+ http_threads = http_threads ,
409443 )
410444 else :
411445 otel_config = _resolve_otel_config ()
@@ -879,6 +913,223 @@ def _resolve_otel_config() -> Any:
879913 )
880914
881915
916+ # ---------------------------------------------------------------------------
917+ # Pre-fork support: rebuilding the app in a worker process
918+ #
919+ # waitress takes the app *object* we hand it. A pre-fork server (granian) takes
920+ # an import path and builds the app inside each forked worker, which means the
921+ # configuration cannot travel as Python objects -- only whatever the child
922+ # inherits. The environment is that channel.
923+ #
924+ # This is a private contract between ``vgi-serve`` and its own workers, not a
925+ # user-facing knob, so it is one opaque blob rather than a family of documented
926+ # env vars. Secrets stay out of it: the signing key travels in VGI_SIGNING_KEY
927+ # (see resolve_shared_signing_key), and auth/OAuth/OTel settings are already
928+ # resolved from the environment by their own helpers, so a child reconstructs
929+ # them exactly as the parent did.
930+ # ---------------------------------------------------------------------------
931+
932+ SERVE_CONFIG_ENV = "VGI_SERVE_CONFIG"
933+
934+
935+ def export_serve_config (
936+ * ,
937+ worker_ref : str ,
938+ prefix : str ,
939+ cors_origins : str ,
940+ describe : bool ,
941+ log_level : int ,
942+ max_stream_response_bytes : int | None ,
943+ ) -> None :
944+ """Publish the parent's serve configuration for worker processes to read.
945+
946+ Args:
947+ worker_ref: The worker reference string the parent was given
948+ (``module:Class``, ``module``, or ``./file.py``). Passed rather
949+ than the resolved class because a child has to import it itself.
950+ prefix: URL prefix for RPC endpoints.
951+ cors_origins: Allowed CORS origins.
952+ describe: Whether to enable the worker + API description pages.
953+ log_level: Logging level for the worker instance.
954+ max_stream_response_bytes: Producer-stream response budget, or None.
955+
956+ """
957+ os .environ [SERVE_CONFIG_ENV ] = json .dumps (
958+ {
959+ "worker_ref" : worker_ref ,
960+ "prefix" : prefix ,
961+ "cors_origins" : cors_origins ,
962+ "describe" : describe ,
963+ "log_level" : log_level ,
964+ "max_stream_response_bytes" : max_stream_response_bytes ,
965+ }
966+ )
967+
968+
969+ def wsgi_app_factory () -> Any :
970+ """Build the WSGI app from the environment — the pre-fork worker entry point.
971+
972+ Granian is pointed at ``vgi.serve:wsgi_app_factory`` and calls this once
973+ per worker process. Everything it needs was published by
974+ :func:`export_serve_config` in the parent, plus ``VGI_SIGNING_KEY``, which
975+ the parent minted and exported so every worker seals state tokens with the
976+ same key.
977+
978+ Returns:
979+ The Falcon WSGI application.
980+
981+ Raises:
982+ RuntimeError: Called without a parent having exported the config,
983+ which means this was invoked directly rather than by ``vgi-serve``.
984+
985+ """
986+ raw = os .environ .get (SERVE_CONFIG_ENV )
987+ if not raw :
988+ msg = (
989+ f"{ SERVE_CONFIG_ENV } is not set. vgi.serve:wsgi_app_factory is the worker entry point "
990+ f"for `vgi-serve --http --server granian`; it cannot be started on its own."
991+ )
992+ raise RuntimeError (msg )
993+ config = json .loads (raw )
994+
995+ # Resolved here, not inherited as objects: each worker re-reads the same
996+ # environment the parent did and arrives at the same answer.
997+ return create_app (
998+ load_worker_class (config ["worker_ref" ]),
999+ prefix = config ["prefix" ],
1000+ cors_origins = config ["cors_origins" ],
1001+ describe = config ["describe" ],
1002+ signing_key = None , # picked up from VGI_SIGNING_KEY
1003+ log_level = config ["log_level" ],
1004+ authenticate = _resolve_authenticate (),
1005+ oauth_resource_metadata = _resolve_oauth_resource_metadata (),
1006+ otel_config = _resolve_otel_config (),
1007+ max_stream_response_bytes = config ["max_stream_response_bytes" ],
1008+ )
1009+
1010+
1011+ def _resolve_http_port (host : str , port : int | None ) -> int :
1012+ """Resolve the listen port: explicit ``--port`` > ``$PORT`` > 8080.
1013+
1014+ ``--port 0`` binds a throwaway socket to let the OS pick, then reports the
1015+ concrete number — process managers and test harnesses read it off stdout.
1016+ """
1017+ import socket
1018+
1019+ if port is None :
1020+ env_port = os .environ .get ("PORT" )
1021+ port = int (env_port ) if env_port else 8080
1022+ if port == 0 :
1023+ with socket .socket (socket .AF_INET , socket .SOCK_STREAM ) as sock :
1024+ sock .bind ((host , 0 ))
1025+ port = int (sock .getsockname ()[1 ])
1026+ return port
1027+
1028+
1029+ def _serve_http_granian (
1030+ worker_cls : type [Worker ],
1031+ * ,
1032+ host : str ,
1033+ port : int ,
1034+ prefix : str ,
1035+ cors_origins : str ,
1036+ describe : bool ,
1037+ effective_level : int ,
1038+ max_stream_response_bytes : int | None ,
1039+ worker_ref : str | None ,
1040+ http_workers : int | None ,
1041+ http_threads : int | None ,
1042+ ) -> None :
1043+ """Serve via granian, which forks workers that import the app themselves.
1044+
1045+ Granian moves HTTP parsing and socket I/O into Rust, off the GIL. On this
1046+ workload that measured 1.45x over waitress at a single worker -- the win
1047+ comes from the I/O leaving the interpreter, not from process fan-out, so
1048+ it does not require paying for extra interpreters.
1049+
1050+ Two things follow from granian building the app per worker rather than
1051+ accepting ours:
1052+
1053+ - The configuration has to travel through the environment. See
1054+ :func:`export_serve_config`.
1055+ - Every worker must seal state tokens with the *same* key, or a client
1056+ that reconnects mid-stream gets an intermittent 400. The parent mints
1057+ and exports one here; :func:`resolve_shared_signing_key` explains why.
1058+
1059+ Args:
1060+ worker_cls: The resolved worker class -- used only for the startup
1061+ banner; each granian worker imports its own from ``worker_ref``.
1062+ host: Bind address.
1063+ port: Bind port, already resolved.
1064+ prefix: URL prefix for RPC endpoints.
1065+ cors_origins: Allowed CORS origins.
1066+ describe: Whether to enable the description pages.
1067+ effective_level: Logging level for the worker instance.
1068+ max_stream_response_bytes: Producer-stream response budget, or None.
1069+ worker_ref: The worker reference string, required here because the
1070+ children import it rather than inheriting a class object.
1071+ http_workers: Number of worker processes; ``None`` means 1. These are
1072+ separate interpreters, so memory scales with the count.
1073+ http_threads: Python threads per worker; ``None`` means 1.
1074+
1075+ """
1076+ try :
1077+ from granian import Granian
1078+ from granian .constants import Interfaces
1079+ except ImportError :
1080+ sys .stderr .write (
1081+ "Error: granian not installed.\n Install with: pip install 'vgi[granian]' (or: uv sync --extra granian)\n "
1082+ )
1083+ sys .exit (1 )
1084+
1085+ if not worker_ref :
1086+ msg = "granian requires the worker reference string so each worker process can import it"
1087+ raise RuntimeError (msg )
1088+
1089+ workers = http_workers or 1
1090+ # Mint before starting children so they inherit one key rather than each
1091+ # generating its own. Only meaningful when we actually fork.
1092+ _key , is_ephemeral = resolve_shared_signing_key (propagate_to_children = workers > 1 )
1093+ _warn_if_ephemeral_signing_key (is_ephemeral = is_ephemeral , multiprocess = workers > 1 )
1094+
1095+ export_serve_config (
1096+ worker_ref = worker_ref ,
1097+ prefix = prefix ,
1098+ cors_origins = cors_origins ,
1099+ describe = describe ,
1100+ log_level = effective_level ,
1101+ max_stream_response_bytes = max_stream_response_bytes ,
1102+ )
1103+
1104+ print (f"PORT:{ port } " , flush = True )
1105+ _logger .info (
1106+ "http_server_starting server=granian host=%s port=%d prefix=%s workers=%d" ,
1107+ host ,
1108+ port ,
1109+ prefix ,
1110+ workers ,
1111+ )
1112+ sys .stderr .write (f"Serving { worker_cls .__name__ } on http://{ host } :{ port } { prefix } (granian, { workers } worker(s))\n " )
1113+ sys .stderr .flush ()
1114+
1115+ Granian (
1116+ "vgi.serve:wsgi_app_factory" ,
1117+ address = host ,
1118+ port = port ,
1119+ interface = Interfaces .WSGI ,
1120+ workers = workers ,
1121+ # Python threads inside each worker, sharing that worker's GIL --
1122+ # distinct from `workers`, which are separate interpreters. Defaults
1123+ # to 1 because granian already runs socket I/O in Rust, so the only
1124+ # thing these threads overlap is app code, and for a CPU-bound
1125+ # request path that is pure contention: measured 1387 turns/s at one
1126+ # thread against 851 at two and 362 at eight. Functions that block on
1127+ # external I/O want more; --http-threads raises it.
1128+ blocking_threads = http_threads or 1 ,
1129+ factory = True ,
1130+ ).serve ()
1131+
1132+
8821133def _serve_http (
8831134 worker_cls : type [Worker ],
8841135 * ,
@@ -893,9 +1144,29 @@ def _serve_http(
8931144 oauth_resource_metadata : Any = None ,
8941145 otel_config : Any = None ,
8951146 max_stream_response_bytes : int | None = None ,
1147+ server : str = "waitress" ,
1148+ worker_ref : str | None = None ,
1149+ http_workers : int | None = None ,
1150+ http_threads : int | None = None ,
8961151) -> None :
8971152 """Start the worker as an HTTP server."""
898- import socket
1153+ port = _resolve_http_port (host , port )
1154+
1155+ if server == "granian" :
1156+ _serve_http_granian (
1157+ worker_cls ,
1158+ host = host ,
1159+ port = port ,
1160+ prefix = prefix ,
1161+ cors_origins = cors_origins ,
1162+ describe = describe ,
1163+ effective_level = effective_level ,
1164+ max_stream_response_bytes = max_stream_response_bytes ,
1165+ worker_ref = worker_ref ,
1166+ http_workers = http_workers ,
1167+ http_threads = http_threads ,
1168+ )
1169+ return
8991170
9001171 try :
9011172 import waitress # type: ignore[import-untyped]
@@ -905,16 +1176,6 @@ def _serve_http(
9051176 )
9061177 sys .exit (1 )
9071178
908- # Port resolution: explicit --port > $PORT env var > 8080
909- if port is None :
910- env_port = os .environ .get ("PORT" )
911- port = int (env_port ) if env_port else 8080
912-
913- if port == 0 :
914- with socket .socket (socket .AF_INET , socket .SOCK_STREAM ) as s :
915- s .bind ((host , 0 ))
916- port = int (s .getsockname ()[1 ])
917-
9181179 wsgi_app = create_app (
9191180 worker_cls ,
9201181 prefix = prefix ,
@@ -934,7 +1195,17 @@ def _serve_http(
9341195 sys .stderr .write (f"Serving { worker_cls .__name__ } on http://{ host } :{ port } { prefix } \n " )
9351196 sys .stderr .flush ()
9361197
937- waitress .serve (wsgi_app , host = host , port = port , _quiet = True )
1198+ # Same Arrow-sized buffers the other two serve paths use. Without them
1199+ # waitress spools multi-MiB bodies through temp files and reads sockets in
1200+ # 8 KiB chunks; measured on this path, tuning is worth ~2x.
1201+ from vgi_rpc .http .server ._serve import _tame_queue_depth_logger , waitress_arrow_tuning
1202+
1203+ _tame_queue_depth_logger ()
1204+ serve_kwargs : dict [str , Any ] = {"host" : host , "port" : port , "_quiet" : True }
1205+ serve_kwargs .update (waitress_arrow_tuning ())
1206+ if http_threads is not None :
1207+ serve_kwargs ["threads" ] = http_threads
1208+ waitress .serve (wsgi_app , ** serve_kwargs )
9381209
9391210
9401211if __name__ == "__main__" :
0 commit comments