Rust HTTP frontend behind --rust-frontend (RFC #130, Step 3) - #169
Rust HTTP frontend behind --rust-frontend (RFC #130, Step 3)#169npuichigo wants to merge 15 commits into
Conversation
37d8313 to
99490a6
Compare
99490a6 to
96d10db
Compare
22e2841 to
e9294b6
Compare
|
Rebased directly onto |
NSagan271
left a comment
NSagan271
left a comment
There was a problem hiding this comment.
First pass; will do a more detailed pass and test later today.
|
@npuichigo I just tried running our benchmark against The Rust Repro (against a # urlencoded — what `requests.post(url, data=...)` sends when there are no files
curl -s -o /dev/null -w "%{http_code}\n" -X POST http://0.0.0.0:8000/generate \
--data-urlencode "text=hi" --data-urlencode 'model_kwargs={"think_mode":true}'
# => 400
# multipart — works
curl -s -o /dev/null -w "%{http_code}\n" -X POST http://0.0.0.0:8000/generate \
-F "text=hi" -F 'model_kwargs={"think_mode":true}'
# => 200This breaks e.g. Root cause: async fn generate(State(st): State<AppState>, mut mp: Multipart) -> Response {axum's Suggested fix: take the raw Note the same |
|
Fixed in 8409ded.
The shared tail is factored into On your note about |
stephen-dwq
left a comment
stephen-dwq
left a comment
There was a problem hiding this comment.
Besides /v1/videos/generation for Cosmos, the frontend works. Ran ./test/qwen3-omni/launch_server.sh with
CUDA_VISIBLE_DEVICES=$DEVICES mstar serve qwen3_omni \
--config configs/qwen3omni_thinker_tp2.yaml \
--cache-dir $QWEN3OMNI_CACHE_DIR \
--socket-path-prefix /tmp/mstar_${WHO}/ \
--upload-dir /tmp/mstar_uploads_${WHO}/ \
--host ${HOST:-0.0.0.0} --port $PORT \
--tensor-comm-protocol $TENSOR_PROTOCOL \
--rust-frontend
and all requests in ./test/qwen3-omni succeed. I have done no concurrent testing of frontend routing under load.
…tep 3) rust/server/ is the axum HTTP surface — routes (/generate + the OpenAI endpoints), request translation, SSE/NDJSON streaming, uploads, media fetch, CORS — built on the vendored transport crate (the rlib target). It speaks a flattened msgpack protocol to a small Python bridge (mstar/api_server/rust_frontend.py) that drives the EXISTING data plane: APIServer.submit_request / iter_result_chunks / abort_request. The submit message is submit_request's signature verbatim; preprocessing, tokenization, and the conductor protocol are untouched. - Opt-in: mstar-serve --rust-frontend replaces only the uvicorn.run call; everything else in the process is identical. Binary resolution: --rust-frontend-bin / MSTAR_SERVER_BIN / $PATH / the in-repo build. - The bridge dogfoods Step 1 (RustZMQCommunicator with a msgpack Codec) in a private socket dir, so entity names never touch the worker mesh. - Frontend tokenization stays off with this backend (the preprocess worker owns it); the bridge rejects pre-tokenized ingest explicitly. - Wire-contract test: the real binary + the real bridge against a stubbed APIServer — chat roundtrip, ingest-failure -> 500 (and the server keeps serving), health gate. CI builds the server crate and now runs the whole test/rust/ suite (cpu torch), closing the gap where the arena tests were not exercised. - Env knobs documented in docs/environment_variables.rst; setup in docs/installation.rst. msgpack added to dependencies.
The one-command CLI builds an explicit argv for the low-level mstar-serve, so the flag added there was unreachable from 'mstar serve <model>'. Forward it (and --rust-frontend-bin), and show the primary form in the installation docs.
Production-envelope hardening, all failure-behavior: - /health now round-trips a ping through the bridge to the Python backend (2s budget). A dead bridge loop or conductor turns the LB check red instead of letting traffic queue into the 600s timeout. Mock mode (no bridge) stays shallow. - SIGTERM/SIGINT stop accepting and drain in-flight requests, with a hard 30s cap so a wedged stream cannot hold the process forever. - MSTAR_MAX_CONCURRENT_REQUESTS (default 256) bounds in-flight generation work; past the cap clients get an immediate 503 (overloaded_error) instead of unbounded queueing. /health and /v1/models bypass it. MSTAR_MAX_BODY_MB (default 128) replaces axum's 2MB default body cap, which multipart media uploads exceed. Tests: health-goes-red when the bridge dies, SIGTERM exits cleanly, saturation returns 503 while /health stays green.
The bridge's serve loop woke every 2 ms to check for frontend messages, so an idle request paid up to 2 ms (avg ~1 ms) just to be noticed — measured as 2.13 ms p50 end-to-end vs 0.63 ms for the uvicorn server on an identical stubbed data plane. Use the transport's blocking receive (get_all_new_messages(blocking=True, timeout_s=...)) on an executor thread instead: the wait happens inside the transport with the GIL released, an arriving submit ends it immediately, and an idle bridge costs zero wakeups (was ~500/s). The event loop stays free for the per-request relay tasks; sends from those tasks are safe concurrently (the transport serializes its push sockets internally). The timeout only bounds how fast stop() is noticed. Measured on the same A/B (stubbed data plane, /v1/chat/completions, localhost): sequential p50 2.13 -> 0.42 ms (now below uvicorn's 0.63); 100-chunk SSE total p50 2.67 -> 1.71 ms; c=32 throughput unchanged. Under an adversarial GIL load (a never-yielding spin thread in the backend process), throughput and streaming still improve (101 -> 108 rps, stream total 504 -> 398 ms) while single-request p50 degrades (15 -> 30 ms: the executor handoff is one more GIL acquisition against a thread that never yields); with a real serving loop, which releases the GIL on device waits, that handoff is cheap.
SSE/NDJSON streaming writes many small chunks; with Nagle enabled each flush can stall ~40 ms against the peer's delayed ACK. Measured on the frontend A/B (100-chunk SSE, stubbed data plane, localhost): total p50/p99 1.9/46 ms -> 0.68/0.74 ms, TTFB 0.36 -> 0.17 ms.
From NSagan271's first-pass review of the Rust frontend: Security (path traversal → arbitrary write, amplified to arbitrary delete by the component-based cleanup filter): - Multipart upload filenames and the client-controlled audio `format` extension are now sanitized to a single path component (Path::file_name / alphanumerics only) on both the Rust side (main.rs, media.rs) and the Python source of the same bug (entrypoint.py os.path.basename, media_io.py ext filter). Correctness: - Admission control now bounds STREAMING requests: the permit was dropped when the Response was built, freeing the slot while the stream was still being served. It's now moved into the response body so it drops at end-of-body, uniformly for streaming and non-streaming. - Multipart parse errors return 400 instead of being swallowed into a submitted request with missing files (both multipart loops). - A failed bridge send now pushes a terminal StreamItem::Error to the client instead of hanging it until request_timeout (600 s). - MSTAR_ALLOW_REMOTE defaults to off on the Python end too (was True), gated through resolve_media_ref; matches the Rust frontend's default. Ergonomics: - --host is forwarded to the Rust frontend (MSTAR_SERVER_HOST); it was hardcoded to 127.0.0.1, breaking multi-node/container binds. - MSTAR_TOKENIZER set with this backend fails fast at launch (the bridge can't ingest pre-tokenized input) instead of 500-ing every request. - The bridge entity is renamed "conductor" -> "bridge" on both sides, to stop shadowing M*'s real Conductor when reading the code. - The frontend's mkdtemp bridge_dir is cleaned up on shutdown. Deferred (noted on the PR): the batched chunk relay — it changes the bridge↔frontend wire on both sides for a concurrency-conditional win, so it belongs in its own change with a benchmark. New tests: tokenizer fail-fast, upload-filename traversal containment.
…ssion) The FastAPI /generate uses Form(...) fields, which Starlette parses from either multipart/form-data OR application/x-www-form-urlencoded — the latter is what `requests.post(url, data=...)` sends when there are no files. axum's Multipart extractor rejects any non-multipart body with a 400 before the handler runs, so file-less clients (e.g. test/qwen3-omni/text_request_single.py) regressed to 400 Bad Request. /generate now dispatches on Content-Type: the existing multipart loop for multipart/form-data, and a new branch that parses urlencoded bodies (form_urlencoded) into the same fields. The urlencoded path carries no file uploads, matching how those clients already behave (files force multipart). The shared tail is factored into generate_finish. An unrecognized Content-Type now gets an explicit 400 naming both accepted types. /v1/images/edits stays Multipart-only — it exists to take an image upload, so multipart is correct there. Test drives both content types through the real binary and asserts each lands as a submit with text + model_kwargs parsed.
Cosmos3 was added to the adapter registry after this PR opened, so the Rust frontend was missing the video surface it serves. Mirror the Python Cosmos3Adapter: a Cosmos3 adapter variant (cosmos3 / cosmos3_droid / cosmos3_super), supports_images + supports_videos, image_to_request, and video_to_request (text-to-video, plus image/video-to-video conditioning via resolve_media_ref; num_frames/fps first-class, other knobs through extra_body). Add the Surface::Videos arm, the VideoGenerationRequest protocol struct, and the videos_generations handler + route returning the mp4 as b64_json. The audio track mstar muxes into the mp4 is a follow-up (no mp4/PCM muxer on the Rust side yet); we return video-only, the same shape mstar degrades to when muxing fails. Covered by a Cosmos3 videos-generations round-trip test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
8409ded to
8b5bc4a
Compare
NSagan271
left a comment
NSagan271
left a comment
There was a problem hiding this comment.
@npuichigo Apologies for the delay in re-reviewing; I've accumulated a backlog of PRs.
Overall, I've mainly found parity issues between the Rust and Python versions, and a few error handling comments/issues.
I'm also running some benchmarking of qwen3-omni, orpheus, and bagel on Rust vs. fastapi. They're still running; I'll post the results when they complete.
cc. @merceod to take a look if you'd like.
Parity with the Python api_server:
- /generate errors now use FastAPI's {"detail": ...} shape (entrypoint.py
raises HTTPException), not the OpenAI {"error": {...}} envelope the /v1
endpoints use.
- streaming/tokenize form fields parse the pydantic bool set (true/t/yes/
y/on/1 + falsy), case-insensitive, 400 on anything else — `streaming=TRUE`
no longer silently downgrades to a non-streaming body.
- ImageGenerationRequest honors `n`: n engine submits with the seed+i
contract (image 0 bit-identical to n=1), images aggregated, matching
serving_images.create_images.
- base64 decode is lenient (Python's validate=False): strip interior
whitespace so line-wrapped data URIs / input_audio decode instead of 400.
- stream accepts an explicit null (Python `bool | None`): stream is now
Option<bool>.
- .avi recognized as a video extension (Python accepts it).
- /v1/audio/speech rejects unsupported response_format with 400 instead of
silently returning WAV bytes under the requested container's name (the
frontend has no compressed-audio encoder).
- Python router: chat/speech/images translate ValueError/TypeError to 400
(was 500), matching the videos handler and the Rust frontend.
Error handling / robustness:
- chat SSE no longer appends a clean finish_reason:"stop" + [DONE] after a
mid-stream backend error — a failed request must not look completed.
- multipart field read failures (body-limit overrun, mid-upload disconnect,
non-UTF-8) propagate as 400 instead of being flattened to empty values
(/generate + /v1/images/edits).
- Python bridge loop wraps per-message dispatch in try/except so a bad
abort/ping/decode can't propagate out of run() and take the process down
(restores FastAPI's per-request isolation).
- Cosmos3 video: dropped audio chunks now log a warning; documented the
compressed-audio and Cosmos3-sound limitations in the docs.
Tests: video image/video conditioning branches + both-set→400 + images on
Cosmos3; admission test rewritten to hold a streaming request's permit
through its body (cap=1) so it actually exercises permit lifetime.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
@npuichigo The benchmark finished, results below. All cells are
Text output:
¹ 1 ms absolute difference at the timer's resolution; the percentage is a quantization artifact, not a signal. The same caveat applies to any ITL cell in the single-digit-milliseconds range. Speech output:
Image output (bagel_cfg_parallel i2i bs=1)" req/s 0.140 → 0.140 (+0.0%), E2E p50 6.982 → 6.910 s (−1.0%). Verdict on the deferred qwen3_omni audio→speech was also benchmarked (bs 1/8/32 × 3 rounds, both before and after the rebase) and shows no regression, but I've left it out of the tables: nothing pins the generated audio length, so the two arms never synthesize equal amounts of audio and every figure needs a work-imbalance caveat to be read honestly. Overall, no performance regression detected (and good improvements on some cells). |
Rust frontend parity:
- Malformed / schema-invalid JSON bodies on the /v1 endpoints now return
FastAPI's 422 + {"detail": [{...}]} (a json_422 helper via
Result<Json<T>, JsonRejection>), not axum's default 400 + text/plain.
- ChatMessage.role is required (Python's pydantic requires it; a missing
role is now a 422, not silently defaulted).
- flatten_messages builds input_modalities in first-encounter order (a
mod_order vec) instead of the BTreeMap's sorted keys, so the list order
matches Python's insertion-ordered dict for downstream walk construction.
- --rust-frontend-bin now implies --rust-frontend (naming a binary but
staying on uvicorn was a footgun).
Robustness:
- bridge.rs logs an undecodable message instead of dropping it silently
(a systematic mismatch otherwise looks like a hang).
- serving.rs drops a malformed (non-4-byte) "token" chunk with a log
instead of forwarding it as an unknown modality.
Revert out-of-scope changes to the existing Python frontend (belong in
their own PR, per review):
- router.py: the chat/speech/images ValueError->400 mapping is reverted to
the original 500 (the nuanced status mapping — AttributeError, propagated
vs adapter-raised ValueError, NotImplementedError->404, non-500 codes vs
hardcoded "server_error" — is a separate PR).
- adapters.py flatten_messages: allow_remote default restored to True, so
the existing Python chat frontend keeps accepting http(s) media URLs
(the round-1 flip to False was a behavior regression for Qwen/Bagel). The
Rust frontend keeps its own off-by-default SSRF gate (MSTAR_ALLOW_REMOTE).
Test: malformed JSON + missing-role both 422 with a detail list.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Thanks for running this so thoroughly — 3 rounds each way, and I appreciate the honest caveats (the 1 ms quantization footnotes, and dropping qwen audio→speech since nothing pins the generated length so the arms never synthesize equal audio). Agreed on the verdict: no regression, and the wins are where they should be — orpheus bs=16 (highest chunk rate) is the best Rust cell (+7.8% audio sec/s, RTF mean −7.4%, TTFA p99 −71.9%), which is exactly the evidence that per-chunk IPC isn't the bottleneck. So I'm dropping the For context on where the branch stands after your and @stephen-dwq's latest passes (all pushed, Two open items still want a maintainer call, since you and stephen lean different ways:
And one deferred follow-up worth its own PR: carrying an error-kind on the bridge wire so a bad-input data-plane failure maps to 400 instead of collapsing to 500. |
merceod
left a comment
merceod
left a comment
There was a problem hiding this comment.
I looked at the PR and overall it looks good. There is just one thing I would recommend fixing before we merge (I left a comment about that).
@NSagan271 @stephen-dwq what do you think?
…200) merceod found that a request failing in the backend came back as HTTP 200 through the Rust frontend. The data plane delivers a failure in-band as a terminal chunk with modality "error" (message in data, HTTP status in metadata.status) that ends cleanly — it does not raise. The Python non-streaming handlers go through collect_results, which re-raises it as an HTTPException; the Rust side never special-cased the "error" modality, so: - /generate non-streaming returned 200 with the error base64'd in outputs["error"] - chat/speech/images/videos non-streaming returned 200 with the error chunk silently filtered (empty content/pcm/data) Fix (merceod's suggestion): the bridge maps a modality=="error" chunk to a terminal StreamItem::Error carrying status (from metadata.status, default 500) + message. Out::Error and collect() thread the status through, so the handlers return the real code — error(status, ..) on the /v1 surfaces, detail_error(status, ..) on /generate. Because the status rides in-band, this also delivers the 400-vs-500 distinction the deferred "error-kind" follow-up was for, with no new wire messages. Chat SSE already ends on an error event via the `errored` flag, so a failed stream no longer looks cleanly completed there either. Test: a non-streaming request whose backend emits an error chunk returns 500 (chat envelope + /generate detail), and a 400-status chunk returns 400. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
@npuichigo all of your changes look good to me. I did some final stress testing before merging; the benchmark only ever exercises the happy path, so I stood both frontends up against the same stubbed data plane and fired ~120 malformed / oddly-shaped requests at each, plus disconnect / timeout / concurrency-cap / large-upload cases, then repeated the interesting ones against a live qwen3_omni server under each frontend in turn. Pretty much everything is at parity, including SSE framing, WAV containers, NDJSON framing.
Confirmed against the live server: Anything that round-trips parameters through a form, an env var, a YAML config or a loosely-typed SDK sends strings, and Everything else is non-blocking — could you open an issue to track it? Nothing below needs to hold up this PR, and a couple of them are Python-side bugs that want their own change:
Happy to hand over the test harness if it's useful to fold into the repo. Also apologies for the extended back-and-forth; just trying to be thorough with these larger-scope Rust PRs. This should be the last thing from my end. |
NSagan's pre-merge stress test found serde is strict where pydantic coerces: a client sending a number or bool as a string (what forms, env vars, YAML configs, and loosely-typed SDKs produce) got a hard 422 under the Rust frontend where FastAPI accepts it — `"stream":"true"`, `"temperature":"0.5"`, `"max_tokens":"16"`/`16.0`, `"n":"2"`, `"fps":"12"`. Since these are declared fields (not passthrough), it's a real narrowing of accepted input. Add a `flex` module of `deserialize_with` helpers (opt_i64 / opt_f64 / opt_bool) that accept the string and integral-float forms alongside the native types, and annotate every declared numeric/bool field across the four request models. Absent stays None (via serde default), explicit null maps to None, and genuinely invalid values still error. The passthrough `extra` map is untouched. Unit tests: string/integral-float coercion, native types + null + absent, and rejection of non-numeric / non-integral / non-bool values. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Thank you — and no need to apologize, the thoroughness caught real things. Blocking item fixed in c06fead. Added a Non-blocking items → tracked in #211, grouped as you framed them: the three Python-side bugs ( On CORS credentials specifically — good catch that And yes please — folding your stress-test harness into the repo would be genuinely useful; happy to take it however is easiest for you. |
NSagan271
left a comment
NSagan271
left a comment
There was a problem hiding this comment.
Looks good from my end! I will make a separate PR to this repo with the testing harness, and also the harness from testing the SHM PR.
4 participants
Note
Decoupled from #168 (rebased directly onto
main): this step needs only the merged #164 transport — the frontend/bridge never touches the SHM arena — so it can review and merge independently of #168's stress-testing timeline.Rust HTTP frontend behind
--rust-frontend(RFC #130, Step 3)rust/server/is an axum server owning the full HTTP surface —/generate, the OpenAI-compatible endpoints (chat/audio/images incl. multipart uploads), SSE/NDJSON streaming, media fetch, CORS — built on the vendored transport crate's library target. It speaks a flattened msgpack protocol to a small Python bridge (mstar/api_server/rust_frontend.py) that drives the existing data plane:APIServer.submit_request/iter_result_chunks/abort_request. The submit message issubmit_request's signature verbatim — preprocessing, tokenization, and the conductor protocol are untouched.Scope
mstar-serve --rust-frontendreplaces exactly theuvicorn.runcall; everything else in the api_server process is identical. Without the flag, nothing changes.RustZMQCommunicatorwith a msgpackCodec, in a private socket dir — the bridge mesh never touches the conductor/worker mesh's entity names.MSTAR_TOKENIZER) stays unset with this backend, and the bridge rejects pre-tokenized ingest with a clear error. The tokenize-override hook from the plan discussion remains available for a later step.What replaces what
--rust-frontendmstar/api_server/openai/*request/response translationrust/server/src/adapters.rs(same three model families)Performance (frontend layer A/B)
Setup: both frontends drive an identical stubbed data plane (the same
APIServerobject),/v1/chat/completions, localhost. The load client is Rust (one keep-alive connection per thread; its own ceiling, measured against/health, is ~89k rps — far above anything below). Each backend stack is pinned withtasksetto a fixed CPU budget — for the Rust side that budget covers both its processes (axum binary + Python bridge) — and the client runs on disjoint cores. uvicorn is the single-process configurationmstar-serveactually runs (multi-worker uvicorn is not deployable for this server: request state lives in-process).Equal budget: 1 CPU each (tokio sizes its pool from the affinity mask → 1 worker thread; the purest runtime-vs-runtime comparison)
Equal budget: 4 CPUs each
GIL-loaded backend, 4 CPUs (a never-yielding pure-Python spin thread in the backend process — a worst-case stand-in for the co-resident serving loop; a real loop releases the GIL on device waits, so production sits between this and the tables above)
Two commits fell out of running this A/B:
get_all_new_messages(blocking=True, timeout_s=…), the blocking pack from the Step-1 review round — its first in-tree consumer) on an executor thread: an arriving submit is picked up immediately, and an idle bridge costs zero wakeups (was ~500/s).TCP_NODELAYon accepted connections — SSE writes many small chunks, and Nagle + delayed ACK turned flushes into ~40 ms stalls (100-chunk stream total was 1.9/46 ms p50/p99 before; 0.72/0.80 ms after, on an unrestricted-CPU run).Tests
test/rust/test_rust_frontend.pyruns the real binary + the real bridge against a stubbedAPIServer: chat roundtrip through the whole HTTP→msgpack→ingest→chunks→SSE path, ingest-failure → clean 500 with the server still serving afterwards, health gate. CI builds the server crate and now runs the entiretest/rust/suite (with CPU torch — this also closes a gap where #168's arena tests were not exercised by CI).Env knobs are documented in
docs/environment_variables.rst; build/opt-in steps indocs/installation.rst.🤖 Generated with Claude Code
Behavior differences from the Python (uvicorn) frontend
Follow-ups and the full list are tracked in #211. Highlights:
image_urlfetches are gated behindMSTAR_ALLOW_REMOTE(off by default — SSRF surface). The Python chat path fetches unconditionally; that's the bug, and Rust frontend (#169) follow-ups: Python-side parity bugs + documented behavior differences #211 tracks closing the Python hole so the two agree./v1/audio/speechonly produceswav/pcm(no compressed encoder); otherresponse_formatvalues are a 400 with the workaround named. Note OpenAI's default ismp3./generaterequires a formContent-Type(FastAPI parsed a JSON body as an empty form)./healthis a deep check (pings the backend bridge), so it can go red where Python's never did.MSTAR_MAX_CONCURRENT_REQUESTS(default 256) returns 503 past the cap; the Python frontend has no counterpart and queues.422+{"detail": [...]}; only the message text differs.