diff --git a/Cargo.lock b/Cargo.lock index 22e04fb..6c71338 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -94,6 +94,28 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[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 = "async-trait" version = "0.1.89" @@ -172,6 +194,12 @@ dependencies = [ "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" @@ -224,6 +252,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd16c4719339c4530435d38e511904438d07cce7950afa3718a84ac36c10e89e" + [[package]] name = "chrono" version = "0.4.44" @@ -277,6 +311,15 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" +[[package]] +name = "clipboard-win" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bde03770d3df201d4fb868f2c9c59e66a3e4e2bd06692a0fe701e7103c7e84d4" +dependencies = [ + "error-code", +] + [[package]] name = "colorchoice" version = "1.0.5" @@ -293,6 +336,19 @@ dependencies = [ "windows-sys 0.59.0", ] +[[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 0.2.2", + "windows-sys 0.59.0", +] + [[package]] name = "console" version = "0.16.3" @@ -415,6 +471,41 @@ dependencies = [ "memchr", ] +[[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 = "der" version = "0.8.0" @@ -434,6 +525,37 @@ dependencies = [ "powerfmt", ] +[[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" @@ -517,6 +639,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "endian-type" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d" + [[package]] name = "equivalent" version = "1.0.2" @@ -533,6 +661,21 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "error-code" +version = "3.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59" + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + [[package]] name = "fallible-iterator" version = "0.3.0" @@ -551,6 +694,17 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +[[package]] +name = "fd-lock" +version = "4.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce92ff622d6dadf7349484f42c93271a0d49b7cc4d466a936405bacbe10aa78" +dependencies = [ + "cfg-if", + "rustix", + "windows-sys 0.59.0", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -800,12 +954,12 @@ dependencies = [ "dirs", "futures", "http", - "indicatif", + "indicatif 0.18.4", "libc", "log", "native-tls", "num_cpus", - "rand", + "rand 0.9.4", "reqwest", "serde", "serde_json", @@ -815,6 +969,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "http" version = "1.4.0" @@ -919,7 +1082,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -1050,6 +1213,12 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" +[[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" @@ -1083,13 +1252,26 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console 0.15.11", + "number_prefix", + "portable-atomic", + "unicode-width 0.2.2", + "web-time", +] + [[package]] name = "indicatif" version = "0.18.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" dependencies = [ - "console", + "console 0.16.3", "portable-atomic", "unicode-width 0.2.2", "unit-prefix", @@ -1129,6 +1311,24 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1218,6 +1418,22 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[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" @@ -1245,6 +1461,12 @@ 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" @@ -1266,6 +1488,28 @@ dependencies = [ "windows-sys 0.61.2", ] +[[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 = "native-tls" version = "0.2.18" @@ -1283,6 +1527,37 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nibble_vec" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43" +dependencies = [ + "smallvec", +] + +[[package]] +name = "nix" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab2156c4fce2f8df6c499cc1c763e4394b7482525bf2a9701c9d79d215f519e4" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + +[[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 = "ntapi" version = "0.4.3" @@ -1326,6 +1601,12 @@ dependencies = [ "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" @@ -1338,6 +1619,28 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[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 = "openssl" version = "0.10.78" @@ -1411,6 +1714,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pem-rfc7468" version = "1.0.0" @@ -1525,6 +1834,7 @@ dependencies = [ name = "puma" version = "0.0.3" dependencies = [ + "async-stream", "axum", "chrono", "clap", @@ -1532,17 +1842,21 @@ dependencies = [ "dirs", "futures", "hf-hub", - "indicatif", + "indicatif 0.18.4", "prettytable-rs", "regex", "reqwest", "rusqlite", "rusqlite_migration", + "rustyline", + "rustyline-derive", "serde", "serde_derive", "serde_json", "sysinfo", "tempfile", + "thiserror 2.0.18", + "tokenizers", "tokio", "tokio-stream", "tower 0.4.13", @@ -1573,14 +1887,45 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "radix_trie" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd" +dependencies = [ + "endian-type", + "nibble_vec", +] + +[[package]] +name = "rand" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -1590,7 +1935,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -1612,6 +1966,17 @@ dependencies = [ "rayon-core", ] +[[package]] +name = "rayon-cond" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "059f538b55efd2309c9794130bc149c6a553db90e9d99c2030785c82f0bd7df9" +dependencies = [ + "either", + "itertools 0.11.0", + "rayon", +] + [[package]] name = "rayon-core" version = "1.13.0" @@ -1688,7 +2053,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -1817,6 +2182,39 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +[[package]] +name = "rustyline" +version = "14.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7803e8936da37efd9b6d4478277f4b2b9bb5cdb37a113e8d63222e58da647e63" +dependencies = [ + "bitflags", + "cfg-if", + "clipboard-win", + "fd-lock", + "home", + "libc", + "log", + "memchr", + "nix", + "radix_trie", + "unicode-segmentation", + "unicode-width 0.1.14", + "utf8parse", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustyline-derive" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5af959c8bf6af1aff6d2b463a57f71aae53d1332da58419e30ad8dc7011d951" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "ryu" version = "1.0.23" @@ -1997,6 +2395,18 @@ dependencies = [ "winapi", ] +[[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" @@ -2195,6 +2605,38 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tokenizers" +version = "0.20.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b08cc37428a476fc9e20ac850132a513a2e1ce32b6a31addf2b74fa7033b905" +dependencies = [ + "aho-corasick", + "derive_builder", + "esaxx-rs", + "getrandom 0.2.17", + "indicatif 0.17.11", + "itertools 0.12.1", + "lazy_static", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.8.7", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 1.0.69", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.52.1" @@ -2419,6 +2861,21 @@ 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.1.14" @@ -2437,6 +2894,12 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "unit-prefix" version = "0.5.2" @@ -2455,7 +2918,7 @@ version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" dependencies = [ - "base64", + "base64 0.22.1", "cookie_store", "der", "flate2", @@ -2479,7 +2942,7 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" dependencies = [ - "base64", + "base64 0.22.1", "http", "httparse", "log", diff --git a/Cargo.toml b/Cargo.toml index 622000e..40cd30b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ name = "puma" version = "0.0.3" edition = "2021" -description = "A lightweight, high-performance inference engine for local AI." +description = "A lightweight, high-performance model engine for local AI." license = "Apache-2.0" repository = "https://github.com/InftyAI/PUMA" homepage = "https://github.com/InftyAI/PUMA" @@ -31,6 +31,7 @@ sysinfo = "0.32" rusqlite = { version = "0.32", features = ["bundled"] } rusqlite_migration = "1.3" regex = "1.11" +thiserror = "2.0" # Web server axum = "0.7" @@ -39,6 +40,10 @@ tower-http = { version = "0.5", features = ["cors", "trace"] } uuid = { version = "1.0", features = ["v4", "serde"] } futures = "0.3" tokio-stream = "0.1" +rustyline = "14.0" +rustyline-derive = "0.10" +tokenizers = "0.20" +async-stream = "0.3" [dev-dependencies] tempfile = "3.12" diff --git a/Makefile b/Makefile index 0e72459..d48a626 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,5 @@ build: + rm -rf ./puma cargo build && cp target/debug/puma ./puma test: diff --git a/README.md b/README.md index 846ad16..b2558c2 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ PUMA Logo -**A lightweight, high-performance inference engine for local AI** +**A lightweight, high-performance model engine for local AI** [![Stability: Active](https://img.shields.io/badge/stability-active-brightgreen.svg)](https://github.com/InftyAI/PUMA) [![Latest Release](https://img.shields.io/github/v/release/InftyAI/PUMA)](https://github.com/InftyAI/PUMA/releases) diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md new file mode 100644 index 0000000..2ddccf5 --- /dev/null +++ b/docs/fsm_architecture.md @@ -0,0 +1,240 @@ +# FSM Architecture for PUMA + +Finite State Machine pattern for sequence lifecycle management. The FSM states +and events are defined in `src/fsm/`; the transition logic lives in the +`SequenceManager` (`src/sequence_manager/mod.rs`), which the `Scheduler` drives. + +## Core Concepts + +### State Machine per Sequence + +Each sequence is an independent state machine. Multiple sequences can be in different states simultaneously. + +**Implementation:** `src/fsm/states.rs` + +**States:** +- `Empty` - Pre-birth placeholder; the `from` state of `Create`, never persisted +- `Waiting` - In queue, no resources allocated +- `Scheduling` - Being scheduled (transitional) +- `Prefilling` - Processing prompt +- `Decoding` - Generating tokens +- `Preempted` - Temporarily suspended (blocks freed) +- `Finished` - Completed successfully +- `Aborted` - Error or cancelled + +### States Own Resources + +Each state owns the resources it needs: +- `WaitingState` has no blocks (not allocated yet) +- `DecodingState` owns blocks via `Arc>` +- `PreemptedState` has no blocks (freed) + +**Benefit:** Type system enforces resource invariants - can't access blocks that don't exist. + +### Events as State Transformations + +Events transform states: `Event + OldState → NewState` + +**Implementation:** `src/fsm/events.rs` + +**Events:** +- `Create` - Empty → Waiting (birth a new sequence) +- `Schedule` - Waiting → Prefilling (allocate blocks) +- `AppendTokens` - Prefilling → Decoding or Decoding → Decoding/Finished +- `Preempt` - Decoding → Preempted (free blocks on OOM) +- `Resume` - Preempted → Waiting (re-queue) +- `Complete` - * → Finished (free blocks) +- `Abort` - * → Aborted (cleanup) + +Invalid transitions return `Error::InvalidTransition`. + +Forking (Decoding → (Decoding, Decoding) for beam search) is **not** an event: +it is a dedicated `SequenceManager::fork` method, since it produces a second +sequence rather than transforming one state into another. + +### FSM Integration with SequenceManager + +The FSM logic lives in the **`SequenceManager`** (the *mechanism* layer), which +owns `BlockManager`. The `Scheduler` (the *policy* layer) drives it with events +and never touches `BlockManager` or transition logic directly. + +**Public API:** `SequenceManager::advance(seq_id, event)` - Event-based abstraction +- Single entry point for all state transitions +- Easy to add logging, metrics, debugging +- Backs up state and restores it on error (no block leaks) + +**Internal Implementation:** Private `transition()` / `transition_*()` methods +- Type-safe helpers that enforce correct state types +- Direct access to `self.block_manager` - no parameter passing +- Called by `advance()` after looking up the current state + +```rust +impl SequenceManager { + pub fn advance(&mut self, seq_id: SequenceId, event: Event) -> Result<&SequenceState> { + // Remove current state, run transition(), reinsert on success + // or restore the backup on error. + } + + fn transition(&mut self, state: SequenceState, event: Event) -> Result { + let seq_id = state.seq_id(); + match (state, event) { + (SequenceState::Waiting(s), Event::Schedule{..}) => + self.transition_schedule(s), + // ... dispatches to internal methods + } + } + + fn transition_schedule(&mut self, state: WaitingState) -> Result { + // Direct access to self.block_manager, self.tokens_per_block + } +} +``` + +## State Diagram + +``` + AddRequest + ↓ + ┌─────────────→ Waiting + │ ↓ ScheduleEvent + │ Prefilling + │ ↓ AppendTokens + │ Decoding + │ ├─→ AppendTokens (continue) + │ │ ├─ Decoding (more tokens) + │ │ └─ Finished (max_tokens) + │ │ + │ ├─→ fork() method (beam search, not an event) + │ │ ├─ Parent: Decoding + │ │ └─ Child: Decoding + │ │ + │ ├─→ PreemptEvent (OOM) + │ │ ↓ Preempted + │ │ ↓ ResumeEvent + │ └───┘ + │ + └─ CompleteEvent → Finished +``` + +## Scheduler / SequenceManager Integration + +The `SequenceManager` owns the FSM and executes all state transitions; the +`Scheduler` owns the queues/batches and decides *which* events to fire. + +**Implementation:** `src/sequence_manager/mod.rs` (mechanism), +`src/scheduler/core.rs` (policy) + +**Key features:** +- Synchronous API (LLMEngine handles async coordination) +- `SequenceManager` maintains `HashMap` +- `SequenceManager` owns `BlockManager` directly - no parameter passing needed +- Clones state before transitions (prevents loss on failure) +- Restores backup on error (no block leaks) +- Event-based `advance()` method for maintainability + +**Tests:** `cargo test --lib` + +## Arc Optimization + +States use `Arc>` for efficient cloning: + +- **State cloning** (for backup/fork): **O(1)** - just increment ref count +- **Modifying blocks**: Copy-on-write via `Arc::make_mut()` - only copies if shared +- **Fork sequences**: Share blocks via Arc, diverge only when modified + +**Benefit:** Hot path performance - state cloning costs ~50 bytes instead of O(n) vector copy. + +## Error Handling + +All state transitions are wrapped with backup/restore: + +1. Remove state from HashMap +2. Clone state as backup +3. Apply event transformation +4. On success: insert new state +5. On failure: restore backup, log warning + +**Prevents:** +- Sequence loss on transition failure +- Block leaks (blocks remain owned by restored state) +- Duplicate sequences (check before insert) + +## Type Safety + +**Runtime checks:** +- Invalid transitions caught by pattern matching +- Return `Error::InvalidTransition` with clear message + +**Resource ownership:** +- States own their blocks via Arc +- Type system prevents accessing freed blocks +- Drop trait ensures cleanup + +## Benefits + +1. **Type Safety** - States enforce resource invariants, invalid transitions caught explicitly +2. **Clear Ownership** - Resources belong to states, automatic cleanup on drop +3. **Testability** - Events are pure functions, test state transitions in isolation +4. **Explicitness** - Every transition is explicit, state diagram maps to code +5. **Performance** - Arc enables O(1) cloning, copy-on-write for efficiency +6. **Reliability** - Backup/restore pattern prevents data loss on errors + +## Two-Level Event Architecture + +PUMA uses a two-level event system (inspired by TokenSpeed): + +**Level 1: FSM Events (Internal)** - `src/fsm/events.rs` +- Created by the Scheduler internally +- Include scheduler context (tokens_per_block, etc.) +- Applied via `SequenceManager::advance(seq_id, event)` +- Examples: `Event::Schedule`, `Event::AppendTokens` + +**Level 2: Scheduler Events (External)** - `src/scheduler/events.rs` +- Created by external components (clients, GPU workers) +- Simple data, no resource pointers +- LLMEngine translates these to FSM events +- Examples: `SchedulerEvent::AddRequest`, `SchedulerEvent::CancelRequest` + +**Why two levels?** External components cannot create FSM events because: +- They don't own `BlockManager` +- They don't know scheduler policies (tokens_per_block) +- FSM events require scheduler context + +## Key Files + +- `src/fsm/states.rs` - State definitions and helper methods +- `src/fsm/events.rs` - FSM event enum (internal transitions) +- `src/sequence_manager/mod.rs` - Owns memory + FSM transition logic (`advance()`) +- `src/scheduler/core.rs` - Scheduling policy that drives the SequenceManager +- `src/scheduler/events.rs` - External scheduler events +- `src/backend/llm_engine.rs` - Event coordinator + +## Usage Example + +```rust +// External component sends a high-level event +scheduler_tx.send(SchedulerEvent::AddRequest { + seq_id: SequenceId(1), + token_ids: vec![/* tokenized prompt */], + max_tokens: 100, + response_tx, // channel the result is delivered on +})?; + +// The Scheduler translates it to an FSM event and drives the SequenceManager, +// which owns BlockManager and performs the transition. +impl Scheduler { + fn schedule_prefill(&mut self) { + while let Some(seq_id) = self.waiting_queue.pop_front() { + // Build the FSM event with scheduler context + let event = Event::Schedule { + tokens_per_block: self.tokens_per_block, + }; + + match self.sequences.advance(seq_id, event) { + Ok(_) => self.prefill_batch.push(seq_id), + Err(e) => { /* apply OOM / queueing policy */ } + } + } + } +} +``` diff --git a/src/api/chat.rs b/src/api/chat.rs index 0dfcf4f..c87f1f3 100644 --- a/src/api/chat.rs +++ b/src/api/chat.rs @@ -7,7 +7,6 @@ use axum::{ Json, }; use futures::stream::StreamExt; -use std::sync::Arc; use tokio_stream::wrappers::ReceiverStream; use uuid::Uuid; @@ -16,11 +15,11 @@ use crate::api::types::{ ChatChoice, ChatChoiceDelta, ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ChatMessage, ChatMessageDelta, ErrorResponse, Usage, }; -use crate::backend::InferenceEngine; +use crate::backend::EngineHandle; /// Main handler for chat completions -pub async fn chat_completions( - State(state): State>, +pub async fn chat_completions( + State(state): State, Json(req): Json, ) -> Response { let engine = state.engine; @@ -83,8 +82,8 @@ pub async fn chat_completions( } /// Non-streaming chat completion -async fn chat_completions_non_stream( - engine: Arc, +async fn chat_completions_non_stream( + engine: EngineHandle, req: ChatCompletionRequest, ) -> Result> { let id = format!("chatcmpl-{}", Uuid::new_v4()); @@ -125,8 +124,8 @@ async fn chat_completions_non_stream( } /// Streaming chat completion -async fn chat_completions_stream( - engine: Arc, +async fn chat_completions_stream( + engine: EngineHandle, req: ChatCompletionRequest, ) -> Sse>> { let id = format!("chatcmpl-{}", Uuid::new_v4()); @@ -236,19 +235,11 @@ async fn chat_completions_stream( Sse::new(ReceiverStream::new(rx)).keep_alive(KeepAlive::default()) } -/// Format chat messages into a prompt +/// Format chat messages into a prompt via the shared prompt formatter. fn format_chat_messages(messages: &[ChatMessage]) -> String { - messages - .iter() - .map(|m| { - if m.role == "system" { - format!("System: {}", m.content) - } else if m.role == "user" { - format!("User: {}", m.content) - } else { - format!("Assistant: {}", m.content) - } - }) - .collect::>() - .join("\n") + crate::utils::prompt::format_conversation( + messages + .iter() + .map(|m| (m.role.as_str(), m.content.as_str())), + ) } diff --git a/src/api/completions.rs b/src/api/completions.rs index 497736b..c9bedf8 100644 --- a/src/api/completions.rs +++ b/src/api/completions.rs @@ -5,11 +5,10 @@ use crate::api::routes::AppState; use crate::api::types::{ CompletionChoice, CompletionRequest, CompletionResponse, ErrorResponse, Usage, }; -use crate::backend::InferenceEngine; /// Handler for legacy text completions -pub async fn completions( - State(state): State>, +pub async fn completions( + State(state): State, Json(req): Json, ) -> impl IntoResponse { let engine = state.engine; diff --git a/src/api/models.rs b/src/api/models.rs index c8304e1..00b3be6 100644 --- a/src/api/models.rs +++ b/src/api/models.rs @@ -6,12 +6,9 @@ use axum::{ use crate::api::routes::AppState; use crate::api::types::{ErrorResponse, Model, ModelList}; -use crate::backend::InferenceEngine; /// List all available models -pub async fn list_models( - State(state): State>, -) -> impl IntoResponse { +pub async fn list_models(State(state): State) -> impl IntoResponse { let registry = state.registry; match registry.load_models(None) { Ok(models) => { @@ -43,8 +40,8 @@ pub async fn list_models( } /// Get a specific model by ID -pub async fn get_model( - State(state): State>, +pub async fn get_model( + State(state): State, Path(model_id): Path, ) -> impl IntoResponse { let registry = state.registry; diff --git a/src/api/routes.rs b/src/api/routes.rs index dae6c0e..15875be 100644 --- a/src/api/routes.rs +++ b/src/api/routes.rs @@ -10,33 +10,30 @@ use tower_http::{ LatencyUnit, }; -use crate::backend::InferenceEngine; +use crate::backend::EngineHandle; use crate::registry::model_registry::ModelRegistry; use super::{chat, completions, models}; /// Shared application state #[derive(Clone)] -pub struct AppState { - pub engine: Arc, +pub struct AppState { + pub engine: EngineHandle, pub registry: Arc, } /// Create the API router with all endpoints -pub fn create_router( - engine: Arc, - registry: Arc, -) -> Router { +pub fn create_router(engine: EngineHandle, registry: Arc) -> Router { let state = AppState { engine, registry }; Router::new() // Chat completions (most important) - .route("/v1/chat/completions", post(chat::chat_completions::)) + .route("/v1/chat/completions", post(chat::chat_completions)) // Legacy completions - .route("/v1/completions", post(completions::completions::)) + .route("/v1/completions", post(completions::completions)) // Models - .route("/v1/models", get(models::list_models::)) - .route("/v1/models/:model", get(models::get_model::)) + .route("/v1/models", get(models::list_models)) + .route("/v1/models/:model", get(models::get_model)) // Health check .route("/health", get(health_check)) // Pass state diff --git a/src/api/tests.rs b/src/api/tests.rs index fba05ca..fd8939e 100644 --- a/src/api/tests.rs +++ b/src/api/tests.rs @@ -13,13 +13,21 @@ use tempfile::TempDir; use tower::util::ServiceExt; // for `oneshot` and `ready` use super::routes::create_router; +use crate::backend::engine; use crate::backend::mock::MockEngine; use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry}; /// Helper to create test app with a pre-registered test model /// Returns the router and the temp directory (which must be kept alive) fn create_test_app() -> (axum::Router, TempDir) { - let engine = Arc::new(MockEngine::new()); + // Build the engine and spawn its runner; the handle drives the router + let (handle, runner) = engine( + MockEngine::new(), + create_test_tokenizer(), + "test-model".to_string(), + ); + tokio::spawn(runner.serve()); + let temp_dir = TempDir::new().unwrap(); let registry = Arc::new(ModelRegistry::new(Some(temp_dir.path().to_path_buf()))); @@ -49,7 +57,13 @@ fn create_test_app() -> (axum::Router, TempDir) { .register_model(test_model) .expect("failed to register test model"); - (create_router(engine, registry), temp_dir) + (create_router(handle, registry), temp_dir) +} + +/// Simple BPE tokenizer for tests +fn create_test_tokenizer() -> tokenizers::Tokenizer { + use tokenizers::models::bpe::BPE; + tokenizers::Tokenizer::new(BPE::default()) } /// Helper to make a JSON request diff --git a/src/backend/engine.rs b/src/backend/engine.rs index 2d464a5..a3bb6f2 100644 --- a/src/backend/engine.rs +++ b/src/backend/engine.rs @@ -1,34 +1,30 @@ +use crate::block_manager::types::TokenId; use std::io; use std::pin::Pin; use tokio_stream::Stream; -/// Inference engine trait -pub trait InferenceEngine: Send + Sync { - /// Generate text completion +/// Backend trait - low-level inference that works with token IDs +/// +/// LLMEngine handles tokenization (text → tokens) +/// Backend handles inference (tokens → tokens) +pub trait Backend: Send + Sync { + /// Generate tokens from input token_ids + /// Returns generated token IDs fn generate( &self, - model: &str, - prompt: &str, + token_ids: Vec, max_tokens: usize, temperature: f32, - ) -> impl std::future::Future> + Send; + ) -> impl std::future::Future, io::Error>> + Send; - /// Generate text with streaming + /// Generate tokens with streaming + /// Returns stream of token IDs as they're generated fn generate_stream( &self, - model: &str, - prompt: &str, + token_ids: Vec, max_tokens: usize, temperature: f32, ) -> impl std::future::Future< - Output = Result + Send>>, io::Error>, + Output = Result + Send>>, io::Error>, > + Send; } - -/// Generation response -#[derive(Debug, Clone)] -pub struct GenerateResponse { - pub text: String, - pub prompt_tokens: usize, - pub completion_tokens: usize, -} diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs new file mode 100644 index 0000000..21e71bc --- /dev/null +++ b/src/backend/llm_engine.rs @@ -0,0 +1,434 @@ +use std::io; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use tokenizers::Tokenizer; +use tokio::sync::mpsc; +use tokio_stream::StreamExt; + +use super::engine::Backend; +use crate::block_manager::allocator::CpuAllocator; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::{SequenceId, TokenId}; +use crate::fsm::FinishReason; +use crate::scheduler::core::Scheduler; +use crate::scheduler::events::{ResponseSender, SchedulerEvent}; + +/// User-facing response for generate() +#[derive(Debug, Clone)] +pub struct GenerateResponse { + pub text: String, + pub prompt_tokens: usize, + pub completion_tokens: usize, +} + +/// Send-side handle to the engine (cheap to clone). +/// +/// Holds only what's needed to translate a user request into a +/// `SchedulerEvent` and push it onto the event channel: +/// - tokenizer (text → tokens) +/// - event sender (fire the event) +/// - sequence id counter (assign ids) +/// +/// The actual work happens in [`EngineRunner::serve`], which owns the scheduler +/// and backend and consumes these events. This handle is `Clone` so the API +/// layer and CLI can share it freely across tasks. +#[derive(Clone)] +pub struct EngineHandle { + tokenizer: Arc, + event_tx: mpsc::UnboundedSender, + seq_id_counter: Arc, + model: String, +} + +impl EngineHandle { + fn next_seq_id(&self) -> SequenceId { + SequenceId(self.seq_id_counter.fetch_add(1, Ordering::Relaxed)) + } + + /// Tokenize prompt + fn tokenize(&self, prompt: &str) -> Result, io::Error> { + self.tokenizer + .encode(prompt, false) + .map(|encoding| encoding.get_ids().to_vec()) + .map_err(|e| io::Error::other(format!("Tokenization failed: {}", e))) + } + + /// Send an AddRequest event to the engine loop + fn send_request( + &self, + token_ids: Vec, + max_tokens: usize, + response_tx: ResponseSender, + ) -> Result<(), io::Error> { + let seq_id = self.next_seq_id(); + self.event_tx + .send(SchedulerEvent::AddRequest { + seq_id, + token_ids, + max_tokens, + response_tx, + }) + .map_err(|e| io::Error::other(format!("Engine send failed: {}", e))) + } + + /// Generate text completion (single response) + pub async fn generate( + &self, + _model: &str, + prompt: &str, + max_tokens: usize, + _temperature: f32, + ) -> Result { + let token_ids = self.tokenize(prompt)?; + let prompt_tokens = token_ids.len(); + + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + self.send_request(token_ids, max_tokens, ResponseSender::Single(response_tx))?; + + let text = response_rx + .await + .map_err(|e| io::Error::other(format!("Response channel closed: {}", e)))? + .map_err(|e| io::Error::other(format!("Scheduler error: {:?}", e)))?; + + Ok(GenerateResponse { + text, + prompt_tokens, + completion_tokens: 0, + }) + } + + /// Generate with streaming + pub async fn generate_stream( + &self, + _model: &str, + prompt: &str, + max_tokens: usize, + _temperature: f32, + ) -> Result + Send>>, io::Error> { + let token_ids = self.tokenize(prompt)?; + + let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel(); + self.send_request(token_ids, max_tokens, ResponseSender::Stream(response_tx))?; + + let stream = async_stream::stream! { + while let Some(result) = response_rx.recv().await { + match result { + Ok(token) => yield token, + Err(e) => { + tracing::error!("Stream error: {:?}", e); + break; + } + } + } + }; + + Ok(Box::pin(stream)) + } + + /// Model name this handle serves + pub fn model(&self) -> &str { + &self.model + } +} + +/// The engine itself: owns the scheduler and backend, drives the event loop. +/// +/// Drains the event channel and runs inference. Spawn its +/// [`serve`](Self::serve) on a task; hold the paired [`EngineHandle`] everywhere +/// else to submit work. +pub struct EngineRunner { + backend: B, + scheduler: Scheduler, + tokenizer: Arc, +} + +impl EngineRunner { + /// Decode token IDs to text + fn decode_tokens(&self, token_ids: &[u32]) -> Result { + self.tokenizer + .decode(token_ids, false) + .map_err(|e| io::Error::other(format!("Detokenization failed: {}", e))) + } + + /// Run one-shot inference for a non-streaming request. + /// + /// Calls [`Backend::generate`], decodes the completion, advances the FSM, + /// and delivers the full text to the client's single-response channel. + async fn forward_single( + &mut self, + seq_id: SequenceId, + token_ids: Arc>, + max_tokens: usize, + ) { + let prompt_len = token_ids.len(); + let completion_tokens = match self + .backend + .generate((*token_ids).clone(), max_tokens, 0.0) + .await + { + Ok(tokens) => tokens, + Err(e) => { + tracing::error!("Backend inference failed for {:?}: {}", seq_id, e); + self.scheduler + .abort_request(seq_id, format!("Backend inference failed: {}", e)); + return; + } + }; + + // The backend returns only the newly generated completion tokens. + let num_completion = completion_tokens.len(); + let text = self.decode_tokens(&completion_tokens).unwrap_or_default(); + + // Advance FSM state (prefill → decode → finished) so blocks are freed, + // then deliver the result to the client channel. + self.scheduler.append_tokens(seq_id, prompt_len); + self.scheduler + .complete_sequence(seq_id, FinishReason::Stop, text); + tracing::debug!( + "Completed {:?}: {} completion tokens", + seq_id, + num_completion + ); + } + + /// Run streaming inference for a streaming request. + /// + /// Calls [`Backend::generate_stream`] and forwards each decoded text chunk to + /// the client via [`send_chunk`](Scheduler::send_chunk) as it arrives, then + /// advances the FSM and closes the stream once generation finishes. + async fn forward_stream( + &mut self, + seq_id: SequenceId, + token_ids: Arc>, + max_tokens: usize, + ) { + let prompt_len = token_ids.len(); + let mut stream = match self + .backend + .generate_stream((*token_ids).clone(), max_tokens, 0.0) + .await + { + Ok(stream) => stream, + Err(e) => { + tracing::error!("Backend stream failed for {:?}: {}", seq_id, e); + self.scheduler + .abort_request(seq_id, format!("Backend stream failed: {}", e)); + return; + } + }; + + // The backend streams only the newly generated completion tokens. + // + // Decode incrementally rather than one token at a time: subword/byte-BPE + // tokenizers split multi-byte characters (emoji, CJK) across several + // tokens, so a lone token often decodes to `�` or "". We keep the running + // list of completion tokens and decode the whole prefix each step, which + // is always well-formed. When the decode ends in the Unicode replacement + // character the trailing multi-byte char is still incomplete, so we hold + // the output back until the next token completes it; otherwise we emit the + // newly-decoded suffix. This is the streaming detokenization scheme vLLM + // and TGI use. + let mut completion_ids: Vec = Vec::new(); + let mut sent_len = 0usize; // bytes of the decoded completion already sent + while let Some(token_id) = stream.next().await { + completion_ids.push(token_id); + // Advance decoding state and allocate KV blocks as needed. + self.scheduler.append_tokens(seq_id, 1); + + let decoded = self.decode_tokens(&completion_ids).unwrap_or_default(); + if let Some(delta) = stream_delta(&decoded, sent_len) { + sent_len = decoded.len(); + self.scheduler.send_chunk(seq_id, delta.to_string()); + } + } + + // Flush any tail held back by the incomplete-multibyte guard: if the + // stream ended while the last decode still had bytes beyond `sent_len` + // (e.g. a trailing `�` from a truncated multi-byte char), emit them now + // so the final output is never silently dropped. + let decoded = self.decode_tokens(&completion_ids).unwrap_or_default(); + if let Some(delta) = stream_flush(&decoded, sent_len) { + self.scheduler.send_chunk(seq_id, delta.to_string()); + } + + let num_completion = completion_ids.len(); + + // Advance FSM state (prefill → decode → finished) so blocks are freed, + // then close the stream channel. + self.scheduler.append_tokens(seq_id, prompt_len); + self.scheduler + .complete_sequence(seq_id, FinishReason::Stop, String::new()); + tracing::debug!( + "Streamed {:?}: {} completion tokens", + seq_id, + num_completion + ); + } + + /// Main loop - drives scheduling and inference + /// + /// Pattern + /// 1. Handle events (add/cancel requests) + /// 2. Schedule (decide what to run) + /// 3. Forward (run inference if work exists) + /// 4. Process outputs (update scheduler state) + pub async fn serve(mut self) { + tracing::info!("EngineRunner started"); + + loop { + // 1. Handle events (non-blocking, drain all). The scheduler owns the + // event channel; the loop drains it and dispatches each event. + while let Ok(event) = self.scheduler.event_rx.try_recv() { + self.scheduler.handle_event(event); + } + + // 2. Schedule (always, not event-driven) + let has_work = self.scheduler.schedule(); + + // 3. Forward: run inference on the scheduled batch and report + // results back to the scheduler, which fulfills client channels. + // Streaming requests use the backend's streaming path (token by + // token); single requests use the one-shot path. + // + // TODO: this processes one sequence at a time, awaiting each backend + // call in turn. For real batch inference, partition the batch by mode + // (see the `PrefillBatch` design) and hand the non-streaming bucket to + // a single `Backend::generate_batch` call instead of looping. + if has_work { + for work in self.scheduler.take_prefill_batch() { + let (seq_id, token_ids, max_tokens, streaming) = work; + if streaming { + self.forward_stream(seq_id, token_ids, max_tokens).await; + } else { + self.forward_single(seq_id, token_ids, max_tokens).await; + } + } + } + + // 4. Small yield to prevent busy loop + tokio::task::yield_now().await; + } + } +} + +/// Incremental streaming detokenization: given the full text decoded from the +/// completion prefix so far and how many bytes were already sent, return the +/// next delta to emit, or `None` to hold output back. +/// +/// A trailing replacement char (`�`) means the last multi-byte character is +/// still incomplete, so we wait for the next token rather than emit a partial +/// char. This mirrors the scheme vLLM and TGI use. +fn stream_delta(decoded: &str, sent_len: usize) -> Option<&str> { + if decoded.ends_with('\u{FFFD}') { + return None; + } + stream_flush(decoded, sent_len) +} + +/// Final flush: emit whatever bytes remain past `sent_len`, including any tail +/// the `�` guard in [`stream_delta`] held back when the stream ended. Returns +/// `None` when nothing is left to send. +fn stream_flush(decoded: &str, sent_len: usize) -> Option<&str> { + if decoded.len() > sent_len { + Some(&decoded[sent_len..]) + } else { + None + } +} + +/// Construct a paired [`EngineHandle`] and [`EngineRunner`]. +/// +/// Spawn `runner.serve()` on a task and share the returned handle with the API / +/// CLI. All request submission goes through events, so the handle never +/// touches the scheduler directly. +pub fn engine( + backend: B, + tokenizer: Tokenizer, + model: String, +) -> (EngineHandle, EngineRunner) { + // Create block manager (100MB memory pool, 512 bytes per block) + let allocator = Box::new(CpuAllocator::new(1024 * 1024 * 100)); + let block_manager = BlockManager::new(allocator, 512); + + // Event channel: the handle produces events, the scheduler consumes them. + let (event_tx, event_rx) = mpsc::unbounded_channel(); + + // Create scheduler (max 32 batch size, 16 tokens per block); it owns the + // event receiver. + let scheduler = Scheduler::new(block_manager, event_rx, 32, 16); + + let tokenizer = Arc::new(tokenizer); + + let handle = EngineHandle { + tokenizer: tokenizer.clone(), + event_tx, + seq_id_counter: Arc::new(AtomicU64::new(1)), + model, + }; + + let runner = EngineRunner { + backend, + scheduler, + tokenizer, + }; + + (handle, runner) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::mock::MockEngine; + + fn create_test_tokenizer() -> Tokenizer { + use tokenizers::models::bpe::BPE; + use tokenizers::Tokenizer as TokenizerBuilder; + + let bpe = BPE::default(); + TokenizerBuilder::new(bpe) + } + + #[tokio::test] + async fn test_llm_engine() { + let backend = MockEngine::new(); + let tokenizer = create_test_tokenizer(); + let (handle, runner) = engine(backend, tokenizer, "test-model".to_string()); + + tokio::spawn(runner.serve()); + + let result = handle.generate("test-model", "Hello world", 100, 0.7).await; + assert!(result.is_ok()); + } + + #[test] + fn test_stream_delta_emits_complete_suffix() { + // Nothing sent yet -> emit the whole thing. + assert_eq!(stream_delta("hello", 0), Some("hello")); + // Some already sent -> emit only the new bytes. + assert_eq!(stream_delta("hello world", 5), Some(" world")); + // Nothing new -> hold. + assert_eq!(stream_delta("hello", 5), None); + } + + #[test] + fn test_stream_delta_holds_incomplete_multibyte() { + // Decode ended on a replacement char: the multi-byte char is not yet + // complete, so we must not emit anything this step. + assert_eq!(stream_delta("hi\u{FFFD}", 2), None); + } + + #[test] + fn test_stream_flush_emits_held_back_tail() { + // Simulates the bug scenario: the loop held back a trailing `�` and the + // stream then ended. The flush must still emit those bytes rather than + // drop them. + let decoded = "hi\u{FFFD}"; + assert_eq!(stream_delta(decoded, 2), None); // loop held it back + assert_eq!(stream_flush(decoded, 2), Some("\u{FFFD}")); // flush emits it + } + + #[test] + fn test_stream_flush_none_when_all_sent() { + assert_eq!(stream_flush("done", 4), None); + } +} diff --git a/src/backend/mock.rs b/src/backend/mock.rs index 92f056e..3639aa7 100644 --- a/src/backend/mock.rs +++ b/src/backend/mock.rs @@ -3,66 +3,100 @@ use std::io; use std::pin::Pin; use tokio_stream::Stream; -use super::engine::{GenerateResponse, InferenceEngine}; +use super::engine::Backend; +use crate::block_manager::types::TokenId; -/// Mock engine for testing (replace with MLX later) +/// Default vocab size the mock samples completion tokens from. +/// +/// Small enough that ids land in the low end of any real tokenizer's vocab +/// (so they decode to real text), large enough for varied output. +const DEFAULT_VOCAB_SIZE: u32 = 1000; + +/// Mock inference engine that behaves like a real autoregressive model. +/// +/// Unlike a naive echo, this consumes the prompt as *context* and emits only +/// **new** completion tokens — never the prompt back. Tokens are produced by a +/// deterministic pseudo-random walk seeded from the context, so: +/// - the same prompt always yields the same completion (reproducible tests), +/// - different prompts yield different completions (input-dependent), +/// - ids stay within the vocab range, so they decode to real text. +/// +/// Generation stops at `max_tokens` (the mock has no EOS concept yet). #[derive(Clone)] -pub struct MockEngine; +pub struct MockEngine { + vocab_size: u32, +} impl MockEngine { pub fn new() -> Self { - Self + Self { + vocab_size: DEFAULT_VOCAB_SIZE, + } + } + + /// Construct with an explicit vocab size (useful for tests pinning output). + pub fn with_vocab_size(vocab_size: u32) -> Self { + Self { + vocab_size: vocab_size.max(1), + } + } + + /// Deterministically generate `max_tokens` completion token ids from the + /// prompt context, mimicking an autoregressive decode. + /// + /// Each step hashes the running context (prompt + tokens produced so far) + /// into the vocab range and appends the result — exactly the shape of real + /// inference (next token conditioned on all prior tokens), just with a hash + /// standing in for a learned distribution. + fn generate_tokens(&self, prompt: &[TokenId], max_tokens: usize) -> Vec { + let mut context: Vec = prompt.to_vec(); + let mut completion = Vec::with_capacity(max_tokens); + for _ in 0..max_tokens { + let next = (hash_tokens(&context) % self.vocab_size as u64) as TokenId; + completion.push(next); + context.push(next); + } + completion } } -impl InferenceEngine for MockEngine { +/// FNV-1a hash over a token sequence — order-sensitive and fast, so each +/// distinct context maps to a distinct next token deterministically. +fn hash_tokens(tokens: &[TokenId]) -> u64 { + let mut hash: u64 = 0xcbf29ce484222325; + for &tok in tokens { + for byte in tok.to_le_bytes() { + hash ^= byte as u64; + hash = hash.wrapping_mul(0x100000001b3); + } + } + hash +} + +impl Backend for MockEngine { async fn generate( &self, - model: &str, - prompt: &str, + token_ids: Vec, max_tokens: usize, _temperature: f32, - ) -> Result { - // Mock response for testing - let response_text = format!( - "This is a mock response from model '{}' for prompt: '{}' (max_tokens: {})", - model, - prompt.chars().take(50).collect::(), - max_tokens - ); - - Ok(GenerateResponse { - text: response_text, - prompt_tokens: prompt.split_whitespace().count(), - completion_tokens: 20, - }) + ) -> Result, io::Error> { + // Return ONLY the new completion tokens, as a real model does. + Ok(self.generate_tokens(&token_ids, max_tokens)) } async fn generate_stream( &self, - model: &str, - _prompt: &str, + token_ids: Vec, max_tokens: usize, _temperature: f32, - ) -> Result + Send>>, io::Error> { - // Mock streaming response - let tokens = vec![ - "This ".to_string(), - "is ".to_string(), - "a ".to_string(), - "mock ".to_string(), - "streaming ".to_string(), - "response ".to_string(), - format!("from model '{}' ", model), - format!("(max_tokens: {}).", max_tokens), - ]; - - // Simulate delay between tokens - let stream = stream::iter(tokens).then(|token| async move { - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - token + ) -> Result + Send>>, io::Error> { + // Same completion, but yielded one token at a time with a small delay to + // mimic a real GPU's inter-token latency. + let completion = self.generate_tokens(&token_ids, max_tokens); + let stream = stream::iter(completion).then(|token_id| async move { + tokio::time::sleep(tokio::time::Duration::from_millis(20)).await; + token_id }); - Ok(Box::pin(stream)) } } @@ -72,3 +106,58 @@ impl Default for MockEngine { Self::new() } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn returns_only_completion_tokens() { + let engine = MockEngine::new(); + let prompt = vec![5, 9, 2]; + let out = engine.generate(prompt.clone(), 4, 0.0).await.unwrap(); + // Completion-only: exactly max_tokens, and it does not start with the + // prompt (a real model returns a continuation, not an echo). + assert_eq!(out.len(), 4); + assert_ne!(&out[..prompt.len().min(out.len())], &prompt[..]); + } + + #[tokio::test] + async fn is_deterministic() { + let engine = MockEngine::new(); + let a = engine.generate(vec![1, 2, 3], 8, 0.0).await.unwrap(); + let b = engine.generate(vec![1, 2, 3], 8, 0.0).await.unwrap(); + assert_eq!(a, b, "same prompt must yield same completion"); + } + + #[tokio::test] + async fn is_input_dependent() { + let engine = MockEngine::new(); + let a = engine.generate(vec![1, 2, 3], 8, 0.0).await.unwrap(); + let b = engine.generate(vec![3, 2, 1], 8, 0.0).await.unwrap(); + assert_ne!(a, b, "different prompts should yield different completions"); + } + + #[tokio::test] + async fn tokens_stay_within_vocab() { + let engine = MockEngine::with_vocab_size(50); + let out = engine.generate(vec![7, 7, 7], 32, 0.0).await.unwrap(); + assert!( + out.iter().all(|&t| t < 50), + "ids must be within vocab range" + ); + } + + #[tokio::test] + async fn stream_matches_generate() { + let engine = MockEngine::new(); + let prompt = vec![10, 20, 30]; + let batched = engine.generate(prompt.clone(), 6, 0.0).await.unwrap(); + let mut streamed = Vec::new(); + let mut s = engine.generate_stream(prompt, 6, 0.0).await.unwrap(); + while let Some(tok) = s.next().await { + streamed.push(tok); + } + assert_eq!(batched, streamed, "streaming and batched output must agree"); + } +} diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 1631447..68e2dab 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -1,4 +1,5 @@ pub mod engine; +pub mod llm_engine; pub mod mock; -pub use engine::*; +pub use llm_engine::{engine, EngineHandle}; diff --git a/src/block_manager/allocator.rs b/src/block_manager/allocator.rs new file mode 100644 index 0000000..a9a7c35 --- /dev/null +++ b/src/block_manager/allocator.rs @@ -0,0 +1,153 @@ +use super::types::*; + +/// Memory allocator trait +pub trait MemoryAllocator: Send + Sync { + fn allocate(&mut self, size_bytes: usize) -> Result; + fn free(&mut self, addr: MemoryAddress) -> Result<()>; + fn get_total_memory(&self) -> usize; + fn get_available_memory(&self) -> usize; +} + +/// Simple CPU memory allocator (for testing/CPU inference) +pub struct CpuAllocator { + total_memory: usize, + used_memory: usize, +} + +impl CpuAllocator { + pub fn new(total_memory: usize) -> Self { + Self { + total_memory, + used_memory: 0, + } + } +} + +impl MemoryAllocator for CpuAllocator { + fn allocate(&mut self, size_bytes: usize) -> Result { + if self.used_memory + size_bytes > self.total_memory { + return Err(Error::OutOfMemory); + } + + // Allocate aligned memory + let layout = std::alloc::Layout::from_size_align(size_bytes, 64) + .map_err(|e| Error::AllocationFailed(e.to_string()))?; + + let ptr = unsafe { std::alloc::alloc(layout) }; + + if ptr.is_null() { + return Err(Error::OutOfMemory); + } + + self.used_memory += size_bytes; + + Ok(MemoryAddress { + ptr, + size: size_bytes, + }) + } + + fn free(&mut self, addr: MemoryAddress) -> Result<()> { + if addr.size > self.used_memory { + return Err(Error::FreeFailed( + "free() called with size larger than used_memory".to_string(), + )); + } + + let layout = std::alloc::Layout::from_size_align(addr.size, 64) + .map_err(|e| Error::FreeFailed(e.to_string()))?; + + unsafe { + std::alloc::dealloc(addr.ptr, layout); + } + + self.used_memory -= addr.size; + + Ok(()) + } + + fn get_total_memory(&self) -> usize { + self.total_memory + } + + fn get_available_memory(&self) -> usize { + self.total_memory - self.used_memory + } +} + +// TODO: Implement CudaAllocator later +// pub struct CudaAllocator { ... } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_allocate_tracks_used_memory() { + let mut alloc = CpuAllocator::new(10_000); + assert_eq!(alloc.get_total_memory(), 10_000); + assert_eq!(alloc.get_available_memory(), 10_000); + + let addr = alloc.allocate(1024).unwrap(); + assert_eq!(addr.size, 1024); + assert!(!addr.ptr.is_null()); + assert_eq!(alloc.get_available_memory(), 10_000 - 1024); + + alloc.free(addr).unwrap(); + assert_eq!(alloc.get_available_memory(), 10_000); + } + + #[test] + fn test_allocate_out_of_memory() { + let mut alloc = CpuAllocator::new(1024); + + // First allocation fits exactly. + let addr = alloc.allocate(1024).unwrap(); + assert_eq!(alloc.get_available_memory(), 0); + + // Any further allocation exceeds capacity. + assert!(matches!(alloc.allocate(1), Err(Error::OutOfMemory))); + + alloc.free(addr).unwrap(); + } + + #[test] + fn test_allocate_exact_boundary() { + let mut alloc = CpuAllocator::new(2048); + let a = alloc.allocate(1024).unwrap(); + let b = alloc.allocate(1024).unwrap(); + assert_eq!(alloc.get_available_memory(), 0); + alloc.free(a).unwrap(); + alloc.free(b).unwrap(); + assert_eq!(alloc.get_available_memory(), 2048); + } + + #[test] + fn test_free_more_than_used_errors() { + let mut alloc = CpuAllocator::new(10_000); + let addr = alloc.allocate(512).unwrap(); + + // Fabricate an address claiming a larger size than is actually used. + let bogus = MemoryAddress { + ptr: addr.ptr, + size: 4096, + }; + assert!(matches!(alloc.free(bogus), Err(Error::FreeFailed(_)))); + + // The real allocation is still accounted for and can be freed. + assert_eq!(alloc.get_available_memory(), 10_000 - 512); + alloc.free(addr).unwrap(); + assert_eq!(alloc.get_available_memory(), 10_000); + } + + #[test] + fn test_reuse_after_free() { + let mut alloc = CpuAllocator::new(1024); + for _ in 0..3 { + let addr = alloc.allocate(1024).unwrap(); + assert_eq!(alloc.get_available_memory(), 0); + alloc.free(addr).unwrap(); + assert_eq!(alloc.get_available_memory(), 1024); + } + } +} diff --git a/src/block_manager/manager.rs b/src/block_manager/manager.rs new file mode 100644 index 0000000..b8176e8 --- /dev/null +++ b/src/block_manager/manager.rs @@ -0,0 +1,356 @@ +use super::allocator::*; +use super::types::*; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Manages memory blocks with reference counting and pooling. +/// +/// Features: +/// - Reference counting: blocks can be shared, freed only when ref_count = 0 +/// - Free pools: reuse freed blocks instead of allocating new memory +/// - Type-specific blocks: different block types with different sizes +/// - Thread-safe ID generation: atomic counter for BlockId +pub struct BlockManager { + /// Underlying memory allocator (CPU or GPU) + allocator: Box, + + /// All blocks (allocated + free) - tracks metadata and physical memory + block_table: HashMap, + + /// Free block pools per type - blocks ready for reuse (ref_count = 0) + /// Keyed by BlockType, values are BlockIds ready to be allocated + free_pools: HashMap>, + + /// Configuration for each block type (size, etc.) + block_configs: HashMap, + + /// Default block type for allocate() calls + default_block_type: BlockType, + + /// Atomic counter for generating unique BlockIds + next_block_id: AtomicU64, +} + +impl BlockManager { + pub fn new(allocator: Box, default_block_size: usize) -> Self { + let mut manager = Self { + allocator, + block_table: HashMap::new(), + free_pools: HashMap::new(), + block_configs: HashMap::new(), + default_block_type: BlockType::StandardKV, + next_block_id: AtomicU64::new(0), + }; + + // Register default type + manager.register_block_type(BlockConfig { + block_type: BlockType::StandardKV, + size_bytes: default_block_size, + }); + + manager + } + + /// Register a new block type with specific configuration. + pub fn register_block_type(&mut self, config: BlockConfig) { + self.free_pools.insert(config.block_type, Vec::new()); + self.block_configs.insert(config.block_type, config); + } + + /// Allocate a block of specific type. + /// + /// Flow: + /// 1. Check free pool first (reuse if available) + /// 2. If pool empty, allocate new physical memory + /// 3. Return BlockId with ref_count = 1 + pub fn allocate_typed(&mut self, block_type: &BlockType) -> Result { + // Try free pool first - reuse freed blocks + if let Some(block_id) = self + .free_pools + .get_mut(block_type) + .and_then(|pool| pool.pop()) + { + let block = self.block_table.get_mut(&block_id).unwrap(); + block.ref_count = 1; + return Ok(block_id); + } + + // No free blocks - allocate new physical memory + let config = self + .block_configs + .get(block_type) + .ok_or_else(|| Error::UnknownBlockType(block_type.as_str().to_string()))?; + + let mem_addr = self.allocator.allocate(config.size_bytes)?; + + let block_id = BlockId(self.next_block_id.fetch_add(1, Ordering::Relaxed)); + + self.block_table.insert( + block_id, + Block { + block_id, + block_type: *block_type, + mem_addr, + ref_count: 1, + }, + ); + + Ok(block_id) + } + + /// Allocate a block using the default type. + pub fn allocate(&mut self) -> Result { + let default_type = self.default_block_type; + self.allocate_typed(&default_type) + } + + /// Decrement reference count. When ref_count reaches 0, block goes to free pool. + /// + /// Note: Physical memory is NOT freed - block is just marked as reusable. + /// This enables fast reallocation without syscalls. + pub fn free(&mut self, block_id: BlockId) -> Result<()> { + let block = self + .block_table + .get_mut(&block_id) + .ok_or(Error::InvalidBlockId(block_id))?; + + if block.ref_count == 0 { + return Err(Error::DoubleFree(block_id)); + } + + block.ref_count -= 1; + + // When ref_count hits 0, move to free pool for reuse + if block.ref_count == 0 { + self.free_pools + .get_mut(&block.block_type) + .unwrap() + .push(block_id); + } + + Ok(()) + } + + /// Increment reference count (for sharing blocks between sequences). + /// + /// Used for fork/beam search where multiple sequences share the same blocks. + pub fn add_ref(&mut self, block_id: BlockId) -> Result<()> { + let block = self + .block_table + .get_mut(&block_id) + .ok_or(Error::InvalidBlockId(block_id))?; + block.ref_count += 1; + Ok(()) + } + + /// Get the physical memory address for a block. + pub fn get_memory_address(&self, block_id: BlockId) -> Result { + self.block_table + .get(&block_id) + .map(|b| b.mem_addr) + .ok_or(Error::InvalidBlockId(block_id)) + } + + /// Get the block type for a block. + pub fn get_block_type(&self, block_id: BlockId) -> Result<&BlockType> { + self.block_table + .get(&block_id) + .map(|b| &b.block_type) + .ok_or(Error::InvalidBlockId(block_id)) + } + + /// Get statistics for all block types (allocated, free, total memory, etc.). + pub fn get_stats(&self) -> HashMap { + let mut stats = HashMap::new(); + + for (type_id, pool) in &self.free_pools { + let config = &self.block_configs[type_id]; + let free_count = pool.len(); + + let allocated_count = self + .block_table + .values() + .filter(|b| b.block_type == *type_id && b.ref_count > 0) + .count(); + + stats.insert( + *type_id, + BlockStats { + total_blocks: free_count + allocated_count, + allocated_blocks: allocated_count, + free_blocks: free_count, + block_size: config.size_bytes, + total_memory: (free_count + allocated_count) * config.size_bytes, + }, + ); + } + + stats + } + + /// Check if a block can be allocated (either from free pool or new memory). + pub fn can_allocate(&self, block_type: &BlockType) -> bool { + // Check free pool first + if let Some(pool) = self.free_pools.get(block_type) { + if !pool.is_empty() { + return true; + } + } + + // Check if allocator has enough memory + if let Some(config) = self.block_configs.get(block_type) { + return self.allocator.get_available_memory() >= config.size_bytes; + } + + false + } +} + +impl Drop for BlockManager { + /// Free all physical memory when BlockManager is dropped. + /// + /// This ensures no memory leaks even if blocks weren't explicitly freed. + fn drop(&mut self) { + // Free all backing memory allocations + for block in self.block_table.values() { + if let Err(e) = self.allocator.free(block.mem_addr) { + eprintln!( + "Warning: failed to free block {:?} memory: {:?}", + block.block_id, e + ); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + + #[test] + fn test_allocate_and_free() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + // Allocate a block + let block_id = manager.allocate().unwrap(); + assert_eq!(block_id, BlockId(0)); + + // Free it + manager.free(block_id).unwrap(); + + // Should be back in free pool + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.free_blocks, 1); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_ref_counting() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + + // Add ref + manager.add_ref(block_id).unwrap(); + + // Free once - should still be allocated (ref_count = 1) + manager.free(block_id).unwrap(); + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 1); + assert_eq!(block_stats.free_blocks, 0); + + // Free again - now should be in free pool (ref_count = 0) + manager.free(block_id).unwrap(); + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.free_blocks, 1); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_double_free_error() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + manager.free(block_id).unwrap(); + + // Should fail on double free + let result = manager.free(block_id); + assert!(matches!(result, Err(Error::DoubleFree(_)))); + } + + #[test] + fn test_oom() { + // Small allocator - only enough for 2 blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut manager = BlockManager::new(allocator, 1024); + + // First block succeeds + let block1 = manager.allocate().unwrap(); + assert_eq!(block1, BlockId(0)); + + // Second block succeeds + let block2 = manager.allocate().unwrap(); + assert_eq!(block2, BlockId(1)); + + // Third block should fail (OOM) + let result = manager.allocate(); + assert!(matches!(result, Err(Error::OutOfMemory))); + } + + #[test] + fn test_free_pool_reuse() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + // Allocate and free + let block1 = manager.allocate().unwrap(); + manager.free(block1).unwrap(); + + // Next allocation should reuse from free pool (same ID) + let block2 = manager.allocate().unwrap(); + assert_eq!(block1, block2); + } + + #[test] + fn test_get_memory_address() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + let addr = manager.get_memory_address(block_id); + assert!(addr.is_ok()); + } + + #[test] + fn test_invalid_block_id() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let manager = BlockManager::new(allocator, 1024); + + let invalid_id = BlockId(999); + let result = manager.get_memory_address(invalid_id); + assert!(matches!(result, Err(Error::InvalidBlockId(_)))); + } + + #[test] + fn test_can_allocate() { + let allocator = Box::new(CpuAllocator::new(2048)); + let mut manager = BlockManager::new(allocator, 1024); + + // Should be able to allocate + assert!(manager.can_allocate(&BlockType::StandardKV)); + + // Allocate both blocks + manager.allocate().unwrap(); + manager.allocate().unwrap(); + + // Should not be able to allocate more + assert!(!manager.can_allocate(&BlockType::StandardKV)); + } +} diff --git a/src/block_manager/mod.rs b/src/block_manager/mod.rs new file mode 100644 index 0000000..a4468f5 --- /dev/null +++ b/src/block_manager/mod.rs @@ -0,0 +1,5 @@ +pub mod allocator; +pub mod manager; +pub mod types; + +// Re-export for convenience diff --git a/src/block_manager/types.rs b/src/block_manager/types.rs new file mode 100644 index 0000000..820ed1f --- /dev/null +++ b/src/block_manager/types.rs @@ -0,0 +1,111 @@ +/// Block type - defines the kind of memory block +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum BlockType { + /// Standard KV cache blocks (for normal attention) + StandardKV, + + /// Compressed KV cache blocks (for DeepSeek-style compression) + CompressedKV, + + /// Compressor state blocks (for DeepSeek RNN compressor) + CompressorState, +} + +impl BlockType { + pub fn as_str(&self) -> &'static str { + match self { + BlockType::StandardKV => "standard_kv", + BlockType::CompressedKV => "compressed_kv", + BlockType::CompressorState => "compressor_state", + } + } +} + +impl std::fmt::Display for BlockType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct BlockId(pub u64); + +impl BlockId { + pub fn new(id: u64) -> Self { + Self(id) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct SequenceId(pub u64); + +impl SequenceId { + pub fn new(id: u64) -> Self { + Self(id) + } +} + +/// Token ID type - used for tokenized text +pub type TokenId = u32; + +#[derive(Debug, Clone, Copy)] +pub struct MemoryAddress { + pub ptr: *mut u8, + pub size: usize, +} + +unsafe impl Send for MemoryAddress {} +unsafe impl Sync for MemoryAddress {} + +pub struct BlockConfig { + pub block_type: BlockType, + pub size_bytes: usize, +} + +pub struct Block { + pub block_id: BlockId, + pub block_type: BlockType, + pub mem_addr: MemoryAddress, + pub ref_count: usize, +} + +#[derive(Debug, Clone)] +pub struct BlockStats { + pub total_blocks: usize, + pub allocated_blocks: usize, + pub free_blocks: usize, + pub block_size: usize, + pub total_memory: usize, +} + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("Unknown block type: {0}")] + UnknownBlockType(String), + + #[error("Invalid block ID: {0:?}")] + InvalidBlockId(BlockId), + + #[error("Double free detected for block: {0:?}")] + DoubleFree(BlockId), + + #[error("Unknown sequence: {0:?}")] + UnknownSequence(SequenceId), + + #[error("Out of memory")] + OutOfMemory, + + #[error("Allocation failed: {0}")] + AllocationFailed(String), + + #[error("Free failed: {0}")] + FreeFailed(String), + + #[error("Invalid state transition: {0}")] + InvalidTransition(&'static str), + + #[error("Sequence aborted: {0}")] + Aborted(String), +} + +pub type Result = std::result::Result; diff --git a/src/cli/chat.rs b/src/cli/chat.rs new file mode 100644 index 0000000..087251e --- /dev/null +++ b/src/cli/chat.rs @@ -0,0 +1,117 @@ +use rustyline::error::ReadlineError; +use rustyline::hint::{Hint, Hinter}; +use rustyline::{Context, Editor}; +use rustyline_derive::{Completer, Helper, Highlighter, Validator}; +use std::io::{self, Write}; +use tokio_stream::StreamExt; + +use crate::backend::EngineHandle; + +#[derive(Clone)] +struct PlaceholderHint { + display: String, +} + +impl Hint for PlaceholderHint { + fn display(&self) -> &str { + &self.display + } + + fn completion(&self) -> Option<&str> { + None + } +} + +/// Hint helper that shows placeholder when input is empty +#[derive(Helper, Completer, Highlighter, Validator)] +struct PlaceholderHinter { + placeholder: String, +} + +impl Hinter for PlaceholderHinter { + type Hint = PlaceholderHint; + + fn hint(&self, line: &str, _pos: usize, _ctx: &Context<'_>) -> Option { + if line.is_empty() { + // Grey/dimmed color: \x1b[2m ... \x1b[0m + Some(PlaceholderHint { + display: format!("\x1b[2m{}\x1b[0m", self.placeholder), + }) + } else { + None + } + } +} + +/// Interactive chat loop for puma run +pub async fn interactive_chat(engine: &EngineHandle, model: &str) -> Result<(), io::Error> { + let mut conversation_history: Vec<(String, String)> = Vec::new(); + + // Setup editor with placeholder hinter + let helper = PlaceholderHinter { + placeholder: "Send a message (Ctrl-C or 'exit' to quit)".to_string(), + }; + let mut rl = Editor::::new() + .map_err(io::Error::other)?; + rl.set_helper(Some(helper)); + + loop { + let readline = rl.readline("> "); + + let input = match readline { + Ok(line) => line.trim().to_string(), + Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => { + break; + } + Err(err) => { + return Err(io::Error::other(err)); + } + }; + + // Exit commands + if input.is_empty() { + continue; + } + if input == "exit" { + break; + } + + // Add user message to history (structured role/content turns) + conversation_history.push(("user".to_string(), input.clone())); + + // Build prompt from conversation history via the shared formatter + let prompt = crate::utils::prompt::format_conversation( + conversation_history + .iter() + .map(|(role, content)| (role.as_str(), content.as_str())), + ); + + // Empty line before response + println!(); + + // Generate response with streaming + match engine.generate_stream(model, &prompt, 512, 0.7).await { + Ok(mut stream) => { + let mut full_response = String::new(); + + // Display tokens as they arrive + while let Some(token) = stream.next().await { + print!("{}", token); + io::stdout().flush()?; + full_response.push_str(&token); + } + + println!("\n"); // Double newline after response + + // Add assistant response to history + conversation_history + .push(("assistant".to_string(), full_response.trim().to_string())); + } + Err(e) => { + eprintln!("Error: {}\n", e); + } + } + } + + Ok(()) +} diff --git a/src/cli/commands.rs b/src/cli/commands.rs index 8542ca7..c1ab516 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -2,7 +2,11 @@ use clap::{Parser, Subcommand}; use colored::Colorize; use prettytable::{format, row, Table}; -use crate::cli::{inspect, ls, rm}; +use tokenizers::Tokenizer; + +use crate::backend::engine; +use crate::backend::mock::MockEngine; +use crate::cli::{chat, inspect, ls, rm}; use crate::downloader::{self, Provider}; use crate::registry::model_registry::ModelRegistry; use crate::system::system_info::SystemInfo; @@ -16,6 +20,16 @@ pub struct Cli { command: Commands, } +impl Cli { + /// Whether this command should emit logs by default. + /// + /// Only the long-running server (`serve`) logs unprompted; one-shot and + /// interactive CLI commands stay quiet unless the user sets `RUST_LOG`. + pub fn wants_default_logging(&self) -> bool { + matches!(self.command, Commands::SERVE(_)) + } +} + #[derive(Subcommand)] #[allow(clippy::upper_case_acronyms)] enum Commands { @@ -188,16 +202,9 @@ pub async fn run(cli: Cli) { Commands::RUN(args) => { let registry = ModelRegistry::new(None); - // Check if model exists + // Ensure model exists locally, download if needed match registry.get_model(&args.model) { - Ok(Some(_)) => { - // Model exists, proceed - println!("Running model: {}", args.model); - // TODO: Implement actual model execution - println!("Model execution not yet implemented"); - } Ok(None) => { - // Model not found, download it first println!( "Model {} not found locally. Downloading...", args.model.cyan().bold() @@ -207,16 +214,47 @@ pub async fn run(cli: Cli) { eprintln!("❌ Error: {}", e); std::process::exit(1); } - - // Now run the model - println!("Running model: {}", args.model.cyan().bold()); - // TODO: Implement actual model execution - println!("Model execution not yet implemented"); } Err(e) => { eprintln!("❌ Error checking model: {}", e); std::process::exit(1); } + Ok(Some(_)) => {} + } + + // Load tokenizer from model directory + let model_info = registry.get_model(&args.model).unwrap().unwrap(); + let tokenizer_path = format!( + "{}/snapshots/{}/tokenizer.json", + model_info.metadata.cache.path, model_info.metadata.cache.revision + ); + + let tokenizer = match Tokenizer::from_file(&tokenizer_path) { + Ok(tok) => tok, + Err(e) => { + eprintln!( + "Warning: Could not load tokenizer from {}: {}.", + tokenizer_path, e + ); + std::process::exit(1); + } + }; + + // Load inference backend + // TODO: Replace MockEngine with real backend that loads model files + // Real backend will use: registry.get_model(&args.model)?.metadata.cache.path + let backend = MockEngine::new(); + + // Create engine: cheap send-side handle + runner that owns the scheduler + let (handle, runner) = engine(backend, tokenizer, args.model.clone()); + + // Spawn the runner's event loop; the handle submits work via events + tokio::spawn(runner.serve()); + + // Start interactive chat + if let Err(e) = chat::interactive_chat(&handle, &args.model).await { + eprintln!("❌ Chat error: {}", e); + std::process::exit(1); } } diff --git a/src/cli/mod.rs b/src/cli/mod.rs index fbfc92f..2e09052 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -1,3 +1,4 @@ +pub mod chat; pub mod commands; pub mod inspect; pub mod ls; diff --git a/src/cli/serve.rs b/src/cli/serve.rs index c9f7a31..55aab77 100644 --- a/src/cli/serve.rs +++ b/src/cli/serve.rs @@ -1,8 +1,11 @@ use colored::Colorize; use std::sync::Arc; +use tokenizers::models::bpe::BPE; +use tokenizers::Tokenizer; use tracing::{debug, info}; use crate::api::routes::create_router; +use crate::backend::engine; use crate::backend::mock::MockEngine; use crate::registry::model_registry::ModelRegistry; @@ -30,16 +33,25 @@ pub async fn execute( info!("Starting PUMA to serve model: {}", model_name); // Initialize backend (MockEngine for now, replace with MLX later) - let engine = Arc::new(MockEngine::new()); - info!("Inference engine initialized"); + let backend = MockEngine::new(); debug!("Using MockEngine backend"); + // TODO: Load the model's real tokenizer; placeholder BPE for now + let tokenizer = Tokenizer::new(BPE::default()); + + // Create engine: cheap send-side handle + runner that owns the scheduler + let (handle, runner) = engine(backend, tokenizer, model_name.to_string()); + + // Spawn the runner's event loop; the handle submits work via events + tokio::spawn(runner.serve()); + info!("Inference engine initialized"); + // Initialize model registry let registry = Arc::new(ModelRegistry::new(None)); info!("Model registry loaded"); // Create router - let app = create_router(engine, registry); + let app = create_router(handle, registry); // Bind address let addr = format!("{}:{}", host, port); @@ -54,7 +66,7 @@ pub async fn execute( info!(" GET /health"); // Start server - debug!("Starting axum server"); + debug!("Starting server"); axum::serve(listener, app).await?; info!("Server shutdown"); diff --git a/src/downloader/mod.rs b/src/downloader/mod.rs index 0067561..5ecc922 100644 --- a/src/downloader/mod.rs +++ b/src/downloader/mod.rs @@ -27,7 +27,10 @@ impl fmt::Display for DownloadError { } pub trait Downloader { - async fn download_model(&self, name: &str) -> Result<(), DownloadError>; + fn download_model( + &self, + name: &str, + ) -> impl std::future::Future> + Send; } /// Provider for downloading models diff --git a/src/fsm/events.rs b/src/fsm/events.rs new file mode 100644 index 0000000..69518cd --- /dev/null +++ b/src/fsm/events.rs @@ -0,0 +1,37 @@ +use super::states::FinishReason; +use crate::block_manager::types::TokenId; + +/// FSM Events - pure data that triggers state transitions +/// +/// These events represent state machine transitions and are created +/// internally by the scheduler. External components should use +/// SchedulerEvent instead. +#[derive(Debug, Clone)] +pub enum Event { + /// Create a new sequence: Empty → Waiting + Create { + token_ids: Vec, + max_tokens: usize, + }, + + /// Allocate blocks and start prefilling + Schedule { tokens_per_block: usize }, + + /// Append tokens to a running sequence + AppendTokens { + num_tokens: usize, + tokens_per_block: usize, + }, + + /// Mark sequence as complete + Complete { reason: FinishReason }, + + /// Preempt a running sequence (free blocks) + Preempt, + + /// Resume a preempted sequence + Resume, + + /// Abort a sequence with error + Abort { reason: String }, +} diff --git a/src/fsm/mod.rs b/src/fsm/mod.rs new file mode 100644 index 0000000..cbe4b1a --- /dev/null +++ b/src/fsm/mod.rs @@ -0,0 +1,43 @@ +//! FSM - Finite State Machine for sequence lifecycle +//! +//! This module defines the *data* of the state machine: the states and the +//! events that trigger transitions. The actual transition logic lives in +//! [`SequenceManager`](crate::sequence_manager::SequenceManager) — its private +//! `transition()`/`transition_*()` methods compute the next state and +//! allocate/free blocks along the way. +//! +//! # Architecture +//! +//! - `states`: State definitions (Waiting, Prefilling, Decoding, etc.) +//! - `events`: Event types that trigger transitions +//! +//! # Usage +//! +//! The `Scheduler` (policy) drives the `SequenceManager` (mechanism) with these +//! events; it never touches `BlockManager` or transition logic directly: +//! +//! ```rust,ignore +//! use crate::fsm::Event; +//! +//! // seq_manager: SequenceManager, owned by the scheduler +//! seq_manager.create(seq_id, token_ids, max_tokens)?; +//! let new_state = seq_manager.advance(seq_id, Event::Schedule { tokens_per_block: 16 })?; +//! ``` +//! +//! # Design Rationale +//! +//! Transition logic lives **inside `SequenceManager`** rather than a free +//! `fsm::apply()` because: +//! - The manager owns `BlockManager` - no need to pass it as a parameter +//! - Direct access to context (tokens_per_block, etc.) +//! - Event-based `advance()` gives a single entry point for logging/metrics +//! - Type-safe internal `transition_*()` methods enforce correct state types +//! +//! See `docs/fsm_architecture.md` for detailed documentation. + +pub mod events; +pub mod states; + +// Re-export commonly used types +pub use events::Event; +pub use states::*; diff --git a/src/fsm/states.rs b/src/fsm/states.rs new file mode 100644 index 0000000..17f5ceb --- /dev/null +++ b/src/fsm/states.rs @@ -0,0 +1,315 @@ +use crate::block_manager::types::*; +use std::sync::Arc; + +/// FSM States - each state owns its resources +/// +/// State transitions: +/// Waiting → Scheduling → Prefilling → Decoding → Finished +/// ↓ ↓ +/// └→ Preempted ←┘ + +#[derive(Debug, Clone)] +pub enum SequenceState { + /// Pre-birth placeholder: the "from" state a new sequence transitions out + /// of via `Event::Create`. Carries the `seq_id` so every state variant can + /// answer `seq_id()`. Never persisted in the scheduler's sequence map. + Empty(SequenceId), + + /// Waiting in queue for scheduling + Waiting(WaitingState), + + /// Being scheduled (allocating resources) + Scheduling(SchedulingState), + + /// Running prefill phase + Prefilling(PrefillingState), + + /// Running decode phase (generating tokens) + Decoding(DecodingState), + + /// Preempted due to OOM + Preempted(PreemptedState), + + /// Finished generation + Finished(FinishedState), + + /// Aborted (error or cancelled) + Aborted(AbortedState), +} + +/// Waiting state - no resources allocated yet +#[derive(Debug, Clone)] +pub struct WaitingState { + pub seq_id: SequenceId, + pub token_ids: Arc>, // Arc for O(1) cloning, GPU gets reference + pub max_tokens: usize, +} + +/// Scheduling state - in process of allocating blocks +#[derive(Debug, Clone)] +pub struct SchedulingState { + pub seq_id: SequenceId, + pub prompt_tokens: usize, + pub max_tokens: usize, + // Partially allocated blocks (if any) + pub blocks: Arc>, +} + +/// Prefilling state - processing prompt +#[derive(Debug, Clone)] +pub struct PrefillingState { + pub seq_id: SequenceId, + pub token_ids: Arc>, // Arc for O(1) cloning, needed if preempted + pub blocks: Arc>, // Arc for cheap cloning (O(1)) + pub tokens_filled: usize, + pub tokens_total: usize, + pub max_tokens: usize, +} + +/// Decoding state - generating tokens +#[derive(Debug, Clone)] +pub struct DecodingState { + pub seq_id: SequenceId, + pub token_ids: Arc>, // Arc for O(1) cloning, needed if preempted + pub blocks: Arc>, // Arc for cheap cloning (O(1)) + pub num_tokens: usize, + pub max_tokens: usize, +} + +/// Preempted state - blocks freed, waiting to resume +#[derive(Debug, Clone)] +pub struct PreemptedState { + pub seq_id: SequenceId, + pub token_ids: Arc>, // Arc for O(1) cloning, needed to resume + pub num_tokens: usize, // How many tokens were generated so far + pub max_tokens: usize, + // No blocks - they were freed +} + +/// Finished state - all resources freed +#[derive(Debug, Clone)] +pub struct FinishedState { + pub seq_id: SequenceId, + pub finish_reason: FinishReason, +} + +/// Aborted state - error or cancellation +#[derive(Debug, Clone)] +pub struct AbortedState { + pub seq_id: SequenceId, + pub reason: String, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum FinishReason { + /// Reached max_tokens limit + MaxTokens, + /// Natural completion (EOS token) + Stop, + /// User cancelled + Cancelled, +} + +// Helper methods for SequenceState +impl SequenceState { + pub fn seq_id(&self) -> SequenceId { + match self { + SequenceState::Empty(seq_id) => *seq_id, + SequenceState::Waiting(s) => s.seq_id, + SequenceState::Scheduling(s) => s.seq_id, + SequenceState::Prefilling(s) => s.seq_id, + SequenceState::Decoding(s) => s.seq_id, + SequenceState::Preempted(s) => s.seq_id, + SequenceState::Finished(s) => s.seq_id, + SequenceState::Aborted(s) => s.seq_id, + } + } + + pub fn is_running(&self) -> bool { + matches!( + self, + SequenceState::Prefilling(_) | SequenceState::Decoding(_) + ) + } + + pub fn is_waiting(&self) -> bool { + matches!( + self, + SequenceState::Waiting(_) | SequenceState::Scheduling(_) + ) + } + + pub fn is_finished(&self) -> bool { + matches!(self, SequenceState::Finished(_) | SequenceState::Aborted(_)) + } + + pub fn blocks(&self) -> Option<&[BlockId]> { + match self { + SequenceState::Prefilling(s) => Some(s.blocks.as_ref()), + SequenceState::Decoding(s) => Some(s.blocks.as_ref()), + SequenceState::Scheduling(s) if !s.blocks.is_empty() => Some(s.blocks.as_ref()), + _ => None, + } + } + + pub fn variant_name(&self) -> &'static str { + match self { + SequenceState::Empty(_) => "Empty", + SequenceState::Waiting(_) => "Waiting", + SequenceState::Scheduling(_) => "Scheduling", + SequenceState::Prefilling(_) => "Prefilling", + SequenceState::Decoding(_) => "Decoding", + SequenceState::Preempted(_) => "Preempted", + SequenceState::Finished(_) => "Finished", + SequenceState::Aborted(_) => "Aborted", + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn waiting(seq_id: u64) -> SequenceState { + SequenceState::Waiting(WaitingState { + seq_id: SequenceId(seq_id), + token_ids: Arc::new(vec![1, 2, 3]), + max_tokens: 10, + }) + } + + fn prefilling(seq_id: u64, blocks: Vec) -> SequenceState { + SequenceState::Prefilling(PrefillingState { + seq_id: SequenceId(seq_id), + token_ids: Arc::new(vec![1, 2, 3]), + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: 3, + max_tokens: 10, + }) + } + + fn decoding(seq_id: u64, blocks: Vec) -> SequenceState { + SequenceState::Decoding(DecodingState { + seq_id: SequenceId(seq_id), + token_ids: Arc::new(vec![1, 2, 3]), + blocks: Arc::new(blocks), + num_tokens: 3, + max_tokens: 10, + }) + } + + #[test] + fn test_seq_id_for_every_variant() { + assert_eq!(SequenceState::Empty(SequenceId(7)).seq_id(), SequenceId(7)); + assert_eq!(waiting(1).seq_id(), SequenceId(1)); + assert_eq!(prefilling(2, vec![]).seq_id(), SequenceId(2)); + assert_eq!(decoding(3, vec![]).seq_id(), SequenceId(3)); + + let scheduling = SequenceState::Scheduling(SchedulingState { + seq_id: SequenceId(4), + prompt_tokens: 3, + max_tokens: 10, + blocks: Arc::new(vec![]), + }); + assert_eq!(scheduling.seq_id(), SequenceId(4)); + + let preempted = SequenceState::Preempted(PreemptedState { + seq_id: SequenceId(5), + token_ids: Arc::new(vec![1]), + num_tokens: 1, + max_tokens: 10, + }); + assert_eq!(preempted.seq_id(), SequenceId(5)); + + let finished = SequenceState::Finished(FinishedState { + seq_id: SequenceId(6), + finish_reason: FinishReason::Stop, + }); + assert_eq!(finished.seq_id(), SequenceId(6)); + + let aborted = SequenceState::Aborted(AbortedState { + seq_id: SequenceId(8), + reason: "boom".to_string(), + }); + assert_eq!(aborted.seq_id(), SequenceId(8)); + } + + #[test] + fn test_is_running() { + assert!(prefilling(1, vec![]).is_running()); + assert!(decoding(1, vec![]).is_running()); + assert!(!waiting(1).is_running()); + assert!(!SequenceState::Empty(SequenceId(1)).is_running()); + } + + #[test] + fn test_is_waiting() { + assert!(waiting(1).is_waiting()); + let scheduling = SequenceState::Scheduling(SchedulingState { + seq_id: SequenceId(1), + prompt_tokens: 3, + max_tokens: 10, + blocks: Arc::new(vec![]), + }); + assert!(scheduling.is_waiting()); + assert!(!decoding(1, vec![]).is_waiting()); + } + + #[test] + fn test_is_finished() { + let finished = SequenceState::Finished(FinishedState { + seq_id: SequenceId(1), + finish_reason: FinishReason::MaxTokens, + }); + let aborted = SequenceState::Aborted(AbortedState { + seq_id: SequenceId(1), + reason: "x".to_string(), + }); + assert!(finished.is_finished()); + assert!(aborted.is_finished()); + assert!(!decoding(1, vec![]).is_finished()); + } + + #[test] + fn test_blocks_exposed_only_for_resource_holding_states() { + // Prefilling/Decoding always expose their blocks. + assert_eq!( + prefilling(1, vec![BlockId(0), BlockId(1)]).blocks(), + Some(&[BlockId(0), BlockId(1)][..]) + ); + assert_eq!( + decoding(1, vec![BlockId(2)]).blocks(), + Some(&[BlockId(2)][..]) + ); + + // Scheduling exposes blocks only when it has partially allocated some. + let scheduling_with = SequenceState::Scheduling(SchedulingState { + seq_id: SequenceId(1), + prompt_tokens: 3, + max_tokens: 10, + blocks: Arc::new(vec![BlockId(9)]), + }); + assert_eq!(scheduling_with.blocks(), Some(&[BlockId(9)][..])); + + let scheduling_empty = SequenceState::Scheduling(SchedulingState { + seq_id: SequenceId(1), + prompt_tokens: 3, + max_tokens: 10, + blocks: Arc::new(vec![]), + }); + assert_eq!(scheduling_empty.blocks(), None); + + // States without KV blocks return None. + assert_eq!(waiting(1).blocks(), None); + assert_eq!(SequenceState::Empty(SequenceId(1)).blocks(), None); + } + + #[test] + fn test_variant_name() { + assert_eq!(SequenceState::Empty(SequenceId(1)).variant_name(), "Empty"); + assert_eq!(waiting(1).variant_name(), "Waiting"); + assert_eq!(prefilling(1, vec![]).variant_name(), "Prefilling"); + assert_eq!(decoding(1, vec![]).variant_name(), "Decoding"); + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..909efec --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,14 @@ +#![allow(dead_code)] + +pub mod api; +pub mod backend; +pub mod block_manager; +pub mod cli; +pub mod downloader; +pub mod fsm; +pub mod registry; +pub mod scheduler; +pub mod sequence_manager; +pub mod storage; +pub mod system; +pub mod utils; diff --git a/src/main.rs b/src/main.rs index 02929ea..a15b72d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,14 @@ +#![allow(dead_code)] + mod api; mod backend; +mod block_manager; mod cli; mod downloader; +mod fsm; mod registry; +mod scheduler; +mod sequence_manager; mod storage; mod system; mod utils; @@ -14,20 +20,28 @@ use crate::cli::commands::{run, Cli}; use crate::utils::file; fn main() { - // Setup tracing subscriber for tower-http TraceLayer + let cli = Cli::parse(); + + // Setup tracing subscriber for tower-http TraceLayer. + // + // An explicit RUST_LOG always wins. Otherwise only the long-running server + // logs by default; one-shot/interactive CLI commands stay quiet so their + // output isn't buried under per-request INFO lines. + let default_filter = if cli.wants_default_logging() { + "info,hf_hub=warn,tower_http=info,rusqlite_migration=warn" + } else { + "off" + }; tracing_subscriber::fmt() .with_env_filter( - tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| { - "info,hf_hub=warn,tower_http=info,rusqlite_migration=warn".into() - }), + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| default_filter.into()), ) .init(); // Create the root folder if it doesn't exist. file::create_folder_if_not_exists(&file::root_home()).unwrap(); - let cli = Cli::parse(); - let runtime = Builder::new_multi_thread() .worker_threads(4) .enable_all() diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs new file mode 100644 index 0000000..252b8c8 --- /dev/null +++ b/src/scheduler/core.rs @@ -0,0 +1,785 @@ +use super::events::{ResponseSender, SchedulerEvent, SchedulerEventReceiver, SchedulerStats}; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use crate::fsm::{Event, FinishReason, SequenceState}; +use crate::sequence_manager::SequenceManager; +use std::collections::{HashMap, VecDeque}; +use std::sync::Arc; +use tracing::{debug, info, warn}; + +/// Scheduler - batch and scheduling policy (synchronous API) +/// +/// Architecture +/// - Pure synchronous methods (no async, no loop) +/// - LLMEngine calls: add_request() → schedule() → process_outputs() +/// - Drives a `SequenceManager` (which owns memory and executes FSM +/// transitions) via `create`/`advance`/`fork` +/// +/// Responsibilities: +/// - This is the *policy* layer; the `SequenceManager` is the *mechanism* +/// layer that owns `BlockManager` and the FSM transition logic +/// - Manages batches (waiting/prefill/decode) +/// - Scheduling policy (what to schedule when) +/// - Owns client response channels and delivers results/errors +pub struct Scheduler { + /// Owns sequence state + block memory and executes FSM transitions. + sequences: SequenceManager, + + /// Response channels to send results back to clients (supports streaming) + response_channels: HashMap, + + /// Waiting sequences (not yet scheduled) + waiting_queue: VecDeque, + + /// Currently running sequences in prefill phase + prefill_batch: Vec, + + /// Currently running sequences in decode phase + decode_batch: Vec, + + /// Receiver for external events (add/cancel/query) from clients. + /// + /// Public so the engine loop can drain it directly and feed each event to + /// [`handle_event`](Self::handle_event). + pub event_rx: SchedulerEventReceiver, + + /// Configuration + max_batch_size: usize, + tokens_per_block: usize, +} + +impl Scheduler { + pub fn new( + block_manager: BlockManager, + event_rx: SchedulerEventReceiver, + max_batch_size: usize, + tokens_per_block: usize, + ) -> Self { + Self { + sequences: SequenceManager::new(block_manager, tokens_per_block), + response_channels: HashMap::new(), + waiting_queue: VecDeque::new(), + prefill_batch: Vec::new(), + decode_batch: Vec::new(), + event_rx, + max_batch_size, + tokens_per_block, + } + } + + /// Dispatch a single external event to its handler. + /// + /// One entry point for every `SchedulerEvent`: mutating events build the + /// corresponding FSM event and apply it; query events reply on their own + /// response channel. + pub fn handle_event(&mut self, event: SchedulerEvent) { + match event { + SchedulerEvent::AddRequest { + seq_id, + token_ids, + max_tokens, + response_tx, + } => self.add_request(seq_id, token_ids, max_tokens, response_tx), + + SchedulerEvent::CancelRequest { seq_id } => { + self.abort_request(seq_id, "User cancelled".to_string()) + } + + SchedulerEvent::GetBlocks { seq_id, response } => { + let _ = response.send(self.get_blocks(seq_id)); + } + + SchedulerEvent::GetStats { response } => { + let _ = response.send(self.get_stats()); + } + } + } + + // ===== Public API (called by LLMEngine) ===== + + pub fn add_request( + &mut self, + seq_id: SequenceId, + token_ids: Vec, + max_tokens: usize, + response_tx: ResponseSender, + ) { + debug!( + "Adding request: seq_id={:?}, num_tokens={}, max_tokens={}", + seq_id, + token_ids.len(), + max_tokens + ); + + // Check duplicate + if self.sequences.contains(seq_id) { + warn!("Duplicate seq_id {:?}", seq_id); + self.send_error( + response_tx, + Error::InvalidTransition("Duplicate sequence ID"), + ); + return; + } + + // Store response channel + self.response_channels.insert(seq_id, response_tx); + + // Birth the sequence (Empty → Waiting) inside the sequence manager. + match self.sequences.create(seq_id, token_ids, max_tokens) { + Ok(()) => { + self.waiting_queue.push_back(seq_id); + } + Err(e) => { + warn!("Failed to create sequence {:?}: {:?}", seq_id, e); + self.response_channels.remove(&seq_id); + } + } + } + + /// Abort a sequence with a caller-supplied reason. + /// + /// The reason is passed through to the FSM `Abort` event, so callers (user + /// cancel, backend failure, …) describe *why* the sequence ended. + pub fn abort_request(&mut self, seq_id: SequenceId, reason: String) { + match self.sequences.advance( + seq_id, + Event::Abort { + reason: reason.clone(), + }, + ) { + Ok(_) => { + info!("Aborted sequence {:?}", seq_id); + + // Remove from queues/batches + self.waiting_queue.retain(|&id| id != seq_id); + self.prefill_batch.retain(|&id| id != seq_id); + self.decode_batch.retain(|&id| id != seq_id); + + // Resolve the client channel so waiters unblock and the entry + // is not leaked. Streaming callers see the error and close; + // single callers get an `Err` instead of hanging forever. + if let Some(response_tx) = self.response_channels.remove(&seq_id) { + self.send_error(response_tx, Error::Aborted(reason)); + } + } + Err(e) => { + warn!("Failed to abort {:?}: {:?}", seq_id, e); + } + } + } + + /// Schedule sequences - returns whether work was scheduled + /// + /// Called by LLMEngine every iteration to: + /// 1. Move waiting → prefilling + /// 2. Allocate blocks + /// 3. Build batches + pub fn schedule(&mut self) -> bool { + let before = self.prefill_batch.len(); + self.schedule_prefill(); + self.schedule_decode(); + self.prefill_batch.len() > before || !self.decode_batch.is_empty() + } + + /// Drain the current prefill batch as inference work for the engine loop. + /// + /// Clears the prefill batch and returns `(seq_id, token_ids, max_tokens, + /// streaming)` per scheduled sequence. The engine loop runs the backend on + /// these and reports results back via [`send_chunk`](Self::send_chunk) and + /// [`complete_sequence`](Self::complete_sequence). The `streaming` flag tells + /// the loop whether to call the backend's streaming or single-shot path. + pub fn take_prefill_batch(&mut self) -> Vec<(SequenceId, Arc>, usize, bool)> { + let batch = std::mem::take(&mut self.prefill_batch); + batch + .into_iter() + .filter_map(|seq_id| match self.sequences.get(seq_id) { + Some(SequenceState::Prefilling(s)) => { + let streaming = self + .response_channels + .get(&seq_id) + .is_some_and(|tx| tx.is_streaming()); + Some((seq_id, Arc::clone(&s.token_ids), s.max_tokens, streaming)) + } + _ => None, + }) + .collect() + } + + // ===== Internal State Management (called by LLMEngine after inference) ===== + + pub fn append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { + let event = Event::AppendTokens { + num_tokens, + tokens_per_block: self.tokens_per_block, + }; + + // On error the sequence manager restores the prior state; the scheduler + // only applies OOM policy. + match self.sequences.advance(seq_id, event) { + Ok(_) => {} + Err(Error::OutOfMemory) => { + warn!("OOM while appending tokens to {:?}", seq_id); + self.handle_oom(); + } + Err(e) => { + warn!("Failed to append tokens to {:?}: {:?}", seq_id, e); + } + } + } + + pub fn complete_sequence(&mut self, seq_id: SequenceId, reason: FinishReason, result: String) { + match self.sequences.advance(seq_id, Event::Complete { reason }) { + Ok(_) => { + info!("Completed sequence {:?}", seq_id); + + // Remove from batches + self.prefill_batch.retain(|&id| id != seq_id); + self.decode_batch.retain(|&id| id != seq_id); + + // Send result back to client + self.send_result(seq_id, result); + } + Err(e) => { + warn!("Failed to complete {:?}: {:?}", seq_id, e); + + // Send error back to client + if let Some(response_tx) = self.response_channels.remove(&seq_id) { + self.send_error(response_tx, e); + } + } + } + } + + pub fn fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { + let _ = self.sequences.fork(parent_id, child_id); + } + + fn preempt_sequence(&mut self, seq_id: SequenceId) { + match self.sequences.advance(seq_id, Event::Preempt) { + Ok(_) => { + info!("Preempted sequence {:?}", seq_id); + + // Remove from batches + self.prefill_batch.retain(|&id| id != seq_id); + self.decode_batch.retain(|&id| id != seq_id); + } + Err(e) => { + warn!("Failed to preempt {:?}: {:?}", seq_id, e); + } + } + } + + fn resume_sequence(&mut self, seq_id: SequenceId) { + match self.sequences.advance(seq_id, Event::Resume) { + Ok(_) => { + debug!("Resumed sequence {:?}", seq_id); + self.waiting_queue.push_back(seq_id); + } + Err(e) => { + warn!("Failed to resume {:?}: {:?}", seq_id, e); + } + } + } + + // ===== Scheduling Logic ===== + + fn try_schedule(&mut self) { + self.schedule_prefill(); + self.schedule_decode(); + } + + fn schedule_prefill(&mut self) { + while let Some(seq_id) = self.waiting_queue.pop_front() { + // Check batch limit + if self.prefill_batch.len() >= self.max_batch_size { + self.waiting_queue.push_front(seq_id); + debug!("Prefill batch full"); + break; + } + + let event = Event::Schedule { + tokens_per_block: self.tokens_per_block, + }; + + // The sequence manager restores prior state on any error, so the + // scheduler only applies queueing policy. + match self.sequences.advance(seq_id, event) { + Ok(_) => { + self.prefill_batch.push(seq_id); + info!("Scheduled {:?} for prefill", seq_id); + } + Err(Error::OutOfMemory) => { + // Put back at the front and stop scheduling this round. + self.waiting_queue.push_front(seq_id); + debug!("OOM - cannot schedule {:?}", seq_id); + break; + } + Err(Error::InvalidTransition(_)) => { + warn!("Cannot schedule {:?} - not in Waiting state", seq_id); + } + Err(e) => { + warn!("Failed to schedule {:?}: {:?}", seq_id, e); + } + } + } + } + + fn schedule_decode(&mut self) { + // TODO: Continuous batching + // - Move completed prefill → decode batch + // - Remove finished sequences + // - Respect max_batch_size + + debug!("Decode batch size: {}", self.decode_batch.len()); + } + + fn handle_oom(&mut self) { + warn!("Handling OOM"); + // TODO: Preemption policy + // Find lowest priority sequence + // Preempt it + // Retry scheduling + } + + // ===== Query Methods ===== + + pub fn get_blocks(&self, seq_id: SequenceId) -> Result> { + self.sequences.blocks(seq_id) + } + + pub fn get_stats(&self) -> SchedulerStats { + SchedulerStats { + num_sequences: self.sequences.len(), + num_running: self.sequences.num_running(), + num_waiting: self.waiting_queue.len(), + num_preempted: self.sequences.num_preempted(), + block_stats: self.sequences.block_stats(), + } + } + + // ===== Response Channel Helpers ===== + + /// Send a decoded text chunk to a streaming client. + /// + /// The chunk is an incremental piece of detokenized output (a delta), not a + /// single token — see the engine's incremental decode. No-op for + /// non-streaming clients, which receive the full text via + /// [`complete_sequence`](Self::complete_sequence). + pub fn send_chunk(&mut self, seq_id: SequenceId, chunk: String) { + if let Some(response_tx) = self.response_channels.get(&seq_id) { + match response_tx { + ResponseSender::Stream(tx) => { + let _ = tx.send(Ok(chunk)); + } + ResponseSender::Single(_) => { + // Single response - streamed chunks are dropped; + // the full result is delivered at completion. + } + } + } + } + + /// Send error to client + fn send_error(&self, response_tx: ResponseSender, error: Error) { + match response_tx { + ResponseSender::Single(tx) => { + let _ = tx.send(Err(error)); + } + ResponseSender::Stream(tx) => { + let _ = tx.send(Err(error)); + } + } + } + + /// Send final result to client + fn send_result(&mut self, seq_id: SequenceId, result: String) { + if let Some(response_tx) = self.response_channels.remove(&seq_id) { + match response_tx { + ResponseSender::Single(tx) => { + let _ = tx.send(Ok(result)); + } + ResponseSender::Stream(_tx) => { + // Already streamed tokens, just close the channel + // (drop _tx automatically closes it) + } + } + } + } +} + +/// Create scheduler event channel +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + use crate::scheduler::events::ResponseSender; + + fn create_test_scheduler() -> Scheduler { + let allocator = Box::new(CpuAllocator::new(100_000)); + let block_manager = BlockManager::new(allocator, 1024); + let (_event_tx, event_rx) = tokio::sync::mpsc::unbounded_channel(); + Scheduler::new(block_manager, event_rx, 10, 16) + } + + fn create_dummy_response() -> ResponseSender { + let (tx, _rx) = tokio::sync::oneshot::channel(); + ResponseSender::Single(tx) + } + + #[test] + fn test_add_request() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + + assert_eq!(scheduler.sequences.len(), 1); + assert_eq!(scheduler.waiting_queue.len(), 1); + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Waiting(_)) + )); + } + + #[test] + fn test_add_duplicate_request() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); // Duplicate + + // Should not add duplicate + assert_eq!(scheduler.sequences.len(), 1); + assert_eq!(scheduler.waiting_queue.len(), 1); + } + + #[test] + fn test_cancel_request_waiting() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.abort_request(SequenceId(1), "User cancelled".to_string()); + + // Should transition to Aborted + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Aborted(_)) + )); + assert_eq!(scheduler.waiting_queue.len(), 0); + } + + #[test] + fn test_abort_resolves_single_channel_and_removes_entry() { + let mut scheduler = create_test_scheduler(); + + // Keep the receiver so we can observe what the client sees. + let (tx, rx) = tokio::sync::oneshot::channel(); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, ResponseSender::Single(tx)); + assert_eq!(scheduler.response_channels.len(), 1); + + scheduler.abort_request(SequenceId(1), "User cancelled".to_string()); + + // The waiter is unblocked with an Aborted error (not left hanging)... + match rx.blocking_recv() { + Ok(Err(Error::Aborted(reason))) => assert_eq!(reason, "User cancelled"), + other => panic!("expected Ok(Err(Aborted)), got {:?}", other), + } + // ...and the channel entry is not leaked. + assert!(scheduler.response_channels.is_empty()); + } + + #[test] + fn test_abort_closes_streaming_channel() { + let mut scheduler = create_test_scheduler(); + + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, ResponseSender::Stream(tx)); + + scheduler.abort_request(SequenceId(1), "boom".to_string()); + + // Streaming client receives the error, then the channel closes. + match rx.blocking_recv() { + Some(Err(Error::Aborted(reason))) => assert_eq!(reason, "boom"), + other => panic!("expected Some(Err(Aborted)), got {:?}", other), + } + assert!(rx.blocking_recv().is_none(), "channel should be closed"); + assert!(scheduler.response_channels.is_empty()); + } + + #[test] + fn test_cancel_request_not_found() { + let mut scheduler = create_test_scheduler(); + + // Should not panic + scheduler.abort_request(SequenceId(999), "User cancelled".to_string()); + assert_eq!(scheduler.sequences.len(), 0); + } + + #[test] + fn test_schedule_prefill_success() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + + // Should move to prefilling + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Prefilling(_)) + )); + assert_eq!(scheduler.prefill_batch.len(), 1); + assert_eq!(scheduler.waiting_queue.len(), 0); + } + + #[test] + fn test_schedule_prefill_multiple() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.add_request(SequenceId(2), vec![0; 50], 100, create_dummy_response()); + scheduler.schedule_prefill(); + + // Both should be scheduled + assert_eq!(scheduler.prefill_batch.len(), 2); + assert_eq!(scheduler.waiting_queue.len(), 0); + } + + #[test] + fn test_schedule_prefill_batch_limit() { + let mut scheduler = create_test_scheduler(); + + // Add more than max_batch_size (10) + for i in 0..15 { + scheduler.add_request(SequenceId(i), vec![0; 100], 200, create_dummy_response()); + } + scheduler.schedule_prefill(); + + // Only 10 should be scheduled + assert_eq!(scheduler.prefill_batch.len(), 10); + assert_eq!(scheduler.waiting_queue.len(), 5); + } + + #[test] + fn test_schedule_prefill_oom() { + // Small allocator - only 2 blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let block_manager = BlockManager::new(allocator, 1024); + let (_event_tx, event_rx) = tokio::sync::mpsc::unbounded_channel(); + let mut scheduler = Scheduler::new(block_manager, event_rx, 10, 16); + + // First request needs 7 blocks (100 tokens / 16 tokens_per_block) + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + + // Should fail - not enough memory + assert_eq!(scheduler.prefill_batch.len(), 0); + assert_eq!(scheduler.waiting_queue.len(), 1); + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Waiting(_)) + )); + } + + #[test] + fn test_append_tokens_prefilling() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + + // Append all tokens + scheduler.append_tokens(SequenceId(1), 100); + + // Should transition to Decoding + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Decoding(_)) + )); + } + + #[test] + fn test_append_tokens_decoding() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + // Append more tokens + scheduler.append_tokens(SequenceId(1), 10); + + if let Some(SequenceState::Decoding(state)) = scheduler.sequences.get(SequenceId(1)) { + assert_eq!(state.num_tokens, 110); + } else { + panic!("Expected Decoding state"); + } + } + + #[test] + fn test_append_tokens_reaches_max() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + // Reach max_tokens + scheduler.append_tokens(SequenceId(1), 100); // 100 + 100 = 200 = max_tokens + + // Should transition to Finished + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Finished(_)) + )); + } + + #[test] + fn test_complete_sequence() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + scheduler.complete_sequence(SequenceId(1), FinishReason::Stop, "test result".to_string()); + + // Should be Finished + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Finished(_)) + )); + assert_eq!(scheduler.prefill_batch.len(), 0); + } + + #[test] + fn test_fork_sequence() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + scheduler.fork_sequence(SequenceId(1), SequenceId(2)); + + // Both should exist + assert!(scheduler.sequences.contains(SequenceId(1))); + assert!(scheduler.sequences.contains(SequenceId(2))); + assert!(matches!( + scheduler.sequences.get(SequenceId(2)), + Some(SequenceState::Decoding(_)) + )); + } + + #[test] + fn test_fork_duplicate_child_id() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); + + scheduler.add_request(SequenceId(2), vec![0; 50], 100, create_dummy_response()); // Child ID exists + scheduler.fork_sequence(SequenceId(1), SequenceId(2)); // Should fail + + // Original seq 2 should be unchanged (Waiting) + assert!(matches!( + scheduler.sequences.get(SequenceId(2)), + Some(SequenceState::Waiting(_)) + )); + } + + #[test] + fn test_preempt_sequence() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + scheduler.preempt_sequence(SequenceId(1)); + + // Should be Preempted + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Preempted(_)) + )); + assert_eq!(scheduler.decode_batch.len(), 0); + } + + #[test] + fn test_resume_sequence() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); + scheduler.preempt_sequence(SequenceId(1)); + + scheduler.resume_sequence(SequenceId(1)); + + // Should be back to Waiting + assert!(matches!( + scheduler.sequences.get(SequenceId(1)), + Some(SequenceState::Waiting(_)) + )); + assert_eq!(scheduler.waiting_queue.len(), 1); + } + + #[test] + fn test_get_blocks() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + + let blocks = scheduler.get_blocks(SequenceId(1)).unwrap(); + assert!(!blocks.is_empty()); + } + + #[test] + fn test_get_blocks_not_found() { + let scheduler = create_test_scheduler(); + + let result = scheduler.get_blocks(SequenceId(999)); + assert!(matches!(result, Err(Error::UnknownSequence(_)))); + } + + #[test] + fn test_get_stats() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.add_request(SequenceId(2), vec![0; 50], 100, create_dummy_response()); + scheduler.schedule_prefill(); + + let stats = scheduler.get_stats(); + assert_eq!(stats.num_sequences, 2); + assert_eq!(stats.num_running, 2); // Both in prefill + assert_eq!(stats.num_waiting, 0); + } + + #[test] + fn test_try_schedule() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.try_schedule(); + + // Should schedule automatically + assert_eq!(scheduler.prefill_batch.len(), 1); + assert_eq!(scheduler.waiting_queue.len(), 0); + } + + #[test] + fn test_cancel_removes_from_batches() { + let mut scheduler = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.schedule_prefill(); + assert_eq!(scheduler.prefill_batch.len(), 1); + + scheduler.abort_request(SequenceId(1), "User cancelled".to_string()); + + // Should remove from batch + assert_eq!(scheduler.prefill_batch.len(), 0); + } +} diff --git a/src/scheduler/events.rs b/src/scheduler/events.rs new file mode 100644 index 0000000..5a88c06 --- /dev/null +++ b/src/scheduler/events.rs @@ -0,0 +1,103 @@ +//! Scheduler Events - External API +//! +//! # Two-Level Event Architecture +//! +//! PUMA follows TokenSpeed's pattern with two distinct event levels: +//! +//! ## Level 1: FSM Events (Internal - `crate::fsm::Event`) +//! - Pure state machine transitions +//! - Created by scheduler internally +//! - Includes scheduler resources (BlockManager, allocators) +//! - Examples: Schedule, AppendTokens, Complete +//! +//! ## Level 2: Scheduler Events (External - `SchedulerEvent`) +//! - API boundary for external components +//! - Simple data, no resource pointers +//! - Can be created by clients, GPU, cache systems +//! - Examples: AddRequest, CancelRequest, TokensGenerated +//! +//! ## Why Two Levels? +//! +//! External components (GPU, clients) **cannot create** FSM events because: +//! - FSM events need `&mut BlockManager` (owned by scheduler) +//! - FSM events include policy decisions (scheduler computes) +//! - FSM events track internal state (scheduler maintains) +//! +//! The scheduler translates external events → FSM events, adding +//! the necessary context and resources. + +use crate::block_manager::types::*; +use std::collections::HashMap; +use tokio::sync::mpsc; + +/// External events that clients/backend send to the scheduler +/// +/// Only includes events that come from outside the scheduler: +/// - AddRequest: Client starts a new request +/// - CancelRequest: Client cancels an ongoing request +/// - GetBlocks: Backend queries block IDs for inference +/// - GetStats: Monitoring/observability queries +/// +/// Response sender - supports both single response and streaming +#[derive(Debug)] +pub enum ResponseSender { + /// Single response (non-streaming) + Single(tokio::sync::oneshot::Sender>), + /// Streaming response (multiple tokens) + Stream(tokio::sync::mpsc::UnboundedSender>), +} + +impl ResponseSender { + /// Whether this sequence expects token-by-token streaming. + pub fn is_streaming(&self) -> bool { + matches!(self, ResponseSender::Stream(_)) + } +} + +/// Internal operations (token generation, state transitions, OOM handling) +/// are NOT events - the scheduler handles them directly in its run loop. +#[derive(Debug)] +pub enum SchedulerEvent { + /// Client starts a new inference request + AddRequest { + seq_id: SequenceId, + token_ids: Vec, // Tokenized by LLMEngine + max_tokens: usize, + response_tx: ResponseSender, // Send result(s) back + }, + + /// Client cancels a request (user stop, disconnect, timeout) + CancelRequest { seq_id: SequenceId }, + + /// Backend queries block IDs for a sequence (for KV cache addressing) + GetBlocks { + seq_id: SequenceId, + response: tokio::sync::oneshot::Sender>>, + }, + + /// Query scheduler statistics + GetStats { + response: tokio::sync::oneshot::Sender, + }, +} + +#[derive(Debug, Clone)] +pub struct SchedulerStats { + pub num_sequences: usize, + pub num_running: usize, + pub num_waiting: usize, + pub num_preempted: usize, + pub block_stats: HashMap, +} + +pub type SchedulerEventSender = mpsc::UnboundedSender; +pub type SchedulerEventReceiver = mpsc::UnboundedReceiver; + +/// Create event channel for scheduler +/// +/// TODO: Use bounded channel to prevent OOM when producers outpace the event loop. +/// Capacity should be calculated based on available memory and average event size. +/// Consider: capacity = (available_memory * 0.1) / sizeof(SchedulerEvent) +pub fn create_event_channel() -> (SchedulerEventSender, SchedulerEventReceiver) { + mpsc::unbounded_channel() +} diff --git a/src/scheduler/id_generator.rs b/src/scheduler/id_generator.rs new file mode 100644 index 0000000..ee9e490 --- /dev/null +++ b/src/scheduler/id_generator.rs @@ -0,0 +1,70 @@ +use crate::block_manager::types::SequenceId; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Sequential ID generator for sequences +/// Thread-safe, starts from 0 +pub struct SequenceIdGenerator { + next_id: AtomicU64, +} + +impl SequenceIdGenerator { + /// Create a new ID generator starting from 0 + pub fn new() -> Self { + Self { + next_id: AtomicU64::new(0), + } + } + + /// Create an ID generator starting from a specific value + pub fn with_start(start: u64) -> Self { + Self { + next_id: AtomicU64::new(start), + } + } + + /// Generate the next sequential ID + /// Thread-safe, can be called concurrently + pub fn next(&self) -> SequenceId { + SequenceId(self.next_id.fetch_add(1, Ordering::Relaxed)) + } + + /// Peek at the next ID without consuming it + pub fn peek(&self) -> SequenceId { + SequenceId(self.next_id.load(Ordering::Relaxed)) + } +} + +impl Default for SequenceIdGenerator { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sequential_generation() { + let gen = SequenceIdGenerator::new(); + assert_eq!(gen.next().0, 0); + assert_eq!(gen.next().0, 1); + assert_eq!(gen.next().0, 2); + } + + #[test] + fn test_with_start() { + let gen = SequenceIdGenerator::with_start(100); + assert_eq!(gen.next().0, 100); + assert_eq!(gen.next().0, 101); + } + + #[test] + fn test_peek() { + let gen = SequenceIdGenerator::new(); + assert_eq!(gen.peek().0, 0); + assert_eq!(gen.peek().0, 0); // Doesn't consume + assert_eq!(gen.next().0, 0); // Now consume + assert_eq!(gen.peek().0, 1); + } +} diff --git a/src/scheduler/mod.rs b/src/scheduler/mod.rs new file mode 100644 index 0000000..f463eba --- /dev/null +++ b/src/scheduler/mod.rs @@ -0,0 +1,5 @@ +pub mod core; +pub mod events; +pub mod id_generator; + +// Re-export for convenience diff --git a/src/sequence_manager/mod.rs b/src/sequence_manager/mod.rs new file mode 100644 index 0000000..dffb7f5 --- /dev/null +++ b/src/sequence_manager/mod.rs @@ -0,0 +1,773 @@ +//! SequenceManager - owns sequence state and the memory those sequences use. +//! +//! This is the *mechanism* layer: it holds the `sequences` map and the +//! `BlockManager`, and executes FSM transitions (which allocate/free blocks). +//! The `Scheduler` is the *policy* layer — it decides which events fire and +//! owns the queues/batches — and drives this manager without ever touching a +//! `BlockManager` or a `BlockId` directly. +//! +//! Keeping memory here (rather than in the scheduler) localizes future changes +//! like multi-GPU block pools or CPU/GPU swap: the scheduler keeps calling +//! `advance`/`create`/`can_allocate`, and the device details stay inside. + +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use crate::fsm::{ + AbortedState, DecodingState, Event, FinishReason, FinishedState, PreemptedState, + PrefillingState, SequenceState, WaitingState, +}; +use std::collections::HashMap; +use std::sync::Arc; +use tracing::{debug, warn}; + +/// Owns sequence state + block memory and executes FSM transitions. +pub struct SequenceManager { + /// Block manager for memory allocation + block_manager: BlockManager, + + /// All sequences and their states + sequences: HashMap, + + /// Tokens per block (used when computing block allocations) + tokens_per_block: usize, +} + +impl SequenceManager { + pub fn new(block_manager: BlockManager, tokens_per_block: usize) -> Self { + assert!(tokens_per_block > 0, "tokens_per_block must be > 0"); + Self { + block_manager, + sequences: HashMap::new(), + tokens_per_block, + } + } + + // ===== Sequence map access (used by the scheduler for policy) ===== + + /// Whether a sequence with this id exists. + pub fn contains(&self, seq_id: SequenceId) -> bool { + self.sequences.contains_key(&seq_id) + } + + /// Look up a sequence's current state. + pub fn get(&self, seq_id: SequenceId) -> Option<&SequenceState> { + self.sequences.get(&seq_id) + } + + /// Total number of tracked sequences. + pub fn len(&self) -> usize { + self.sequences.len() + } + + pub fn is_empty(&self) -> bool { + self.sequences.is_empty() + } + + /// Count sequences currently running (prefill or decode). + pub fn num_running(&self) -> usize { + self.sequences.values().filter(|s| s.is_running()).count() + } + + /// Count sequences currently preempted. + pub fn num_preempted(&self) -> usize { + self.sequences + .values() + .filter(|s| matches!(s, SequenceState::Preempted(_))) + .count() + } + + /// Block ids held by a sequence (for KV-cache addressing). + pub fn blocks(&self, seq_id: SequenceId) -> Result> { + match self.sequences.get(&seq_id) { + Some(state) => Ok(state.blocks().map(|b| b.to_vec()).unwrap_or_default()), + None => Err(Error::UnknownSequence(seq_id)), + } + } + + /// Per-block-type memory statistics. + pub fn block_stats(&self) -> HashMap { + self.block_manager.get_stats() + } + + // ===== Transition entry points (used by the scheduler) ===== + + /// Create a new sequence in the Waiting state and store it. + /// + /// Returns `Err` if the id already exists. + pub fn create( + &mut self, + seq_id: SequenceId, + token_ids: Vec, + max_tokens: usize, + ) -> Result<()> { + if self.sequences.contains_key(&seq_id) { + return Err(Error::InvalidTransition("Duplicate sequence ID")); + } + let event = Event::Create { + token_ids, + max_tokens, + }; + let state = self.transition(SequenceState::Empty(seq_id), event)?; + self.sequences.insert(seq_id, state); + Ok(()) + } + + /// Apply an FSM event to an existing sequence, updating the map in place. + /// + /// On success the sequence's state is replaced with the new state. On + /// failure the original state is restored and the error is returned, so the + /// caller can apply rollback/OOM policy. + pub fn advance(&mut self, seq_id: SequenceId, event: Event) -> Result<&SequenceState> { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => return Err(Error::UnknownSequence(seq_id)), + }; + let backup = state.clone(); + + match self.transition(state, event) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + Ok(&self.sequences[&seq_id]) + } + Err(e) => { + // Restore prior state so the map is never left missing a seq. + self.sequences.insert(seq_id, backup); + Err(e) + } + } + } + + /// Fork a decoding sequence into a child via copy-on-write block sharing. + pub fn fork(&mut self, parent_id: SequenceId, child_id: SequenceId) -> Result<()> { + if self.sequences.contains_key(&child_id) { + warn!("Child ID {:?} already exists", child_id); + return Err(Error::InvalidTransition("Duplicate child ID")); + } + + let parent_state = match self.sequences.remove(&parent_id) { + Some(s) => s, + None => { + warn!("Parent sequence {:?} not found", parent_id); + return Err(Error::UnknownSequence(parent_id)); + } + }; + + let backup = parent_state.clone(); + + if let SequenceState::Decoding(s) = parent_state { + // Copy-on-write: increment ref counts on shared blocks. + let mut reffed = Vec::with_capacity(s.blocks.len()); + for &block_id in s.blocks.iter() { + if let Err(e) = self.block_manager.add_ref(block_id) { + warn!("Failed to fork {:?}: {:?}", parent_id, e); + // Roll back the refs we already bumped so those blocks can + // still return to the free pool later (otherwise leak). + for &done in &reffed { + if let Err(e) = self.block_manager.free(done) { + warn!("Failed to roll back ref for {:?}: {:?}", done, e); + } + } + self.sequences.insert(parent_id, backup); + return Err(e); + } + reffed.push(block_id); + } + + let child = DecodingState { + seq_id: child_id, + token_ids: Arc::clone(&s.token_ids), + blocks: Arc::clone(&s.blocks), + num_tokens: s.num_tokens, + max_tokens: s.max_tokens, + }; + + self.sequences.insert(parent_id, SequenceState::Decoding(s)); + self.sequences + .insert(child_id, SequenceState::Decoding(child)); + debug!("Forked sequence {:?} → {:?}", parent_id, child_id); + Ok(()) + } else { + warn!("Cannot fork {:?} - not in Decoding state", parent_id); + self.sequences.insert(parent_id, backup); + Err(Error::InvalidTransition("Can only fork Decoding sequences")) + } + } + + // ===== FSM Transitions ===== + + /// Compute the next state for `(state, event)`, allocating/freeing blocks as + /// needed. Pure with respect to the sequence map — does not read or write + /// `self.sequences`; the caller stores the result. + fn transition(&mut self, state: SequenceState, event: Event) -> Result { + let seq_id = state.seq_id(); + match (state, event) { + // Empty → Waiting (new sequence) + ( + SequenceState::Empty(_), + Event::Create { + token_ids, + max_tokens, + }, + ) => self.transition_create(seq_id, token_ids, max_tokens), + + // Waiting → Prefilling + (SequenceState::Waiting(s), Event::Schedule { .. }) => self.transition_schedule(s), + + // Prefilling → Prefilling/Decoding + (SequenceState::Prefilling(s), Event::AppendTokens { num_tokens, .. }) => { + self.transition_append_tokens_prefilling(s, num_tokens) + } + + // Decoding → Decoding/Finished + (SequenceState::Decoding(s), Event::AppendTokens { num_tokens, .. }) => { + self.transition_append_tokens_decoding(s, num_tokens) + } + + // Any running state → Finished + (state, Event::Complete { reason }) => self.transition_complete(state, reason), + + // Running → Preempted + ( + state @ (SequenceState::Prefilling(_) | SequenceState::Decoding(_)), + Event::Preempt, + ) => self.transition_preempt(state), + + // Preempted → Waiting + (SequenceState::Preempted(s), Event::Resume) => self.transition_resume(s), + + // Any → Aborted + (state, Event::Abort { reason }) => self.transition_abort(state, reason), + + // Invalid transitions + _ => Err(Error::InvalidTransition("Invalid state transition")), + } + } + + /// Transition: Empty → Waiting (construct a new sequence's initial state) + fn transition_create( + &mut self, + seq_id: SequenceId, + token_ids: Vec, + max_tokens: usize, + ) -> Result { + Ok(SequenceState::Waiting(WaitingState { + seq_id, + token_ids: Arc::new(token_ids), + max_tokens, + })) + } + + /// Transition: Waiting → Prefilling (allocate blocks for tokenized prompt) + fn transition_schedule(&mut self, state: WaitingState) -> Result { + let prompt_tokens = state.token_ids.len(); + let blocks_needed = prompt_tokens.div_ceil(self.tokens_per_block); + + let mut blocks = Vec::new(); + for _ in 0..blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => blocks.push(block_id), + Err(Error::OutOfMemory) => { + // Cleanup on OOM + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(Error::OutOfMemory); + } + Err(e) => { + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(e); + } + } + } + + debug!( + "Scheduled seq {:?}: allocated {} blocks for {} tokens", + state.seq_id, + blocks.len(), + prompt_tokens + ); + + Ok(SequenceState::Prefilling(PrefillingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: prompt_tokens, + max_tokens: state.max_tokens, + })) + } + + /// Transition: Prefilling → Decoding or Prefilling (append tokens) + fn transition_append_tokens_prefilling( + &mut self, + mut state: PrefillingState, + num_tokens: usize, + ) -> Result { + state.tokens_filled += num_tokens; + + debug!( + "Prefilling seq {:?}: {}/{} tokens", + state.seq_id, state.tokens_filled, state.tokens_total + ); + + if state.tokens_filled >= state.tokens_total { + debug!("Seq {:?} transitioning to Decoding", state.seq_id); + Ok(SequenceState::Decoding(DecodingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + blocks: state.blocks, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) + } else { + Ok(SequenceState::Prefilling(state)) + } + } + + /// Transition: Decoding → Decoding or Finished (append tokens, maybe allocate more blocks) + fn transition_append_tokens_decoding( + &mut self, + mut state: DecodingState, + num_tokens: usize, + ) -> Result { + state.num_tokens += num_tokens; + + // Check if finished + if state.num_tokens >= state.max_tokens { + // Free blocks + for &block_id in state.blocks.iter() { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!( + "Seq {:?} finished: reached max_tokens ({})", + state.seq_id, state.max_tokens + ); + + return Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: FinishReason::MaxTokens, + })); + } + + // Check if need more blocks + let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); + let initial_block_count = state.blocks.len(); + + if state.blocks.len() < blocks_needed { + let blocks = Arc::make_mut(&mut state.blocks); + + while blocks.len() < blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => { + blocks.push(block_id); + debug!( + "Seq {:?}: allocated block, total blocks: {}", + state.seq_id, + blocks.len() + ); + } + Err(e) => { + warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); + for block_id in blocks.drain(initial_block_count..) { + let _ = self.block_manager.free(block_id); + } + return Err(e); + } + } + } + } + + Ok(SequenceState::Decoding(state)) + } + + /// Transition: Prefilling/Decoding → Finished + fn transition_complete( + &mut self, + state: SequenceState, + reason: FinishReason, + ) -> Result { + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + debug!("Completed seq {:?}: {:?}", seq_id, reason); + + Ok(SequenceState::Finished(FinishedState { + seq_id, + finish_reason: reason, + })) + } + + /// Transition: Prefilling/Decoding → Preempted (free blocks) + fn transition_preempt(&mut self, state: SequenceState) -> Result { + let seq_id = state.seq_id(); + let (token_ids, num_tokens, max_tokens) = match &state { + SequenceState::Prefilling(s) => { + (Arc::clone(&s.token_ids), s.tokens_filled, s.max_tokens) + } + SequenceState::Decoding(s) => (Arc::clone(&s.token_ids), s.num_tokens, s.max_tokens), + _ => { + return Err(Error::InvalidTransition( + "Can only preempt Prefilling/Decoding", + )) + } + }; + + // Free blocks + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + debug!("Preempted seq {:?}", seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id, + token_ids, + num_tokens, + max_tokens, + })) + } + + /// Transition: Preempted → Waiting + fn transition_resume(&mut self, state: PreemptedState) -> Result { + debug!("Resuming seq {:?}", state.seq_id); + + Ok(SequenceState::Waiting(WaitingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + max_tokens: state.max_tokens, + })) + } + + /// Transition: Any → Aborted (free blocks, cleanup) + fn transition_abort(&mut self, state: SequenceState, reason: String) -> Result { + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + warn!("Aborted seq {:?}: {}", seq_id, reason); + + Ok(SequenceState::Aborted(AbortedState { seq_id, reason })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + use crate::fsm::FinishReason; + + /// Build a manager with `max_blocks` worth of memory and the given + /// `tokens_per_block`. Block byte-size is fixed at 1024; capping allocator + /// memory is how we force OOM in tests. + fn manager(max_blocks: usize, tokens_per_block: usize) -> SequenceManager { + let block_bytes = 1024; + let allocator = Box::new(CpuAllocator::new(max_blocks * block_bytes)); + let block_manager = BlockManager::new(allocator, block_bytes); + SequenceManager::new(block_manager, tokens_per_block) + } + + fn schedule_event() -> Event { + Event::Schedule { + tokens_per_block: 4, + } + } + + fn append(num_tokens: usize) -> Event { + Event::AppendTokens { + num_tokens, + tokens_per_block: 4, + } + } + + #[test] + fn test_create_inserts_waiting_sequence() { + let mut mgr = manager(10, 4); + let id = SequenceId(1); + + mgr.create(id, vec![1, 2, 3], 20).unwrap(); + + assert!(mgr.contains(id)); + assert_eq!(mgr.len(), 1); + assert!(!mgr.is_empty()); + assert_eq!(mgr.get(id).unwrap().variant_name(), "Waiting"); + } + + #[test] + fn test_create_duplicate_id_errors() { + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3], 20).unwrap(); + + let err = mgr.create(id, vec![4, 5], 20).unwrap_err(); + assert!(matches!(err, Error::InvalidTransition(_))); + // Original sequence is untouched. + assert_eq!(mgr.len(), 1); + } + + #[test] + fn test_advance_unknown_sequence_errors() { + let mut mgr = manager(10, 4); + let err = mgr.advance(SequenceId(99), schedule_event()).unwrap_err(); + assert!(matches!(err, Error::UnknownSequence(SequenceId(99)))); + } + + #[test] + fn test_schedule_allocates_blocks() { + // 6 tokens, 4 tokens/block => 2 blocks needed. + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3, 4, 5, 6], 20).unwrap(); + + let state = mgr.advance(id, schedule_event()).unwrap(); + assert_eq!(state.variant_name(), "Prefilling"); + assert_eq!(mgr.blocks(id).unwrap().len(), 2); + } + + #[test] + fn test_schedule_oom_rolls_back_to_waiting() { + // Needs 2 blocks but only 1 is available -> OOM, no blocks leaked. + let mut mgr = manager(1, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3, 4, 5, 6], 20).unwrap(); + + let err = mgr.advance(id, schedule_event()).unwrap_err(); + assert!(matches!(err, Error::OutOfMemory)); + + // State restored to Waiting and no blocks are held. + assert_eq!(mgr.get(id).unwrap().variant_name(), "Waiting"); + assert!(mgr.blocks(id).unwrap().is_empty()); + + // The partially-allocated block was freed: a 1-block prompt now fits. + let id2 = SequenceId(2); + mgr.create(id2, vec![1, 2], 20).unwrap(); + assert!(mgr.advance(id2, schedule_event()).is_ok()); + } + + #[test] + fn test_prefill_to_decoding_transition() { + // 4 tokens, 4/block => 1 block, prefill completes in one append. + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(id, schedule_event()).unwrap(); + + // Partial prefill stays in Prefilling. + let state = mgr.advance(id, append(2)).unwrap(); + assert_eq!(state.variant_name(), "Prefilling"); + + // Reaching tokens_total flips to Decoding. + let state = mgr.advance(id, append(2)).unwrap(); + assert_eq!(state.variant_name(), "Decoding"); + assert_eq!(mgr.num_running(), 1); + } + + #[test] + fn test_decoding_reaches_max_tokens_finishes_and_frees() { + let mut mgr = manager(10, 4); + let id = SequenceId(1); + // prompt of 4 tokens, max_tokens 5. + mgr.create(id, vec![1, 2, 3, 4], 5).unwrap(); + mgr.advance(id, schedule_event()).unwrap(); + mgr.advance(id, append(4)).unwrap(); // -> Decoding, num_tokens=4 + + let state = mgr.advance(id, append(1)).unwrap(); // hits max_tokens=5 + assert_eq!(state.variant_name(), "Finished"); + // Blocks freed on finish. + assert!(mgr.blocks(id).unwrap().is_empty()); + assert_eq!(mgr.num_running(), 0); + } + + #[test] + fn test_preempt_frees_blocks_then_resume_returns_to_waiting() { + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(id, schedule_event()).unwrap(); + mgr.advance(id, append(4)).unwrap(); // Decoding + + let state = mgr.advance(id, Event::Preempt).unwrap(); + assert_eq!(state.variant_name(), "Preempted"); + assert_eq!(mgr.num_preempted(), 1); + assert!(mgr.blocks(id).unwrap().is_empty()); + + let state = mgr.advance(id, Event::Resume).unwrap(); + assert_eq!(state.variant_name(), "Waiting"); + assert_eq!(mgr.num_preempted(), 0); + } + + #[test] + fn test_complete_and_abort_finish_the_sequence() { + let mut mgr = manager(10, 4); + + let a = SequenceId(1); + mgr.create(a, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(a, schedule_event()).unwrap(); + let state = mgr + .advance( + a, + Event::Complete { + reason: FinishReason::Stop, + }, + ) + .unwrap(); + assert_eq!(state.variant_name(), "Finished"); + assert!(mgr.blocks(a).unwrap().is_empty()); + + let b = SequenceId(2); + mgr.create(b, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(b, schedule_event()).unwrap(); + let state = mgr + .advance( + b, + Event::Abort { + reason: "cancelled".to_string(), + }, + ) + .unwrap(); + assert_eq!(state.variant_name(), "Aborted"); + assert!(mgr.blocks(b).unwrap().is_empty()); + } + + #[test] + fn test_invalid_transition_preserves_state() { + // Resume on a Waiting sequence is invalid. + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3], 20).unwrap(); + + let err = mgr.advance(id, Event::Resume).unwrap_err(); + assert!(matches!(err, Error::InvalidTransition(_))); + // Sequence remains present and in its original state. + assert!(mgr.contains(id)); + assert_eq!(mgr.get(id).unwrap().variant_name(), "Waiting"); + } + + #[test] + fn test_fork_shares_blocks_of_decoding_parent() { + let mut mgr = manager(10, 4); + let parent = SequenceId(1); + let child = SequenceId(2); + mgr.create(parent, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(parent, schedule_event()).unwrap(); + mgr.advance(parent, append(4)).unwrap(); // Decoding + + let parent_blocks = mgr.blocks(parent).unwrap(); + mgr.fork(parent, child).unwrap(); + + // Child exists, is Decoding, and shares the parent's blocks. + assert_eq!(mgr.get(child).unwrap().variant_name(), "Decoding"); + assert_eq!(mgr.blocks(child).unwrap(), parent_blocks); + assert_eq!(mgr.num_running(), 2); + } + + #[test] + fn test_fork_non_decoding_parent_errors() { + let mut mgr = manager(10, 4); + let parent = SequenceId(1); + mgr.create(parent, vec![1, 2, 3], 20).unwrap(); // Waiting + + let err = mgr.fork(parent, SequenceId(2)).unwrap_err(); + assert!(matches!(err, Error::InvalidTransition(_))); + // Parent is restored and child was never inserted. + assert!(mgr.contains(parent)); + assert!(!mgr.contains(SequenceId(2))); + } + + /// Sum of free blocks across all block types. + fn free_blocks(mgr: &SequenceManager) -> usize { + mgr.block_stats().values().map(|s| s.free_blocks).sum() + } + + #[test] + fn test_fork_rolls_back_refs_on_partial_failure() { + let mut mgr = manager(10, 4); + let parent = SequenceId(1); + mgr.create(parent, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(parent, schedule_event()).unwrap(); + mgr.advance(parent, append(4)).unwrap(); // Decoding, holds 1 real block + + let real_block = mgr.blocks(parent).unwrap()[0]; + + // Inject a bogus second block so add_ref succeeds on the real block and + // then fails on the bogus one, exercising the mid-loop rollback path. + match mgr.sequences.get_mut(&parent) { + Some(SequenceState::Decoding(s)) => { + s.blocks = Arc::new(vec![real_block, BlockId(9999)]); + } + _ => panic!("parent should be Decoding"), + } + + let err = mgr.fork(parent, SequenceId(2)).unwrap_err(); + assert!(matches!(err, Error::InvalidBlockId(_))); + // Child was never created and the parent is restored. + assert!(!mgr.contains(SequenceId(2))); + assert!(mgr.contains(parent)); + + // The real block's ref was rolled back to 1 (parent only), so aborting + // the parent drops it to 0 and returns it to the free pool. Without the + // rollback its ref would be 2 and it would leak (never pooled). + let free_before = free_blocks(&mgr); + mgr.advance( + parent, + Event::Abort { + reason: "x".to_string(), + }, + ) + .unwrap(); + assert_eq!( + free_blocks(&mgr), + free_before + 1, + "real block should return to the free pool after abort" + ); + } + + #[test] + fn test_fork_duplicate_child_errors() { + let mut mgr = manager(10, 4); + let parent = SequenceId(1); + let child = SequenceId(2); + mgr.create(parent, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(parent, schedule_event()).unwrap(); + mgr.advance(parent, append(4)).unwrap(); + + mgr.create(child, vec![9], 20).unwrap(); // child id already taken + let err = mgr.fork(parent, child).unwrap_err(); + assert!(matches!(err, Error::InvalidTransition(_))); + } +} diff --git a/src/system/system_info.rs b/src/system/system_info.rs index ff62ebc..2cc2fcb 100644 --- a/src/system/system_info.rs +++ b/src/system/system_info.rs @@ -280,3 +280,82 @@ impl SystemInfo { println!(" Running Models: {}", self.running_models); } } + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + #[test] + fn test_collect_populates_basic_fields() { + let info = SystemInfo::collect(); + + // Version comes from Cargo and must be non-empty. + assert_eq!(info.version, env!("CARGO_PKG_VERSION")); + // Any real host reports at least one CPU core. + assert!(info.cpu_cores >= 1); + // format_size never yields an empty string. + assert!(!info.total_memory.is_empty()); + assert!(!info.cache_size.is_empty()); + // running_models tracking is not implemented yet. + assert_eq!(info.running_models, 0); + } + + #[test] + fn test_serde_roundtrip() { + let info = SystemInfo { + version: "1.2.3".to_string(), + os: "TestOS".to_string(), + architecture: "x86_64".to_string(), + cpu_cores: 8, + total_memory: "16.00 GB".to_string(), + gpu_info: vec![GpuInfo { + name: "Test GPU".to_string(), + backend: "Metal".to_string(), + memory: Some("10 GPU cores".to_string()), + }], + cache_dir: "/tmp/cache".to_string(), + cache_size: "0 B".to_string(), + models_count: 2, + running_models: 0, + }; + + let json = serde_json::to_string(&info).unwrap(); + let back: SystemInfo = serde_json::from_str(&json).unwrap(); + + assert_eq!(back.version, info.version); + assert_eq!(back.cpu_cores, info.cpu_cores); + assert_eq!(back.gpu_info.len(), 1); + assert_eq!(back.gpu_info[0].backend, "Metal"); + assert_eq!(back.gpu_info[0].memory.as_deref(), Some("10 GPU cores")); + } + + #[test] + fn test_calculate_cache_size_missing_dir_is_zero() { + let missing = PathBuf::from("/nonexistent/puma/cache/path/xyz"); + assert_eq!(SystemInfo::calculate_cache_size(&missing), 0); + } + + #[test] + fn test_cache_size_sums_files_and_subdirs() { + let tmp = tempfile::tempdir().unwrap(); + let root = tmp.path().to_path_buf(); + + // A file at the top level and one nested in a subdirectory. + fs::write(root.join("a.bin"), vec![0u8; 4096]).unwrap(); + let sub = root.join("sub"); + fs::create_dir(&sub).unwrap(); + fs::write(sub.join("b.bin"), vec![0u8; 4096]).unwrap(); + + // Disk usage is block-based (blocks * 512); an empty dir contributes 0 + // files, so total should reflect both files' on-disk size (> 0) and + // recurse into the subdirectory. + let total = SystemInfo::calculate_cache_size(&root); + let sub_only = SystemInfo::dir_size(&sub); + assert!(sub_only > 0, "nested file should contribute size"); + assert!( + total >= sub_only, + "top-level total {total} should include nested {sub_only}" + ); + } +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 73f5099..04e8ed2 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,2 +1,3 @@ pub mod file; pub mod format; +pub mod prompt; diff --git a/src/utils/prompt.rs b/src/utils/prompt.rs new file mode 100644 index 0000000..551706d --- /dev/null +++ b/src/utils/prompt.rs @@ -0,0 +1,62 @@ +//! Prompt formatting — the seam between the chat API's role/content language +//! and the flat text the tokenizer consumes. +//! +//! Roles ("system"/"user"/"assistant") are an *API-level* concept. The +//! inference engine only ever sees token ids; it has no notion of turns. This +//! module serializes structured turns into one prompt string, which the engine +//! handle then tokenizes. +//! +//! Today it uses a simple `Role: content` convention with a trailing +//! `Assistant:` cue. When a real model is wired in, this is the single place to +//! swap in that model's chat template (e.g. `<|im_start|>` markers) — nothing +//! below the tokenizer changes. + +/// Format conversation turns into a single prompt string. +/// +/// Each turn is `(role, content)`. The result labels every turn as +/// `Role: content` (one per line) and ends with an empty `Assistant:` cue so +/// the model continues as the assistant. +pub fn format_conversation<'a, I>(turns: I) -> String +where + I: IntoIterator, +{ + let mut prompt = String::new(); + for (role, content) in turns { + prompt.push_str(&format!("{}: {}\n", display_role(role), content)); + } + prompt.push_str("Assistant:"); + prompt +} + +/// Capitalize a role label for display (`user` → `User`), defaulting unknown +/// roles to `Assistant` to match prior behavior. +fn display_role(role: &str) -> &'static str { + match role { + "system" => "System", + "user" => "User", + _ => "Assistant", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn formats_turns_with_assistant_cue() { + let prompt = format_conversation([("system", "Be helpful."), ("user", "Hello")]); + assert_eq!(prompt, "System: Be helpful.\nUser: Hello\nAssistant:"); + } + + #[test] + fn unknown_role_defaults_to_assistant() { + let prompt = format_conversation([("tool", "result")]); + assert_eq!(prompt, "Assistant: result\nAssistant:"); + } + + #[test] + fn empty_conversation_is_just_the_cue() { + let prompt = format_conversation([]); + assert_eq!(prompt, "Assistant:"); + } +}