From acbe20cdfd94ff1db668108bbbe7ff82c18cc03a Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Wed, 29 Jul 2026 20:16:19 -0700 Subject: [PATCH] add the sglang-server rust implementation (api server, runtime, tokenizer manager, detokenizer) --- rust/sglang-server/Cargo.lock | 2997 +++++++++++++++++ rust/sglang-server/src/api_server.rs | 97 + rust/sglang-server/src/api_server/common.rs | 226 ++ rust/sglang-server/src/api_server/frame.rs | 865 +++++ rust/sglang-server/src/api_server/guard.rs | 130 + rust/sglang-server/src/api_server/log.rs | 41 + .../src/api_server/native_api.rs | 602 ++++ rust/sglang-server/src/api_server/openai.rs | 55 + rust/sglang-server/src/api_server/submit.rs | 126 + rust/sglang-server/src/detokenizer.rs | 597 ++++ rust/sglang-server/src/environ.rs | 1 - rust/sglang-server/src/error.rs | 1 - rust/sglang-server/src/lib.rs | 225 +- rust/sglang-server/src/message.rs | 10 +- rust/sglang-server/src/runtime.rs | 492 +++ rust/sglang-server/src/runtime/threads.rs | 142 + rust/sglang-server/src/tokenizer.rs | 162 + rust/sglang-server/src/tokenizer_manager.rs | 6 + .../src/tokenizer_manager/egress.rs | 217 ++ .../src/tokenizer_manager/ingress.rs | 1210 +++++++ 20 files changed, 8196 insertions(+), 6 deletions(-) create mode 100644 rust/sglang-server/Cargo.lock create mode 100644 rust/sglang-server/src/api_server.rs create mode 100644 rust/sglang-server/src/api_server/common.rs create mode 100644 rust/sglang-server/src/api_server/frame.rs create mode 100644 rust/sglang-server/src/api_server/guard.rs create mode 100644 rust/sglang-server/src/api_server/log.rs create mode 100644 rust/sglang-server/src/api_server/native_api.rs create mode 100644 rust/sglang-server/src/api_server/openai.rs create mode 100644 rust/sglang-server/src/api_server/submit.rs create mode 100644 rust/sglang-server/src/detokenizer.rs create mode 100644 rust/sglang-server/src/runtime/threads.rs create mode 100644 rust/sglang-server/src/tokenizer.rs create mode 100644 rust/sglang-server/src/tokenizer_manager/egress.rs create mode 100644 rust/sglang-server/src/tokenizer_manager/ingress.rs diff --git a/rust/sglang-server/Cargo.lock b/rust/sglang-server/Cargo.lock new file mode 100644 index 000000000000..3c924db68087 --- /dev/null +++ b/rust/sglang-server/Cargo.lock @@ -0,0 +1,2997 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "form_urlencoded", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec 0.6.3", +] + +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec 0.8.0", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures", +] + +[[package]] +name = "bstr" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cee35f73844aa3014bb606320a6c1f010249dbdf43342fe54b5a4f6a8ed4b79" +dependencies = [ + "memchr", + "regex-automata", + "serde_core", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.2.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core 0.10.1", +] + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "console" +version = "0.15.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width", + "windows-sys 0.59.0", +] + +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "daachorse" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f55d7153ba3b507595872a3874803f07a8a81d1e888abed8e5db7da0597d6e2" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn", +] + +[[package]] +name = "dirs" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" +dependencies = [ + "libc", + "option-ext", + "redox_users", + "windows-sys 0.61.2", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dynamo-tokenizers" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8618ea35c99831eb8f41a1d239b6f5c009e5f647f118f9c20267c51b7415dab8" +dependencies = [ + "aho-corasick", + "anyhow", + "base64 0.22.1", + "blake3", + "fastokens", + "moka", + "rayon", + "rustc-hash 1.1.0", + "serde", + "serde_json", + "strum", + "tiktoken-rs", + "tokenizers", + "tracing", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + +[[package]] +name = "fancy-regex" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" +dependencies = [ + "bit-set 0.5.3", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "fancy-regex" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" +dependencies = [ + "bit-set 0.8.0", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "fastokens" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8728655e193e0d08d7a95d63cf1fdb9b768d282cab0a112ecb006615bae9f067" +dependencies = [ + "daachorse", + "fancy-regex 0.17.0", + "hf-hub 0.4.3", + "icu_normalizer", + "memchr", + "pcre2", + "rayon", + "serde", + "serde_json", + "strum", + "thiserror", +] + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + +[[package]] +name = "flume" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" +dependencies = [ + "fastrand", + "futures-core", + "futures-sink", + "spin", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 5.3.0", + "wasip2", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 6.0.0", + "rand_core 0.10.1", + "wasm-bindgen", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + +[[package]] +name = "hf-hub" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" +dependencies = [ + "dirs", + "http", + "indicatif", + "libc", + "log", + "rand 0.9.5", + "reqwest", + "serde", + "serde_json", + "thiserror", + "ureq", + "windows-sys 0.60.2", +] + +[[package]] +name = "hf-hub" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aef3982638978efa195ff11b305f51f1f22f4f0a6cabee7af79b383ebee6a213" +dependencies = [ + "dirs", + "log", +] + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "tokio", + "tokio-rustls", + "tower-service", + "webpki-roots 1.0.8", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "utf16_iter", + "utf8_iter", + "write16", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console", + "number_prefix", + "portable-atomic", + "unicode-width", + "web-time", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libredox" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" +dependencies = [ + "libc", +] + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "moka" +version = "0.12.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "957228ad12042ee839f93c8f257b62b4c0ab5eaae1d4fa60de53b27c9d7c5046" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "option-ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pcre2" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e970b0fcce0c7ee6ef662744ff711f21ccd6f11b7cf03cd187a80e89797fc67" +dependencies = [ + "libc", + "log", + "pcre2-sys", +] + +[[package]] +name = "pcre2-sys" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18b9073c1a2549bd409bf4a32c94d903bb1a09bf845bc306ae148897fa0760a4" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash 2.1.3", + "rustls", + "socket2", + "thiserror", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" +dependencies = [ + "bytes", + "getrandom 0.4.3", + "lru-slab", + "rand 0.10.2", + "rand_pcg", + "ring", + "rustc-hash 2.1.3", + "rustls", + "rustls-pki-types", + "slab", + "thiserror", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.61.2", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core 0.9.5", +] + +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "redox_users" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" +dependencies = [ + "getrandom 0.2.17", + "libredox", + "thiserror", +] + +[[package]] +name = "regex" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a0e75113e14dc5acb068cd0786884f214f1312650a3d36d269f5c4f3cdee8a2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f388202e4b80542a0921078cc23b6333bcf1409c1e3f86404cae4766a6131db" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", + "webpki-roots 1.0.8", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rmp" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ba8be72d372b2c9b35542551678538b562e7cf86c3315773cae48dfbfe7790c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "rmp-serde" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f81bee8c8ef9b577d1681a70ebbc962c232461e397b22c208c43c04b67a155" +dependencies = [ + "rmp", + "serde", +] + +[[package]] +name = "rmpv" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a4e1d4b9b938a26d2996af33229f0ca0956c652c1375067f0b45291c1df8417" +dependencies = [ + "rmp", + "serde", + "serde_bytes", +] + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustls" +version = "0.23.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" +dependencies = [ + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_bytes" +version = "0.11.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" +dependencies = [ + "serde", + "serde_core", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "sglang-server" +version = "0.1.0" +dependencies = [ + "async-stream", + "axum", + "bytes", + "core_affinity", + "dynamo-tokenizers", + "flume", + "futures", + "hf-hub 0.5.0", + "pyo3", + "regex-syntax", + "rmp-serde", + "rmpv", + "serde", + "serde_json", + "thiserror", + "tokio", + "tracing", + "tracing-appender", + "tracing-subscriber", + "uuid", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", +] + +[[package]] +name = "spin" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +dependencies = [ + "lock_api", +] + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "strum" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "symlink" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7973cce6668464ea31f176d85b13c7ab3bba2cb3b77a2ed26abd7801688010a" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tiktoken-rs" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a19830747d9034cd9da43a60eaa8e552dfda7712424aebf187b7a60126bae0d" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bstr", + "fancy-regex 0.13.0", + "lazy_static", + "regex", + "rustc-hash 1.1.0", +] + +[[package]] +name = "time" +version = "0.3.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18dfaaeddcb932337b5e7866ee7d0ce9b76d2fd092997146f187ec09b4558a50" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" + +[[package]] +name = "time-macros" +version = "0.2.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c431b87111666e491a90baa837f914fb45cd5dc3c268591b0220ff5057f2085f" +dependencies = [ + "num-conv", + "time-core", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokenizers" +version = "0.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a620b996116a59e184c2fa2dfd8251ea34a36d0a514758c6f966386bd2e03476" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "hf-hub 0.4.3", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.9.5", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-appender" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "050686193eb999b4bb3bc2acfa891a13da00f79734704c4b8b4ef1a10b368a3c" +dependencies = [ + "crossbeam-channel", + "symlink", + "thiserror", + "time", + "tracing-subscriber", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "ureq" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "once_cell", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "socks", + "url", + "webpki-roots 0.26.11", +] + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf16_iter" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "uuid" +version = "1.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea5fab0d6c3c01ae70085a09cb03d4c7a1d6314e2b3e075392783396d724ca0a" +dependencies = [ + "getrandom 0.4.3", + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.8", +] + +[[package]] +name = "webpki-roots" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "write16" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/sglang-server/src/api_server.rs b/rust/sglang-server/src/api_server.rs new file mode 100644 index 000000000000..911989b297db --- /dev/null +++ b/rust/sglang-server/src/api_server.rs @@ -0,0 +1,97 @@ +//! API server (axum / tokio). I/O-bound; own pinned multi-thread runtime. Only +//! this module knows HTTP, so other protocols can mount the same `AppState`. +//! `/generate` submits a `Request` then awaits one `Done` (unary) or relays SSE +//! frames (`data: {json}` … `[DONE]`), byte-compatible with Python +//! `http_server.generate_request`; `/server_info` reuses it for one control result. +mod common; +mod frame; +mod guard; +mod log; +mod native_api; +mod openai; +mod submit; + +use std::sync::Arc; + +use axum::Router; + +use crate::runtime::ServerArgs; +use crate::tokenizer_manager::ActivityCounter; +use crate::tokenizer_manager::Senders; + +/// Shared handler state: the submit machinery (`senders`, `egress_buf`) +/// + shared tokenizer. +#[derive(Clone)] +struct AppState { + senders: Senders, + egress_buf: usize, + server_args: Arc, + /// Egress heartbeat (bumped per drained ring frame). + egress_activity: ActivityCounter, +} + +pub async fn serve( + listener: std::net::TcpListener, + senders: Senders, + egress_buf: usize, + server_args: Arc, + egress_activity: ActivityCounter, + // The SAME set ingress releases from — see `Ingress::on_abort`. Constructing a + // local one here would leave the api server admitting rids that nothing ever + // releases. + shutdown: flume::Receiver<()>, +) { + let state = AppState { + senders, + egress_buf, + server_args: server_args.clone(), + egress_activity, + }; + // Each endpoint module registers its own routes and merges here. + let app = Router::new() + .merge(common::routes()) + .merge(native_api::routes()) + .merge(openai::routes()) + // TODO(auth): no API-key boundary yet. Python gates every route (except + // /health*, /metrics*, OPTIONS) via `add_api_key_middleware`; until ported, + // a configured `api_key` does NOT protect these routes. + // + // No body limit, matching the Python server. + .layer(axum::extract::DefaultBodyLimit::disable()) + .with_state(state); + let app = log::apply(app, &server_args); + + // The listener was already bound synchronously in `runtime::start` (so a port + // conflict fails startup); adopt it into the tokio reactor here. + let listener = match tokio::net::TcpListener::from_std(listener) { + Ok(l) => l, + Err(e) => { + tracing::error!(error = %e, "failed to adopt pre-bound listener"); + return; + } + }; + if let Ok(addr) = listener.local_addr() { + tracing::info!(%addr, "sglang-server api listening"); + } + // Non-graceful shutdown: on the signal, stop accepting and RETURN without + // waiting for in-flight handlers (a `/generate` blocked on egress would wedge + // the join). Returning unwinds `block_on` in `runtime::start` → the api tokio + // runtime drops → detached handlers cancel → their `AbortGuard`s fire, release + // `Senders` clones → tok/detok channels close → workers exit. Full drain is + // deferred (see `request_shutdown`). + // `with_connect_info` exposes the peer address to the access-log middleware. + let serve = axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ); + tokio::select! { + r = serve => { + if let Err(e) = r { + tracing::error!(error = %e, "axum serve exited"); + } + } + _ = shutdown.recv_async() => { + tracing::info!("shutdown: stopping accepts, aborting in-flight handlers"); + } + } +} diff --git a/rust/sglang-server/src/api_server/common.rs b/rust/sglang-server/src/api_server/common.rs new file mode 100644 index 000000000000..70a7b73232c8 --- /dev/null +++ b/rust/sglang-server/src/api_server/common.rs @@ -0,0 +1,226 @@ +//! Common control-plane endpoints — `/server_info`, `/get_model_info` +//! (+ `/model_info` alias), plus the control-request submission path +//! (`await_control_result`, on the shared `submit`). Data-plane endpoints (incl. `/health*`, +//! which round-trips a generate probe) live in the sibling `native_api` and +//! `openai` modules; the shared `AppState` lives in the parent +//! `api_server` module. + +use axum::{ + Router, + extract::State, + http::StatusCode, + response::{IntoResponse, Response}, + routing::get, +}; + +use super::AppState; +use super::guard::AbortGuard; +use super::submit::submit; +use crate::message::{ControlRequest, EgressItem, GetInternalStateReq, RequestKind}; +use crate::runtime::ServerArgs; + +/// The routes this module owns, mounted by `api_server::serve`. +pub(super) fn routes() -> Router { + Router::new() + // Control-plane: reuses the ingress FSM (no tokenization), returns one + // non-streamed JSON result. Adding one = a route line + its struct tag. + .route("/server_info", get(server_info)) + // Static config, no scheduler round-trip. `/get_model_info` (+ `/model_info` + // alias). + .route("/get_model_info", get(model_info)) + .route("/model_info", get(model_info)) +} + +/// Submit a control request through the ingress FSM (no tokenization) and await the +/// scheduler's single msgpack result (a `structs.asdict` named map). Returns the +/// raw bytes, or an error `Response` to return as-is. +async fn await_control_result( + state: &AppState, + control: ControlRequest, +) -> Result { + let (rid, mut rx) = submit(state, RequestKind::Control(Box::new(control)), false).await?; + // Control requests register a detok entry like any other, and only + // `handle_result` removes it — so a request that never produces one (a stalled + // scheduler, a client that hangs up mid-await) leaves the entry behind. A + // monitor polling `/server_info` then leaks one `DetokState` per poll, forever. + // The guard deregisters on drop; it is disarmed below when the result lands. + let mut guard = AbortGuard::new(state.senders.clone(), rid.clone()); + let received = rx.recv().await; + if received.is_some() { + guard.disarm(&rid); // completed normally — nothing to abort + } + match received { + Some(EgressItem::Control(bytes)) => Ok(bytes), + Some(EgressItem::Error(e)) => { + let code = + StatusCode::from_u16(e.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + Err((code, e.to_string()).into_response()) + } + // A control request never receives generation frames. + Some(EgressItem::Frame(_)) | Some(EgressItem::Done(_)) => Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "unexpected generation output for control request", + ) + .into_response()), + None => Err((StatusCode::from_u16(499).unwrap(), "request aborted").into_response()), + } +} + +/// `GET /get_model_info` (+ `/model_info` alias) — static model metadata from +/// `server_args` (no scheduler round-trip); `is_generation` always true. +async fn model_info(State(state): State) -> Response { + let sa = &state.server_args; + let body = serde_json::json!({ + "model_path": sa.model_path, + "tokenizer_path": sa.tokenizer_path, + "is_generation": true, + // Python's `TokenizerManager` merges this into every request + // (`{**preferred, **client}`); this server has no equivalent yet, so + // `RustServer.launch` REFUSES to start when it is set. It can therefore + // only be null here — echoing it keeps the field's shape. + "preferred_sampling_params": sa.preferred_sampling_params, + "weight_version": serde_json::Value::Null, + }); + ( + StatusCode::OK, + [("content-type", "application/json")], + serde_json::to_vec(&body).unwrap_or_default(), + ) + .into_response() +} + +/// `GET /server_info` — surface only an allowlist ([`INTERNAL_STATE_ALLOWLIST`] + +/// curated [`ServerArgs`] accessors), never the raw server-args dump (embeds +/// `api_key`/`admin_api_key`; see [`shape_server_info`]). +/// +/// TODO(server_info): Python also includes `kv_events`; add once plumbed. +async fn server_info(State(state): State) -> Response { + let bytes = match await_control_result( + &state, + ControlRequest::GetInternalStateReq(GetInternalStateReq::new( + crate::ids::Rid::new().to_string(), + )), + ) + .await + { + Ok(b) => b, + Err(resp) => return resp, + }; + match shape_server_info(&bytes, &state.server_args) { + Ok(json) => (StatusCode::OK, [("content-type", "application/json")], json).into_response(), + Err(e) => { + tracing::error!(error = %e, "server_info: shaping failed"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + "bad server_info response", + ) + .into_response() + } + } +} + +/// Runtime-metric keys `get_internal_state` adds atop the server-args dump. We copy +/// ONLY these out of `internal_state` (an allowlist), so the co-mingled +/// `api_key`/`admin_api_key` can never reach the response. +const INTERNAL_STATE_ALLOWLIST: &[&str] = &[ + "last_gen_throughput", + "memory_usage", + "effective_max_running_requests_per_dp", + "avg_spec_accept_length", + "step_time_dict", +]; + +fn shape_server_info(msgpack: &[u8], server_args: &ServerArgs) -> Result, String> { + // GetInternalStateReqOutput asdict → `{ "internal_state": { server-args dump + + // metrics }, ... }`. Pull that inner map out (it is NOT safe to expose whole). + let mut obj: serde_json::Map = + rmp_serde::from_slice(msgpack).map_err(|e| e.to_string())?; + let internal = match obj.remove("internal_state") { + Some(serde_json::Value::Object(m)) => m, + _ => serde_json::Map::new(), + }; + + // Copy only the allowlisted runtime metrics — never the raw server-args dump. + let mut state_out = serde_json::Map::new(); + for &k in INTERNAL_STATE_ALLOWLIST { + match internal.get(k) { + Some(v) if !v.is_null() => { + state_out.insert(k.to_string(), v.clone()); + } + _ => {} + } + } + + // Top-level non-secret config from typed accessors (structurally can't surface + // a key field, unlike the raw dump). + let response = serde_json::json!({ + "model_path": server_args.model_path, + "served_model_name": server_args.served_model_name, + "tokenizer_path": server_args.tokenizer_path, + "max_context_length": server_args.model_config.context_len, + "max_total_num_tokens": server_args.max_total_num_tokens, + "version": server_args.version, + "internal_states": [serde_json::Value::Object(state_out)], + }); + serde_json::to_vec(&response).map_err(|e| e.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The scheduler's `internal_state` embeds the full server-args dump (incl. + /// `api_key`/`admin_api_key`). `/server_info` must surface only the allowlisted + /// runtime metrics + curated config — never the secrets — and must not re-nest + /// the dump under `internal_states[].internal_state`. + #[test] + fn shape_server_info_excludes_secrets_and_dump() { + // GetInternalStateReqOutput.asdict → { "internal_state": { …dump+metrics… } }. + let internal = rmpv::Value::Map(vec![ + ( + rmpv::Value::from("api_key"), + rmpv::Value::from("secret-token"), + ), + ( + rmpv::Value::from("admin_api_key"), + rmpv::Value::from("admin-token"), + ), + (rmpv::Value::from("model_path"), rmpv::Value::from("/m")), + ( + rmpv::Value::from("last_gen_throughput"), + rmpv::Value::from(1.5), + ), + ( + rmpv::Value::from("effective_max_running_requests_per_dp"), + rmpv::Value::from(32), + ), + ]); + let outer = rmpv::Value::Map(vec![(rmpv::Value::from("internal_state"), internal)]); + let mut msgpack = Vec::new(); + rmpv::encode::write_value(&mut msgpack, &outer).unwrap(); + + let sa = + ServerArgs::from_json(r#"{"model_path": "/m", "api_key": "secret-token"}"#).unwrap(); + let out = shape_server_info(&msgpack, &sa).unwrap(); + let text = String::from_utf8(out.clone()).unwrap(); + // No secret leaks anywhere in the serialized response. + assert!(!text.contains("secret-token"), "api_key leaked: {text}"); + assert!( + !text.contains("admin-token"), + "admin_api_key leaked: {text}" + ); + + let v: serde_json::Value = serde_json::from_slice(&out).unwrap(); + // Allowlisted metric surfaced; the whole dump did not. + let state0 = &v["internal_states"][0]; + assert_eq!(state0["last_gen_throughput"], 1.5); + assert_eq!(state0["effective_max_running_requests_per_dp"], 32); + assert!( + state0.get("internal_state").is_none(), + "must not re-nest the dump under internal_state" + ); + assert!(state0.get("api_key").is_none()); + // Curated top-level config comes from typed accessors, not the dump. + assert_eq!(v["model_path"], "/m"); + } +} diff --git a/rust/sglang-server/src/api_server/frame.rs b/rust/sglang-server/src/api_server/frame.rs new file mode 100644 index 000000000000..53fe0d8220d6 --- /dev/null +++ b/rust/sglang-server/src/api_server/frame.rs @@ -0,0 +1,865 @@ +//! Frame shaping for the native `/generate` protocol: the cumulative +//! [`OutputAccumulator`] plus the functions that render [`ChunkEvent`]s / +//! accumulated state into wire JSON (`meta_info`, logprob tuples, error and +//! abort frames). No HTTP here — the sibling `native_api` module owns the handlers +//! and streams; it calls these per frame. + +use crate::message::{ChunkEvent, ChunkExtras}; + +/// The text slot of a `[logprob, token_id, text]` tuple: the decoded token when +/// `return_text_in_logprobs` supplied a text buffer, else `null`. +fn text_slot(texts: Option<&[String]>, j: usize) -> serde_json::Value { + texts + .and_then(|t| t.get(j)) + .map(|s| serde_json::json!(s)) + .unwrap_or(serde_json::Value::Null) +} + +/// A decoded-text column becomes the tuples' text source only when populated +/// (`return_text_in_logprobs`); empty → `None` → null text slots. +fn opt_texts(t: &[String]) -> Option<&[String]> { + (!t.is_empty()).then_some(t) +} + +/// The logprob slot of a tuple: a finite value, or `null` for the `NaN` sentinel. +fn lp_value(v: f32) -> serde_json::Value { + if v.is_nan() { + serde_json::Value::Null + } else { + serde_json::json!(v) + } +} + +/// SGLang logprob shape: a list of `[logprob, token_id, text]` tuples. `texts` +/// (parallel to `idxs`) fills the text slot when set, else `null`. +fn logprob_tuples(vals: &[f32], idxs: &[i32], texts: Option<&[String]>) -> serde_json::Value { + let tuples: Vec = vals + .iter() + .zip(idxs.iter()) + .enumerate() + .map(|(j, (&v, &tid))| serde_json::json!([lp_value(v), tid, text_slot(texts, j)])) + .collect(); + serde_json::Value::Array(tuples) +} + +/// Ragged top-k / token-ids shape: one entry per position — a list of +/// `[logprob, token_id, text]` tuples, or `null` when `lens[p] == 0` (mirrors +/// `detokenize_top_logprobs_tokens`). `texts` is parallel to `vals`/`idxs`. +fn ragged_logprob_tuples( + vals: &[f32], + idxs: &[i32], + lens: &[u32], + texts: Option<&[String]>, +) -> serde_json::Value { + let mut positions = Vec::with_capacity(lens.len()); + let mut off = 0usize; + for &l in lens { + let l = l as usize; + if l == 0 { + positions.push(serde_json::Value::Null); + } else { + // Bounds-checked like `hidden_states_rows`: a header whose `lens` run + // past the value buffer would otherwise panic the api thread on an + // out-of-range index. + let tuples: Vec = (off..off + l) + .filter_map(|j| { + Some(serde_json::json!([ + lp_value(*vals.get(j)?), + *idxs.get(j)?, + text_slot(texts, j) + ])) + }) + .collect(); + positions.push(serde_json::Value::Array(tuples)); + } + off += l; + } + serde_json::Value::Array(positions) +} + +/// Append a flat family's `[logprob, token_id, text]` tuples to `dst`, comma +/// separated and WITHOUT the enclosing brackets, so a cumulative frame can +/// concatenate each delta instead of re-rendering every accumulated position. +/// +/// Byte-identical to [`logprob_tuples`]'s serialization: `serde_json` writes an +/// array as `[`, elements joined by `,`, `]` with no spaces, and each element here +/// is rendered by the same `Value` Display. +fn push_logprob_tuples(dst: &mut String, vals: &[f32], idxs: &[i32], texts: Option<&[String]>) { + use std::fmt::Write; + for (j, (&v, &tid)) in vals.iter().zip(idxs.iter()).enumerate() { + if !dst.is_empty() { + dst.push(','); + } + let _ = write!(dst, "[{},{tid},{}]", lp_value(v), text_slot(texts, j)); + } +} + +/// Ragged counterpart of [`push_logprob_tuples`] — one entry per position, `null` +/// where `lens[p] == 0`. Mirrors [`ragged_logprob_tuples`] including its +/// bounds-checked skip, so a header whose `lens` run past the value buffer yields +/// the same (shortened) row rather than panicking. +fn push_ragged_tuples( + dst: &mut String, + vals: &[f32], + idxs: &[i32], + lens: &[u32], + texts: Option<&[String]>, +) { + use std::fmt::Write; + let mut off = 0usize; + for &l in lens { + let l = l as usize; + if !dst.is_empty() { + dst.push(','); + } + if l == 0 { + dst.push_str("null"); + } else { + dst.push('['); + let mut first = true; + for j in off..off + l { + let (Some(&v), Some(&tid)) = (vals.get(j), idxs.get(j)) else { + continue; + }; + if !first { + dst.push(','); + } + first = false; + let _ = write!(dst, "[{},{tid},{}]", lp_value(v), text_slot(texts, j)); + } + dst.push(']'); + } + off += l; + } +} + +/// Render one set-once family (they ride the prefill or the final chunk) into a +/// standalone JSON array, so it is serialized when it arrives rather than on every +/// subsequent frame. +fn ragged_array_json(vals: &[f32], idxs: &[i32], lens: &[u32], texts: Option<&[String]>) -> String { + let mut body = String::new(); + push_ragged_tuples(&mut body, vals, idxs, lens, texts); + format!("[{body}]") +} + +/// Reshape flat hidden-state f32s + per-row lengths into `meta_info`'s nested +/// `list[list[float]]` (one row per output position). +fn hidden_states_rows(vals: &[f32], lens: &[u32]) -> serde_json::Value { + let mut rows = Vec::with_capacity(lens.len()); + let mut off = 0usize; + for &l in lens { + let l = l as usize; + // `get`, not a clamped index: clamping only the END leaves `off` past + // `vals.len()` after one over-long row, making the next range reversed + // (`start > end`) — which panics on the api thread rather than yielding + // an empty row. Same reasoning as the egress decoder's `take_f32`. + rows.push(serde_json::json!(vals.get(off..off + l).unwrap_or(&[]))); + off += l; + } + serde_json::Value::Array(rows) +} + +/// The `{ "error": { message, code } }` object every error path emits (an SSE +/// event's data, a unary body, or one entry of a batch array). +pub(super) fn error_value(code: u16, message: &str) -> serde_json::Value { + serde_json::json!({ "error": { "message": message, "code": code } }) +} + +/// Format a decoded [`ChunkEvent`] as one SGLang `/generate` frame's JSON. `rid` +/// (response `meta_info.id`) is passed as a string; the event's numeric `rid` is +/// just the shard routing key. +pub(super) fn frame_value(out: &ChunkEvent, rid: &str) -> serde_json::Value { + let mut v = serde_json::json!({ + "text": out.text, + "meta_info": { + "id": rid, + "prompt_tokens": out.prompt_tokens, + "completion_tokens": out.completion_tokens, + // Full dict (type + matched + message + status_code + …), or null. + "finish_reason": out.finish_reason, + }, + }); + if !out.token_ids.is_empty() { + v["output_ids"] = serde_json::json!(out.token_ids); + } + // Logprobs + hidden states ride behind the boxed extras (absent for a plain + // token/text frame). `[logprob, token_id, text|null]` tuples; text + // (`return_text_in_logprobs`) was decoded on the detok shard into `*_txt`. + let Some(ex) = out.extras.as_deref() else { + return v; + }; + if !ex.out_lp_val.is_empty() { + v["meta_info"]["output_token_logprobs"] = + logprob_tuples(&ex.out_lp_val, &ex.out_lp_idx, opt_texts(&ex.out_lp_txt)); + } + if !ex.in_lp_val.is_empty() { + v["meta_info"]["input_token_logprobs"] = + logprob_tuples(&ex.in_lp_val, &ex.in_lp_idx, opt_texts(&ex.in_lp_txt)); + } + if !ex.out_top_lens.is_empty() { + v["meta_info"]["output_top_logprobs"] = ragged_logprob_tuples( + &ex.out_top_val, + &ex.out_top_idx, + &ex.out_top_lens, + opt_texts(&ex.out_top_txt), + ); + } + if !ex.in_top_lens.is_empty() { + v["meta_info"]["input_top_logprobs"] = ragged_logprob_tuples( + &ex.in_top_val, + &ex.in_top_idx, + &ex.in_top_lens, + opt_texts(&ex.in_top_txt), + ); + } + if !ex.out_tid_lens.is_empty() { + v["meta_info"]["output_token_ids_logprobs"] = ragged_logprob_tuples( + &ex.out_tid_val, + &ex.out_tid_idx, + &ex.out_tid_lens, + opt_texts(&ex.out_tid_txt), + ); + } + if !ex.in_tid_lens.is_empty() { + v["meta_info"]["input_token_ids_logprobs"] = ragged_logprob_tuples( + &ex.in_tid_val, + &ex.in_tid_idx, + &ex.in_tid_lens, + opt_texts(&ex.in_tid_txt), + ); + } + if !ex.hidden_lens.is_empty() { + v["meta_info"]["hidden_states"] = hidden_states_rows(&ex.hidden_val, &ex.hidden_lens); + } + v +} + +/// Cumulative frame JSON from the accumulator's memoized parts — O(1) in the +/// accumulated length, where rebuilding the `Value` is O(T) per frame and so O(T²) +/// per request. +/// +/// Byte-identical to `frame_value(..).to_string()`, which requires emitting +/// `meta_info`'s keys in the alphabetical order `serde_json`'s `BTreeMap` gives +/// them — pinned for both the plain and the logprob shapes by +/// `cumulative_frame_json_matches_serde`. `None` only when the extras memo was +/// invalidated (see `extras_memo_broken`), leaving the `Value` path as the +/// fallback. +pub(super) fn cumulative_frame_json( + acc: &OutputAccumulator, + rid: &str, + index: Option, +) -> Option { + use std::fmt::Write; + + if acc.extras_memo_broken { + return None; + } + let o = acc.snapshot(); + // Through `Value`, not `to_string` on the struct: the `Value` path sorts the + // finish reason's own keys via `BTreeMap`, and this must match it byte for byte. + let finish = serde_json::to_value(&o.finish_reason).ok()?.to_string(); + + // KEEP ALPHABETICAL: `serde_json::Map` is a `BTreeMap` here, so this is the + // order the `Value` path produces and the order the equivalence test asserts. + let mut m = String::new(); + let _ = write!(m, "{{\"completion_tokens\":{}", o.completion_tokens); + let _ = write!(m, ",\"finish_reason\":{finish}"); + if let Some(h) = &acc.hidden_json { + let _ = write!(m, ",\"hidden_states\":{h}"); + } + let _ = write!(m, ",\"id\":{}", serde_json::Value::String(rid.to_string())); + if let Some(v) = &acc.in_tid_json { + let _ = write!(m, ",\"input_token_ids_logprobs\":{v}"); + } + if let Some(v) = &acc.in_lp_json { + let _ = write!(m, ",\"input_token_logprobs\":{v}"); + } + if let Some(v) = &acc.in_top_json { + let _ = write!(m, ",\"input_top_logprobs\":{v}"); + } + // The `Value` path keys these off the source columns being non-empty; an empty + // source renders to an empty body, so the two guards coincide. + if !acc.out_tid_json.is_empty() { + let _ = write!(m, ",\"output_token_ids_logprobs\":[{}]", acc.out_tid_json); + } + if !acc.out_lp_json.is_empty() { + let _ = write!(m, ",\"output_token_logprobs\":[{}]", acc.out_lp_json); + } + if !acc.out_top_json.is_empty() { + let _ = write!(m, ",\"output_top_logprobs\":[{}]", acc.out_top_json); + } + let _ = write!(m, ",\"prompt_tokens\":{}}}", o.prompt_tokens); + + let mut s = String::with_capacity(acc.text_json.len() + acc.ids_json.len() + m.len() + 40); + s.push('{'); + if let Some(i) = index { + let _ = write!(s, "\"index\":{i},"); + } + s.push_str("\"meta_info\":"); + s.push_str(&m); + if !acc.ids_json.is_empty() { + s.push_str(",\"output_ids\":["); + s.push_str(&acc.ids_json); + s.push(']'); + } + s.push_str(",\"text\":\""); + s.push_str(&acc.text_json); + s.push_str("\"}"); + Some(s) +} + +/// Attach the batch `index` (batch streams only) and render to the SSE `data` text. +pub(super) fn tag_value(mut v: serde_json::Value, index: Option) -> String { + if let Some(i) = index { + v["index"] = serde_json::json!(i); + } + v.to_string() +} + +/// One streaming frame's JSON: cumulative ignores `delta`, incremental ships it. +pub(super) fn stream_frame_string( + delta: ChunkEvent, + acc: &OutputAccumulator, + incremental: bool, + rid_str: &str, + index: Option, +) -> String { + if !incremental { + return cumulative_frame_string(acc, rid_str, index); + } + tag_value(stream_frame_value(delta, acc, true, rid_str), index) +} + +/// A cumulative frame's JSON, built purely from the accumulator (which is why a +/// backlog can coalesce to its last); falls back to the `Value` builder on extras. +pub(super) fn cumulative_frame_string( + acc: &OutputAccumulator, + rid_str: &str, + index: Option, +) -> String { + cumulative_frame_json(acc, rid_str, index) + .unwrap_or_else(|| tag_value(frame_value(acc.snapshot(), rid_str), index)) +} + +/// Format one streaming frame: the accumulator's cumulative view (default), or this +/// step's delta with the cumulative token count in `meta_info` (matching Python). +pub(super) fn stream_frame_value( + delta: ChunkEvent, + acc: &OutputAccumulator, + incremental: bool, + rid_str: &str, +) -> serde_json::Value { + if incremental { + let mut d = delta; + d.completion_tokens = acc.snapshot().completion_tokens; + frame_value(&d, rid_str) + } else { + frame_value(acc.snapshot(), rid_str) + } +} + +/// Folds per-chunk [`ChunkEvent`] deltas into a cumulative view — used by the drain +/// loops needing cumulative output (every unary response + the cumulative SGLang +/// stream; OpenAI streaming forwards deltas and skips this). Holds a single +/// [`ChunkEvent`] so `snapshot` hands back a **borrow** per frame — no per-frame +/// clone of the growing buffers (that added O(T²) atop the wire's inherent O(T²)). +/// Shared with the [`openai`] submodule. +#[derive(Default)] +pub(super) struct OutputAccumulator { + out: ChunkEvent, + /// Serialized cumulative `output_ids` body (`"1,2,3"`, no brackets), appended per + /// delta so a frame memcpy's it instead of rebuilding the array — O(T), not O(T²). + ids_json: String, + /// JSON-escaped cumulative text, without the surrounding quotes. Escaping is + /// per-character, so `escape(a + b) == escape(a) + escape(b)` and deltas append. + text_json: String, + /// Memoized bodies (no enclosing brackets) of the three CUMULATIVE logprob + /// families, appended per delta — the same O(T) trick `ids_json` uses, extended + /// to the families that made a cumulative stream with logprobs O(T²). Cumulative + /// is SGLang's default, so that path re-rendered every accumulated position on + /// every frame: measured 117 ms for one 500-token top-5 request, versus 1.2 ms + /// incremental. + out_lp_json: String, + out_top_json: String, + out_tid_json: String, + /// Set-once families — they ride the prefill or the final chunk, so they are + /// rendered when they arrive rather than on every frame after. + in_lp_json: Option, + in_top_json: Option, + in_tid_json: Option, + hidden_json: Option, + /// Set once a family's text column falls out of lockstep with its values, at + /// which point the memo is abandoned for the `Value` path. + /// + /// Appending per delta assumes `text_slot(accumulated, global_j)` equals + /// `text_slot(delta, local_j)`, which holds only while every delta supplies + /// either a text per value or none at all. That is what a real request does — + /// `return_text_in_logprobs` is per-request, so the detok shard fills `*_txt` + /// for all deltas or none — but a mixed sequence would silently diverge from + /// `frame_value`, so it is detected rather than assumed. + extras_memo_broken: bool, +} + +/// Append `s` JSON-escaped (no surrounding quotes) — `serde_json` quotes it, and the +/// quotes are the first and last bytes of a string encoding. +fn push_escaped(dst: &mut String, s: &str) { + if s.is_empty() { + return; + } + let quoted = serde_json::to_string(s).expect("str-to-json should never fail"); + dst.push_str("ed[1..quoted.len() - 1]); +} + +impl OutputAccumulator { + /// Fold one delta frame in. Output families concatenate; input families and + /// hidden states are set-once / last-writer-wins (they ride the prefill/final + /// chunk), matching the Python `meta_info` assignment. + pub(super) fn fold(&mut self, d: &ChunkEvent) { + use std::fmt::Write; + + // Grow the memoized serializations alongside the raw cumulative buffers. + push_escaped(&mut self.text_json, &d.text); + for &id in &d.token_ids { + if !self.ids_json.is_empty() { + self.ids_json.push(','); + } + let _ = write!(self.ids_json, "{id}"); + } + + let o = &mut self.out; + o.rid.clone_from(&d.rid); // constant across the request; keeps the accumulated view coherent + o.text.push_str(&d.text); + o.token_ids.extend_from_slice(&d.token_ids); // token_ids doubles as output_ids + o.completion_tokens += d.completion_tokens; + o.prompt_tokens = d.prompt_tokens; // constant across the request + if d.finish_reason.is_some() { + o.finish_reason = d.finish_reason.clone(); + } + // Logprobs/hidden ride behind the boxed extras — most frames have none, so + // only allocate the accumulator's box once a delta actually carries some. + let Some(de) = d.extras.as_deref() else { + return; + }; + let oe = o + .extras + .get_or_insert_with(|| Box::new(ChunkExtras::default())); + oe.out_lp_val.extend_from_slice(&de.out_lp_val); + oe.out_lp_idx.extend_from_slice(&de.out_lp_idx); + oe.out_top_val.extend_from_slice(&de.out_top_val); + oe.out_top_idx.extend_from_slice(&de.out_top_idx); + oe.out_top_lens.extend_from_slice(&de.out_top_lens); + oe.out_tid_val.extend_from_slice(&de.out_tid_val); + oe.out_tid_idx.extend_from_slice(&de.out_tid_idx); + oe.out_tid_lens.extend_from_slice(&de.out_tid_lens); + oe.out_lp_txt.extend_from_slice(&de.out_lp_txt); + oe.out_top_txt.extend_from_slice(&de.out_top_txt); + oe.out_tid_txt.extend_from_slice(&de.out_tid_txt); + // Append THIS delta's tuples, indexed within the delta — equivalent to + // indexing the accumulated arrays only while texts stay in lockstep, which + // the guard below verifies. + push_logprob_tuples( + &mut self.out_lp_json, + &de.out_lp_val, + &de.out_lp_idx, + opt_texts(&de.out_lp_txt), + ); + push_ragged_tuples( + &mut self.out_top_json, + &de.out_top_val, + &de.out_top_idx, + &de.out_top_lens, + opt_texts(&de.out_top_txt), + ); + push_ragged_tuples( + &mut self.out_tid_json, + &de.out_tid_val, + &de.out_tid_idx, + &de.out_tid_lens, + opt_texts(&de.out_tid_txt), + ); + let lockstep = |txt: &Vec, val: &Vec| txt.is_empty() || txt.len() == val.len(); + if !lockstep(&oe.out_lp_txt, &oe.out_lp_val) + || !lockstep(&oe.out_top_txt, &oe.out_top_val) + || !lockstep(&oe.out_tid_txt, &oe.out_tid_val) + { + self.extras_memo_broken = true; + } + if !de.in_lp_val.is_empty() { + oe.in_lp_val = de.in_lp_val.clone(); + oe.in_lp_idx = de.in_lp_idx.clone(); + oe.in_lp_txt = de.in_lp_txt.clone(); + let mut body = String::new(); + push_logprob_tuples( + &mut body, + &oe.in_lp_val, + &oe.in_lp_idx, + opt_texts(&oe.in_lp_txt), + ); + self.in_lp_json = Some(format!("[{body}]")); + } + // Input families ride once (prefill); `lens` non-empty marks their arrival. + if !de.in_top_lens.is_empty() { + oe.in_top_val = de.in_top_val.clone(); + oe.in_top_idx = de.in_top_idx.clone(); + oe.in_top_lens = de.in_top_lens.clone(); + oe.in_top_txt = de.in_top_txt.clone(); + self.in_top_json = Some(ragged_array_json( + &oe.in_top_val, + &oe.in_top_idx, + &oe.in_top_lens, + opt_texts(&oe.in_top_txt), + )); + } + if !de.in_tid_lens.is_empty() { + oe.in_tid_val = de.in_tid_val.clone(); + oe.in_tid_idx = de.in_tid_idx.clone(); + oe.in_tid_lens = de.in_tid_lens.clone(); + oe.in_tid_txt = de.in_tid_txt.clone(); + self.in_tid_json = Some(ragged_array_json( + &oe.in_tid_val, + &oe.in_tid_idx, + &oe.in_tid_lens, + opt_texts(&oe.in_tid_txt), + )); + } + // Hidden states are non-cumulative: the latest non-empty set wins. + if !de.hidden_lens.is_empty() { + oe.hidden_val = de.hidden_val.clone(); + oe.hidden_lens = de.hidden_lens.clone(); + self.hidden_json = + Some(hidden_states_rows(&oe.hidden_val, &oe.hidden_lens).to_string()); + } + } + + /// Borrow the cumulative output for an intermediate streaming frame. + pub(super) fn snapshot(&self) -> &ChunkEvent { + &self.out + } + + /// Consume into the final cumulative output. + pub(super) fn into_output(self) -> ChunkEvent { + self.out + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn flat_logprob_tuples_shape() { + let v = logprob_tuples(&[-0.5, -1.5], &[10, 20], None); + assert_eq!( + v, + serde_json::json!([ + [-0.5f32, 10, serde_json::Value::Null], + [-1.5f32, 20, serde_json::Value::Null] + ]) + ); + } + + /// With a text buffer, the tuple's third slot carries the decoded token. + #[test] + fn flat_logprob_tuples_with_text() { + let texts = vec!["a".to_string(), "b".to_string()]; + let v = logprob_tuples(&[-0.5, -1.5], &[10, 20], Some(&texts)); + assert_eq!( + v, + serde_json::json!([[-0.5f32, 10, "a"], [-1.5f32, 20, "b"]]) + ); + } + + /// Ragged reshape restores null positions (len 0) — mirrors + /// detokenize_top_logprobs_tokens emitting None for empty positions. + #[test] + fn ragged_logprob_tuples_restores_null_positions() { + // 2 positions: first null (len 0), second k=1. + let v = ragged_logprob_tuples(&[-0.3], &[9], &[0, 1], None); + assert_eq!( + v, + serde_json::json!([ + serde_json::Value::Null, + [[-0.3f32, 9, serde_json::Value::Null]] + ]) + ); + } + + /// The `NaN` sentinel (the Python `None` logprob for the first prompt token) + /// becomes a JSON `null` logprob, while its token id in the parallel `idx` + /// column is preserved. Guards the scheduler-killing prompt-logprob crash. + #[test] + fn nan_sentinel_becomes_null_logprob() { + // Flat (input/output logprobs): first value absent, second present. + let flat = logprob_tuples(&[f32::NAN, -0.5], &[10, 20], None); + assert_eq!( + flat, + serde_json::json!([ + [serde_json::Value::Null, 10, serde_json::Value::Null], + [-0.5f32, 20, serde_json::Value::Null], + ]) + ); + // Ragged (top-k / token-ids logprobs): a NaN inside a position → null. + let ragged = ragged_logprob_tuples(&[f32::NAN], &[7], &[1], None); + assert_eq!( + ragged, + serde_json::json!([[[serde_json::Value::Null, 7, serde_json::Value::Null]]]) + ); + } + + /// End-to-end: a `ChunkEvent` carrying a prompt-logprob request (first input + /// logprob is the `NaN` sentinel) formats without panicking and emits + /// `input_token_logprobs` with a leading `[null, token_id, text]`. + #[test] + fn prompt_logprob_frame_emits_null_first() { + let out = ChunkEvent { + extras: Some(Box::new(ChunkExtras { + in_lp_val: vec![f32::NAN, -0.5], + in_lp_idx: vec![10, 20], + in_lp_txt: vec!["".into(), "hi".into()], + ..Default::default() + })), + ..Default::default() + }; + let frame = frame_value(&out, "1"); + assert_eq!( + frame["meta_info"]["input_token_logprobs"], + serde_json::json!([[serde_json::Value::Null, 10, ""], [-0.5f32, 20, "hi"]]) + ); + } + + /// The accumulator folds deltas cumulatively and `snapshot` borrows the + /// running state (no per-frame clone); `into_output` moves the same state. + #[test] + fn accumulator_snapshot_is_cumulative() { + let mut acc = OutputAccumulator::default(); + acc.fold(&ChunkEvent { + text: "he".into(), + token_ids: vec![1, 2], + completion_tokens: 2, + ..Default::default() + }); + { + let s = acc.snapshot(); + assert_eq!(s.text, "he"); + assert_eq!(s.token_ids, vec![1, 2]); + } + acc.fold(&ChunkEvent { + text: "llo".into(), + token_ids: vec![3], + completion_tokens: 1, + ..Default::default() + }); + { + let s = acc.snapshot(); + assert_eq!(s.text, "hello"); // cumulative + assert_eq!(s.token_ids, vec![1, 2, 3]); + assert_eq!(s.completion_tokens, 3); + } + let out = acc.into_output(); + assert_eq!(out.text, "hello"); + } + + /// A populated text column (decoded on the detok shard) → `Some`; empty + /// (`return_text_in_logprobs` off) → `None` → null text slots. + #[test] + fn opt_texts_gates_on_population() { + assert!(opt_texts(&[]).is_none()); + let t = vec!["x".to_string()]; + assert_eq!(opt_texts(&t), Some(t.as_slice())); + } + + /// The memoized cumulative fast path must emit **byte-identical** JSON to the + /// `serde_json::Value` builder it replaces — same keys, same alphabetical order + /// (`Map` is a `BTreeMap`; no `preserve_order`), same escaping. Covers unicode + /// and control chars, an empty-ids first frame, a finish_reason, and the batch + /// `index`. Guards the O(T) rewrite of the O(T²) `output_ids` serialization. + #[test] + fn cumulative_frame_json_matches_serde() { + let deltas = [ + ChunkEvent { + rid: "7".into(), + text: String::new(), + token_ids: vec![], + completion_tokens: 0, + prompt_tokens: 128, + ..Default::default() + }, + ChunkEvent { + rid: "7".into(), + text: "He\"llo\n\t".into(), + token_ids: vec![1000], + completion_tokens: 1, + prompt_tokens: 128, + ..Default::default() + }, + ChunkEvent { + rid: "7".into(), + text: " 世界 🌍 \\".into(), + token_ids: vec![-2, 3], + completion_tokens: 2, + prompt_tokens: 128, + ..Default::default() + }, + ChunkEvent { + rid: "7".into(), + text: "!".into(), + token_ids: vec![9], + completion_tokens: 1, + prompt_tokens: 128, + finish_reason: serde_json::from_value( + serde_json::json!({"type": "stop", "matched": 9}), + ) + .expect("finish reason must parse"), + ..Default::default() + }, + ]; + + for index in [None, Some(3usize)] { + let mut acc = OutputAccumulator::default(); + for d in &deltas { + acc.fold(d); + let fast = cumulative_frame_json(&acc, "7", index).expect("no extras → fast path"); + let slow = tag_value(frame_value(acc.snapshot(), "7"), index); + assert_eq!(fast, slow, "index={index:?} text={:?}", acc.snapshot().text); + } + } + } + + /// The same equivalence, for the shape that made cumulative streaming O(T²): + /// every logprob family at once, across several deltas, with and without + /// `return_text_in_logprobs` texts and with a null ragged position. + /// + /// This is the guard on the memoization. The fast path hand-writes + /// `meta_info`'s keys, so it has to reproduce `serde_json`'s alphabetical + /// `BTreeMap` order and every family's exact tuple encoding; asserting equality + /// against the `Value` builder after each fold is what makes that safe to + /// maintain. + #[test] + fn cumulative_frame_json_matches_serde_with_logprobs() { + for with_texts in [false, true] { + let txt = |v: &[&str]| -> Vec { + if with_texts { + v.iter().map(|s| (*s).to_string()).collect() + } else { + Vec::new() + } + }; + let deltas = [ + // Prefill: the set-once input families and a null top-k position. + ChunkEvent { + rid: "9".into(), + prompt_tokens: 4, + extras: Some(Box::new(ChunkExtras { + in_lp_val: vec![f32::NAN, -1.5], + in_lp_idx: vec![10, 11], + in_lp_txt: txt(&["a", "b"]), + in_top_val: vec![-0.25], + in_top_idx: vec![12], + in_top_lens: vec![0, 1], + in_top_txt: txt(&["c"]), + in_tid_val: vec![-2.0], + in_tid_idx: vec![13], + in_tid_lens: vec![1], + in_tid_txt: txt(&["d"]), + ..Default::default() + })), + ..Default::default() + }, + ChunkEvent { + rid: "9".into(), + text: "He\"llo".into(), + token_ids: vec![100], + completion_tokens: 1, + prompt_tokens: 4, + extras: Some(Box::new(ChunkExtras { + out_lp_val: vec![-0.5], + out_lp_idx: vec![100], + out_lp_txt: txt(&["He\"llo"]), + out_top_val: vec![-0.5, -3.0], + out_top_idx: vec![100, 7], + out_top_lens: vec![2], + out_top_txt: txt(&["He\"llo", "x"]), + out_tid_val: vec![-0.5], + out_tid_idx: vec![100], + out_tid_lens: vec![1], + out_tid_txt: txt(&["He\"llo"]), + ..Default::default() + })), + ..Default::default() + }, + ChunkEvent { + rid: "9".into(), + text: " 世界".into(), + token_ids: vec![-2, 3], + completion_tokens: 2, + prompt_tokens: 4, + finish_reason: serde_json::from_value( + serde_json::json!({"type": "stop", "matched": 3}), + ) + .expect("finish reason must parse"), + extras: Some(Box::new(ChunkExtras { + out_lp_val: vec![f32::NAN, -0.125], + out_lp_idx: vec![-2, 3], + out_lp_txt: txt(&[" 世", "界"]), + // A zero-length position must render as `null`, not `[]`. + out_top_val: vec![-0.125], + out_top_idx: vec![3], + out_top_lens: vec![0, 1], + out_top_txt: txt(&["界"]), + out_tid_val: vec![], + out_tid_idx: vec![], + out_tid_lens: vec![0, 0], + out_tid_txt: txt(&[]), + hidden_val: vec![0.5, -0.25, 1.0], + hidden_lens: vec![2, 1], + ..Default::default() + })), + }, + ]; + + for index in [None, Some(2usize)] { + let mut acc = OutputAccumulator::default(); + for d in &deltas { + acc.fold(d); + let fast = cumulative_frame_json(&acc, "9", index) + .expect("the extras memo must stay valid for a well-formed request"); + let slow = tag_value(frame_value(acc.snapshot(), "9"), index); + assert_eq!(fast, slow, "with_texts={with_texts} index={index:?}"); + } + } + } + } + + /// A delta sequence that supplies texts for some values and not others breaks + /// the append-equivalence the memo rests on (`text_slot` is indexed globally, + /// so a gap shifts every later text). The accumulator must notice and defer to + /// the `Value` builder rather than emit a frame that disagrees with it. + #[test] + fn mismatched_logprob_texts_fall_back_to_the_value_path() { + let mut acc = OutputAccumulator::default(); + acc.fold(&ChunkEvent { + rid: "1".into(), + extras: Some(Box::new(ChunkExtras { + out_lp_val: vec![-0.5], + out_lp_idx: vec![5], + ..Default::default() // no texts + })), + ..Default::default() + }); + assert!(cumulative_frame_json(&acc, "1", None).is_some()); + acc.fold(&ChunkEvent { + rid: "1".into(), + extras: Some(Box::new(ChunkExtras { + out_lp_val: vec![-0.25], + out_lp_idx: vec![6], + out_lp_txt: vec!["b".into()], // …now texts: out of lockstep + ..Default::default() + })), + ..Default::default() + }); + assert!( + cumulative_frame_json(&acc, "1", None).is_none(), + "a text column out of lockstep must invalidate the memo" + ); + } +} diff --git a/rust/sglang-server/src/api_server/guard.rs b/rust/sglang-server/src/api_server/guard.rs new file mode 100644 index 000000000000..728d03ac5df4 --- /dev/null +++ b/rust/sglang-server/src/api_server/guard.rs @@ -0,0 +1,130 @@ +//! Abort-on-disconnect guard for in-flight requests. Handlers arm a guard per +//! submitted rid; axum dropping the handler/SSE stream (client disconnected) +//! drops the guard, which aborts whatever wasn't disarmed (mirrors Python's +//! `is_disconnected` abort). + +use std::collections::HashSet; + +use crate::ids::Rid; +use crate::tokenizer_manager::{AbortSource, Senders}; + +/// Aborts still-in-flight rids on drop. Each rid is disarmed on natural finish; +/// whatever remains at drop is aborted. +pub(super) struct AbortGuard { + senders: Senders, + /// Rids still in flight. `Rid` carries its own partition key, so there is no + /// separate routing value to keep alongside it. + /// + /// A set, not a `Vec`: `disarm` runs once per request that finishes, and over a + /// batch a linear scan makes the guard quadratic in the batch size — measured + /// 13.3 ms for a 4096-item batch, more than all of that batch's real transform + /// work combined. `Rid`'s identity is its id string, so set membership is the + /// same relation `retain` was testing. The cost is two hashes of a ~40-byte + /// string on the single-request path (~80 ns against a ~40 µs request), which + /// is why the trade is worth making rather than threading slot indices in from + /// the batch call sites. + rids: HashSet, +} + +impl AbortGuard { + pub(super) fn new(senders: Senders, rid: Rid) -> Self { + Self { + senders, + rids: HashSet::from([rid]), + } + } + + /// Guard covering no rids yet — a batch arms each as it's submitted so a + /// mid-fan-out disconnect aborts every request already handed to the scheduler. + pub(super) fn new_empty(senders: Senders) -> Self { + Self { + senders, + rids: HashSet::new(), + } + } + + /// Track a request for abort-on-drop. + pub(super) fn arm(&mut self, rid: Rid) { + self.rids.insert(rid); + } + + /// Request finished naturally — don't abort it on drop. + pub(super) fn disarm(&mut self, rid: &Rid) { + self.rids.remove(rid); + } +} + +impl Drop for AbortGuard { + fn drop(&mut self) { + // Report the abort and nothing more. There is no in-flight rid registry to + // release from: `Rid::from_client` makes each client rid internally unique, + // so a resubmit of the "same" rid is a different `Rid` and cannot be caught + // up in this abort. That removes the ordering hazard split ownership created. + // + // The lane is unbounded, so this send only fails at shutdown, when the loop + // is gone and nothing is generating anyway. + for rid in self.rids.drain() { + let _ = self.senders.abort.send(AbortSource::Guard(rid)); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn senders_with_abort(abort: flume::Sender) -> Senders { + Senders { + tm: flume::unbounded().0, + abort, + tok: flume::unbounded().0, + detok: vec![], + } + } + + /// A batch guard aborts exactly the rids still armed at drop — the ones whose + /// requests never reached a terminal — and leaves the finished ones alone. + #[test] + fn guard_aborts_only_the_rids_still_armed() { + let (abort_tx, abort_rx) = flume::unbounded(); + let done: Rid = "done".into(); + let mut guard = AbortGuard::new(senders_with_abort(abort_tx), done.clone()); + guard.arm("aborted".into()); + guard.disarm(&done); // finished naturally + drop(guard); + + assert!( + matches!(abort_rx.try_recv().unwrap(), AbortSource::Guard(r) if r.as_str() == "aborted") + ); + assert!( + abort_rx.try_recv().is_err(), + "a disarmed rid must not be aborted" + ); + } + + /// An armed guard aborts its rid on drop — exactly the cleanup a busy-skipped + /// `/health_generate` probe relies on. It never sees a terminal frame here, so + /// dropping the guard is the only path that deregisters its detok sink (via the + /// ingress `on_abort`). Regression for the detok-entry leak per health probe. + #[test] + fn armed_guard_aborts_on_drop() { + let (tm_tx, tm_rx) = flume::unbounded(); + drop(AbortGuard::new(senders_with_abort(tm_tx), "r7".into())); + assert!( + matches!(tm_rx.try_recv(), Ok(AbortSource::Guard(rid)) if rid.as_str() == "r7"), + "armed guard must abort its rid on drop", + ); + assert!(tm_rx.try_recv().is_err(), "exactly one abort"); + } + + /// A disarmed rid (finished naturally) is not aborted on drop. + #[test] + fn disarmed_guard_does_not_abort() { + let (tm_tx, tm_rx) = flume::unbounded(); + let id = Rid::from("r9"); + let mut guard = AbortGuard::new(senders_with_abort(tm_tx), "r9".into()); + guard.disarm(&id); + drop(guard); + assert!(tm_rx.try_recv().is_err(), "disarmed rid must not abort"); + } +} diff --git a/rust/sglang-server/src/api_server/log.rs b/rust/sglang-server/src/api_server/log.rs new file mode 100644 index 000000000000..0cf9ab596253 --- /dev/null +++ b/rust/sglang-server/src/api_server/log.rs @@ -0,0 +1,41 @@ +//! HTTP access logging — one INFO line per request, content-matching the +//! Python server's uvicorn access log. Gated exactly like uvicorn's +//! (`--log-level-http warning` turns it off, see +//! `ServerArgs::http_access_log_enabled`); when disabled the middleware is not +//! installed at all — zero cost. + +use axum::{Router, response::Response}; + +use crate::runtime::ServerArgs; + +/// Install the access-log middleware when `server_args` enables it; identity +/// otherwise (the layer is never installed, so disabled stays zero-cost). +pub(super) fn apply(app: Router, server_args: &ServerArgs) -> Router { + if server_args.http_access_log_enabled() { + app.layer(axum::middleware::from_fn(access_log)) + } else { + app + } +} + +/// Access log — one INFO line per request, content-matching the Python server's +/// uvicorn access log (`127.0.0.1:54232 - "GET /model_info HTTP/1.1" 200 OK`). +/// Logged when the response head is ready; for SSE that's stream start, same as +/// uvicorn. +async fn access_log( + axum::extract::ConnectInfo(peer): axum::extract::ConnectInfo, + req: axum::extract::Request, + next: axum::middleware::Next, +) -> Response { + let method = req.method().clone(); + let uri = req.uri().clone(); + let version = req.version(); + let res = next.run(req).await; + let status = res.status(); + tracing::info!( + "{peer} - \"{method} {uri} {version:?}\" {} {}", + status.as_u16(), + status.canonical_reason().unwrap_or("") + ); + res +} diff --git a/rust/sglang-server/src/api_server/native_api.rs b/rust/sglang-server/src/api_server/native_api.rs new file mode 100644 index 000000000000..4b1d0bee1862 --- /dev/null +++ b/rust/sglang-server/src/api_server/native_api.rs @@ -0,0 +1,602 @@ +//! The native SGLang data-plane endpoints: `/generate` (submit a request, then +//! either fold egress frames to one unary JSON response or relay them as SSE +//! `data: {json}` … `[DONE]`, byte-compatible with Python +//! `http_server.generate_request`) and `/health` + `/health_generate` (which +//! round-trip a 1-token generate probe). Frame shaping (`meta_info`, logprob +//! tuples, cumulative vs incremental streams) lives here, as does +//! generate-request submission (`submit`); the shared `AppState` lives in the +//! parent `api_server` module. + +use std::convert::Infallible; + +use axum::{ + Json, Router, + extract::State, + extract::rejection::JsonRejection, + http::StatusCode, + response::{ + IntoResponse, Response, + sse::{Event, Sse}, + }, + routing::{get, post}, +}; +use tokio::sync::mpsc; + +use super::AppState; +use super::frame::{ + OutputAccumulator, cumulative_frame_string, error_value, frame_value, stream_frame_string, + tag_value, +}; +use super::guard::AbortGuard; +use super::submit::{pre_submit_error, submit}; +use crate::environ::env_bool; +use crate::ids::Rid; +use crate::message::{EgressItem, GenerateBody, GenerateRequest, RequestKind, SamplingParams}; + +/// The routes this module owns, mounted by `api_server::serve`. +pub(super) fn routes() -> Router { + Router::new() + .route("/generate", post(generate)) + .merge(health_routes()) +} + +/// `/health` + `/health_generate`. Both env knobs are resolved ONCE here, at +/// router build (server startup) — changing them on a live process needs a +/// restart. The deep-probe handler is built once with +/// `SGLANG_HEALTH_CHECK_TIMEOUT` frozen in and serves `/health_generate` +/// always; `SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION` (default true, mirroring +/// Python) decides whether `/health` shares it or is a plain 200 (routing the +/// request already proves the frontend is up). +fn health_routes() -> Router { + let timeout = + std::time::Duration::from_secs(crate::environ::env_u64("SGLANG_HEALTH_CHECK_TIMEOUT", 20)); + let probe = get(move |state: State| health_generate(state, timeout)); + let health = if env_bool("SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION", true) { + probe.clone() + } else { + get(|| async { StatusCode::OK.into_response() }) + }; + Router::new() + .route("/health", health) + .route("/health_generate", probe) +} + +/// `GET /health_generate` — deep health: confirm the scheduler → detok path is +/// producing output. 200 iff the egress heartbeat advances within `timeout` +/// (from `SGLANG_HEALTH_CHECK_TIMEOUT`, frozen at router build), else 503. +/// (`/health` uses the same handler when its env gate is on.) +/// +/// Fires a pre-tokenized 1-token probe (`input_ids = [0]`, skips the tokenizer) so +/// an idle pipeline produces a frame, then watches the *global* +/// [`AppState::egress_activity`] counter (not the probe's own rid) — so a busy +/// server passes immediately and a backlog never false-503s (the analogue of +/// Python's `last_receive_tstamp`). The `HEALTH_CHECK` skip + `http_worker_ipc` +/// ack are irrelevant here: this single-process server owns the egress ring. +async fn health_generate(State(state): State, timeout: std::time::Duration) -> Response { + let baseline = state + .egress_activity + .load(std::sync::atomic::Ordering::Relaxed); + + // Fire the probe (the heartbeat is the signal, not its own response). A busy + // scheduler skips it with no terminal frame, so its detok registration is + // cleaned up only by the `AbortGuard` below. + let probe = GenerateRequest { + // The `HEALTH_CHECK_` rid form + rid: Rid::new_health_check(), + input_ids: Some(vec![0]), + // One greedy token: the cheapest round-trip that still produces a frame. + sampling_params: SamplingParams { + max_new_tokens: Some(1), + temperature: 0.0, + ..Default::default() + }, + stream: false, + ..Default::default() + }; + let (rid, _keepalive) = + match submit(&state, RequestKind::Generate(Box::new(probe)), false).await { + // Hold the receiver so the probe's sink stays open until it completes. + Ok(v) => v, + Err(resp) => return resp, + }; + // Deregister on drop (never disarmed): a busy-skipped probe has no terminal + // frame, so without this abort it leaks one detok entry per call. + let _abort_guard = AbortGuard::new(state.senders.clone(), rid); + + // Watch the heartbeat advance (timeout frozen at router build, default 20s). + let deadline = tokio::time::Instant::now() + timeout; + loop { + if state + .egress_activity + .load(std::sync::atomic::Ordering::Relaxed) + != baseline + { + return StatusCode::OK.into_response(); + } + if tokio::time::Instant::now() >= deadline { + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } +} + +/// `POST /generate` — the native generation endpoint. Splits the body +/// into per-request payloads (a scalar body → one, a list body → a batch) and +/// dispatches to the single or batch path; a malformed body is a 400 before +/// anything reaches the scheduler. +/// +/// The body is extracted as a `Result` so a deserialization failure is answered +/// with **400** (Python's status for a bad request) carrying serde's field-level +/// message, instead of axum's default 422. +async fn generate( + State(state): State, + body: Result, JsonRejection>, +) -> Response { + let body = match body { + Ok(Json(body)) => body, + // A body that fails to parse has no readable `stream` flag, so this one + // can only answer unary — as Python's does (FastAPI rejects before its + // handler runs). + Err(rejection) => { + return pre_submit_error(StatusCode::BAD_REQUEST, &rejection.body_text(), false); + } + }; + let stream = body.stream; + // Fan `text`/`input_ids`/`sampling_params` (scalar or list) into per-request + // payloads. `is_batch` = list form → the response is a JSON array. + let (payloads, is_batch) = match body.into_requests() { + Ok(v) => v, + // The error carries its own status (a bad batch is `Validation` → 400). + Err(e) => { + let code = StatusCode::from_u16(e.http_status()).unwrap_or(StatusCode::BAD_REQUEST); + return pre_submit_error(code, &e.to_string(), stream); + } + }; + if !is_batch { + // `into_requests` guarantees exactly one payload for a non-batch body. + let payload = payloads + .into_iter() + .next() + .expect("into_requests yields >=1 payload"); + generate_single(&state, payload, stream).await + } else { + generate_batch(&state, payloads, stream).await + } +} + +/// Answer an error raised *before* anything was submitted, in the shape the client +/// asked for. +/// +/// A single (non-batched) `/generate`: submit one request, then either stream its +/// SSE frames or fold to one unary response. +async fn generate_single(state: &AppState, req: GenerateRequest, stream: bool) -> Response { + // `return_text_in_logprobs` is decoded on the detok shard into `*_txt`, so + // `frame_value` just reads them — no tokenizer needed here. + let (rid_str, mut rx) = match submit(state, RequestKind::Generate(Box::new(req)), stream).await + { + Ok(v) => v, + Err(resp) => return resp, + }; + // Abort on client disconnect: the guard fires when dropped before the request + // finishes (axum drops the handler/SSE stream). Disarmed on a natural terminal. + // `rid_str` is the response `meta_info.id`, reused for every frame. + let mut guard = AbortGuard::new(state.senders.clone(), rid_str.clone()); + // Cumulative frames (SGLang default) vs per-step deltas. + let incremental = state.server_args.incremental_streaming_output; + + if stream { + // A single request is a 1-element batch without the `index` field — reuse + // the same stream so the frame/abort/truncation logic lives in one place. + use futures::StreamExt; + let s = generation_event_stream(vec![(rid_str, rx)], guard, incremental, false) + .map(|data| Ok::<_, Infallible>(Event::default().data(data))); + Sse::new(s).into_response() + } else { + // Unary: fold to the terminal, respond once. Disarm only on a real terminal + // (a truncation leaves the guard armed so the scheduler work is aborted). + let (status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await; + if terminal { + guard.disarm(&rid_str); + } + (status, Json(value)).into_response() + } +} + +/// Fold a unary request to its terminal → (HTTP status, result/`error` JSON, saw-terminal); +/// `false` = truncation, caller keeps the abort guard armed. Shared by single + batch. +async fn drain_unary( + rx: &mut mpsc::Receiver, + rid_str: &str, +) -> (StatusCode, serde_json::Value, bool) { + let mut acc = OutputAccumulator::default(); + while let Some(item) = rx.recv().await { + match item { + EgressItem::Frame(out) => acc.fold(&out), + EgressItem::Done(out) => { + acc.fold(&out); + let final_out = acc.into_output(); + // A validation abort carries its own HTTP status + diagnostic. + if let Some((code, message)) = final_out + .finish_reason + .as_ref() + .and_then(|f| f.abort_status()) + { + let status = + StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + return (status, error_value(code, message), true); + } + return (StatusCode::OK, frame_value(&final_out, rid_str), true); + } + EgressItem::Error(e) => { + let code = e.http_status(); + let status = + StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + return (status, error_value(code, &e.to_string()), true); + } + EgressItem::Control(_) => continue, // never on `/generate` + } + } + // Sender dropped without a terminal item: the shard dropped this request (a + // truncation — a client disconnect would have dropped the handler future). + ( + StatusCode::INTERNAL_SERVER_ERROR, + error_value(500, "response truncated before completion"), + false, + ) +} + +/// Batch `/generate`: submit all sub-requests first (scheduler runs them together), +/// then either (unary) drain each in order into a JSON array, or (streaming) +/// multiplex their streams into one SSE response, each frame carrying its `index`. +/// One [`AbortGuard`] covers the batch. A failed unary item is its own +/// `{ "error": … }` entry; the batch response is 200. +async fn generate_batch( + state: &AppState, + requests: Vec, + stream: bool, +) -> Response { + // No cross-item rid collision to worry about: `into_requests` rejected duplicate + // rids within this batch, and `Rid::from_client` made each one unique against + // every other in-flight request. + let mut guard = AbortGuard::new_empty(state.senders.clone()); + let mut receivers = Vec::with_capacity(requests.len()); + for req in requests { + match submit(state, RequestKind::Generate(Box::new(req)), stream).await { + Ok((rid, rx)) => { + guard.arm(rid.clone()); + receivers.push((rid, rx)); + } + Err(resp) => return resp, + } + } + + if stream { + // Multiplex the N streams (mirrors the Python `_handle_batch_request` path); + // `guard` moves into the stream so a disconnect aborts what's unfinished. + use futures::StreamExt; + let incremental = state.server_args.incremental_streaming_output; + let s = generation_event_stream(receivers, guard, incremental, true) + .map(|data| Ok::<_, Infallible>(Event::default().data(data))); + Sse::new(s).into_response() + } else { + // Unary: drain each in order (already all submitted, so they run together). + let mut results = Vec::with_capacity(receivers.len()); + for (rid_str, mut rx) in receivers { + let (_status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await; + if terminal { + guard.disarm(&rid_str); + } + results.push(value); + } + (StatusCode::OK, Json(serde_json::Value::Array(results))).into_response() + } +} + +/// Await the next item from `rx`, then drain whatever queued behind it (so the caller +/// can coalesce a backlog, as Python's `state.out_list` does), handing the receiver +/// back for `FuturesUnordered` to re-poll. Empty result = channel closed. +async fn recv_indexed( + index: usize, + mut rx: mpsc::Receiver, +) -> (usize, mpsc::Receiver, Vec) { + let mut items = Vec::new(); + match rx.recv().await { + Some(item) => items.push(item), + None => return (index, rx, items), // closed + } + while let Ok(item) = rx.try_recv() { + items.push(item); + } + (index, rx, items) +} + +/// Multiplex `receivers` (one per request) into SSE `data` strings + a final `[DONE]`; +/// `with_index` tags each frame (batch only), `incremental` = delta vs cumulative, +/// `guard` aborts unfinished on drop. +fn generation_event_stream( + receivers: Vec<(Rid, mpsc::Receiver)>, + mut guard: AbortGuard, + incremental: bool, + with_index: bool, +) -> impl futures::Stream { + async_stream::stream! { + use futures::StreamExt; + + let n = receivers.len(); + let rid_strs: Vec = receivers.iter().map(|(rid, _)| rid.clone()).collect(); + let mut accs: Vec = + (0..n).map(|_| OutputAccumulator::default()).collect(); + + // Batch position, tagged onto every frame (a single request omits it). + let idx = |i: usize| with_index.then_some(i); + + // Poll all receivers concurrently; re-arm a receiver's future after each + // non-terminal frame so its stream keeps flowing. + let mut futs = futures::stream::FuturesUnordered::new(); + for (i, (_, rx)) in receivers.into_iter().enumerate() { + futs.push(recv_indexed(i, rx)); + } + + while let Some((i, rx, items)) = futs.next().await { + if items.is_empty() { + // Channel closed with no terminal → truncation for this item; + // leave its rid armed so the scheduler work is aborted. + yield tag_value(error_value(500, "response truncated before completion"), idx(i)); + continue; + } + + // Cumulative frames supersede one another, so a drained backlog collapses + // to its last (Python's `out_list[-1]`); deltas can't be dropped. + let mut coalesced = false; // a cumulative frame is pending + let mut terminal = None; // (finish_reason) of a `Done` in this batch + let mut failed = None; // an `Error` in this batch + + for item in items { + match item { + EgressItem::Frame(out) => { + accs[i].fold(&out); + if incremental { + yield stream_frame_string(out, &accs[i], true, rid_strs[i].client_facing(), idx(i)); + } else { + coalesced = true; + } + } + EgressItem::Done(out) => { + accs[i].fold(&out); + terminal = Some(out); + } + EgressItem::Error(e) => failed = Some(e), + EgressItem::Control(_) => {} // never on /generate + } + } + + if let Some(e) = failed { + yield tag_value(error_value(e.http_status(), &e.to_string()), idx(i)); + guard.disarm(&rid_strs[i]); + } else if let Some(out) = terminal { + // A validation abort → an error object, not a frame. The final frame + // carries the full cumulative state, so any coalesced ones are moot. + yield match out.finish_reason.as_ref().and_then(|f| f.abort_status()) { + Some((code, message)) => tag_value(error_value(code, message), idx(i)), + None => stream_frame_string(out, &accs[i], incremental, rid_strs[i].client_facing(), idx(i)), + }; + guard.disarm(&rid_strs[i]); // terminal → not re-pushed + } else { + if coalesced { + yield cumulative_frame_string(&accs[i], rid_strs[i].client_facing(), idx(i)); + } + futs.push(recv_indexed(i, rx)); // keep this item flowing + } + } + yield "[DONE]".to_string(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::ChunkEvent; + use crate::tokenizer_manager::Senders; + use futures::StreamExt; + fn senders() -> Senders { + Senders { + tm: flume::unbounded().0, + abort: flume::unbounded().0, + tok: flume::unbounded().0, + detok: vec![], + } + } + + fn frame(rid: u64, text: &str) -> EgressItem { + EgressItem::Frame(ChunkEvent { + rid: Rid::from(rid.to_string()), + text: text.into(), + completion_tokens: 1, + ..Default::default() + }) + } + fn done(rid: u64, text: &str) -> EgressItem { + EgressItem::Done(ChunkEvent { + rid: Rid::from(rid.to_string()), + text: text.into(), + completion_tokens: 1, + // Parsed from the wire map Python emits, not a hand-built enum. + finish_reason: Some( + serde_json::from_value(serde_json::json!({"type": "length", "length": 1})) + .expect("finish reason must parse"), + ), + ..Default::default() + }) + } + fn parse(s: &str) -> serde_json::Value { + serde_json::from_str(s).expect("frame is JSON") + } + + /// Two sub-requests' frames interleave into one stream, each tagged with its + /// batch `index`; text accumulates per item; `[DONE]` comes only after both + /// terminate, then the stream ends. + #[tokio::test] + async fn interleaves_indexes_and_accumulates() { + let (tx0, rx0) = mpsc::channel(8); + let (tx1, rx1) = mpsc::channel(8); + let receivers = vec![("10".into(), rx0), ("11".into(), rx1)]; + let stream = + generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true); + futures::pin_mut!(stream); + + // Drive deterministically: exactly one channel has data before each poll. + tx0.send(frame(10, "a")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["index"], 0); + assert_eq!(v["text"], "a"); + + tx1.send(frame(11, "b")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["index"], 1); + assert_eq!(v["text"], "b"); + + tx0.send(done(10, "!")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["index"], 0); + assert_eq!(v["text"], "a!", "cumulative per item"); + assert_eq!(v["meta_info"]["finish_reason"]["type"], "length"); + + tx1.send(done(11, "?")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["index"], 1); + assert_eq!(v["text"], "b?"); + + assert_eq!(stream.next().await.unwrap(), "[DONE]"); + assert!(stream.next().await.is_none()); + } + + /// A per-item error is surfaced with its `index` and doesn't end the batch; + /// `[DONE]` still waits for the other item. + #[tokio::test] + async fn per_item_error_carries_index() { + let (tx0, rx0) = mpsc::channel(8); + let (tx1, rx1) = mpsc::channel(8); + let receivers = vec![("10".into(), rx0), ("11".into(), rx1)]; + let stream = + generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true); + futures::pin_mut!(stream); + + tx0.send(EgressItem::Error(crate::error::Error::Validation( + "bad".into(), + ))) + .await + .unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["index"], 0); + assert_eq!(v["error"]["code"], 400); + + tx1.send(done(11, "ok")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["index"], 1); + + assert_eq!(stream.next().await.unwrap(), "[DONE]"); + } + + /// `incremental=true`: each frame carries this step's **delta** text/output_ids, + /// but `meta_info.completion_tokens` stays cumulative (matching Python). + #[tokio::test] + async fn incremental_emits_deltas_with_cumulative_count() { + let (tx, rx) = mpsc::channel(8); + let receivers = vec![("10".into(), rx)]; + let stream = + generation_event_stream(receivers, AbortGuard::new_empty(senders()), true, true); + futures::pin_mut!(stream); + + tx.send(frame(10, "Hello")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["text"], "Hello"); + assert_eq!(v["meta_info"]["completion_tokens"], 1); + + tx.send(frame(10, " world")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["text"], " world", "delta, not cumulative 'Hello world'"); + assert_eq!( + v["meta_info"]["completion_tokens"], 2, + "count stays cumulative" + ); + + tx.send(done(10, "!")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["text"], "!"); + assert_eq!(v["meta_info"]["completion_tokens"], 3); + assert_eq!(v["meta_info"]["finish_reason"]["type"], "length"); + + assert_eq!(stream.next().await.unwrap(), "[DONE]"); + } + + /// The single-request shape (`with_index=false`, one receiver) omits the + /// `index` field entirely, and still terminates with `[DONE]`. + #[tokio::test] + async fn single_shape_omits_index() { + let (tx, rx) = mpsc::channel(8); + let receivers = vec![("10".into(), rx)]; + let stream = + generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, false); + futures::pin_mut!(stream); + + tx.send(done(10, "hi")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["text"], "hi"); + assert!(v.get("index").is_none(), "single response has no index"); + + assert_eq!(stream.next().await.unwrap(), "[DONE]"); + } + + /// A backlog of cumulative chunks collapses to a single frame carrying the latest + /// state — each cumulative frame supersedes the last, so emitting the intermediate + /// ones ships the full O(T) payload again for nothing. Mirrors the Python waiter's + /// `out = out_list[-1]`. This is the whole point of draining in `recv_indexed`. + #[tokio::test] + async fn cumulative_backlog_coalesces_to_latest() { + let (tx, rx) = mpsc::channel(8); + let receivers = vec![("10".into(), rx)]; + let stream = + generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, false); + futures::pin_mut!(stream); + + // Three chunks queued before the stream is ever polled (a client falling behind). + tx.send(frame(10, "a")).await.unwrap(); + tx.send(frame(10, "b")).await.unwrap(); + tx.send(frame(10, "c")).await.unwrap(); + + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["text"], "abc", "one frame, full cumulative text"); + assert_eq!(v["meta_info"]["completion_tokens"], 3, "no tokens lost"); + + // The terminal frame still carries everything, and only then does [DONE] land. + tx.send(done(10, "!")).await.unwrap(); + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["text"], "abc!"); + assert_eq!(v["meta_info"]["finish_reason"]["type"], "length"); + assert_eq!(stream.next().await.unwrap(), "[DONE]"); + } + + /// Incremental frames are *deltas*, so a backlog must emit every one — dropping + /// any would silently lose tokens. Only the cumulative protocol may coalesce. + #[tokio::test] + async fn incremental_backlog_emits_every_delta() { + let (tx, rx) = mpsc::channel(8); + let receivers = vec![("10".into(), rx)]; + let stream = + generation_event_stream(receivers, AbortGuard::new_empty(senders()), true, false); + futures::pin_mut!(stream); + + tx.send(frame(10, "a")).await.unwrap(); + tx.send(frame(10, "b")).await.unwrap(); + tx.send(frame(10, "c")).await.unwrap(); + + for (n, expect) in [(1, "a"), (2, "b"), (3, "c")] { + let v = parse(&stream.next().await.unwrap()); + assert_eq!(v["text"], expect, "delta {n} must not be dropped"); + assert_eq!( + v["meta_info"]["completion_tokens"], n, + "count stays cumulative" + ); + } + } +} diff --git a/rust/sglang-server/src/api_server/openai.rs b/rust/sglang-server/src/api_server/openai.rs new file mode 100644 index 000000000000..22678f6d335c --- /dev/null +++ b/rust/sglang-server/src/api_server/openai.rs @@ -0,0 +1,55 @@ +//! OpenAI-compatible endpoints: `/v1/completions`, `/v1/chat/completions`, and +//! `/v1/models`. Each runs the same tokenize→generate→detok pipeline as +//! `/generate` and shapes the neutral [`ChunkEvent`] delta into OpenAI types +//! (`dynamo-protocols`), with chat-template rendering (`dynamo-renderer`) and +//! reasoning / tool-call parsing (`dynamo-parsers`). +//! +//! Mounted on the shared [`AppState`](super::AppState) by the parent +//! `api_server` module; the submit machinery and control plane live there. + +use axum::{ + extract::State, + http::StatusCode, + response::{IntoResponse, Response}, +}; + +use axum::{Router, routing::get}; + +use super::AppState; + +/// The routes this module owns, mounted by `api_server::serve`. +pub(super) fn routes() -> Router { + // `/v1/models` is OpenAI-compatible; completions/chat land here too. + Router::new().route("/v1/models", get(available_models)) +} + +/// `GET /v1/models` — OpenAI-compatible model list. Served from `server_args`; +/// no scheduler round-trip. Mirrors `http_server.available_models`. +/// +/// TODO(v1/models): when `--enable-lora`, append a `ModelCard` per loaded LoRA +/// adapter (`id=lora_name, root=lora_path, parent=served_model_name, +/// max_model_len=None`). Adapters load/unload at runtime, so that part needs a +/// control-request query to the scheduler's LoRA registry. +async fn available_models(State(state): State) -> Response { + let created = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + let name = &state.server_args.served_model_name; + let base = serde_json::json!({ + "id": name, + "object": "model", + "created": created, + "owned_by": "sglang", + "root": name, + "parent": serde_json::Value::Null, + "max_model_len": state.server_args.model_config.context_len, + }); + let list = serde_json::json!({ "object": "list", "data": [base] }); + ( + StatusCode::OK, + [("content-type", "application/json")], + serde_json::to_vec(&list).unwrap_or_default(), + ) + .into_response() +} diff --git a/rust/sglang-server/src/api_server/submit.rs b/rust/sglang-server/src/api_server/submit.rs new file mode 100644 index 000000000000..80c85c04839a --- /dev/null +++ b/rust/sglang-server/src/api_server/submit.rs @@ -0,0 +1,126 @@ +//! Request submission into the ingress pipeline, shared by every endpoint +//! module: mint the client-visible rid (uuid hex, Python-parity), build the +//! `Request`, and hand it to the TM with an egress receiver for the response. + +use std::convert::Infallible; + +use axum::{ + Json, + http::StatusCode, + response::{ + IntoResponse, Response, + sse::{Event, Sse}, + }, +}; +use tokio::sync::mpsc; + +use super::AppState; +use super::frame::error_value; +use crate::fsm::RequestState; +use crate::ids::Rid; +use crate::message::{EgressItem, EgressSink, Request, RequestKind}; +use crate::tokenizer_manager::TmEvent; + +/// Submit one request; returns the rid, its hashed routing key, and the egress +/// receiver. Every request arrives with its final rid — a generate request from +/// `into_requests` (or the `HEALTH_CHECK_` the health probe sets), a +/// control request from its constructor — so this only echoes it back. +pub(super) async fn submit( + state: &AppState, + kind: RequestKind, + // `stream`: the client is reading an SSE stream, so it expects 200 plus an + // error frame rather than a 4xx — same rule `pre_submit_error` applies + // everywhere else. + stream: bool, +) -> Result<(Rid, mpsc::Receiver), Response> { + let rid = match &kind { + // Generate rids are already final: `GenerateBody::into_requests` normalized the + // client's, or minted one. Control requests have no client-facing rid. + RequestKind::Generate(g) => g.rid.clone(), + RequestKind::Control(c) => c.rid().into(), + }; + // Two in-flight requests can name the same client rid, but they cannot share a + // `Rid`: `into_requests` built each through `Rid::from_client`, which appends a + // uniquifier. So nothing here needs to check for a collision — the detok table + // key is unique by construction, and `client_facing` restores what the client + // sent for `meta_info.id`. + // Async-aware send so a full TM inbox yields (backpressure) instead of parking + // a thread; Err only when the inbox is closed (shutdown). + let (tx, rx) = mpsc::channel::(state.egress_buf); + let request = Request { + rid: rid.clone(), + state: RequestState::Received, + sink: EgressSink::Local(tx), + kind, + }; + match state.senders.tm.send_async(TmEvent::Ingress(request)).await { + Ok(()) => Ok((rid, rx)), + // `SendError` has a single meaning — the channel is disconnected. + Err(_) => { + tracing::error!(%rid, "tm inbox closed; request rejected"); + // Return 503 so the client can retry. + Err(pre_submit_error( + StatusCode::SERVICE_UNAVAILABLE, + "service unavailable", + stream, + )) + } + } +} + +/// Shape an error that occurs *before* (or instead of) a successful submit into a +/// client response. Two parity points with Python's `generate_request`. The body +/// is the same `{"error": {...}}` object every other path emits — not bare text, +/// which a client parsing JSON chokes on. And a streaming request gets 200 plus +/// one SSE error frame and `[DONE]`, not a 4xx: the client has already committed +/// to reading a stream, and Python answers it inside `stream_results()`. +pub(super) fn pre_submit_error(code: StatusCode, message: &str, stream: bool) -> Response { + let body = error_value(code.as_u16(), message); + if !stream { + return (code, Json(body)).into_response(); + } + let frames = [body.to_string(), "[DONE]".to_string()]; + Sse::new(futures::stream::iter( + frames.map(|data| Ok::<_, Infallible>(Event::default().data(data))), + )) + .into_response() +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Unary pre-submit errors are a 4xx/5xx with a JSON `{"error":...}` body; + /// streaming ones are 200 + an SSE error frame + `[DONE]`, because Python + /// answers from inside `stream_results()` once the stream is committed. + #[tokio::test] + async fn pre_submit_errors_match_python_shape() { + let unary = pre_submit_error(StatusCode::BAD_REQUEST, "bad input", false); + assert_eq!(unary.status(), StatusCode::BAD_REQUEST); + let body = axum::body::to_bytes(unary.into_body(), 64 * 1024) + .await + .unwrap(); + let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON body"); + assert_eq!(v["error"]["message"], "bad input"); + assert_eq!(v["error"]["code"], 400); + + let streamed = pre_submit_error(StatusCode::BAD_REQUEST, "bad input", true); + assert_eq!( + streamed.status(), + StatusCode::OK, + "the stream itself is 200" + ); + let body = axum::body::to_bytes(streamed.into_body(), 64 * 1024) + .await + .unwrap(); + let text = String::from_utf8(body.to_vec()).unwrap(); + assert!( + text.contains(r#""code":400"#), + "carries the status in-band: {text}" + ); + assert!( + text.trim_end().ends_with("data: [DONE]"), + "terminated: {text}" + ); + } +} diff --git a/rust/sglang-server/src/detokenizer.rs b/rust/sglang-server/src/detokenizer.rs new file mode 100644 index 000000000000..d0a8009946aa --- /dev/null +++ b/rust/sglang-server/src/detokenizer.rs @@ -0,0 +1,597 @@ +//! Detokenizer shards — CPU-bound, one pinned thread per shard. +//! +//! Each shard owns a *local* `rid -> DetokState` map. There is no lock: a given +//! rid is routed to exactly one shard (by `Rid::shard`) for both its +//! `Register` and all its `Chunk`s, so the map has a single accessor. +//! +//! The hash PARTITIONS, the rid IDENTIFIES. Keying the map by the hash meant two +//! distinct rids that happened to collide became one entry: `Register` evicted the +//! first client's sink and their tokens were then written to the second client's +//! connection. Chunks carry the rid string (moved out of the frame header, which +//! owns it and would otherwise drop it), so a collision now only co-locates. +//! +//! Real detokenization uses dynamo-tokenizers' `DecodeStream`, a stateful +//! incremental decoder (TGI/vLLM-style: it buffers partial UTF-8 / byte-fallback +//! tokens and only emits text once a valid boundary is reached). Each request +//! gets its own `DecodeStream`. When no tokenizer is configured (or +//! `skip_tokenizer_init` is set) the backend is `Skip`: no decoding, the raw +//! `output_ids` are emitted instead of text. +//! +//! Per-chunk egress flow (no FSM state change inside Streaming): +//! ChunkEvent{finish:None} -> step ids -> delta -> Server frame +//! ChunkEvent{finish:Some} -> step ids -> delta -> final frame + +use std::collections::HashMap; + +use crate::error::Error; +use crate::fsm::{Event, RequestState}; +use crate::ids::Rid; +use crate::message::DetokMsg; +use crate::message::{ChunkEvent, EgressItem, EgressSink, Matched, SinkError, TokenIds}; +use crate::runtime::Runnable; +use crate::tokenizer_manager::AbortSource; + +/// Default for `skip_special_tokens` (SGLang's SamplingParams default). The +/// per-request value isn't available on the egress side yet; see the note in +/// `DetokenizerBackend::new_decoder`. +const SKIP_SPECIAL_TOKENS: bool = true; + +/// Per-request incremental decoder. `step` feeds the new token ids for one chunk +/// and returns the newly decoded text delta (empty if the ids only produced a +/// partial/incomplete multi-byte sequence that needs more tokens). +pub trait StreamDecoder: Send { + fn step(&mut self, token_ids: &[i32]) -> Result; +} + +/// Real decoder wrapping a dynamo-tokenizers `DecodeStream`. +struct DynamoDecoder { + stream: dynamo_tokenizers::DecodeStream, +} + +impl StreamDecoder for DynamoDecoder { + fn step(&mut self, token_ids: &[i32]) -> Result { + let mut out = String::new(); + for &id in token_ids { + if let Some(chunk) = self + .stream + .step(id as u32) + .map_err(|e| Error::Detokenize(e.to_string()))? + { + out.push_str(&chunk); + } + } + Ok(out) + } +} + +/// Shard-wide detok backend. Cloned per shard; mints a fresh per-request decoder +/// on each `Register`. +#[derive(Clone)] +pub enum DetokenizerBackend { + Dynamo(dynamo_tokenizers::Tokenizer), + /// No decoding at all — the shard emits each chunk's raw output token ids as + /// `output_ids` (no `DecodeStream`, no accumulation). Used for + /// `skip_tokenizer_init` and when no tokenizer is configured. + Skip, +} + +impl DetokenizerBackend { + /// Mint a per-request decoder, or `None` in skip mode (the shard passes the + /// token ids through untouched instead of decoding text). + fn new_decoder(&self) -> Option> { + match self { + // NOTE: the stream is seeded with an empty prompt context, which is + // correct for the common case. Seeding with the prompt's trailing + // tokens (for perfect first-token spacing) would require Register to + // carry input_ids — deferred. + DetokenizerBackend::Dynamo(t) => Some(Box::new(DynamoDecoder { + stream: t.decode_stream(&[], SKIP_SPECIAL_TOKENS), + })), + DetokenizerBackend::Skip => None, + } + } + + /// Decode each logprob token id to its own text (one id at a time, matching + /// Python's `batch_decode([[id] for id in ids])`). Runs on this CPU-bound + /// shard, not the api-server I/O threads. `Skip` mode (no tokenizer) yields + /// no text, so the `[logprob, token_id, text]` tuple's text slot stays null. + fn decode_logprob_texts(&self, idxs: &[i32]) -> Vec { + match self { + DetokenizerBackend::Dynamo(t) => idxs + .iter() + .map(|&id| { + t.decode(&[id as u32], false) + .map(String::from) + .unwrap_or_default() + }) + .collect(), + DetokenizerBackend::Skip => Vec::new(), + } + } +} + +struct DetokState { + sink: EgressSink, + /// `return_text_in_logprobs`: whether to decode this request's logprob token + /// ids to text (in this shard) for the `[logprob, token_id, text]` tuples. + decode_logprob_text: bool, + /// `SamplingParams.no_stop_trim`: keep the matched stop in the output. Default + /// (`false`) trims it off the final chunk (see [`trim_stop_str`]). + no_stop_trim: bool, + /// Per-request incremental decoder; `None` in `skip_tokenizer_init` mode. + /// This is the *only* per-request accumulation the shard keeps: the decoder's + /// internal byte/UTF-8 buffer. Decoded **text deltas** are emitted per chunk + /// (no cumulative buffer here) — the api-server's drain loop reassembles the + /// cumulative view where a consumer needs it (every unary response and the + /// cumulative SGLang `/generate` stream); OpenAI streaming forwards deltas. + decoder: Option>, + /// Egress half of the lifecycle FSM. Lives here because the ingress + /// `Request` (and its FSM) was handed to the scheduler when queued; the + /// shard is the sole owner of the request's egress state, so no lock. + fsm: RequestState, +} + +/// One detokenizer shard: owns a *local* `rid -> DetokState` map (single accessor, +/// no lock) and the egress backend. Spawned (pinned) per shard as a [`Runnable`]; +/// a given rid is routed to exactly one shard. +pub struct DetokenizerWorker { + shard: usize, + rx: flume::Receiver, + backend: DetokenizerBackend, + /// Unbounded abort lane, used to abort a request the shard had to drop + /// (client backpressure) so the scheduler stops generating for it. + abort: flume::Sender, +} + +impl DetokenizerWorker { + pub fn new( + shard: usize, + rx: flume::Receiver, + backend: DetokenizerBackend, + abort: flume::Sender, + ) -> Self { + Self { + shard, + rx, + backend, + abort, + } + } +} + +impl Runnable for DetokenizerWorker { + fn run(self) { + let mut table: HashMap = HashMap::new(); + tracing::debug!(shard = self.shard, "detokenizer worker started"); + + // Plain `recv`: exits when the `DetokMsg` channel closes (every `Senders` + // clone gone). On shutdown that happens once the API runtime drop cancels + // in-flight handlers (their `AbortGuard`s release the last clones) and + // tm-ingress/tm-egress exit — no shutdown signal needed here. + while let Ok(msg) = self.rx.recv() { + match msg { + DetokMsg::Register { + rid, + sink, + decode_logprob_text, + no_stop_trim, + } => { + table.insert( + rid.clone(), + DetokState { + sink, + decode_logprob_text, + no_stop_trim, + decoder: self.backend.new_decoder(), + // Registered == handed to the scheduler == Queued. + fsm: RequestState::Queued, + }, + ); + } + // One decode step's chunks for this shard, batched by tm-egress. + DetokMsg::Chunks(evs) => { + for ev in evs { + handle_chunk(&mut table, ev, &self.backend, &self.abort); + } + } + DetokMsg::Result { rid, payload } => handle_result(&mut table, &rid, payload), + DetokMsg::Fail { rid, message } => { + handle_fail(&mut table, &rid, message, &self.abort) + } + DetokMsg::Deregister { rid } => { + table.remove(&rid); + } + } + } + } +} + +/// Control-request result: deliver the JSON payload to the sink verbatim as a +/// single `Done` frame — no detokenization, no streaming. +fn handle_result(table: &mut HashMap, rid: &Rid, payload: bytes::Bytes) { + if let Some(mut st) = table.remove(rid) { + let _ = st.sink.try_send(EgressItem::Control(payload)); + // Egress FSM: a control request goes straight to Completed (no Streaming + // / Finalizing states — single response, never streamed). + st.fsm = RequestState::Completed; + } +} + +/// Terminal per-request failure (bad request header): send an `Error` to the sink +/// (the api-server turns it into an HTTP 400) and drop the request. +/// Terminal per-request failure. `Internal` (500), not `Validation` (400): the +/// producers of this message are server faults — a malformed scheduler output +/// frame — not bad client input. Also aborts the request on the scheduler, which +/// otherwise keeps generating tokens for a connection that will never read them. +fn handle_fail( + table: &mut HashMap, + rid: &Rid, + message: String, + abort: &flume::Sender, +) { + if let Some(mut st) = table.remove(rid) { + // Abort first: `try_send` on the sink can release the handler, which frees + // the rid for reuse (same ordering hazard as the disconnect path). + let _ = abort.send(AbortSource::Detok(rid.clone())); + let _ = st + .sink + .try_send(EgressItem::Error(Error::Internal(message))); + st.fsm = RequestState::Completed; + } +} + +fn handle_chunk( + table: &mut HashMap, + mut ev: ChunkEvent, + backend: &DetokenizerBackend, + abort: &flume::Sender, +) { + // Copied once: `ev` is moved into the sink below, but the rid is still + // needed to look the request up and to remove it. + let rid = ev.rid.clone(); + + let Some(st) = table.get_mut(&rid) else { + // Late chunk after completion/abort — drop. + return; + }; + let decode_logprob_text = st.decode_logprob_text; + let no_stop_trim = st.no_stop_trim; + + // Queued → Streaming on the first chunk (the scheduler picked it). + if matches!(st.fsm, RequestState::Queued) { + let _ = st.fsm.apply(Event::SchedulerPicked); + } + + let finished = ev.finish_reason.is_some(); + // Matched-stop trim (Python `trim_matched_stop`): the final chunk's finish + // reason names the stop it matched — a stop STRING or a stop TOKEN id. By + // default that stop is removed from the output; `no_stop_trim` keeps it. + let matched = finished + .then(|| ev.finish_reason.as_ref().and_then(|fr| fr.matched())) + .flatten() + .cloned(); + // Count generated tokens (incl. a matched stop token) *before* trimming. + let n_tok = ev.token_ids.len() as u64; + // Stop TOKEN: drop it before decode, so it reaches neither `text` nor `output_ids`. + trim_stop_token(&mut ev.token_ids, &matched, no_stop_trim); + + // Fully incremental: decode just this chunk's delta. `token_ids` stays in the + // event — it's ALSO surfaced as the `/generate` response's `output_ids` (the + // Python server returns them by default alongside `text`), in both normal and + // `skip_tokenizer_init` mode. Nothing cumulative is kept here — the api-server's + // drain loop reassembles it where needed. + let mut delta_text = match &mut st.decoder { + Some(decoder) => match decoder.step(&ev.token_ids) { + Ok(delta) => delta, + Err(e) => { + // Abort too: this is terminal for the request, and without it the + // scheduler keeps generating for a connection that is already gone + // — the other two terminal paths (disconnect, fail) both abort. + let _ = st.fsm.apply(Event::Error(e.clone())); + let _ = abort.send(AbortSource::Detok(rid.clone())); + let _ = st.sink.try_send(EgressItem::Error(e)); + table.remove(&rid); + return; + } + }, + // skip_tokenizer_init: no decode; the token ids pass through in `ev`. + None => String::new(), + }; + // Stop STRING: trim it (and anything after) from the decoded delta's tail. + if let Some(Matched::Str(stop)) = &matched { + trim_stop_str(&mut delta_text, stop, no_stop_trim); + } + + // Streaming → Streaming (finish:false) or Streaming → Finalizing (finish:true). + let _ = st.fsm.apply(Event::Chunk { finish: finished }); + + // `return_text_in_logprobs`: decode each logprob token id to text HERE (this + // CPU-bound shard) rather than on the api-server I/O threads. Flat text columns + // stay parallel to the `idx` buffers, so `sglang_frame` just reads them. Only the + // logprob-carrying frames have an `extras` box; a plain token frame skips this. + if decode_logprob_text && let Some(ex) = ev.extras.as_deref_mut() { + ex.out_lp_txt = backend.decode_logprob_texts(&ex.out_lp_idx); + ex.in_lp_txt = backend.decode_logprob_texts(&ex.in_lp_idx); + ex.out_top_txt = backend.decode_logprob_texts(&ex.out_top_idx); + ex.in_top_txt = backend.decode_logprob_texts(&ex.in_top_idx); + ex.out_tid_txt = backend.decode_logprob_texts(&ex.out_tid_idx); + ex.in_tid_txt = backend.decode_logprob_texts(&ex.in_tid_idx); + } + + // Fill the decode outputs in place; the pre-decode columns (boxed logprobs/hidden, + // token_ids, prompt_tokens, finish_reason) already ride in `ev`. The API handler + // formats this delta (and accumulates for the cumulative view). + ev.text = delta_text; + ev.completion_tokens = n_tok; + + if finished { + // The Done frame *is* the final frame: Finalizing → Completed. + let sent = st.sink.try_send(EgressItem::Done(ev)).is_ok(); + let _ = st.fsm.apply(if sent { + Event::FinalFrameSent + } else { + Event::Disconnect + }); + table.remove(&rid); + } else { + // Every intermediate chunk emits its delta frame. A failed send means the + // client can't receive it — `Closed` (gone) or `Full` (backpressure: not + // reading fast enough). Either way we can't buffer unboundedly, and + // silently dropping the frame would truncate the response and still look + // like success at EOS. So treat both as terminal: drop the request AND + // abort scheduler work for it. + if let Err(e) = st.sink.try_send(EgressItem::Frame(ev)) { + match e { + SinkError::Full => { + tracing::warn!( + rid = %rid, + "detok: sink full; aborting (client backpressure)" + ) + } + SinkError::Closed => { + tracing::debug!(rid = %rid, "detok: sink closed; aborting (client gone)") + } + } + let _ = st.fsm.apply(Event::Disconnect); + // Abort ONLY when the sink is full. `Closed` means the handler future is + // already gone, so its `AbortGuard` has run: it aborted and released the + // rid. A second abort from here is unordered with respect to that + // release, so it lands after a resubmit of the same rid has registered + // and deregisters the NEW request — the cross-wiring the rid registry + // exists to prevent, reached through the one abort producer that + // bypasses the guard's ordering. + if matches!(e, SinkError::Full) { + let _ = abort.send(AbortSource::Detok(rid.clone())); + } + table.remove(&rid); + } + } +} + +/// Drop a matched stop TOKEN from the final chunk (Python `trim_matched_stop`, +/// token branch); `no_stop_trim` / non-token match keeps it. +fn trim_stop_token(token_ids: &mut TokenIds, matched: &Option, no_stop_trim: bool) { + // Token id 0 is NOT a match: Python guards with `if not matched`, and 0 is + // falsy there, so it trims nothing. Trimming on 0 drops a real generated token + // for any model whose stop id happens to be 0. + if !no_stop_trim && matches!(matched, Some(Matched::Token(t)) if *t != 0) { + token_ids.pop(); + } +} + +/// Remove the matched stop string from the decoded final chunk (Python +/// `trim_matched_stop`, string branch); `no_stop_trim` keeps it. Truncates at +/// the FIRST occurrence. +fn trim_stop_str(text: &mut String, stop: &str, no_stop_trim: bool) { + if stop.is_empty() { + return; + } + if let Some(pos) = text.find(stop) { + text.truncate(if no_stop_trim { pos + stop.len() } else { pos }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::mpsc; + + /// A non-terminal chunk that can't be delivered (sink full → client + /// backpressure) drops the request AND aborts scheduler work — it does not + /// silently keep state, which would later read as a clean completion at EOS. + #[test] + fn full_sink_drops_request_and_aborts_scheduler() { + // Capacity-1 sink, pre-filled so the next send hits `Full`. + let (tx, _rx) = mpsc::channel::(1); + tx.try_send(EgressItem::Frame(ChunkEvent::default())) + .unwrap(); + + let mut table = HashMap::new(); + table.insert( + Rid::from("1"), + DetokState { + sink: EgressSink::Local(tx), + decode_logprob_text: false, + no_stop_trim: false, + decoder: None, + fsm: RequestState::Queued, + }, + ); + + let (tm_tx, tm_rx) = flume::unbounded::(); + let ev = ChunkEvent { + rid: Rid::from("1"), + token_ids: vec![5], + ..Default::default() // finish_reason None → non-terminal + }; + handle_chunk(&mut table, ev, &DetokenizerBackend::Skip, &tm_tx); + + // Request removed (no lingering state to be mistaken for success)... + assert!(!table.contains_key(&Rid::from("1"))); + // ...and the scheduler was told to abort it. + assert!(matches!( + tm_rx.try_recv(), + Ok(AbortSource::Detok(rid)) if rid == Rid::from("1") + )); + } + + /// `trim_stop_str` reproduces the base's stop-string semantics: `stop: "3"` on + /// output " 1, 2, 3" yields " 1, 2, " by default and " 1, 2, 3" with + /// `no_stop_trim`. + #[test] + fn trim_stop_str_matches_base() { + let mut t = " 1, 2, 3".to_string(); + trim_stop_str(&mut t, "3", false); + assert_eq!(t, " 1, 2, "); + + let mut t = " 1, 2, 3".to_string(); + trim_stop_str(&mut t, "3", true); + assert_eq!(t, " 1, 2, 3"); + + // Empty / absent stop is a no-op. + let mut t = "abc".to_string(); + trim_stop_str(&mut t, "", false); + assert_eq!(t, "abc"); + + // The stop can occur twice in the final chunk + let mut t = "a STOP b STOP".to_string(); + trim_stop_str(&mut t, "STOP", false); + assert_eq!(t, "a "); + + let mut t = "a STOP b STOP".to_string(); + trim_stop_str(&mut t, "STOP", true); + assert_eq!(t, "a STOP"); + } + + /// Two requests on the SAME shard keep separate entries. This is what a + /// A shard-hash collision now degrades to: the hash partitions, the rid + /// identifies. Keying the table by the hash made colliding rids one entry, so + /// `Register` evicted the first client's sink and their tokens were written to + /// the second client's connection. A single shard forces co-location + /// deterministically, without needing to find a real 64-bit collision. + #[test] + fn co_located_requests_keep_their_own_sinks() { + let (tx_a, mut rx_a) = mpsc::channel::(4); + let (tx_b, mut rx_b) = mpsc::channel::(4); + let mut table = HashMap::new(); + let state = |tx| DetokState { + sink: EgressSink::Local(tx), + decode_logprob_text: false, + no_stop_trim: false, + decoder: None, + fsm: RequestState::Queued, + }; + table.insert(Rid::from("alice"), state(tx_a)); + table.insert(Rid::from("bob"), state(tx_b)); + let (tm_tx, _tm_rx) = flume::unbounded::(); + + let chunk = |rid: &str, id: i32| ChunkEvent { + rid: Rid::from(rid.to_string()), + token_ids: vec![id], + ..Default::default() + }; + handle_chunk( + &mut table, + chunk("alice", 11), + &DetokenizerBackend::Skip, + &tm_tx, + ); + handle_chunk( + &mut table, + chunk("bob", 22), + &DetokenizerBackend::Skip, + &tm_tx, + ); + + let ids = |rx: &mut mpsc::Receiver| match rx.try_recv() { + Ok(EgressItem::Frame(ev)) => ev.token_ids, + other => panic!("expected a frame, got {other:?}"), + }; + assert_eq!( + ids(&mut rx_a), + vec![11], + "alice must not receive bob's tokens" + ); + assert_eq!( + ids(&mut rx_b), + vec![22], + "bob must not receive alice's tokens" + ); + assert_eq!(table.len(), 2, "neither registration evicted the other"); + } + + /// Drive a final (`finish_reason`) chunk through `handle_chunk` in skip mode and + /// return the emitted `Done` event. + fn final_chunk( + no_stop_trim: bool, + finish_reason: serde_json::Value, + ids: Vec, + ) -> ChunkEvent { + let (tx, mut rx) = mpsc::channel::(4); + let mut table = HashMap::new(); + table.insert( + Rid::from("1"), + DetokState { + sink: EgressSink::Local(tx), + decode_logprob_text: false, + no_stop_trim, + decoder: None, // skip mode → output_ids passthrough + fsm: RequestState::Queued, + }, + ); + let (tm_tx, _tm_rx) = flume::unbounded::(); + let ev = ChunkEvent { + rid: Rid::from("1"), + token_ids: ids, + // Parsed from the wire map, so the trim paths are driven by the same + // shape Python emits rather than a hand-built enum. + finish_reason: Some( + serde_json::from_value(finish_reason).expect("finish reason must parse"), + ), + ..Default::default() + }; + handle_chunk(&mut table, ev, &DetokenizerBackend::Skip, &tm_tx); + match rx.try_recv() { + Ok(EgressItem::Done(out)) => out, + other => panic!("expected Done, got {other:?}"), + } + } + + /// Token id 0 is not a match: Python's `trim_matched_stop` guards with + /// `if not matched`, and 0 is falsy there. Trimming on it drops a real + /// generated token for any model whose stop id is 0. + #[test] + fn matched_token_zero_does_not_trim() { + let mut ids = vec![1, 2, 0]; + trim_stop_token(&mut ids, &Some(Matched::Token(0)), false); + assert_eq!(ids, vec![1, 2, 0], "id 0 is not a matched stop"); + // A real stop id still trims. + let mut ids = vec![1, 2, 3]; + trim_stop_token(&mut ids, &Some(Matched::Token(3)), false); + assert_eq!(ids, vec![1, 2]); + } + + /// A matched stop TOKEN is dropped from the surfaced `output_ids` by default + /// (but still counted in `completion_tokens`); `no_stop_trim` keeps it. + #[test] + fn stop_token_trimmed_from_output_ids() { + let fr = serde_json::json!({ "type": "stop", "matched": 3 }); + let out = final_chunk(false, fr.clone(), vec![1, 2, 3]); + assert_eq!(out.token_ids, vec![1, 2], "matched stop token dropped"); + assert_eq!( + out.completion_tokens, 3, + "generated count still includes it" + ); + + let out = final_chunk(true, fr, vec![1, 2, 3]); + assert_eq!(out.token_ids, vec![1, 2, 3], "no_stop_trim keeps it"); + } + + /// A non-stop finish (`length`, no `matched`) never trims. + #[test] + fn length_finish_keeps_all_tokens() { + let fr = serde_json::json!({ "type": "length", "length": 3 }); + let out = final_chunk(false, fr, vec![1, 2, 3]); + assert_eq!(out.token_ids, vec![1, 2, 3]); + } +} diff --git a/rust/sglang-server/src/environ.rs b/rust/sglang-server/src/environ.rs index e1a691e5f6ea..afa70ff25fb1 100644 --- a/rust/sglang-server/src/environ.rs +++ b/rust/sglang-server/src/environ.rs @@ -2,7 +2,6 @@ //! unset → default, invalid → warn + default (never an error). One shared //! parser per type — call sites pass their variable name + default instead of //! each hand-rolling a reader. -#![allow(dead_code)] // TODO: remove when the consumer PR lands /// Python `EnvBool.parse`: true = `true/1/yes/y`, false = `false/0/no/n` /// (case-insensitive); anything else is invalid. diff --git a/rust/sglang-server/src/error.rs b/rust/sglang-server/src/error.rs index a1eff40826fa..6b5feb8c14c9 100644 --- a/rust/sglang-server/src/error.rs +++ b/rust/sglang-server/src/error.rs @@ -1,6 +1,5 @@ //! Error type shared by all stages. Kept `Clone` so a single failure can be //! reported to the client stream and logged without moving ownership around. -#![allow(dead_code)] // TODO: remove when the consumer PR lands use thiserror::Error; diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index 1d4e6e8fbbe1..e82bf57a059e 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -9,8 +9,9 @@ //! * `push_result` — Python scheduler thread pushes one control result. //! //! All are non-blocking, so the GIL is never held across a wait. -#![allow(dead_code)] // TODO: remove when the consumer PR lands +mod api_server; +mod detokenizer; mod environ; mod error; mod fsm; @@ -18,12 +19,232 @@ mod ids; mod message; mod ring; mod runtime; +mod tokenizer; mod tokenizer_manager; mod utils; +use std::net::SocketAddr; + use pyo3::prelude::*; +use pyo3::pybacked::PyBackedBytes; +use pyo3::types::PyBytes; + +use crate::runtime::{Runtime, RuntimeConfig}; + +/// Columnar ingress batch handed to Python by [`Server::recv_requests`]. +/// `frozen`: immutable snapshot, so field access never contends on a borrow. +#[pyclass(frozen, get_all)] +struct IngressBatch { + /// One msgpack scalar header per request (`input_ids` omitted). + headers: Vec>, + /// The raw-data plane today just all requests' raw little-endian int64 + /// ids, concatenated; sliced per request via `lengths`. + data: Py, + /// Per-request token count (0 for control requests). + lengths: Vec, +} + +/// Handle owned by the Python scheduler process. Construct once via +/// [`Server::start`], then poll it from the scheduler event loop. +#[pyclass] +struct Server { + rt: Runtime, +} + +#[pymethods] +impl Server { + /// Boot the frontend (spawns all threads) and return immediately. + #[new] + #[pyo3(signature = ( + http_addr = None, + ingress_ring_cap = 8192, + egress_ring_cap = 8192, + channel_cap = 8192, + cores = None, + + server_args_json = "{}", + ))] + // pyo3 `#[new]` constructor: the wide arg list is the Python-facing boot + // surface (all optional overrides), not a call-site ergonomics problem. + #[allow(clippy::too_many_arguments)] + fn start( + http_addr: Option, + ingress_ring_cap: usize, + egress_ring_cap: usize, + channel_cap: usize, + cores: Option>, + server_args_json: &str, + ) -> PyResult { + // Static server metadata (server_args + model_config) dumped by the + // scheduler; parse and validate mandatory fields now so a bad/missing + // field is a boot error, not a request-time 500. + let server_args: runtime::ServerArgs = runtime::ServerArgs::from_json(server_args_json) + .map_err(|e| { + PyErr::new::(format!( + "bad server_args_json: {e}" + )) + })?; + server_args.validate_mandatory().map_err(|e| { + PyErr::new::(format!("server_args: {e}")) + })?; + // The HTTP listen address, tokenizer source/threads/shards all live in the + // `server_args` blob; resolve them from there so the scheduler doesn't + // re-pass them. The explicit params stay as optional overrides for + // standalone callers (tests) that construct a `Server` without a full + // `server_args`. + let http_addr: SocketAddr = http_addr + .unwrap_or_else(|| server_args.bind()) + .parse() + .map_err(|e| { + PyErr::new::(format!("bad http_addr: {e}")) + })?; + + let cfg = RuntimeConfig { + rust_server_args: runtime::RustServerServerArgs { + http_addr, + api_worker_num: server_args.api_worker_num(), + ingress_ring_cap, + egress_ring_cap, + channel_cap, + cores, + }, + server_args: std::sync::Arc::new(server_args), + }; + let rt = runtime::start(cfg).map_err(|e| { + PyErr::new::(format!("runtime start failed: {e}")) + })?; + Ok(Server { rt }) + } + + /// Non-blocking drain of the ingress ring, returned **columnar** as an + /// [`IngressBatch`] so the large `input_ids` tensor never goes through + /// msgpack (see the field docs for the layout). The `ids` cells are copied + /// **directly into the result `bytes`** (one copy, no intermediate buffer). + /// + /// Runs entirely GIL-held, deliberately. `drain` is a `try_recv` loop plus an + /// uncontended stash lock (the Python thread is the only consumer), so it + /// cannot block — there is nothing for a detach to overlap with. And detaching + /// is far from free: reacquiring the GIL waits out the interpreter's switch + /// interval, so a `py.detach` here cost up to 5 ms whenever another Python + /// thread was runnable, to cover ~0.2 µs of work. Held, the whole call is a + /// fraction of a microsecond on an empty ring. + #[pyo3(signature = (max = 256))] + fn recv_requests(&self, py: Python<'_>, max: usize) -> PyResult { + let cols = self.rt.ingress.drain(max); + let headers = cols + .headers + .iter() + .map(|h| PyBytes::new(py, h).unbind()) + .collect(); + // Single pass: copy each raw ids cell straight into the output `bytes`. + let data = PyBytes::new_with(py, cols.ids_total, |buf| { + let mut pos = 0; + for cell in &cols.ids { + let end = pos + cell.len(); + buf[pos..end].copy_from_slice(cell); + pos = end; + } + Ok(()) + })? + .unbind(); + Ok(IngressBatch { + headers, + data, + lengths: cols.lengths, + }) + } + + /// Park up to `timeout_ms` for an incoming request so the idle scheduler loop + /// sleeps instead of spinning at 100% CPU. Returns `True` when a request is + /// ready (the next `recv_requests` includes it). The GIL is released while + /// parked, and `flume` wakes the moment a request is pushed, so this adds no + /// latency to real requests — only the idle wait is bounded by `timeout_ms`. + #[pyo3(signature = (timeout_ms = 1000))] + fn wait_ingress(&self, py: Python<'_>, timeout_ms: u64) -> bool { + py.detach(|| { + self.rt + .ingress + .wait(std::time::Duration::from_millis(timeout_ms)) + }) + } + + /// Push a whole decode batch as ONE frame: a columnar msgpack `header` plus + /// the raw `data_cols` (per-column `bytes`), concatenated here. Blocks for + /// backpressure; `False` only on shutdown. + /// + /// Framed and pushed with the GIL HELD, detaching only if the ring is full. + /// This runs on the scheduler's CUDA-launch thread every decode step, where the + /// unconditional detach was the single worst boundary cost: framing is + /// ~0.1–0.2 µs, but reacquiring the GIL waits out the interpreter's switch + /// interval (5 ms by default) whenever another Python thread is runnable — + /// 17–50% of a 10–30 ms decode step, landing nondeterministically. Held, the + /// whole boundary is ~1.3 µs per step. + /// + /// The slow path keeps its detach because a full ring genuinely parks: the + /// scheduler must feel backpressure rather than drop output it has already + /// committed to. It essentially never fires — measured headroom is ~100×. + fn push_batch(&self, py: Python<'_>, header: &[u8], data_cols: Vec) -> bool { + let cols: Vec<&[u8]> = data_cols.iter().map(|d| d.as_ref()).collect(); + self.push_frame(py, crate::message::frame_egress_batch_cols(header, &cols)) + } + + /// Push a control-request result. Blocks for backpressure; `False` only on + /// shutdown. + fn push_result(&self, py: Python<'_>, rid: &str, payload: &[u8]) -> bool { + self.push_frame(py, crate::message::frame_egress_result(rid, payload)) + } + + /// Route a terminal failure back to request `rid`. Blocks for backpressure; + /// `False` only on shutdown. + fn push_error(&self, py: Python<'_>, rid: &str, message: &str) -> bool { + self.push_frame(py, crate::message::frame_egress_error(rid, message)) + } + + /// Signal all threads to stop (best effort). + fn shutdown(&self) { + self.rt.request_shutdown(); + } +} + +impl Server { + /// Hand one already-framed egress message to the ring: GIL-held when it fits, + /// detaching only to park on a full ring. Shared by every push path — they + /// differ solely in how the frame is built. `false` only on shutdown. + #[inline] + fn push_frame(&self, py: Python<'_>, frame: bytes::Bytes) -> bool { + match self.rt.egress.try_push(frame) { + Ok(()) => true, + // Consumer gone (shutdown): the frame is unavoidably lost. + Err(None) => false, + // Full: the scheduler must block here so backpressure reaches it, and + // blocking is exactly when releasing the GIL pays for itself. + Err(Some(frame)) => py.detach(|| self.rt.egress.push(frame)), + } + } +} + +/// Keeps the non-blocking log writer's background thread alive for the process +/// lifetime (dropping the guard would stop log delivery). +static LOG_GUARD: std::sync::OnceLock = + std::sync::OnceLock::new(); #[pymodule] -fn _core(_m: &Bound<'_, PyModule>) -> PyResult<()> { +fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { + // Initialize tracing once; ignore if already set by the host process. + // Non-blocking writer: emitting threads (axum workers, egress, detok) only + // enqueue; a dedicated thread does the stdout formatting-flush + syscall. + // The queue is bounded and lossy — under extreme pressure log lines are + // dropped instead of stalling request threads. + let (writer, guard) = tracing_appender::non_blocking(std::io::stdout()); + let _ = LOG_GUARD.set(guard); + let _ = tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .with_writer(writer) + .try_init(); + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/rust/sglang-server/src/message.rs b/rust/sglang-server/src/message.rs index 81d91e37ff8d..91fdacf34487 100644 --- a/rust/sglang-server/src/message.rs +++ b/rust/sglang-server/src/message.rs @@ -15,8 +15,14 @@ mod request; mod sampling; mod types; -pub use egress::{ChunkEvent, EgressSink}; -pub use request::{GenerateRequest, RequestKind}; +pub use egress::{ + ChunkEvent, ChunkExtras, EGRESS_TAG_BATCH, EGRESS_TAG_ERROR, EGRESS_TAG_RESULT, EgressItem, + EgressSink, SinkError, for_each_chunk, frame_egress_batch_cols, frame_egress_error, + frame_egress_result, +}; +pub use finish_reason::Matched; +pub(crate) use io_struct::{AbortReq, ControlRequest, GetInternalStateReq}; +pub use request::{GenerateBody, GenerateRequest, RequestKind}; pub(crate) use sampling::{SamplingParams, SamplingParamsInput}; pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds}; diff --git a/rust/sglang-server/src/runtime.rs b/rust/sglang-server/src/runtime.rs index cd9751feda08..7119b0a86999 100644 --- a/rust/sglang-server/src/runtime.rs +++ b/rust/sglang-server/src/runtime.rs @@ -12,7 +12,499 @@ //! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling //! axum's worker threads. +use std::sync::{Arc, Mutex}; +use std::thread::JoinHandle; + mod config; mod runnable; +mod threads; + +pub use config::{RuntimeConfig, RustServerServerArgs, ServerArgs}; + +use crate::message::DetokMsg; +use crate::ring::{ + EgressConsumer, EgressProducer, IngressConsumer, IngressProducer, egress_ring, ingress_ring, +}; +use crate::runtime::threads::{plan_cores, spawn_pool}; +use crate::tokenizer_manager::{Senders, TmEvent}; +use crate::{api_server, detokenizer, tokenizer, tokenizer_manager}; // Re-export so stages keep importing `crate::runtime::Runnable`. +pub use runnable::Runnable; + +/// Live runtime. Held by the pyo3 bridge; the Python boundary reads `ingress` +/// and `egress`. `request_shutdown` (also run on `Drop`) stops every stage. +pub struct Runtime { + pub ingress: IngressConsumer, + pub egress: EgressProducer, + /// Worker join handles, joined by `request_shutdown` / `Drop`. + threads: Mutex>>, + /// The single shutdown sender. + shutdown_tx: Mutex>>, +} + +/// Deadline for joining worker threads on shutdown. Past it we abandon the join +/// so process teardown can't deadlock on a worker that somehow failed to exit. +const SHUTDOWN_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); + +impl Runtime { + /// Stop the runtime and join every worker thread (with a bounded wait). + /// + /// Dropping `shutdown_tx` wakes the tm-ingress/tm-egress selectors (which + /// otherwise never see their inbox close — one self-holds a `tm` sender, the + /// other's inbox is the Python-fed egress ring). Those exit and drop their + /// `Senders` clones; the api thread's `serve` returns non-gracefully, so its + /// `block_on` unwinds and the api tokio runtime is dropped — cancelling + /// in-flight handlers, whose `AbortGuard`s release the remaining clones. With + /// every clone gone the tok/detok channels close and those workers exit. + /// + /// In-flight requests are **aborted**, not drained — this is the hard-stop + /// path (also run on `Drop`). Clients of aborted requests retry. + pub fn request_shutdown(&self) { + drop(self.shutdown_tx.lock().unwrap().take()); + let handles: Vec> = self.threads.lock().unwrap().drain(..).collect(); + if handles.is_empty() { + return; // Idempotent: a `Drop` after an explicit shutdown has nothing to join. + } + // Join off-thread and wait with a deadline: a stuck worker can't wedge exit. + let (done_tx, done_rx) = flume::bounded::<()>(1); + std::thread::spawn(move || { + for h in handles { + let _ = h.join(); + } + let _ = done_tx.send(()); + }); + if done_rx.recv_timeout(SHUTDOWN_JOIN_TIMEOUT).is_err() { + tracing::warn!( + "shutdown: workers did not exit within {SHUTDOWN_JOIN_TIMEOUT:?}; abandoning join" + ); + } + } +} + +impl Drop for Runtime { + fn drop(&mut self) { + self.request_shutdown(); + } +} + +/// Boot the whole frontend. Returns once threads are spawned (non-blocking), +/// so the Python caller regains control of the GIL immediately. `Err` on a +/// startup misconfiguration (e.g. no tokenizer for a non-skip server). +pub fn start(cfg: RuntimeConfig) -> Result { + // Bind the API server port before spawning any thread, so an unavailable + // port (EADDRINUSE) is a hard startup error. + let listener = std::net::TcpListener::bind(cfg.rust_server_args.http_addr) + .map_err(|e| format!("bind {} failed: {e}", cfg.rust_server_args.http_addr))?; + listener + .set_nonblocking(true) + .map_err(|e| format!("listener set_nonblocking failed: {e}"))?; + + let (shutdown_tx, shutdown_rx) = flume::unbounded::<()>(); + let mut threads = Vec::new(); + let plan = plan_cores(&cfg); + + // --- rings (Rust ↔ Python) --- + let (ingress_tx, ingress_rx): (IngressProducer, IngressConsumer) = + ingress_ring(cfg.rust_server_args.ingress_ring_cap); + let (egress_tx, egress_rx): (EgressProducer, EgressConsumer) = + egress_ring(cfg.rust_server_args.egress_ring_cap); + + // --- inter-stage channels --- + let (tm_tx, tm_rx) = flume::bounded::(cfg.rust_server_args.channel_cap); + let (tok_tx, tok_rx) = + flume::bounded::(cfg.rust_server_args.channel_cap); + let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num; + let mut detok_tx = Vec::with_capacity(detokenizer_worker_num); + let mut detok_rx = Vec::with_capacity(detokenizer_worker_num); + for _ in 0..detokenizer_worker_num { + let (tx, rx) = flume::bounded::(cfg.rust_server_args.channel_cap); + detok_tx.push(tx); + detok_rx.push(rx); + } + + // Aborts get their own UNBOUNDED lane: on the bounded inbox they are dropped + // exactly under the overload that makes them necessary (see `Senders::abort`). + let (abort_tx, abort_rx) = flume::unbounded::(); + let senders = Senders { + tm: tm_tx.clone(), + abort: abort_tx.clone(), + tok: tok_tx, + detok: detok_tx, + }; + + // `skip_tokenizer_init`: clients send token ids and receive token ids — no + // tokenizer is loaded, and the egress emits raw `output_ids` (no decode). + let skip_tokenizer_init = cfg.server_args.skip_tokenizer_init; + + // The same instance is shared by the tokenizer pool (encode) and the detok + // shards (decode); `None` only under `skip_tokenizer_init`. + let dyn_tokenizer = tokenizer::load_tokenizer( + // Empty only in minimal standalone blobs (the Python dump always + // resolves it); empty → no tokenizer, allowed only under + // `skip_tokenizer_init`. + (!cfg.server_args.tokenizer_path.is_empty()).then_some(&*cfg.server_args.tokenizer_path), + cfg.server_args.revision.as_deref(), + skip_tokenizer_init, + )?; + + // --- Detokenizer shards (pinned, CPU bound) --- + { + // Default: a real tokenizer decodes to text. `None` (→ `Skip`, raw + // `output_ids`) only happens under `skip_tokenizer_init` — + // `load_tokenizer` rejects a non-skip server with no tokenizer. + let backend = match &dyn_tokenizer { + Some(t) => detokenizer::DetokenizerBackend::Dynamo(t.clone()), + None => detokenizer::DetokenizerBackend::Skip, + }; + let detok_cores = plan.as_ref().map(|p| p.detok.clone()); + // Each shard owns its receiver outright (one consumer per shard), so the + // owned `detok_rx` Vec is moved out element-by-element via the iterator. + let count = detok_rx.len(); + let mut rxs = detok_rx.into_iter(); + spawn_pool("detokenizer", detok_cores, count, &mut threads, |i| { + detokenizer::DetokenizerWorker::new( + i, + rxs.next().unwrap(), + backend.clone(), + abort_tx.clone(), + ) + }); + } + + // --- Tokenizer pool (pinned, CPU bound) --- + // Only spawned when a real tokenizer is loaded; under `skip_tokenizer_init` + // there is none and ingress never routes to the pool, so we skip it. + if let Some(t) = &dyn_tokenizer { + // Reuse the single loaded tokenizer (shared with the detok shards). + let tokenizer: Arc = + Arc::new(tokenizer::DynamoTokenizer::new(t.clone())); + let tok_cores = plan.as_ref().map(|p| p.tok.clone()); + // Workers share the MPMC inbox (`tok_rx`) and the read-only backend, so + // each gets a cheap clone of both. + spawn_pool( + "tokenizer", + tok_cores, + cfg.server_args.tokenizer_worker_num, + &mut threads, + |_i| tokenizer::TokenizerWorker::new(tok_rx.clone(), tm_tx.clone(), tokenizer.clone()), + ); + } + + // Egress heartbeat: bumped per drained frame, watched by `/health_generate`. + let egress_activity: tokenizer_manager::ActivityCounter = + Arc::new(std::sync::atomic::AtomicU64::new(0)); + + // --- Egress dispatcher: drains egress ring → routes chunks to shards --- + { + // First TM core; egress is the hotter router (every output token). One + // worker today via `spawn_pool`, so sharding by `Rid::shard` later (see + // `TM_CORES`) is just a larger count + per-shard receivers. + let cores = plan + .as_ref() + .and_then(|p| p.tm.first().copied()) + .map(|c| vec![c]); + let mut egress_rx = Some(egress_rx); // moved into the single worker + let activity = egress_activity.clone(); + let shutdown_rx = shutdown_rx.clone(); + spawn_pool("tm-egress", cores, 1, &mut threads, |_| { + tokenizer_manager::Egress::new( + egress_rx.take().unwrap(), + senders.clone(), + activity.clone(), + shutdown_rx.clone(), + ) + }); + } + + // --- TokenizerManager ingress loop --- + { + // Second TM core when present, else share the first (1-core / API-set + // fallback) — still off the CPU-bound pool cores either way. + let cores = plan + .as_ref() + .and_then(|p| p.tm.get(1).or_else(|| p.tm.first()).copied()) + .map(|c| vec![c]); + let limits = tokenizer_manager::Limits::try_from(&*cfg.server_args) + .map_err(|e| format!("ingress limits: {e}"))?; + let mut parts = Some((tm_rx, ingress_tx)); // moved into the single worker + let shutdown_rx = shutdown_rx.clone(); + spawn_pool("tm-ingress", cores, 1, &mut threads, |_| { + let (tm_rx, ingress_tx) = parts.take().unwrap(); + tokenizer_manager::Ingress::new( + tm_rx, + abort_rx.clone(), + senders.clone(), + ingress_tx, + limits.clone(), + shutdown_rx.clone(), + ) + }); + } + + // --- API server (tokio, I/O bound) --- + { + let cfg = cfg.clone(); + let api_cores = plan.as_ref().map(|p| p.api.clone()); + let senders = senders.clone(); + let api_activity = egress_activity.clone(); + let shutdown_rx = shutdown_rx.clone(); + let handle = std::thread::Builder::new() + .name("api-runtime".into()) + .spawn(move || { + let mut builder = tokio::runtime::Builder::new_multi_thread(); + builder + .worker_threads(cfg.rust_server_args.api_worker_num) + .enable_all(); + if let Some(cores) = api_cores { + let next = std::sync::atomic::AtomicUsize::new(0); + builder.on_thread_start(move || { + let idx = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if let Some(c) = cores.get(idx % cores.len()) { + core_affinity::set_for_current(*c); + } + }); + } + let rt = builder.build().expect("build api runtime"); + rt.block_on(api_server::serve( + listener, + senders, + cfg.rust_server_args.channel_cap, + cfg.server_args.clone(), + // Egress heartbeat watched by `/health_generate`. + api_activity, + shutdown_rx, + )) + }) + .expect("spawn api runtime"); + threads.push(handle); + } + + Ok(Runtime { + ingress: ingress_rx, + egress: egress_tx, + threads: Mutex::new(threads), + shutdown_tx: Mutex::new(Some(shutdown_tx)), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Minimal boot args. `skip_tokenizer_init` avoids loading a tokenizer/detok + /// model; `model_config` carries the two fields `Limits::from_server_args` + /// requires. They are mandatory at boot, so a fixture without them panics the + /// runtime instead of exercising what these tests are about — `start` does not + /// run `ServerArgs::validate_mandatory` itself, `Server::start` does. + const TEST_SERVER_ARGS: &str = r#"{ + "skip_tokenizer_init": true, + "model_config": {"context_len": 2048, "vocab_size": 1000} + }"#; + + /// Regression: `request_shutdown` must actually stop the API server — it joins + /// the api thread once the listener closes, so the port stops accepting. + /// (Previously it set an unread flag and the port kept accepting.) + #[test] + fn request_shutdown_closes_listener() { + // Pick a free port: bind :0, read the assigned addr, release it. + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + // `skip_tokenizer_init` → no tokenizer/detok model load; minimal boot. + let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + // Bind is synchronous in `start`, so the port is already accepting. + let rt = start(cfg).expect("start runtime"); + assert!( + std::net::TcpStream::connect(addr).is_ok(), + "server not listening on {addr} after start returned", + ); + + // Joins the api thread; the listener is closed by the time it returns. + rt.request_shutdown(); + + assert!( + std::net::TcpStream::connect(addr).is_err(), + "port still accepting connections after shutdown", + ); + } + + /// Regression: shutdown must return promptly even with an in-flight `/generate`. + /// No scheduler drains the ingress ring or feeds the egress ring here, so the + /// handler parks on its egress channel forever. Graceful shutdown would wait + /// for it (deadlock → only the 5s bounded-join fallback returns); the + /// non-graceful path cancels the handler via the api runtime drop, whose + /// `AbortGuard` releases the last `Senders` clone so the workers exit. + #[test] + fn shutdown_returns_with_in_flight_request() { + use std::io::Write; + use std::time::{Duration, Instant}; + + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + let rt = start(cfg).expect("start runtime"); + + // Fire a request that will block (already-tokenized → valid → pushed to the + // ring, then the handler awaits egress frames that never arrive). + let mut conn = std::net::TcpStream::connect(addr).expect("connect"); + let body = r#"{"input_ids":[1,2,3],"stream":false,"sampling_params":{"max_new_tokens":8}}"#; + let req = format!( + "POST /generate HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n\ + Content-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + conn.write_all(req.as_bytes()).unwrap(); + conn.flush().unwrap(); + std::thread::sleep(Duration::from_millis(300)); // reach the blocked state + + let t = Instant::now(); + rt.request_shutdown(); + let elapsed = t.elapsed(); + assert!( + elapsed < Duration::from_secs(3), + "shutdown took {elapsed:?} with an in-flight request (deadlock?)", + ); + drop(conn); + } + + /// Regression: a >2MB body must reach the JSON layer and fail on its + /// *content* (unknown field → 4xx), never on size (413). + #[test] + fn accepts_multi_megabyte_generate_body() { + use std::io::{Read, Write}; + + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + let rt = start(cfg).expect("start runtime"); + + // ~3MB of input_ids plus a `text`, which is mutually exclusive with them: + // the body parses in full and is then rejected by `into_requests` with a + // 400, proving it got past any size limit (a 413 would fire before + // parsing). The rejection must come from OUR validation, not from serde — + // an unknown field used to serve here, but unknown fields are now ignored + // to match Python, so such a body would be accepted, dispatched to a ring + // nobody drains in this test, and hang the connection. + let ids = "1,".repeat(1_500_000); + let body = format!( + r#"{{"input_ids":[{}1],"text":"x","sampling_params":{{"max_new_tokens":1}}}}"#, + ids + ); + assert!(body.len() > 2 * 1024 * 1024, "test body must exceed 2MB"); + + let mut conn = std::net::TcpStream::connect(addr).expect("connect"); + let req = format!( + "POST /generate HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n\ + Content-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + conn.write_all(req.as_bytes()).unwrap(); + conn.flush().unwrap(); + + let mut response = String::new(); + conn.read_to_string(&mut response).unwrap(); + let status_line = response.lines().next().unwrap_or(""); + let code: u16 = status_line + .split_whitespace() + .nth(1) + .and_then(|c| c.parse().ok()) + .unwrap_or(0); + // A 400 from the mutually-exclusive-inputs check proves the body was read + // and parsed in full; 413 would mean it was rejected on size beforehand. + assert!( + (400..500).contains(&code) && code != 413, + "expected a JSON-layer 4xx (not 413), got: {status_line}" + ); + + rt.request_shutdown(); + } + + /// Regression: a port conflict must fail `start` (so the scheduler doesn't + /// advertise ready), not return an `Ok` runtime whose listener never binds. + #[test] + fn start_fails_on_port_conflict() { + // Hold the port so the runtime's bind conflicts (EADDRINUSE). + let hog = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = hog.local_addr().unwrap(); + + let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + let err = match start(cfg) { + Ok(_) => panic!("bind conflict must fail startup, got Ok"), + Err(e) => e, + }; + assert!(err.contains("bind"), "error should mention bind: {err}"); + } + + /// `server_args` missing a mandatory `model_config` field must be a startup + /// ERROR, not a panic. + /// + /// `Limits::try_from` is fallible and the ingress loop is built inside a + /// `spawn_pool` closure, so resolving it there would put the failure on a + /// freshly spawned worker thread — a thread `start` never inspects. The boot + /// would report success and the server would accept connections with no + /// ingress loop behind them, hanging every request instead of refusing to + /// start. Only `Server::start` runs `validate_mandatory`, so `start` cannot + /// assume these fields are present. + #[test] + fn start_fails_when_model_config_is_incomplete() { + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + // Boots fine in every other respect — only `model_config` is absent. + let server_args = ServerArgs::from_json(r#"{"skip_tokenizer_init": true}"#).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + let err = match start(cfg) { + Ok(_) => panic!("an incomplete model_config must not boot, got Ok"), + Err(e) => e, + }; + assert!(err.contains("ingress limits"), "{err}"); + } +} diff --git a/rust/sglang-server/src/runtime/threads.rs b/rust/sglang-server/src/runtime/threads.rs new file mode 100644 index 000000000000..5e56ae4d4e0b --- /dev/null +++ b/rust/sglang-server/src/runtime/threads.rs @@ -0,0 +1,142 @@ +//! Thread-group machinery: CPU-core partitioning and the pinned-thread +//! spawners (for [`Runnable`] stages) used by `runtime::start`. +//! +//! Adding a new thread group (encoder, weight loader, KV-cache offloader, …) is +//! three small steps and no spawn boilerplate: +//! 1. a struct implementing [`Runnable`]; +//! 2. a core set for it (a field on [`CorePlan`] + a slice in [`plan_cores`]); +//! 3. one [`spawn_pool`] (N pinned workers) or [`spawn_stage`] (singleton) call. + +use std::thread::JoinHandle; + +use core_affinity::CoreId; + +use super::{Runnable, RuntimeConfig}; + +/// Cores reserved for the two TokenizerManager router threads (`tm-ingress`, +/// `tm-egress`) — light, latency-sensitive channel routers, so one core each. +/// +/// TODO(tm-scaling): both TM threads are single-consumer serialization points, +/// each with its own ceiling. `tm-ingress` runs validate + `normalize_sampling_params` +/// for *every* request before fanning out to the (pooled) tokenizer workers, so a +/// high request-arrival / short-request workload is bounded by that one thread's +/// per-request cost (kept O(fields), see `sampling::normalize_sampling_params`). +/// Sharding ingress by rid — like the tokenizer/detok pools — lifts that ceiling. +/// +/// `tm-egress` is a head-of-line ceiling of a different kind — it +/// does a *blocking* send per chunk to the owning detok shard, so one slow shard +/// stalls the dispatcher and thus every shard (see `Egress::route`). Sharding the +/// dispatcher alone doesn't fix it: each egress-ring frame is a whole batch fanned +/// to *all* shards, so any dispatcher still blocks on the slow one. The real fix +/// is a per-shard egress ring (the scheduler pushing each request's output to its +/// shard's ring), each drained by its own dispatcher — at which point this needs +/// one core per ingress/egress shard rather than a fixed 2. +const TM_CORES: usize = 2; + +/// Partition the machine's cores into four disjoint sets: the I/O-bound API +/// pool, the CPU-bound tokenizer and detokenizer pools, and the two TM router +/// threads. Falls back to no pinning if affinity isn't available or there aren't +/// enough cores for the (CPU-bound) pools. A new thread group adds a field here +/// and a slice in [`plan_cores`]. +pub(super) struct CorePlan { + pub(super) api: Vec, + pub(super) tok: Vec, + pub(super) detok: Vec, + pub(super) tm: Vec, +} + +pub(super) fn plan_cores(cfg: &RuntimeConfig) -> Option { + // `cores` carries the pinning decision: `None`/empty → run unpinned. The + // caller (Python `_partition_cores`) passes this rank's NUMA-local cores + // minus the scheduler's reserved launch cores. + let cores: Vec = match &cfg.rust_server_args.cores { + Some(ids) if !ids.is_empty() => ids.iter().map(|&id| CoreId { id }).collect(), + _ => return None, + }; + if cores.len() + < cfg.rust_server_args.api_worker_num + + cfg.server_args.tokenizer_worker_num + + cfg.server_args.detokenizer_worker_num + { + tracing::warn!( + available = cores.len(), + "not enough cores to pin all pools; running unpinned" + ); + return None; + } + let mut it = cores.into_iter(); + let api: Vec = it + .by_ref() + .take(cfg.rust_server_args.api_worker_num) + .collect(); + let tok = it + .by_ref() + .take(cfg.server_args.tokenizer_worker_num) + .collect(); + let detok = it + .by_ref() + .take(cfg.server_args.detokenizer_worker_num) + .collect(); + // The two TM router threads get up to `TM_CORES` leftover cores; when none + // are spare they fall back to the API set so they never float onto the + // CPU-bound tokenizer/detok cores. + let mut tm: Vec = it.by_ref().take(TM_CORES).collect(); + if tm.is_empty() { + tm = api.clone(); + } + Some(CorePlan { + api, + tok, + detok, + tm, + }) +} + +/// Pin the calling thread to `core` if one was assigned (no-op otherwise). +fn pin_current(core: Option) { + if let Some(c) = core { + core_affinity::set_for_current(c); + } +} + +/// Pick the pinned core for worker `i` from an optional pool core set. +pub(super) fn pool_core(cores: &Option>, i: usize) -> Option { + cores.as_ref().and_then(|c| c.get(i).copied()) +} + +/// Spawn a single [`Runnable`] stage on a named thread, optionally pinned. +/// Used by [`spawn_pool`]; every group goes through the pool spawner now. +fn spawn_stage( + name: &str, + core: Option, + stage: impl Runnable, + threads: &mut Vec>, +) { + let handle = std::thread::Builder::new() + .name(name.to_string()) + .spawn(move || { + pin_current(core); + stage.run(); + }) + .expect("spawn stage"); + threads.push(handle); +} + +/// Spawn a pool of `count` [`Runnable`] workers, each pinned to `cores[i]` (when +/// available) and named `{name}-{i}`. `build(i)` constructs worker `i` — cloning +/// shared handles, or moving a per-worker resource out of a captured iterator. +pub(super) fn spawn_pool( + name: &str, + cores: Option>, + count: usize, + threads: &mut Vec>, + mut build: F, +) where + R: Runnable, + F: FnMut(usize) -> R, +{ + for i in 0..count { + let core = pool_core(&cores, i); + spawn_stage(&format!("{name}-{i}"), core, build(i), threads); + } +} diff --git a/rust/sglang-server/src/tokenizer.rs b/rust/sglang-server/src/tokenizer.rs new file mode 100644 index 000000000000..ee90d605a764 --- /dev/null +++ b/rust/sglang-server/src/tokenizer.rs @@ -0,0 +1,162 @@ +//! Tokenizer pool — CPU-bound, runs on pinned OS threads (off the async +//! executor). Each worker pulls a `Request` from the shared `flume` receiver, +//! fills `input_ids`, and moves the request back to the TokenizerManager inbox. +//! +//! The text→ids step is behind [`TextTokenizer`], implemented by +//! [`DynamoTokenizer`] (dynamo-tokenizers: HuggingFace / tiktoken / fastokens). +//! A non-skip server requires a real tokenizer (enforced at startup); under +//! `skip_tokenizer_init` the pool isn't spawned at all. +//! +//! Mirrors the Python `_tokenize_one_request` text path: when the request +//! already carries `input_ids` it skips tokenization (handled upstream in the +//! TokenizerManager `classify`); otherwise the prompt text is encoded here. + +use std::path::Path; +use std::sync::Arc; + +use crate::error::Error; +use crate::fsm::Event; +use crate::message::{Request, RequestKind, TokenIds}; +use crate::runtime::Runnable; +use crate::tokenizer_manager::TmEvent; + +/// Pluggable text→token-ids backend. `Send + Sync` so one instance is shared +/// (read-only) across all pinned workers. +pub trait TextTokenizer: Send + Sync { + fn encode(&self, text: &str) -> Result; +} + +/// Load the tokenizer shared (Arc-backed) by the encode pool and detok shards. +/// `None` under `skip_tokenizer_init`, else required (missing/failed load → `Err`). +/// `tokenizer_path` is a tokenizer file, a model dir, or an HF Hub repo id +/// (resolved from the local cache — no network). +pub fn load_tokenizer( + tokenizer_path: Option<&str>, + revision: Option<&str>, + skip_tokenizer_init: bool, +) -> Result, String> { + if skip_tokenizer_init { + tracing::info!("skip_tokenizer_init: token ids in and out; no tokenizer/detokenizer"); + return Ok(None); + } + let path = tokenizer_path.ok_or_else(|| { + "no tokenizer configured: set tokenizer_path or enable skip_tokenizer_init".to_string() + })?; + + let file = resolve_model_file(path, revision, "tokenizer.json") + .ok_or_else(|| format!("tokenizer.json not found for '{path}'"))?; + let tokenizer = dynamo_tokenizers::Tokenizer::from_file_with_options( + &file, + dynamo_tokenizers::TokenizerOptions { + add_special_tokens: true, + }, + ) + .map_err(|e| format!("tokenizer load failed ({file}): {e}"))?; + tracing::info!(%path, "loaded tokenizer"); + Ok(Some(tokenizer)) +} + +/// Resolve a model file from the tokenizer source: a dir → `dir/`, a file → +/// its sibling, else an HF Hub repo id → the local cache. `None` if not found. +pub fn resolve_model_file(path: &str, revision: Option<&str>, filename: &str) -> Option { + let p = Path::new(path); + if p.is_dir() { + let f = p.join(filename); + return f.is_file().then(|| f.to_string_lossy().into_owned()); + } + if p.is_file() { + // `path` is a file (e.g. `tokenizer.json`); look for the sibling. + let f = p.parent()?.join(filename); + return f.is_file().then(|| f.to_string_lossy().into_owned()); + } + // Not a local path → HF Hub repo id (offline cache lookup). + resolve_from_hub_cache(path, revision, filename) +} + +/// Locate a file for an HF Hub repo id in the local cache (`HF_HOME`). Offline — +/// the scheduler pre-downloads the model. `None` if not cached. +fn resolve_from_hub_cache(repo_id: &str, revision: Option<&str>, filename: &str) -> Option { + use hf_hub::{Cache, Repo, RepoType}; + + let rev = revision.unwrap_or("main"); + Cache::from_env() + .repo(Repo::with_revision( + repo_id.to_string(), + RepoType::Model, + rev.to_string(), + )) + .get(filename) + .map(|p| p.to_string_lossy().into_owned()) +} + +/// Real tokenizer over an already-loaded dynamo `Tokenizer` (Arc inside). +pub struct DynamoTokenizer { + inner: dynamo_tokenizers::Tokenizer, +} + +impl DynamoTokenizer { + pub fn new(inner: dynamo_tokenizers::Tokenizer) -> Self { + Self { inner } + } +} + +impl TextTokenizer for DynamoTokenizer { + fn encode(&self, text: &str) -> Result { + if text.is_empty() { + // Match Python sglang: reject an empty prompt as a 400 (`Validation`), + // not the misleading 500 a tokenize error would give. + return Err(Error::Validation("prompt cannot be empty".into())); + } + let encoding = self + .inner + .encode(text) + .map_err(|e| Error::Tokenize(e.to_string()))?; + // Vocab ids are non-negative and fit in i32. + Ok(encoding.token_ids().iter().map(|&id| id as i32).collect()) + } +} + +/// One tokenizer worker: pulls a `Request` off the shared inbox, fills +/// `input_ids`, returns it to the TokenizerManager. Pinned; backend shared. +pub struct TokenizerWorker { + rx: flume::Receiver, + tm: flume::Sender, + tokenizer: Arc, +} + +impl TokenizerWorker { + pub fn new( + rx: flume::Receiver, + tm: flume::Sender, + tokenizer: Arc, + ) -> Self { + Self { rx, tm, tokenizer } + } +} + +impl Runnable for TokenizerWorker { + fn run(self) { + while let Ok(mut req) = self.rx.recv() { + // The tokenizer pool only ever receives generate requests. Encode, + // then advance the FSM: `TokenizeDone` on success (→ PreSendValidating). + let event = { + let RequestKind::Generate(g) = &mut req.kind else { + tracing::error!("tokenizer pool received a non-generate request"); + continue; + }; + match self.tokenizer.encode(g.text.as_deref().unwrap_or("")) { + Ok(ids) => { + g.input_ids = Some(ids); + Event::TokenizeDone + } + Err(err) => Event::Error(err), + } + }; + let _ = req.state.apply(event); + if self.tm.send(TmEvent::Tokenized(req)).is_err() { + tracing::error!("tm inbox closed; dropping request"); + break; + } + } + } +} diff --git a/rust/sglang-server/src/tokenizer_manager.rs b/rust/sglang-server/src/tokenizer_manager.rs index a830a7b4ac4f..29bf73e918e2 100644 --- a/rust/sglang-server/src/tokenizer_manager.rs +++ b/rust/sglang-server/src/tokenizer_manager.rs @@ -10,6 +10,12 @@ //! the rest of the pipeline only through `flume` channels: [`TmEvent`] into the //! ingress loop, [`Senders`] fanning out to the pools. +mod egress; +mod ingress; + +pub use egress::{ActivityCounter, Egress}; +pub use ingress::{Ingress, Limits}; + use crate::ids::Rid; use crate::message::{DetokMsg, Request}; diff --git a/rust/sglang-server/src/tokenizer_manager/egress.rs b/rust/sglang-server/src/tokenizer_manager/egress.rs new file mode 100644 index 000000000000..1ad6df0ec900 --- /dev/null +++ b/rust/sglang-server/src/tokenizer_manager/egress.rs @@ -0,0 +1,217 @@ +//! TokenizerManager egress thread — drains the egress ring (scheduler output +//! pushed from Python) and routes each message to the detok shard that owns its +//! `Rid::shard`. Routing is a pure function of the rid, so it matches the shard +//! the request registered with on ingress — no shared map, no lock. +//! +//! The ring carries a 1-byte frame tag: `BATCH` (a whole decode batch, fanned +//! out here into per-request chunks), `RESULT` (a single control-request JSON +//! payload, e.g. `/server_info`), or `ERROR` (a terminal per-request failure the +//! scheduler ingress couldn't decode, routed back as a 400). + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use bytes::Bytes; + +use crate::ids::Rid; +use crate::message::DetokMsg; +use crate::message::{ + ChunkEvent, EGRESS_TAG_BATCH, EGRESS_TAG_ERROR, EGRESS_TAG_RESULT, for_each_chunk, +}; +use crate::ring::EgressConsumer; +use crate::runtime::Runnable; +use crate::tokenizer_manager::{Senders, recv}; + +/// A monotonic counter bumped once per egress-ring frame the dispatcher drains. +/// It's the rust-native equivalent of the Python `TokenizerManager`'s +/// `last_receive_tstamp`: `/health_generate` watches it advance to confirm the +/// scheduler → detok path is alive (the value itself is meaningless). +pub type ActivityCounter = Arc; + +/// Egress dispatcher stage. Owns the egress-ring consumer + the detok-shard +/// senders, so the runtime spawns it as a [`Runnable`]. +pub struct Egress { + egress: EgressConsumer, + senders: Senders, + activity: ActivityCounter, + shutdown: flume::Receiver<()>, +} + +impl Egress { + pub fn new( + egress: EgressConsumer, + senders: Senders, + activity: ActivityCounter, + shutdown: flume::Receiver<()>, + ) -> Self { + Self { + egress, + senders, + activity, + shutdown, + } + } +} + +impl Runnable for Egress { + fn run(self) { + // Reused across frames (`clear` keeps capacity) — steady state allocates nothing. + let shards = self.senders.detok.len(); + let mut buckets: Vec> = (0..shards).map(|_| Vec::new()).collect(); + + while let Some(bytes) = recv(self.egress.receiver(), &self.shutdown) { + let Some((&tag, body)) = bytes.split_first() else { + continue; + }; + match tag { + // A whole decode batch: bucket each request by the shard owning its + // rid, then hand each shard its chunks in one send. + EGRESS_TAG_BATCH => { + for b in buckets.iter_mut() { + b.clear(); + } + let decoded = for_each_chunk(body, |ev| { + // The rid picks the shard; a hash collision only co-locates two + // requests now, it no longer merges them. + buckets[ev.rid.shard(shards)].push(ev); + }); + // Routing only fills the buckets; nothing is delivered until the + // sends below. So dropping them here makes rejection atomic for + // free: a frame whose columns drifted would otherwise deliver the + // requests decoded before the bad one, carrying another request's + // logprobs — the corruption the decoder's bounds checks exist to + // prevent. Better a lost frame than a silently wrong one. + if !decoded.ok { + // Dropping the frame keeps wrong data off the wire, but a + // request whose chunk was in it would otherwise wait forever: + // mid-stream it gets a hole, and if its FINAL chunk was here + // it never sees `Done` and the connection hangs — there is no + // server-side timeout. Fail from the HEADER's rids, not the + // buckets: a frame that fails at request 0 buckets nothing, + // so bucket-driven cleanup would leave every request in it + // hanging. + // Distinguish "failed N requests" from "named nobody": a + // frame whose header would not decode at all yields no rids, + // so nothing downstream fails and every request in it waits + // forever. That is the case worth paging on, and it used to + // log the same line as the recoverable one. + if decoded.rids.is_empty() { + tracing::error!( + "egress: bad batch frame named NO rids; any request in \ + it will hang (header undecodable, or empty rid column)" + ); + } else { + tracing::warn!( + rids = decoded.rids.len(), + "egress: bad batch frame; failing its requests" + ); + } + for b in buckets.iter_mut() { + b.clear(); + } + for rid in decoded.rids { + // 500, not 400: the client's request was fine — the + // scheduler's own output frame was not. + let shard = rid.shard(shards); + let _ = self.senders.detok[shard].send(DetokMsg::Fail { + rid, + message: "internal error: malformed scheduler output frame".into(), + }); + } + continue; + } + for (i, b) in buckets.iter_mut().enumerate() { + if b.is_empty() { + continue; + } + let chunks = DetokMsg::Chunks(std::mem::take(b)); + if self.senders.detok[i].send(chunks).is_err() { + tracing::error!("egress: detok shard closed"); + } + } + // Any frame off the ring = the scheduler produced output → alive. + self.activity.fetch_add(1, Ordering::Relaxed); + } + EGRESS_TAG_RESULT => { + if let Some((rid, msg)) = decode_result(body) { + self.route(&rid, msg); + } + } + EGRESS_TAG_ERROR => { + if let Some((rid, msg)) = decode_error(body) { + self.route(&rid, msg); + } + } + other => tracing::warn!(tag = other, "egress: unknown frame tag"), + } + } + } +} + +impl Egress { + /// Route one message to the shard owning `rid`. HOL ceiling: a slow shard stalls + /// this thread; the fix is a per-shard egress ring (see `threads::TM_CORES`). + #[inline] + fn route(&self, rid: &Rid, msg: DetokMsg) { + if self.senders.detok_for(rid).send(msg).is_err() { + tracing::error!("egress: detok shard closed"); + } + } +} + +/// Control result: `[rid, payload]` → single non-streamed delivery to the sink. +fn decode_result(body: &[u8]) -> Option<(Rid, DetokMsg)> { + let val = rmpv::decode::read_value(&mut &body[..]).ok()?; + let rmpv::Value::Array(arr) = val else { + return None; + }; + let mut items = arr.into_iter(); + let rid = Rid::from(items.next()?.as_str()?); + // The decode already owns the payload buffer — move it out. + let payload = match items.next()? { + rmpv::Value::Binary(b) => Bytes::from(b), + rmpv::Value::String(s) => Bytes::from(s.into_bytes()), + _ => return None, + }; + Some((rid.clone(), DetokMsg::Result { rid, payload })) +} + +/// Per-request failure: `[rid, message]` → terminal `Error` to the sink (→ 400). +fn decode_error(body: &[u8]) -> Option<(Rid, DetokMsg)> { + let val = rmpv::decode::read_value(&mut &body[..]).ok()?; + let rmpv::Value::Array(arr) = val else { + return None; + }; + let mut items = arr.into_iter(); + let rid = Rid::from(items.next()?.as_str()?); + let message = match items.next()? { + rmpv::Value::String(s) => s.into_str()?, + _ => return None, + }; + Some((rid.clone(), DetokMsg::Fail { rid, message })) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::DetokMsg; + use crate::message::frame_egress_error; + + /// A framed error round-trips: `frame_egress_error` → tag stripped → + /// `decode_error` yields the rid + a `Fail` carrying the message. + #[test] + fn error_frame_roundtrips_to_fail() { + let framed = frame_egress_error("42", "invalid request: bad field"); + assert_eq!(framed[0], EGRESS_TAG_ERROR); + let (rid, msg) = decode_error(&framed[1..]).expect("decodes"); + let want = Rid::from("42"); + assert_eq!(rid, want); + match msg { + DetokMsg::Fail { rid, message } => { + assert_eq!(rid.clone(), want); + assert_eq!(message, "invalid request: bad field"); + } + _ => panic!("expected Fail"), + } + } +} diff --git a/rust/sglang-server/src/tokenizer_manager/ingress.rs b/rust/sglang-server/src/tokenizer_manager/ingress.rs new file mode 100644 index 000000000000..57d425a490bd --- /dev/null +++ b/rust/sglang-server/src/tokenizer_manager/ingress.rs @@ -0,0 +1,1210 @@ +//! TokenizerManager — ingress side. +//! +//! [`Ingress`] is a single-consumer stage draining one inbox fed by both the API +//! server (fresh requests) and the Tokenizer pool (returned requests). It owns +//! the request while driving the ingress FSM and hands it off by *moving* it to +//! the next stage; nothing here is shared, so no locks. +//! +//! Edges driven here (from the design table): +//! Received → Validating +//! Validating → Normalizing (generate: sampling-param normalize/verify) +//! Validating → PreSendValidating (control: no tokenize, no sampling params) +//! Normalizing → {Encoding | Tokenizing | PreSendValidating} (by ValidationOutcome) +//! Tokenizing → PreSendValidating (on TokenizeDone, when the request returns) +//! PreSendValidating → Queued (checks needing the tokenized length) +//! Queued → ring (handed to the scheduler) +//! +//! The egress edges (Streaming/Finalizing/Completed) are driven on the egress +//! side (see `egress` + `detokenizer`). + +use bytes::Bytes; + +use crate::error::Error; +use crate::fsm::{Event, RequestState, ValidationOutcome}; + +use crate::message::{ + AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, Request, + RequestKind, +}; +use crate::ring::IngressProducer; +use crate::runtime::{Runnable, ServerArgs}; +use crate::tokenizer_manager::{AbortSource, Senders, TmEvent}; + +/// Ingress FSM dispatcher stage. Owns its inbox + downstream handles, so the +/// runtime spawns it as a [`Runnable`] rather than calling a free `run_*` fn +/// with positional arguments. +pub struct Ingress { + rx: flume::Receiver, + /// Unbounded abort lane (see [`Senders::abort`]). Selected against `rx` so an + /// abort is handled promptly even while the bounded inbox is saturated. + abort_rx: flume::Receiver, + senders: Senders, + ingress: IngressProducer, + limits: Limits, + shutdown: flume::Receiver<()>, +} + +/// Longest client-supplied rid accepted. It keys the detok table and travels on +/// every chunk, so its length is a recurring cost; Python mints 32-byte uuid hex. +const MAX_RID_LEN: usize = 128; + +/// What ingress admits, resolved once at boot from the scheduler's `server_args`. +/// A struct rather than more positional `new` arguments — these grew from two to +/// six, and every one of them is a `u64`/`bool` that would be trivial to swap at +/// a call site. +/// +/// NOT `Default`-able on purpose. `vocab_size` and `context_len` are mandatory, +/// and their zero value is the most restrictive setting there is — a derived +/// `Default` would silently build limits that reject every request rather than +/// failing loudly. Tests construct these explicitly (see `test_limits`). +#[derive(Clone, Debug)] +pub struct Limits { + /// Token-ids-in mode: a generate request must arrive already tokenized. + pub skip_tokenizer_init: bool, + /// `model_config.vocab_size`; bounds client-supplied token ids. Mandatory — + /// [`ServerArgs::validate_mandatory`](crate::runtime::ServerArgs) rejects a + /// boot without it, so ingress can check unconditionally. + pub vocab_size: u64, + /// `model_config.context_len`, the ceiling for input + `max_new_tokens`. + /// Mandatory, as above. + pub context_len: u64, + /// Output slots reserved on top of the input (eagle draft tokens). + pub num_reserved_tokens: u64, + /// Clamp `max_new_tokens` to what fits instead of rejecting the request. + pub allow_auto_truncate: bool, + /// Whether the server can produce hidden states at all. + pub enable_return_hidden_states: bool, +} + +impl TryFrom<&ServerArgs> for Limits { + type Error = Error; + + fn try_from(sa: &ServerArgs) -> Result { + Ok(Self { + skip_tokenizer_init: sa.skip_tokenizer_init, + vocab_size: sa + .model_config + .vocab_size + .ok_or_else(|| Error::Validation("vocab_size missing".into()))?, + context_len: sa + .model_config + .context_len + .ok_or_else(|| Error::Validation("context_len missing".into()))?, + num_reserved_tokens: sa.num_reserved_tokens, + allow_auto_truncate: sa.allow_auto_truncate, + enable_return_hidden_states: sa.enable_return_hidden_states, + }) + } +} + +impl Ingress { + pub fn new( + rx: flume::Receiver, + abort_rx: flume::Receiver, + senders: Senders, + ingress: IngressProducer, + limits: Limits, + shutdown: flume::Receiver<()>, + ) -> Self { + Self { + rx, + abort_rx, + senders, + ingress, + limits, + shutdown, + } + } +} + +/// Which lane produced the next item. +enum Lane { + Abort(AbortSource), + Event(TmEvent), +} + +impl Runnable for Ingress { + fn run(self) { + loop { + // Select, not a drain-then-block: an abort arriving while the inbox is + // idle must still be handled at once. + let next = flume::Selector::new() + .recv(&self.abort_rx, |r| r.ok().map(Lane::Abort)) + .recv(&self.rx, |r| r.ok().map(Lane::Event)) + .recv(&self.shutdown, |_| None) + .wait(); + match next { + Some(Lane::Abort(rid)) => self.on_abort(rid), + // A fresh request and one returning from the tokenizer pool. + Some(Lane::Event(TmEvent::Ingress(req) | TmEvent::Tokenized(req))) => { + self.drive(req) + } + None => { + // Shutdown, or the inbox closed. Drain whatever is still queued + // on the abort lane first: those requests are in flight on the + // scheduler, and the selector may report the closed inbox before + // it ever looks at a pending abort. + while let Ok(source) = self.abort_rx.try_recv() { + self.on_abort(source); + } + return; + } + } + } + } +} + +impl Ingress { + /// Reject a request: → `Failed`, notify the client, deregister (unconditional + /// — a no-op when nothing was registered). + /// `registered` says whether this request ever reached `register_detok`. It + /// must: `Deregister`'s handler is a bare `table.remove(&rid)`, so a + /// request rejected BEFORE registering would evict whatever entry currently + /// holds that key — a concurrent request's sink — leaving that client with no + /// terminal frame and a hung connection. Python cannot hit this because it + /// validates before `rid_to_state[obj.rid] = state`. + fn fail(&self, req: &mut Request, err: Error, registered: bool) { + // Log only server faults (500); 4xx/499/503 are expected and would spam. + if err.http_status() == 500 { + tracing::error!(rid = %req.rid, error = %err, "ingress rejected request"); + } + let _ = req.state.apply(Event::Error(err.clone())); + let _ = req.sink.try_send(EgressItem::Error(err)); // client may be gone + if registered { + let _ = self.senders.detok_for(&req.rid).send(DetokMsg::Deregister { + rid: req.rid.clone(), + }); + } + } + + /// Drive a request through its ingress states until it terminates (failed or + /// pushed to the ring) or is handed to the tokenizer pool (re-entering as a + /// `Tokenized` event). Each arm acts and advances the FSM; the loop + /// re-dispatches. The arms are the design table's states, `Failed` the single + /// reject path. + fn drive(&self, mut req: Request) { + // Flipped once `register_detok` succeeds; `fail` must not deregister before + // that (see `fail`). A pool return re-enters `drive` already registered. + let mut registered = !matches!(req.state, RequestState::Received); + loop { + match req.state.clone() { + // Validate, then register the sink before the request leaves Rust. + // Failures move to `Failed` and fall through to the reject arm. + RequestState::Received => { + if let Err(e) = validate(&mut req, &self.limits) { + let _ = req.state.apply(Event::Error(e)); // → Failed + continue; + } + if !self.register_detok(&req) { + let _ = req + .state + .apply(Event::Error(Error::Internal("detok shard gone".into()))); + continue; + } + registered = true; + // `validate` advanced Received → Validating; keep driving. + } + // Control skips normalization (no sampling params) straight to the + // pre-send checks; generate goes to Normalizing. + RequestState::Validating => match &req.kind { + RequestKind::Control(_) => { + let _ = req + .state + .apply(Event::Validated(ValidationOutcome::AlreadyTokenized)); + } + RequestKind::Generate(_) => { + let _ = req.state.apply(Event::NeedsNormalize); + } + }, + // Normalize + verify sampling params (off the scheduler loop), then + // pick the branch; a bad param becomes `Failed`. + RequestState::Normalizing => { + let outcome = { + let RequestKind::Generate(g) = &mut req.kind else { + // Unreachable (control never reaches here); reject so a + // bug can't leak/hang a registered request. + self.fail( + &mut req, + Error::Internal("non-generate request in Normalizing".into()), + registered, + ); + return; + }; + match g + .sampling_params + .normalize(self.limits.skip_tokenizer_init, self.limits.vocab_size) + { + Err(e) => Err(e), + // Client ids skip the pool; text goes to the tokenizer. + Ok(()) if g.already_tokenized() => { + Ok(ValidationOutcome::AlreadyTokenized) + } + Ok(()) => Ok(ValidationOutcome::NeedsTokenize), + } + }; + match outcome { + Err(e) => { + let _ = req.state.apply(Event::Error(e)); // → Failed + } + Ok(o) => { + // AlreadyTokenized → Queued, NeedsTokenize → Tokenizing. + let _ = req.state.apply(Event::Validated(o)); + } + } + } + // Hand off to the tokenizer pool; it returns the request as a + // `Tokenized` event (PreSendValidating, or Failed on error). + // Doesn't loop. + RequestState::Tokenizing => { + if let Err(err) = self.senders.tok.send(req) { + // Pool gone (workers exited); flume hands the request back. + let mut req = err.into_inner(); + // Past `Received`, so registration happened. + self.fail( + &mut req, + Error::Internal("tokenizer pool gone".into()), + true, + ); + } + return; + } + // The checks that need the final `input_ids`: every branch + // converges here (client ids arrive directly, text arrives from + // the tokenizer pool), so they run once per request regardless of + // how it was tokenized. `validate` runs too early — at `Received` + // a text request has no ids yet. + RequestState::PreSendValidating => { + if let RequestKind::Generate(g) = &mut req.kind + && let Err(e) = check_total_tokens(g, &self.limits) + { + let _ = req.state.apply(Event::Error(e)); // → Failed + continue; + } + let _ = req.state.apply(Event::PreSendValidated); // → Queued + } + // Push the wire message (control frame or generate payload) to the ring. + RequestState::Queued => { + // `matches!` reads the discriminant without holding a borrow, + // so `req` can be moved into the push below. + if matches!(req.kind, RequestKind::Generate(_)) { + self.push_to_ring(req); + } else { + self.push_control_to_ring(req); + } + return; + } + // The single reject path for every post-register failure. + RequestState::Failed(e) => { + self.fail(&mut req, e, registered); + return; + } + // Unreachable (egress states never reach here). Reject via `fail`/ + // return (not apply + continue, which would spin on a terminal state). + other => { + self.fail( + &mut req, + Error::Internal(format!("unexpected ingress state: {other:?}")), + registered, + ); + return; + } + } + } + } + + /// Register the egress sink with the owning detok shard (by id) so the response + /// has a home. Carries the per-request detok flags — `return_text_in_logprobs` + /// (decode logprob text on this shard) and `no_stop_trim` (keep the matched + /// stop in the output) — so the shard needs no back-reference to the request. + /// Returns `false` if the shard is gone. + fn register_detok(&self, req: &Request) -> bool { + let (decode_logprob_text, no_stop_trim) = match &req.kind { + RequestKind::Generate(g) => ( + g.return_text_in_logprobs.unwrap_or(false), + g.sampling_params.no_stop_trim, + ), + RequestKind::Control(_) => (false, false), + }; + self.senders + .detok_for(&req.rid) + .send(DetokMsg::Register { + rid: req.rid.clone(), + sink: req.sink.clone(), + decode_logprob_text, + no_stop_trim, + }) + .is_ok() + } + + /// Push a bare control request (`[tag, rid, nil]`) onto the ingress ring. The + /// scheduler dispatches it (e.g. `GetInternalStateReq`) and replies via the + /// egress ring as a single `Result`. + fn push_control_to_ring(&self, mut req: Request) { + let encode = match &req.kind { + RequestKind::Control(control) => control.encode(), + _ => Err(Error::Internal( + "non-control request reached push_control_to_ring".into(), + )), + }; + let header = match encode { + Ok(b) => b, + Err(e) => { + self.fail(&mut req, e, true); // on the push path: registered + return; + } + }; + // Control requests carry no tensor cell — empty `ids`. + if !self.ingress.try_push(IngressMsg { + header, + ids: Bytes::new(), + }) { + self.fail(&mut req, Error::QueueFull, true); // registered + } + } + + /// Client disconnected (or a detok terminal): deregister the sink, then push an + /// `AbortReq(rid)` so the scheduler stops generating for it. + /// + /// A failed push is logged, not retried: the scheduler keeps generating and the + /// chunks arrive for a rid no longer in the detok table, where they are dropped. + /// That wastes GPU work until the request finishes on its own, but it cannot be + /// misdelivered — the rid is unique to this request for the process's lifetime + /// ([`Rid::from_client`]), so no later request can ever answer to it. + fn on_abort(&self, source: AbortSource) { + let rid = source.rid().clone(); + let _ = self + .senders + .detok_for(&rid) + .send(DetokMsg::Deregister { rid: rid.clone() }); + + // The ring is BOUNDED and drops pushes under exactly the load this matters + // for, so report the miss rather than assuming the scheduler was told. + match ControlRequest::AbortReq(AbortReq::new(rid.as_str().to_string(), false)).encode() { + Ok(header) => { + if !self.ingress.try_push(IngressMsg { + header, + ids: Bytes::new(), + }) { + tracing::error!( + rid = %rid, + "abort dropped: ingress ring full; the scheduler keeps generating \ + for this request until it finishes on its own" + ); + } + } + Err(e) => tracing::error!(rid = %rid, error = %e, "abort encode failed"), + } + } + + /// Serialize the tokenized request to its `TokenizedGenerateReqInput` wire and + /// push it onto the ingress ring for the scheduler. On backpressure, fail it. + fn push_to_ring(&self, mut req: Request) { + // Only generate requests reach here (control uses `push_control_to_ring`). + // Validate + serialize while borrowing `g` immutably; the resulting `Bytes` + // own their data, so the borrow ends before any `fail(&mut req)`. + let serialized = match &req.kind { + RequestKind::Generate(g) if g.already_tokenized() => g + .encode_header() + .map(|header| (header, g.encode_data_buf())), + RequestKind::Generate(_) => Err(Error::Tokenize("empty input_ids".into())), + _ => Err(Error::Internal( + "non-generate request reached push_to_ring".into(), + )), + }; + let (header, ids) = match serialized { + Ok(v) => v, + Err(e) => { + self.fail(&mut req, e, true); // on the push path: registered + return; + } + }; + + if !self.ingress.try_push(IngressMsg { header, ids }) { + self.fail(&mut req, Error::QueueFull, true); // registered + } + // On success the scheduler owns the request (egress arrives by rid); we + // drop our `Request` here — the detok shard holds the sink. + } +} + +/// `Received → Validating` + admissibility check. Under `skip_tokenizer_init` a +/// generate request must already carry token ids (no tokenizer to byte-encode +/// text); control requests carry none and are exempt. +fn validate(req: &mut Request, limits: &Limits) -> Result<(), Error> { + let (skip_tokenizer_init, vocab_size) = (limits.skip_tokenizer_init, limits.vocab_size); + let _ = req + .state + .apply(Event::Validated(ValidationOutcome::NeedsTokenize)); + + // The rid is the request's identity everywhere downstream: it keys the detok + // table, and it rides on EVERY chunk of EVERY decode step. An unbounded + // client-supplied rid is therefore a per-step cost, not a one-off. Python's is + // a 32-byte uuid hex, so this is generous. + // Measured on the CLIENT-facing form: the uniquifier `Rid::from_client` appends + // is this server's own overhead, and charging the client for bytes it did not + // send would reject a rid exactly at the documented limit. + let client_rid_len = req.rid.client_facing().len(); + if client_rid_len > MAX_RID_LEN { + return Err(Error::Validation(format!( + "rid is {client_rid_len} bytes, over the {MAX_RID_LEN}-byte limit" + ))); + } + if skip_tokenizer_init + && matches!(&req.kind, RequestKind::Generate(g) if !g.already_tokenized()) + { + // `Validation` (400), not `Tokenize` (500): the client sent a request this + // server cannot serve, which is their error to fix — Python 400s it too. + return Err(Error::Validation( + "skip_tokenizer_init is set: request must provide input_ids".into(), + )); + } + + // Client-supplied token ids must be in-vocabulary: an out-of-range id + // reaches the embedding lookup and kills the scheduler process, so 400 + // here instead — mirroring the Python `TokenizerManager` validation. + if let RequestKind::Generate(g) = &req.kind { + if let Some(ids) = &g.input_ids { + for &id in ids { + if id < 0 || id as u64 >= vocab_size { + return Err(Error::Validation(format!( + "input_ids contains out-of-vocabulary token id {id}; \ + valid range is [0, {vocab_size})" + ))); + } + } + } + if let Some(ids) = &g.token_ids_logprob { + for &id in ids { + if id < 0 || id as u64 >= vocab_size { + return Err(Error::Validation(format!( + "token_ids_logprob contains out-of-vocabulary token id \ + {id}; valid range is [0, {vocab_size})" + ))); + } + } + } + } + + // The scheduler only computes hidden states when launched for it, so without + // this the request would 200 with `meta_info.hidden_states` silently absent + // (Python `TokenizerManager._validate_one_request`). + if !limits.enable_return_hidden_states + && matches!(&req.kind, RequestKind::Generate(g) if g.return_hidden_states) + { + return Err(Error::Validation( + "The server is not configured to return the hidden states. \ + Please set `--enable-return-hidden-states` to enable this feature." + .into(), + )); + } + + Ok(()) +} + +/// The context-window checks that need the tokenized length, mirroring Python +/// `TokenizerManager._validate_one_request`: the input alone must fit, and then +/// input + `max_new_tokens` must fit. Without them the scheduler silently clamps +/// and the client gets a 200 with a truncated completion instead of an actionable +/// 400. +/// +/// Under `allow_auto_truncate` both clamp instead of rejecting — the launch flag +/// opted into that. +fn check_total_tokens(g: &mut GenerateRequest, limits: &Limits) -> Result<(), Error> { + let max_req_len = limits.context_len; + // Python counts the reserved slots as part of the input, so a request can be + // rejected for them even when the prompt alone fits. + let input_len = + g.input_ids.as_ref().map_or(0, |ids| ids.len()) as u64 + limits.num_reserved_tokens; + + // Input length first, and unconditionally: `max_new_tokens: null` means "no + // cap", which must not disable this. Python's comparison is `>=` — a prompt + // that exactly fills the window leaves no room to generate. + if input_len >= max_req_len { + if !limits.allow_auto_truncate { + return Err(Error::Validation(format!( + "The input ({input_len} tokens) is longer than the model's context \ + length ({max_req_len} tokens)." + ))); + } + if let Some(ids) = &mut g.input_ids { + ids.truncate(max_req_len as usize); + } + } + let input_len = + g.input_ids.as_ref().map_or(0, |ids| ids.len()) as u64 + limits.num_reserved_tokens; + + let Some(max_new_tokens) = g.sampling_params.max_new_tokens else { + return Ok(()); // no cap requested → nothing to add to the input length + }; + let total = input_len.saturating_add(max_new_tokens.max(0) as u64); + if total <= max_req_len { + return Ok(()); + } + if !limits.allow_auto_truncate { + return Err(Error::Validation(format!( + "Requested token count exceeds the model's maximum context length of \ + {max_req_len} tokens. You requested a total of {total} tokens: {input_len} \ + tokens from the input messages and {max_new_tokens} tokens for the \ + completion. Please reduce the number of tokens in the input messages or \ + the completion to fit within the limit." + ))); + } + let clamped = max_req_len.saturating_sub(input_len) as i64; + // Re-check what the clamp can break. `verify` already ran (in Normalizing), so + // lowering `max_new_tokens` here can leave `min_new_tokens > max_new_tokens` — + // and `is_normalized: true` stops the scheduler from re-verifying, so nothing + // downstream would catch it. Python validates before it verifies; we can't + // reorder the FSM, so we re-assert the one invariant the clamp can violate. + if g.sampling_params.min_new_tokens > clamped { + return Err(Error::Validation(format!( + "min_new_tokens must be in [0, max_new_tokens({clamped})], got {}", + g.sampling_params.min_new_tokens + ))); + } + g.sampling_params.max_new_tokens = Some(clamped); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::fsm::RequestState; + use crate::message::{EgressSink, GenerateRequest, SamplingParams}; + use crate::ring::{IngressConsumer, ingress_ring}; + use tokio::sync::mpsc; + + /// An `Ingress` plus its detok-shard receiver, ring consumer (keep alive — + /// dropping it closes the ring → false QueueFull), and tm inbox sender. + fn make_ingress() -> ( + Ingress, + flume::Receiver, + IngressConsumer, + flume::Sender, + ) { + make_ingress_with(test_limits()) + } + + fn make_ingress_with_abort( + abort_rx: flume::Receiver, + ) -> ( + Ingress, + flume::Receiver, + IngressConsumer, + flume::Sender, + ) { + make_ingress_inner(test_limits(), abort_rx) + } + + fn make_ingress_with( + limits: Limits, + ) -> ( + Ingress, + flume::Receiver, + IngressConsumer, + flume::Sender, + ) { + let (abort_tx, abort_rx) = flume::unbounded::(); + std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx + make_ingress_inner(limits, abort_rx) + } + + fn make_ingress_inner( + limits: Limits, + abort_rx: flume::Receiver, + ) -> ( + Ingress, + flume::Receiver, + IngressConsumer, + flume::Sender, + ) { + let (tok_tx, _tok_rx) = flume::unbounded(); + let (detok_tx, detok_rx) = flume::unbounded(); + let senders = Senders { + tm: flume::unbounded().0, + abort: flume::unbounded().0, + tok: tok_tx, + detok: vec![detok_tx], + }; + let (ingress_producer, consumer) = ingress_ring(16); + let (tm_tx, tm_rx) = flume::unbounded(); + // Keep the shutdown sender alive (leak) so its branch never fires — tests + // end `run` by dropping `tm_tx`, not by shutdown. + let (sd_tx, sd_rx) = flume::unbounded::<()>(); + std::mem::forget(sd_tx); + let ingress = Ingress::new(tm_rx, abort_rx, senders, ingress_producer, limits, sd_rx); + (ingress, detok_rx, consumer, tm_tx) + } + + /// Both abort sources do the same two things: drop the detok entry so no + /// further chunk can be delivered, and tell the scheduler to stop generating. + /// + /// Neither releases anything, and nothing needs them to. Release ordering used + /// to be the delicate part here — `AbortGuard::drop` releasing a rid right + /// after enqueuing the abort ordered the SEND, not the EFFECT, so a retry of + /// the same rid could `Register` ahead of the stale abort and be torn down by + /// it. `Rid::from_client` removes the premise: a retry carries a different + /// `Rid`, so no abort in flight can name it. + #[test] + fn every_abort_source_deregisters_and_stops_the_scheduler() { + for source in [ + AbortSource::Guard("x".into()), + AbortSource::Detok("x".into()), + ] { + let (detok_tx, detok_rx) = flume::unbounded::(); + let (ingress_producer, consumer) = ingress_ring(16); + let (sd_tx, sd_rx) = flume::unbounded::<()>(); + std::mem::forget(sd_tx); + let ingress = Ingress::new( + flume::unbounded().1, + flume::unbounded().1, + Senders { + tm: flume::unbounded().0, + abort: flume::unbounded().0, + tok: flume::unbounded().0, + detok: vec![detok_tx], + }, + ingress_producer, + test_limits(), + sd_rx, + ); + + ingress.on_abort(source.clone()); + + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "x"), + "{source:?} must drop the detok entry", + ); + assert_eq!( + consumer.drain(8).headers.len(), + 1, + "{source:?} must push an AbortReq so the scheduler stops", + ); + } + } + + /// A context ceiling high enough that only a test which sets one on purpose + /// can reach it. `context_len` is mandatory now, so "no ceiling" has to be a + /// large number rather than `None`; kept well below `u64::MAX` so the + /// `as i64` in the auto-truncate clamp cannot go negative if a future test + /// does reach this path. + const NO_CONTEXT_CEILING: u64 = 1 << 40; + + /// The default test limits: a real tokenizer, vocab 1000, no context ceiling. + /// Spelled out rather than `..Default::default()` — `Limits` deliberately has + /// no `Default`, because a zero `vocab_size`/`context_len` would reject every + /// request instead of behaving like "unset". + fn test_limits() -> Limits { + Limits { + skip_tokenizer_init: false, + vocab_size: 1000, + context_len: NO_CONTEXT_CEILING, + num_reserved_tokens: 0, + allow_auto_truncate: false, + enable_return_hidden_states: false, + } + } + + fn generate_req(id: u64, sampling_params: SamplingParams) -> Request { + let (tx, _rx) = mpsc::channel(8); + Request { + rid: id.to_string().into(), + state: RequestState::Received, + sink: EgressSink::Local(tx), + kind: RequestKind::Generate(Box::new(GenerateRequest { + rid: id.to_string().into(), + input_ids: Some(vec![1, 2, 3]), + sampling_params, + ..Default::default() + })), + } + } + + /// `input + max_new_tokens` past the context window is an actionable 400, not a + /// silently truncated 200 (Python `TokenizerManager._validate_one_request`). + /// The message names both halves so the client can fix the right one. + #[test] + fn total_tokens_over_context_is_rejected() { + let limits = Limits { + context_len: 10, + ..test_limits() + }; + let mut g = GenerateRequest { + input_ids: Some(vec![1, 2, 3]), + sampling_params: SamplingParams { + max_new_tokens: Some(100), + ..Default::default() + }, + ..Default::default() + }; + let err = check_total_tokens(&mut g, &limits).unwrap_err(); + let msg = err.to_string(); + assert_eq!(err.http_status(), 400); + assert!(msg.contains("total of 103 tokens"), "{msg}"); + assert!(msg.contains("3 tokens from the input"), "{msg}"); + assert!(msg.contains("100 tokens for the completion"), "{msg}"); + // Exactly filling the window is allowed (Python compares with `>`). + g.sampling_params.max_new_tokens = Some(7); + assert!(check_total_tokens(&mut g, &limits).is_ok()); + assert_eq!(g.sampling_params.max_new_tokens, Some(7), "left alone"); + } + + /// The reserved slots (eagle draft tokens) count as input, so a request can be + /// rejected for them even when the prompt alone would fit. + #[test] + fn reserved_tokens_count_toward_the_limit() { + let limits = Limits { + context_len: 10, + num_reserved_tokens: 5, + ..test_limits() + }; + let mut g = GenerateRequest { + input_ids: Some(vec![1, 2, 3]), + sampling_params: SamplingParams { + max_new_tokens: Some(3), // 3 + 3 fits, but 3 + 5 + 3 does not + ..Default::default() + }, + ..Default::default() + }; + let msg = check_total_tokens(&mut g, &limits).unwrap_err().to_string(); + assert!(msg.contains("8 tokens from the input"), "{msg}"); + } + + /// `--allow-auto-truncate` opts into clamping instead of rejecting; with no + /// context length, or no `max_new_tokens` cap, there is nothing to check. + #[test] + fn auto_truncate_clamps_and_unknowns_skip() { + let sp = |max_new_tokens| SamplingParams { + max_new_tokens, + ..Default::default() + }; + let mut g = GenerateRequest { + input_ids: Some(vec![1, 2, 3]), + sampling_params: sp(Some(100)), + ..Default::default() + }; + let truncating = Limits { + context_len: 10, + allow_auto_truncate: true, + ..test_limits() + }; + assert!(check_total_tokens(&mut g, &truncating).is_ok()); + assert_eq!(g.sampling_params.max_new_tokens, Some(7), "clamped to fit"); + + // Unknown context length → no ceiling to enforce. + g.sampling_params = sp(Some(100)); + assert!(check_total_tokens(&mut g, &test_limits()).is_ok()); + assert_eq!(g.sampling_params.max_new_tokens, Some(100), "untouched"); + + // No cap requested → nothing to add to the input length, but the input + // itself is still checked (see `input_length_is_checked_unconditionally`). + g.sampling_params = sp(None); + let roomy = Limits { + context_len: 100, + ..test_limits() + }; + assert!(check_total_tokens(&mut g, &roomy).is_ok()); + } + + /// `max_new_tokens: null` means "no cap", NOT "skip the checks" — the input + /// alone must still fit. Gating the whole function on `max_new_tokens` let an + /// over-long prompt through to the scheduler with no ingress error at all. + /// Python compares with `>=`: a prompt that exactly fills the window leaves no + /// room to generate. + #[test] + fn input_length_is_checked_unconditionally() { + let limits = Limits { + context_len: 3, + ..test_limits() + }; + let req = |max_new_tokens| GenerateRequest { + input_ids: Some(vec![1, 2, 3]), // exactly fills a 3-token window + sampling_params: SamplingParams { + max_new_tokens, + ..Default::default() + }, + ..Default::default() + }; + for max_new_tokens in [None, Some(1)] { + let err = check_total_tokens(&mut req(max_new_tokens), &limits) + .expect_err("input == context_len must be rejected (Python uses >=)"); + assert_eq!(err.http_status(), 400); + assert!(err.to_string().contains("longer than the model's context")); + } + // One token shorter fits, with or without a cap. + let mut g = GenerateRequest { + input_ids: Some(vec![1, 2]), + ..Default::default() + }; + g.sampling_params.max_new_tokens = None; + assert!(check_total_tokens(&mut g, &limits).is_ok()); + + // Under auto-truncate the input is cut to fit instead of rejected. + let truncating = Limits { + allow_auto_truncate: true, + ..limits.clone() + }; + let mut g = req(None); + assert!(check_total_tokens(&mut g, &truncating).is_ok()); + assert_eq!( + g.input_ids.as_deref(), + Some(&[1, 2, 3][..]), + "fits at the cap" + ); + } + + /// The clamp runs AFTER `verify` (which happens in `Normalizing`), so lowering + /// `max_new_tokens` can leave `min_new_tokens > max_new_tokens`. Nothing + /// downstream re-checks — `is_normalized: true` makes the scheduler's own + /// verify early-return — so the clamp has to re-assert it here. + #[test] + fn auto_truncate_cannot_invert_min_and_max_new_tokens() { + let limits = Limits { + context_len: 10, + allow_auto_truncate: true, + ..test_limits() + }; + let mut g = GenerateRequest { + input_ids: Some(vec![1, 2, 3]), // clamps max_new_tokens to 7 + sampling_params: SamplingParams { + max_new_tokens: Some(100), + min_new_tokens: 50, // …which is below min_new_tokens + ..Default::default() + }, + ..Default::default() + }; + let err = check_total_tokens(&mut g, &limits) + .expect_err("a clamp that inverts min/max must 400, not ride the wire"); + assert_eq!(err.http_status(), 400); + assert!(err.to_string().contains("min_new_tokens"), "{err}"); + + // A clamp that keeps the invariant still clamps. + g.sampling_params.min_new_tokens = 2; + g.sampling_params.max_new_tokens = Some(100); + assert!(check_total_tokens(&mut g, &limits).is_ok()); + assert_eq!(g.sampling_params.max_new_tokens, Some(7)); + } + + /// `return_hidden_states` on a server not launched for it is a 400: the + /// scheduler never computes them, so the request would otherwise 200 with + /// `meta_info.hidden_states` silently missing. + #[test] + fn hidden_states_gated_on_server_support() { + let req = |want| { + let mut r = generate_req(31, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut r.kind { + g.return_hidden_states = want; + } + r + }; + let disabled = test_limits(); + let err = validate(&mut req(true), &disabled).unwrap_err(); + assert_eq!(err.http_status(), 400); + assert!( + err.to_string().contains("--enable-return-hidden-states"), + "message must name the flag: {err}" + ); + // Not asking for them (the client sent `false`, or sent nothing and + // `into_requests` resolved the default), or asking on a server that + // supports them, is fine. + assert!(validate(&mut req(false), &disabled).is_ok()); + let enabled = Limits { + enable_return_hidden_states: true, + ..test_limits() + }; + assert!(validate(&mut req(true), &enabled).is_ok()); + } + + /// End-to-end through `drive`: an over-context request is rejected on the way + /// to the ring, after registration — so it must be deregistered, not leaked. + #[test] + fn over_context_request_deregisters_and_never_reaches_the_ring() { + let (ingress, detok_rx, consumer, _tm_tx) = make_ingress_with(Limits { + context_len: 4, + ..test_limits() + }); + ingress.drive(generate_req( + 33, + SamplingParams { + max_new_tokens: Some(64), + ..Default::default() + }, + )); + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "33"), + "registered before the check", + ); + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "33"), + "must deregister on reject", + ); + assert!( + consumer.drain(16).headers.is_empty(), + "must not reach the scheduler" + ); + } + + /// A dropped ring push is survivable, and this pins WHY. The ring is bounded, + /// so under load the scheduler never learns to stop and keeps generating; its + /// chunks then arrive for a rid the detok table no longer holds and are + /// dropped. That wastes GPU work but cannot MISDELIVER, because + /// `Rid::from_client` guarantees no later request ever answers to that rid. + /// The detok entry is dropped either way — that is the half that must not + /// depend on the ring. + /// + /// Ring capacity 1: the first abort pushes, the second finds it full. + #[test] + fn abort_deregisters_even_when_the_ring_push_is_dropped() { + let (tok_tx, _tok_rx) = flume::unbounded(); + let (detok_tx, detok_rx) = flume::unbounded(); + let (abort_tx, abort_rx) = flume::unbounded::(); + let senders = Senders { + tm: flume::unbounded().0, + abort: abort_tx, + tok: tok_tx, + detok: vec![detok_tx], + }; + let (producer, _consumer) = ingress_ring(1); + let (_tm_tx, tm_rx) = flume::unbounded(); + let (sd_tx, sd_rx) = flume::unbounded::<()>(); + std::mem::forget(sd_tx); + let ingress = Ingress::new(tm_rx, abort_rx, senders, producer, test_limits(), sd_rx); + + ingress.on_abort(AbortSource::Guard("pushed".into())); + ingress.on_abort(AbortSource::Guard("dropped".into())); + + // Both deregisters land regardless of whether the ring accepted the push. + for expected in ["pushed", "dropped"] { + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == expected), + "{expected}: the detok entry must be dropped even when the ring is full", + ); + } + } + + /// The rid keys the detok table and rides on every chunk of every decode step, + /// so an unbounded client-supplied one is a recurring cost, not a one-off. + #[test] + fn oversized_rid_is_rejected() { + let mut req = generate_req(51, SamplingParams::default()); + req.rid = "x".repeat(MAX_RID_LEN + 1).into(); + let err = validate(&mut req, &test_limits()).expect_err("must be rejected"); + assert_eq!(err.http_status(), 400); + assert!(err.to_string().contains("over the"), "{err}"); + + // A uuid-sized rid — what Python mints — is nowhere near the cap. + let mut req = generate_req(52, SamplingParams::default()); + req.rid = "0123456789abcdef0123456789abcdef".into(); + assert!(validate(&mut req, &test_limits()).is_ok()); + } + + /// A request rejected BEFORE `register_detok` must not send `Deregister`: the + /// handler is a bare `table.remove(&rid)`, so it would evict whatever entry + /// holds that key — a concurrent request's sink — leaving that client hung with + /// no terminal frame. Python validates before it inserts, so it cannot hit this. + #[test] + fn pre_registration_failure_does_not_deregister() { + // Rejected inside `validate` (out-of-vocab id), which runs before registration. + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let mut req = generate_req(41, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut req.kind { + g.input_ids = Some(vec![2_000_000_000]); + } + ingress.drive(req); + assert!( + detok_rx.try_recv().is_err(), + "a pre-registration reject must send NOTHING to the shard — a Deregister \ + here removes a live request's sink" + ); + + // A post-registration reject still deregisters (the leak fix stays fixed). + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + ingress.drive(generate_req( + 42, + SamplingParams { + top_p: 2.0, // rejected by `normalize`, after registration + ..Default::default() + }, + )); + assert!(matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. }))); + assert!(matches!( + detok_rx.try_recv(), + Ok(DetokMsg::Deregister { .. }) + )); + } + + /// A request rejected at normalization (post-register) must not leak: the shard + /// sees `Register` then `Deregister`. Regression for RSS growth on bad input. + #[test] + fn rejected_request_deregisters_from_shard() { + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + // top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it. + let bad = SamplingParams { + top_p: 2.0, + ..Default::default() + }; + ingress.drive(generate_req(7, bad)); + + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "7"), + "expected Register for rid 7", + ); + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "7"), + "expected Deregister for rid 7 (leak fix)", + ); + assert!( + detok_rx.try_recv().is_err(), + "no further shard messages — registration fully cleaned up", + ); + } + + /// Regression: an out-of-vocabulary client token id must be rejected at + /// ingress with a 400 — passed through, it reaches the embedding lookup + /// and kills the scheduler process (`make_ingress` bounds vocab at 1000). + #[test] + fn out_of_vocab_input_ids_rejected() { + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let mut req = generate_req(21, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut req.kind { + g.input_ids = Some(vec![1, 2_000_000_000]); + } + ingress.drive(req); + // Rejected before registration: the only shard message is nothing at + // all, or a Deregister if registration happened first — never a push. + match detok_rx.try_recv() { + Err(_) => {} + Ok(DetokMsg::Deregister { .. }) => {} + Ok(_) => panic!("out-of-vocab request must not be admitted"), + } + } + + /// Same guard for negative ids and for `token_ids_logprob` entries. + #[test] + fn negative_and_logprob_token_ids_rejected() { + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let mut req = generate_req(22, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut req.kind { + g.input_ids = Some(vec![-1]); + } + ingress.drive(req); + match detok_rx.try_recv() { + Err(_) | Ok(DetokMsg::Deregister { .. }) => {} + Ok(_) => panic!("negative token id must not be admitted"), + } + + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let mut req = generate_req(23, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut req.kind { + g.token_ids_logprob = Some(vec![999_999]); + } + ingress.drive(req); + match detok_rx.try_recv() { + Err(_) | Ok(DetokMsg::Deregister { .. }) => {} + Ok(_) => panic!("out-of-vocab token_ids_logprob must not be admitted"), + } + } + + /// A valid request is registered and handed onward — never deregistered. + #[test] + fn admitted_request_keeps_registration() { + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + // Empty map → all sampling defaults, passes normalization. + ingress.drive(generate_req(9, SamplingParams::default())); + + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "9"), + "expected Register for rid 9", + ); + assert!( + detok_rx.try_recv().is_err(), + "admitted request must not be deregistered", + ); + } + + /// A pool return in `Failed` state (failed encode) is rejected via the same + /// path and deregistered, not leaked. + #[test] + fn tokenize_failure_deregisters_via_ingress() { + let (ingress, detok_rx, _consumer, tm_tx) = make_ingress(); + // The pool marks a failed encode as `Failed(err)` before returning it. + let mut req = generate_req(11, SamplingParams::default()); + let _ = req + .state + .apply(Event::Error(Error::Tokenize("boom".into()))); + tm_tx.send(TmEvent::Tokenized(req)).unwrap(); + // Close the inbox so the run loop returns after draining the one event. + drop(tm_tx); + ingress.run(); + + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "11"), + "tokenize failure must deregister rid 11", + ); + assert!(detok_rx.try_recv().is_err(), "no further shard messages"); + } + + /// An abort deregisters (by the id hashed from the rid string), so a request + /// aborted before any terminal chunk can't leak. + #[test] + fn abort_deregisters_from_shard() { + // Aborts arrive on their own unbounded lane now, not the request inbox. + let (abort_tx, abort_rx) = flume::unbounded::(); + let (ingress, detok_rx, _consumer, tm_tx) = make_ingress_with_abort(abort_rx); + abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap(); + drop(abort_tx); + drop(tm_tx); + ingress.run(); + + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "rid-13"), + "abort must deregister by rid", + ); + assert!(detok_rx.try_recv().is_err(), "no further shard messages"); + } + + /// A successful pool return (Queued, ids filled) is pushed to the ring, not + /// rejected; its registration is untouched. + #[test] + fn tokenized_return_pushes_without_deregister() { + let (ingress, detok_rx, _consumer, tm_tx) = make_ingress(); + let mut req = generate_req(15, SamplingParams::default()); + // Simulate a successful pool return: ids filled, PreSendValidating. + if let RequestKind::Generate(g) = &mut req.kind { + g.input_ids = Some(vec![1, 2, 3]); + } + req.state = RequestState::PreSendValidating; + tm_tx.send(TmEvent::Tokenized(req)).unwrap(); + drop(tm_tx); + ingress.run(); + + // Pushed to the ring; the shard sees nothing. + assert!( + detok_rx.try_recv().is_err(), + "a queued pool-return must be pushed, not touch the shard", + ); + } + + /// If the pool is gone, a request needing tokenization is rejected + + /// deregistered, not silently dropped. + #[test] + fn tokenize_pool_gone_deregisters() { + // `make_ingress` drops the tok receiver, so `tok.send` fails. + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + // No ids → NeedsTokenize → Tokenizing branch. + let mut req = generate_req(21, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut req.kind { + g.input_ids = None; + } + ingress.drive(req); + + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "21"), + "expected Register for rid 21", + ); + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "21"), + "pool-gone hand-off must deregister rid 21", + ); + assert!(detok_rx.try_recv().is_err(), "no further shard messages"); + } +}