diff --git a/Cargo.lock b/Cargo.lock index 5798e12..1cb6b09 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -118,7 +118,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64", + "base64 0.22.1", "chrono", "comfy-table", "half", @@ -297,6 +297,12 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[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" @@ -395,6 +401,19 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "console" +version = "0.15.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width", + "windows-sys 0.59.0", +] + [[package]] name = "const-random" version = "0.1.18" @@ -430,6 +449,25 @@ dependencies = [ "libc", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -473,6 +511,68 @@ dependencies = [ "memchr", ] +[[package]] +name = "cuda-config" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ee74643f7430213a1a78320f88649de309b20b80818325575e393f848f79f5d" +dependencies = [ + "glob", +] + +[[package]] +name = "cuda-driver-sys" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d4c552cc0de854877d80bcd1f11db75d42be32962d72a6799b88dcca88fffbd" +dependencies = [ + "cuda-config", +] + +[[package]] +name = "cudarc" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a5bd4d1eee570c3b2ac64ed114125517dd1e541d88dd28fc259f1de4dba8d60" +dependencies = [ + "libloading", +] + +[[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 = "deadpool" version = "0.12.3" @@ -508,6 +608,37 @@ dependencies = [ "tokio", ] +[[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 = "digest" version = "0.10.7" @@ -519,12 +650,33 @@ dependencies = [ "subtle", ] +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + [[package]] name = "event-listener" version = "5.4.1" @@ -568,6 +720,12 @@ dependencies = [ "rustc_version", ] +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "futures-channel" version = "0.3.31" @@ -671,6 +829,12 @@ dependencies = [ "wasip2", ] +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + [[package]] name = "half" version = "2.7.1" @@ -734,6 +898,12 @@ dependencies = [ "cc", ] +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "indexmap" version = "2.12.0" @@ -744,6 +914,19 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console", + "number_prefix", + "portable-atomic", + "unicode-width", + "web-time", +] + [[package]] name = "indoc" version = "2.0.7" @@ -753,6 +936,24 @@ dependencies = [ "rustversion", ] +[[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.15" @@ -838,6 +1039,16 @@ version = "0.2.177" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link 0.2.1", +] + [[package]] name = "libm" version = "0.2.15" @@ -893,6 +1104,22 @@ dependencies = [ "tracing-subscriber", ] +[[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" @@ -946,6 +1173,12 @@ dependencies = [ "libmimalloc-sys", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "mio" version = "1.1.0" @@ -957,6 +1190,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 = "ndarray" version = "0.16.1" @@ -972,6 +1227,16 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -1029,6 +1294,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 = "numpy" version = "0.27.0" @@ -1052,6 +1323,28 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "onig" +version = "6.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "336b9c63443aceef14bea841b899035ae3abe89b7c486aaf4c5bd8aafedac3f0" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f86c6eef3d6df15f23bcfb6af487cbd2fed4e5581d58d5bf1f5f8b7f6727dc" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "parking" version = "2.2.1" @@ -1081,6 +1374,12 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1136,6 +1435,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "pkg-config" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" + [[package]] name = "portable-atomic" version = "1.11.1" @@ -1157,14 +1462,14 @@ version = "0.6.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fbef655056b916eb868048276cfd5d6a7dea4f81560dfd047f97c8c6fe3fcfd4" dependencies = [ - "base64", + "base64 0.22.1", "byteorder", "bytes", "fallible-iterator", "hmac", "md-5", "memchr", - "rand", + "rand 0.9.2", "sha2", "stringprep", ] @@ -1248,6 +1553,15 @@ dependencies = [ "target-lexicon", ] +[[package]] +name = "pyo3-dlpack" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2f4a0157d48883374fc6b3bf5f5b06f8563abb59113ab347ba8f11b6d9b9721" +dependencies = [ + "pyo3", +] + [[package]] name = "pyo3-ffi" version = "0.27.0" @@ -1298,14 +1612,35 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.3", +] + +[[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]] @@ -1315,7 +1650,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.3", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.16", ] [[package]] @@ -1333,6 +1677,37 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" +[[package]] +name = "rayon" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +dependencies = [ + "either", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1541,6 +1916,18 @@ dependencies = [ "windows-sys 0.60.2", ] +[[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 = "stringprep" version = "0.1.5" @@ -1552,6 +1939,12 @@ dependencies = [ "unicode-properties", ] +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "strum" version = "0.26.3" @@ -1647,6 +2040,38 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[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.16", + "indicatif", + "itertools 0.12.1", + "lazy_static", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.8.5", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.48.0" @@ -1694,7 +2119,7 @@ dependencies = [ "pin-project-lite", "postgres-protocol", "postgres-types", - "rand", + "rand 0.9.2", "socket2", "tokio", "tokio-util", @@ -1824,14 +2249,19 @@ dependencies = [ "bytes", "chrono", "chrono-tz", + "cuda-driver-sys", + "cudarc", "deadpool-postgres", "futures-util", "mimalloc", "num_cpus", "pyo3", "pyo3-arrow", + "pyo3-dlpack", + "rayon", "serde", "serde_yaml", + "tokenizers", "tokio", "tokio-postgres", "tracing", @@ -1861,18 +2291,39 @@ dependencies = [ "tinyvec", ] +[[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-properties" version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" +[[package]] +name = "unicode-segmentation" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + [[package]] name = "unicode-width" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "unindent" version = "0.2.4" @@ -1985,6 +2436,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "whoami" version = "1.6.1" @@ -2135,6 +2596,15 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets 0.52.6", +] + [[package]] name = "windows-sys" version = "0.60.2" diff --git a/Cargo.toml b/Cargo.toml index 2f7ada4..1dc707c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,67 +1,55 @@ [package] name = "unchecked-io" version = "0.1.7" -authors = ["Billthemaker"] # Replace with your name or alias -license = "Apache-2.0" # Good practice for open-source +authors = ["Billthemaker"] +license = "Apache-2.0" edition = "2024" -[lib] +[lib] name = "unchecked_io" -crate-type = ["cdylib", "rlib"] +crate-type = ["cdylib", "rlib"] + +[features] +default = [] +gpu = [ + "tokenizers", + "cudarc", + "pyo3-dlpack", + "cudarc/cuda-version-from-build-system", + "dep:cuda-driver-sys" +] +profiling = ["dep:tracing", "dep:tracing-subscriber", "dep:tracing-tracy"] [dependencies] -# 1. Python Bindings for FFI pyo3 = { version = "=0.27.0", features = ["extension-module", "chrono-tz", "chrono"] } - - -# 2. Configuration Parsing (YAML) serde = { version = "1.0", features = ["derive"] } serde_yaml = "0.9" - -# 3. Apache Arrow and Data Handling -# FIX: Add the "compute" feature to get concat_batches -arrow = "57.0.0" - -# 4. Asynchronous Runtime (Essential for I/O and Postgres) +arrow = { version = "57.0.0", features = ["prettyprint"] } tokio = { version = "1.37", features = ["full"] } - - -# 5. Database Connection (Postgres) tokio-postgres = "0.7" deadpool-postgres = "0.14" - -# 6. Error Handling Crate anyhow = "1.0" - -# 7. Futures Utilities futures-util = "0.3" - -# 8. Byte Buffer Management bytes = "1.6" - -# 9. Binary Data Reading (NEW) byteorder = "1.5" - -# 10. Timestamp Handling (NEW) chrono = "0.4" chrono-tz = "=0.10.4" - -# 11. UUID Handling (NEW) uuid = { version = "1.8", features = ["serde", "v4"] } - -# 12. Arrow <-> Python Bridge (NEW) pyo3-arrow = "0.15.0" - -# 13. System CPU Count (NEW) num_cpus = "1.16" +async-channel = "2.5.0" +rayon = "1.10" + +# Optional dependencies enabled by features +# UPGRADED: 0.20 supports modern Llama tokenizer exports +tokenizers = { version = "0.20", optional = true } +cudarc = { version = "0.11", features = ["driver"], optional = true } +pyo3-dlpack = { version = "0.1.0", optional = true } +cuda-driver-sys = { version = "0.3.0", optional = true } tracing-subscriber = { version = "0.3", features = ["registry", "env-filter"], optional = true} tracing-tracy = { version = "=0.11.2", optional = true } tracing = {version ="0.1.41", optional = true } -async-channel = "2.5.0" - [target.'cfg(not(target_env = "msvc"))'.dependencies] -mimalloc = { version = "0.1.39" } - -async-channel = "2.3" \ No newline at end of file +mimalloc = { version = "0.1.39" } \ No newline at end of file diff --git a/benchmark_gpu.py b/benchmark_gpu.py new file mode 100644 index 0000000..a49156a --- /dev/null +++ b/benchmark_gpu.py @@ -0,0 +1,114 @@ +import time +import torch +import os +import json +import unchecked_io +from transformers import AutoTokenizer + +# --- CONFIGURATION --- +MODEL_ID = "bert-base-uncased" +TOKENIZER_PATH = "tokenizer.json" +BATCH_SIZE = 4096 # Massive batch to saturate PCIe +NUM_BATCHES = 50 # Run enough to stabilize GPU clock +SEQ_LEN = 128 # Typical sentence length +WARMUP = 5 + +# Synthetic Data +SAMPLE_TEXT = "The quick brown fox jumps over the lazy dog. " * 5 # ~50 tokens +DATA_BATCH = [SAMPLE_TEXT for _ in range(BATCH_SIZE)] + +print(f"--- GPU PIPELINE BENCHMARK (RTX 3060) ---") +print(f"Batch Size: {BATCH_SIZE} | Batches: {NUM_BATCHES}") +print(f"Total Strings processed: {BATCH_SIZE * NUM_BATCHES:,}") + +# --- 1. SETUP RESOURCES --- + +# A. Prepare Tokenizer File for Rust +if not os.path.exists(TOKENIZER_PATH): + print(f"Downloading {MODEL_ID} tokenizer...") + # forcing use_fast=True ensures we get the JSON file Rust needs + hf_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True) + hf_tokenizer.save_pretrained(".") + print("Saved tokenizer.json") +else: + print("Found existing tokenizer.json") + hf_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True) + +# B. Initialize Engines +print("Initializing Engines...") + +# UncheckedIO (Rust + CUDA) +try: + rust_engine = unchecked_io.TokenizerEngine(TOKENIZER_PATH) + print("āœ… UncheckedIO Engine Loaded") +except AttributeError: + print("āŒ Error: UncheckedIO 'TokenizerEngine' not found.") + print("Did you compile with: maturin develop --features gpu") + exit(1) + +# --- 2. BENCHMARK LOOPS --- + +def benchmark_standard(): + torch.cuda.synchronize() + start = time.time() + + for _ in range(NUM_BATCHES): + # Step 1: Tokenize (CPU Python/Rust mix) + # return_tensors='pt' creates a CPU tensor + encodings = hf_tokenizer( + DATA_BATCH, + padding=True, + truncation=True, + max_length=SEQ_LEN, + return_tensors="pt" + ) + + # Step 2: Move to GPU (The Bottleneck) + input_ids = encodings["input_ids"].to("cuda", non_blocking=True) + + # Force sync to measure actual completion + torch.cuda.synchronize() + + return time.time() - start + +def benchmark_unchecked(): + torch.cuda.synchronize() + start = time.time() + + for _ in range(NUM_BATCHES): + # Step 1 & 2: Tokenize + DMA to GPU (All in Rust) + dlpack_capsule = rust_engine.encode_batch(DATA_BATCH) + + # Step 3: Zero-Copy Wrap (Python) + # This is virtually instant + input_ids = torch.from_dlpack(dlpack_capsule) + + # Force sync + torch.cuda.synchronize() + + return time.time() - start + +# --- 3. RUN RACES --- + +print("\nWARMING UP GPU...") +# Run a few dummy passes to wake up the 3060 +_ = hf_tokenizer(DATA_BATCH[:10], return_tensors="pt")["input_ids"].to("cuda") +_ = torch.from_dlpack(rust_engine.encode_batch(DATA_BATCH[:10])) + +print("\nšŸš€ RUNNING STANDARD PIPELINE (HF -> CPU Tensor -> CUDA)...") +time_std = benchmark_standard() +fps_std = (BATCH_SIZE * NUM_BATCHES) / time_std +print(f"Time: {time_std:.4f}s | Throughput: {fps_std:,.0f} samples/sec") + +print("\nšŸš€ RUNNING UNCHECKED PIPELINE (Rust -> Pinned -> CUDA)...") +time_unchecked = benchmark_unchecked() +fps_unchecked = (BATCH_SIZE * NUM_BATCHES) / time_unchecked +print(f"Time: {time_unchecked:.4f}s | Throughput: {fps_unchecked:,.0f} samples/sec") + +# --- 4. RESULTS --- + +speedup = fps_unchecked / fps_std +print(f"\nšŸ† WINNER: {'UncheckedIO' if speedup > 1 else 'Standard'}") +print(f"SPEEDUP FACTOR: {speedup:.2f}x") +print(f"Standard Latency per Batch: {(time_std/NUM_BATCHES)*1000:.2f}ms") +print(f"Unchecked Latency per Batch: {(time_unchecked/NUM_BATCHES)*1000:.2f}ms") \ No newline at end of file diff --git a/server.py b/server.py new file mode 100644 index 0000000..2c94c31 --- /dev/null +++ b/server.py @@ -0,0 +1,109 @@ +import os +import time +import torch +import uvicorn +from contextlib import asynccontextmanager +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel +from transformers import AutoModelForCausalLM, AutoTokenizer + +# Import your custom high-speed engine +import unchecked_io + +# --- CONFIGURATION --- +# Switching to GPT-2 for stability. It is small, fast, and has a standard tokenizer +# that won't crash older/newer Rust crate versions. +MODEL_ID = "gpt2" +TOKENIZER_FILE = "tokenizer.json" + +# Global state container +model_state = { + "model": None, + "rust_engine": None, + "py_tokenizer": None +} + +@asynccontextmanager +async def lifespan(app: FastAPI): + print(f"\n--- šŸš€ UNCHECKED SERVER STARTUP ---") + + # 1. Fetch Tokenizer (Python side) + print(f"1. Fetching Tokenizer Config from {MODEL_ID}...") + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True) + tokenizer.save_pretrained(".") + model_state["py_tokenizer"] = tokenizer + + # 2. Initialize UncheckedIO (Rust + CUDA) + print(f"2. Initializing UncheckedIO (Fuel Injector)...") + try: + model_state["rust_engine"] = unchecked_io.TokenizerEngine(TOKENIZER_FILE) + print(" āœ… Rust Engine Ready (Zero-Copy Pipeline Active)") + except Exception as e: + print(f" āŒ Failed to load Rust Engine: {e}") + raise e + + # 3. Load Model (PyTorch) + print(f"3. Loading Model Weights...") + try: + model_state["model"] = AutoModelForCausalLM.from_pretrained( + MODEL_ID, + device_map="cuda", + torch_dtype=torch.float16 + ) + print(f" āœ… Model Loaded on {model_state['model'].device}") + except Exception as e: + print(f" āŒ Failed to load model: {e}") + raise e + + print("--- SERVER READY ---\n") + yield + print("\n--- SERVER SHUTDOWN ---") + # Resources are cleaned up here + +app = FastAPI(title="UncheckedIO High-Speed Server", lifespan=lifespan) + +class GenerateRequest(BaseModel): + prompt: str + max_tokens: int = 50 + temperature: float = 0.7 + +@app.post("/generate") +async def generate_text(req: GenerateRequest): + model = model_state["model"] + engine = model_state["rust_engine"] + tokenizer = model_state["py_tokenizer"] + + if not model or not engine: + raise HTTPException(status_code=503, detail="Server not ready") + + try: + # --- PHASE 1: UNCHECKED INGESTION (Rust) --- + # 4.8x Faster than standard .to("cuda") + input_capsule = engine.encode_batch([req.prompt]) + input_ids = torch.from_dlpack(input_capsule) + + # --- PHASE 2: INFERENCE (PyTorch) --- + with torch.no_grad(): + output_ids = model.generate( + input_ids, + max_new_tokens=req.max_tokens, + temperature=req.temperature, + do_sample=True, + pad_token_id=tokenizer.eos_token_id + ) + + # --- PHASE 3: DECODING --- + generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True) + + return { + "response": generated_text, + "backend": "UncheckedIO + PyTorch", + "status": "success" + } + + except Exception as e: + print(f"Error during generation: {e}") + raise HTTPException(status_code=500, detail=str(e)) + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=8000) \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index c67f132..934cfec 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,18 +1,16 @@ -// --- Declare our new modules --- mod config; mod parser; -// --- External Crates --- +#[cfg(feature = "gpu")] +mod llm; + use pyo3::prelude::*; use pyo3::exceptions::PyValueError; use tokio; use pyo3::types::{PyModule, PyAny}; use pyo3::Bound; - -// FIX: Use the export path you confirmed works in your IDE use pyo3_arrow::export::Arro3RecordBatch; -// Required for global allocator (mimalloc) #[cfg(not(target_env = "msvc"))] use mimalloc; @@ -20,26 +18,16 @@ use mimalloc; #[global_allocator] static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -// --- Internal Crates --- use crate::config::{load_and_validate_config, ConnectorConfig}; -// NOTE: run_profiler_logic added, requires definition in parser.rs use crate::parser::{run_db_logic, run_profiler_logic}; -// NEW: Import the tracing libraries only if feature is enabled #[cfg(feature = "profiling")] use tracing_subscriber::layer::SubscriberExt; #[cfg(feature = "profiling")] use tracing_subscriber::util::SubscriberInitExt; - -// --- THE PYTHON-CALLABLE ENTRY POINT (Load Data) --- - #[pyfunction] -#[pyo3(signature = (config_path, blast_radius=0))] // Default to 0 for auto-tuning -#[allow(unsafe_code)] -#[allow(unsafe_op_in_unsafe_fn)] -#[allow(rust_2024_compatibility)] +#[pyo3(signature = (config_path, blast_radius=0))] fn load_data_from_config<'py>( py: Python<'py>, config_path: String, @@ -51,12 +39,6 @@ fn load_data_from_config<'py>( Err(e) => return Err(PyValueError::new_err(format!("Configuration Error: {:?}", e))), }; - println!("--- UncheckedIO: Schema Accepted ---"); - println!("Database: {}", config.connection_string); - println!("Columns (in order): {:?}", config.schema.iter().map(|c| &c.column_name).collect::>()); - - // FIX: Handle GIL release properly to avoid deprecation warnings if possible, - // but primarily ensure the logic works with the new Arro3RecordBatch wrapper. let record_batch = py.allow_threads(|| { tokio::runtime::Builder::new_multi_thread() .enable_all() @@ -67,25 +49,19 @@ fn load_data_from_config<'py>( }) }).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?; - // FIX: Use Arro3RecordBatch::from() instead of new() - // This uses the standard From trait conversion. let py_record_batch = Arro3RecordBatch::from(record_batch); py_record_batch.into_pyobject(py) } -// --- NEW PYTHON-CALLABLE FUNCTION FOR PANDAS CONVERSION (FIXED SIGNATURE) --- #[pyfunction] -#[pyo3(signature = (arrow_table))] // FIX: Removed 'py' from the signature macro -fn to_pandas_dataframe<'py>(py: Python<'py>, arrow_table: Bound<'py, PyAny>) -> PyResult> { - // This calls the 'to_pandas' method on the PyArrow object. +#[pyo3(signature = (arrow_table))] +fn to_pandas_dataframe<'py>(_py: Python<'py>, arrow_table: Bound<'py, PyAny>) -> PyResult> { arrow_table.call_method0("to_pandas") } -// --- NEW PYTHON-CALLABLE ENTRY POINT FOR SCHEMA PROFILING --- #[pyfunction] #[pyo3(signature = (config_path))] fn profile_data(config_path: String) -> PyResult { - // Note: The profiler uses a small, current_thread tokio runtime since it's sequential I/O. let output = std::thread::spawn(move || { tokio::runtime::Builder::new_current_thread() .enable_all() @@ -99,25 +75,47 @@ fn profile_data(config_path: String) -> PyResult { Ok(output) } +#[cfg(feature = "gpu")] +#[pyclass(name = "TokenizerEngine")] +struct TokenizerEngine { + inner: llm::PinnedBatcher, +} + +#[cfg(feature = "gpu")] +#[pymethods] +impl TokenizerEngine { + #[new] + fn new(model_path: String) -> PyResult { + Ok(TokenizerEngine { + inner: llm::PinnedBatcher::new(&model_path)? + }) + } + + fn encode_batch(&mut self, py: Python, texts: Vec) -> PyResult> { + self.inner.batch_encode_to_gpu(py, texts) + } +} -// --- PYTHON MODULE EXPORT --- #[pymodule] fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { - // NEW: Only initialize Tracy if the feature is enabled #[cfg(feature = "profiling")] { - // Use 'default()' instead of 'new()' to avoid argument mismatch errors - // The .try_init() prevents crashing on module reloads (e.g. Jupyter) let _ = tracing_subscriber::registry() .with(tracing_tracy::TracyLayer::default()) .try_init(); - println!("UncheckedIO: Profiling Mode ENABLED šŸš€"); } m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?; m.add_function(wrap_pyfunction!(to_pandas_dataframe, m)?)?; m.add_function(wrap_pyfunction!(profile_data, m)?)?; + + #[cfg(feature = "gpu")] + { + m.add_class::()?; + println!("UncheckedIO: GPU Acceleration ENABLED āš”"); + } + Ok(()) } \ No newline at end of file diff --git a/src/llm.rs b/src/llm.rs new file mode 100644 index 0000000..402e93d --- /dev/null +++ b/src/llm.rs @@ -0,0 +1,199 @@ +// src/llm.rs +#![cfg(feature = "gpu")] + +use std::sync::Arc; +use tokenizers::Tokenizer; +use cudarc::driver::{CudaDevice, DeviceSlice, CudaSlice, DevicePtr}; +use pyo3::prelude::*; +use pyo3::exceptions::PyValueError; +use pyo3::types::PyCapsule; +use std::ffi::CStr; +use rayon::prelude::*; + +// Import raw CUDA symbols. Note the _v2 suffixes which are required by the sys crate. +use cuda_driver_sys::{ + cuMemHostRegister_v2, + cuMemHostUnregister, + cuMemcpyHtoD_v2, + cudaError_enum +}; + +// --- DLPack v0.4 ABI Definitions --- +#[repr(C)] +struct DLDataType { code: u8, bits: u8, lanes: u16 } +#[repr(C)] +struct DLDevice { device_type: i32, device_id: i32 } +#[repr(C)] +struct DLTensor { + data: *mut std::ffi::c_void, + device: DLDevice, + ndim: i32, + dtype: DLDataType, + shape: *mut i64, + strides: *mut i64, + byte_offset: u64, +} +#[repr(C)] +struct DLManagedTensor { + dl_tensor: DLTensor, + manager_ctx: *mut std::ffi::c_void, + deleter: Option, +} + +const CAPSULE_NAME: &CStr = unsafe { CStr::from_bytes_with_nul_unchecked(b"dltensor\0") }; + +struct PinnedHostBuffer { + data: Vec, +} + +impl PinnedHostBuffer { + fn new(capacity: usize) -> PyResult { + let data = vec![0i64; capacity]; + let ptr = data.as_ptr() as *mut std::ffi::c_void; + let bytes = capacity * std::mem::size_of::(); + + unsafe { + // Use _v2 variant + let result = cuMemHostRegister_v2(ptr, bytes, 0); + if result != cudaError_enum::CUDA_SUCCESS { + return Err(PyValueError::new_err(format!("Failed to Pin Memory: {:?}", result))); + } + } + + Ok(Self { data }) + } +} + +impl Drop for PinnedHostBuffer { + fn drop(&mut self) { + unsafe { + let ptr = self.data.as_ptr() as *mut std::ffi::c_void; + cuMemHostUnregister(ptr); + } + } +} + +pub struct PinnedBatcher { + tokenizer: Tokenizer, + device: Arc, + gpu_buffer: CudaSlice, + host_buffer: PinnedHostBuffer, + max_elements: usize, +} + +impl PinnedBatcher { + pub fn new(model_path: &str) -> PyResult { + let tokenizer = Tokenizer::from_file(model_path) + .map_err(|e| PyValueError::new_err(format!("Failed to load tokenizer: {}", e)))?; + + let device = CudaDevice::new(0) + .map_err(|e| PyValueError::new_err(format!("No CUDA GPU found: {:?}", e)))?; + + let max_elements = 4096 * 512; + + let gpu_buffer = device.alloc_zeros::(max_elements) + .map_err(|e| PyValueError::new_err(format!("GPU Alloc Failed: {:?}", e)))?; + + let host_buffer = PinnedHostBuffer::new(max_elements)?; + + Ok(PinnedBatcher { + tokenizer, + device, + gpu_buffer, + host_buffer, + max_elements, + }) + } + + pub fn batch_encode_to_gpu(&mut self, py: Python, texts: Vec) -> PyResult> { + let encodings = self.tokenizer.encode_batch(texts, true) + .map_err(|e| PyValueError::new_err(format!("Tokenization failed: {}", e)))?; + + if encodings.is_empty() { + return Err(PyValueError::new_err("Empty batch provided")); + } + + let batch_size = encodings.len(); + let seq_len = encodings[0].len(); + let total_elements = batch_size * seq_len; + + if total_elements > self.max_elements { + return Err(PyValueError::new_err(format!( + "Batch too large: {} tokens (Max: {})", + total_elements, self.max_elements + ))); + } + + // Parallel Write into Pinned Memory + // Explicit types added to closure to satisfy type inference + self.host_buffer.data[..total_elements].par_chunks_mut(seq_len) + .zip(encodings.par_iter()) + .for_each(|(dest, enc): (&mut [i64], &tokenizers::Encoding)| { + let ids = enc.get_ids(); + let len = ids.len().min(dest.len()); + for i in 0..len { + dest[i] = ids[i] as i64; + } + }); + + // Direct DMA Copy (Pinned Host -> Device) + unsafe { + let src_ptr = self.host_buffer.data.as_ptr() as *const std::ffi::c_void; + let dst_ptr = *self.gpu_buffer.device_ptr(); + let bytes = total_elements * std::mem::size_of::(); + + let result = cuMemcpyHtoD_v2(dst_ptr, src_ptr, bytes); + + if result != cudaError_enum::CUDA_SUCCESS { + return Err(PyValueError::new_err(format!("Direct GPU Copy Failed: {:?}", result))); + } + } + + // Zero-Copy Handover + let ptr_address = *self.gpu_buffer.device_ptr() as *mut std::ffi::c_void; + let shape = Box::into_raw(vec![batch_size as i64, seq_len as i64].into_boxed_slice()) as *mut i64; + + let dl_tensor = DLTensor { + data: ptr_address, + device: DLDevice { device_type: 2, device_id: 0 }, + ndim: 2, + dtype: DLDataType { code: 0, bits: 64, lanes: 1 }, + shape, + strides: std::ptr::null_mut(), + byte_offset: 0, + }; + + let managed_tensor = Box::new(DLManagedTensor { + dl_tensor, + manager_ctx: std::ptr::null_mut(), + deleter: Some(dlpack_deleter), + }); + + let managed_ptr = Box::into_raw(managed_tensor); + + unsafe { + let capsule_ptr = pyo3::ffi::PyCapsule_New( + managed_ptr as *mut _, + CAPSULE_NAME.as_ptr(), + None + ); + + if capsule_ptr.is_null() { + return Err(PyValueError::new_err("Failed to create PyCapsule")); + } + + Ok(Bound::from_owned_ptr(py, capsule_ptr).into_any().unbind()) + } + } +} + +unsafe extern "C" fn dlpack_deleter(managed_ptr: *mut DLManagedTensor) { + if managed_ptr.is_null() { return; } + let managed = Box::from_raw(managed_ptr); + if !managed.dl_tensor.shape.is_null() { + let _ = Box::from_raw(std::slice::from_raw_parts_mut( + managed.dl_tensor.shape, + managed.dl_tensor.ndim as usize + )); + } +} \ No newline at end of file