made zero-copy JSON - shared memory logic
This commit is contained in:
484
Cargo.lock
generated
484
Cargo.lock
generated
@@ -118,7 +118,7 @@ dependencies = [
|
|||||||
"arrow-schema",
|
"arrow-schema",
|
||||||
"arrow-select",
|
"arrow-select",
|
||||||
"atoi",
|
"atoi",
|
||||||
"base64",
|
"base64 0.22.1",
|
||||||
"chrono",
|
"chrono",
|
||||||
"comfy-table",
|
"comfy-table",
|
||||||
"half",
|
"half",
|
||||||
@@ -297,6 +297,12 @@ version = "1.5.0"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
|
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "base64"
|
||||||
|
version = "0.13.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "base64"
|
name = "base64"
|
||||||
version = "0.22.1"
|
version = "0.22.1"
|
||||||
@@ -395,6 +401,19 @@ dependencies = [
|
|||||||
"crossbeam-utils",
|
"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]]
|
[[package]]
|
||||||
name = "const-random"
|
name = "const-random"
|
||||||
version = "0.1.18"
|
version = "0.1.18"
|
||||||
@@ -430,6 +449,25 @@ dependencies = [
|
|||||||
"libc",
|
"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]]
|
[[package]]
|
||||||
name = "crossbeam-utils"
|
name = "crossbeam-utils"
|
||||||
version = "0.8.21"
|
version = "0.8.21"
|
||||||
@@ -473,6 +511,68 @@ dependencies = [
|
|||||||
"memchr",
|
"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]]
|
[[package]]
|
||||||
name = "deadpool"
|
name = "deadpool"
|
||||||
version = "0.12.3"
|
version = "0.12.3"
|
||||||
@@ -508,6 +608,37 @@ dependencies = [
|
|||||||
"tokio",
|
"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]]
|
[[package]]
|
||||||
name = "digest"
|
name = "digest"
|
||||||
version = "0.10.7"
|
version = "0.10.7"
|
||||||
@@ -519,12 +650,33 @@ dependencies = [
|
|||||||
"subtle",
|
"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]]
|
[[package]]
|
||||||
name = "equivalent"
|
name = "equivalent"
|
||||||
version = "1.0.2"
|
version = "1.0.2"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
|
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]]
|
[[package]]
|
||||||
name = "event-listener"
|
name = "event-listener"
|
||||||
version = "5.4.1"
|
version = "5.4.1"
|
||||||
@@ -568,6 +720,12 @@ dependencies = [
|
|||||||
"rustc_version",
|
"rustc_version",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "fnv"
|
||||||
|
version = "1.0.7"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "futures-channel"
|
name = "futures-channel"
|
||||||
version = "0.3.31"
|
version = "0.3.31"
|
||||||
@@ -671,6 +829,12 @@ dependencies = [
|
|||||||
"wasip2",
|
"wasip2",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "glob"
|
||||||
|
version = "0.3.3"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "half"
|
name = "half"
|
||||||
version = "2.7.1"
|
version = "2.7.1"
|
||||||
@@ -734,6 +898,12 @@ dependencies = [
|
|||||||
"cc",
|
"cc",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "ident_case"
|
||||||
|
version = "1.0.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "indexmap"
|
name = "indexmap"
|
||||||
version = "2.12.0"
|
version = "2.12.0"
|
||||||
@@ -744,6 +914,19 @@ dependencies = [
|
|||||||
"hashbrown",
|
"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]]
|
[[package]]
|
||||||
name = "indoc"
|
name = "indoc"
|
||||||
version = "2.0.7"
|
version = "2.0.7"
|
||||||
@@ -753,6 +936,24 @@ dependencies = [
|
|||||||
"rustversion",
|
"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]]
|
[[package]]
|
||||||
name = "itoa"
|
name = "itoa"
|
||||||
version = "1.0.15"
|
version = "1.0.15"
|
||||||
@@ -838,6 +1039,16 @@ version = "0.2.177"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976"
|
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]]
|
[[package]]
|
||||||
name = "libm"
|
name = "libm"
|
||||||
version = "0.2.15"
|
version = "0.2.15"
|
||||||
@@ -893,6 +1104,22 @@ dependencies = [
|
|||||||
"tracing-subscriber",
|
"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]]
|
[[package]]
|
||||||
name = "matchers"
|
name = "matchers"
|
||||||
version = "0.2.0"
|
version = "0.2.0"
|
||||||
@@ -946,6 +1173,12 @@ dependencies = [
|
|||||||
"libmimalloc-sys",
|
"libmimalloc-sys",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "minimal-lexical"
|
||||||
|
version = "0.2.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "mio"
|
name = "mio"
|
||||||
version = "1.1.0"
|
version = "1.1.0"
|
||||||
@@ -957,6 +1190,28 @@ dependencies = [
|
|||||||
"windows-sys 0.61.2",
|
"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]]
|
[[package]]
|
||||||
name = "ndarray"
|
name = "ndarray"
|
||||||
version = "0.16.1"
|
version = "0.16.1"
|
||||||
@@ -972,6 +1227,16 @@ dependencies = [
|
|||||||
"rawpointer",
|
"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]]
|
[[package]]
|
||||||
name = "nu-ansi-term"
|
name = "nu-ansi-term"
|
||||||
version = "0.50.3"
|
version = "0.50.3"
|
||||||
@@ -1029,6 +1294,12 @@ dependencies = [
|
|||||||
"libc",
|
"libc",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "number_prefix"
|
||||||
|
version = "0.4.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "numpy"
|
name = "numpy"
|
||||||
version = "0.27.0"
|
version = "0.27.0"
|
||||||
@@ -1052,6 +1323,28 @@ version = "1.21.3"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
|
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]]
|
[[package]]
|
||||||
name = "parking"
|
name = "parking"
|
||||||
version = "2.2.1"
|
version = "2.2.1"
|
||||||
@@ -1081,6 +1374,12 @@ dependencies = [
|
|||||||
"windows-link 0.2.1",
|
"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]]
|
[[package]]
|
||||||
name = "percent-encoding"
|
name = "percent-encoding"
|
||||||
version = "2.3.2"
|
version = "2.3.2"
|
||||||
@@ -1136,6 +1435,12 @@ version = "0.1.0"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
|
checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "pkg-config"
|
||||||
|
version = "0.3.32"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "portable-atomic"
|
name = "portable-atomic"
|
||||||
version = "1.11.1"
|
version = "1.11.1"
|
||||||
@@ -1157,14 +1462,14 @@ version = "0.6.9"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "fbef655056b916eb868048276cfd5d6a7dea4f81560dfd047f97c8c6fe3fcfd4"
|
checksum = "fbef655056b916eb868048276cfd5d6a7dea4f81560dfd047f97c8c6fe3fcfd4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"base64",
|
"base64 0.22.1",
|
||||||
"byteorder",
|
"byteorder",
|
||||||
"bytes",
|
"bytes",
|
||||||
"fallible-iterator",
|
"fallible-iterator",
|
||||||
"hmac",
|
"hmac",
|
||||||
"md-5",
|
"md-5",
|
||||||
"memchr",
|
"memchr",
|
||||||
"rand",
|
"rand 0.9.2",
|
||||||
"sha2",
|
"sha2",
|
||||||
"stringprep",
|
"stringprep",
|
||||||
]
|
]
|
||||||
@@ -1248,6 +1553,15 @@ dependencies = [
|
|||||||
"target-lexicon",
|
"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]]
|
[[package]]
|
||||||
name = "pyo3-ffi"
|
name = "pyo3-ffi"
|
||||||
version = "0.27.0"
|
version = "0.27.0"
|
||||||
@@ -1298,14 +1612,35 @@ version = "5.3.0"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
|
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]]
|
[[package]]
|
||||||
name = "rand"
|
name = "rand"
|
||||||
version = "0.9.2"
|
version = "0.9.2"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1"
|
checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"rand_chacha",
|
"rand_chacha 0.9.0",
|
||||||
"rand_core",
|
"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]]
|
[[package]]
|
||||||
@@ -1315,7 +1650,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||||||
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
|
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"ppv-lite86",
|
"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]]
|
[[package]]
|
||||||
@@ -1333,6 +1677,37 @@ version = "0.2.1"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
|
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]]
|
[[package]]
|
||||||
name = "redox_syscall"
|
name = "redox_syscall"
|
||||||
version = "0.5.18"
|
version = "0.5.18"
|
||||||
@@ -1541,6 +1916,18 @@ dependencies = [
|
|||||||
"windows-sys 0.60.2",
|
"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]]
|
[[package]]
|
||||||
name = "stringprep"
|
name = "stringprep"
|
||||||
version = "0.1.5"
|
version = "0.1.5"
|
||||||
@@ -1552,6 +1939,12 @@ dependencies = [
|
|||||||
"unicode-properties",
|
"unicode-properties",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "strsim"
|
||||||
|
version = "0.11.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "strum"
|
name = "strum"
|
||||||
version = "0.26.3"
|
version = "0.26.3"
|
||||||
@@ -1647,6 +2040,38 @@ version = "0.1.1"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
|
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]]
|
[[package]]
|
||||||
name = "tokio"
|
name = "tokio"
|
||||||
version = "1.48.0"
|
version = "1.48.0"
|
||||||
@@ -1694,7 +2119,7 @@ dependencies = [
|
|||||||
"pin-project-lite",
|
"pin-project-lite",
|
||||||
"postgres-protocol",
|
"postgres-protocol",
|
||||||
"postgres-types",
|
"postgres-types",
|
||||||
"rand",
|
"rand 0.9.2",
|
||||||
"socket2",
|
"socket2",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tokio-util",
|
"tokio-util",
|
||||||
@@ -1824,14 +2249,19 @@ dependencies = [
|
|||||||
"bytes",
|
"bytes",
|
||||||
"chrono",
|
"chrono",
|
||||||
"chrono-tz",
|
"chrono-tz",
|
||||||
|
"cuda-driver-sys",
|
||||||
|
"cudarc",
|
||||||
"deadpool-postgres",
|
"deadpool-postgres",
|
||||||
"futures-util",
|
"futures-util",
|
||||||
"mimalloc",
|
"mimalloc",
|
||||||
"num_cpus",
|
"num_cpus",
|
||||||
"pyo3",
|
"pyo3",
|
||||||
"pyo3-arrow",
|
"pyo3-arrow",
|
||||||
|
"pyo3-dlpack",
|
||||||
|
"rayon",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_yaml",
|
"serde_yaml",
|
||||||
|
"tokenizers",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tokio-postgres",
|
"tokio-postgres",
|
||||||
"tracing",
|
"tracing",
|
||||||
@@ -1861,18 +2291,39 @@ dependencies = [
|
|||||||
"tinyvec",
|
"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]]
|
[[package]]
|
||||||
name = "unicode-properties"
|
name = "unicode-properties"
|
||||||
version = "0.1.4"
|
version = "0.1.4"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d"
|
checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "unicode-segmentation"
|
||||||
|
version = "1.12.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "unicode-width"
|
name = "unicode-width"
|
||||||
version = "0.2.2"
|
version = "0.2.2"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
|
checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "unicode_categories"
|
||||||
|
version = "0.1.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "unindent"
|
name = "unindent"
|
||||||
version = "0.2.4"
|
version = "0.2.4"
|
||||||
@@ -1985,6 +2436,16 @@ dependencies = [
|
|||||||
"wasm-bindgen",
|
"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]]
|
[[package]]
|
||||||
name = "whoami"
|
name = "whoami"
|
||||||
version = "1.6.1"
|
version = "1.6.1"
|
||||||
@@ -2135,6 +2596,15 @@ dependencies = [
|
|||||||
"windows-link 0.2.1",
|
"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]]
|
[[package]]
|
||||||
name = "windows-sys"
|
name = "windows-sys"
|
||||||
version = "0.60.2"
|
version = "0.60.2"
|
||||||
|
|||||||
58
Cargo.toml
58
Cargo.toml
@@ -1,67 +1,55 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "unchecked-io"
|
name = "unchecked-io"
|
||||||
version = "0.1.7"
|
version = "0.1.7"
|
||||||
authors = ["Billthemaker"] # Replace with your name or alias
|
authors = ["Billthemaker"]
|
||||||
license = "Apache-2.0" # Good practice for open-source
|
license = "Apache-2.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
|
||||||
[lib]
|
[lib]
|
||||||
name = "unchecked_io"
|
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]
|
[dependencies]
|
||||||
# 1. Python Bindings for FFI
|
|
||||||
pyo3 = { version = "=0.27.0", features = ["extension-module", "chrono-tz", "chrono"] }
|
pyo3 = { version = "=0.27.0", features = ["extension-module", "chrono-tz", "chrono"] }
|
||||||
|
|
||||||
|
|
||||||
# 2. Configuration Parsing (YAML)
|
|
||||||
serde = { version = "1.0", features = ["derive"] }
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
serde_yaml = "0.9"
|
serde_yaml = "0.9"
|
||||||
|
arrow = { version = "57.0.0", features = ["prettyprint"] }
|
||||||
# 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)
|
|
||||||
tokio = { version = "1.37", features = ["full"] }
|
tokio = { version = "1.37", features = ["full"] }
|
||||||
|
|
||||||
|
|
||||||
# 5. Database Connection (Postgres)
|
|
||||||
tokio-postgres = "0.7"
|
tokio-postgres = "0.7"
|
||||||
deadpool-postgres = "0.14"
|
deadpool-postgres = "0.14"
|
||||||
|
|
||||||
# 6. Error Handling Crate
|
|
||||||
anyhow = "1.0"
|
anyhow = "1.0"
|
||||||
|
|
||||||
# 7. Futures Utilities
|
|
||||||
futures-util = "0.3"
|
futures-util = "0.3"
|
||||||
|
|
||||||
# 8. Byte Buffer Management
|
|
||||||
bytes = "1.6"
|
bytes = "1.6"
|
||||||
|
|
||||||
# 9. Binary Data Reading (NEW)
|
|
||||||
byteorder = "1.5"
|
byteorder = "1.5"
|
||||||
|
|
||||||
# 10. Timestamp Handling (NEW)
|
|
||||||
chrono = "0.4"
|
chrono = "0.4"
|
||||||
chrono-tz = "=0.10.4"
|
chrono-tz = "=0.10.4"
|
||||||
|
|
||||||
# 11. UUID Handling (NEW)
|
|
||||||
uuid = { version = "1.8", features = ["serde", "v4"] }
|
uuid = { version = "1.8", features = ["serde", "v4"] }
|
||||||
|
|
||||||
# 12. Arrow <-> Python Bridge (NEW)
|
|
||||||
pyo3-arrow = "0.15.0"
|
pyo3-arrow = "0.15.0"
|
||||||
|
|
||||||
# 13. System CPU Count (NEW)
|
|
||||||
num_cpus = "1.16"
|
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-subscriber = { version = "0.3", features = ["registry", "env-filter"], optional = true}
|
||||||
tracing-tracy = { version = "=0.11.2", optional = true }
|
tracing-tracy = { version = "=0.11.2", optional = true }
|
||||||
tracing = {version ="0.1.41", optional = true }
|
tracing = {version ="0.1.41", optional = true }
|
||||||
|
|
||||||
async-channel = "2.5.0"
|
|
||||||
|
|
||||||
[target.'cfg(not(target_env = "msvc"))'.dependencies]
|
[target.'cfg(not(target_env = "msvc"))'.dependencies]
|
||||||
mimalloc = { version = "0.1.39" }
|
mimalloc = { version = "0.1.39" }
|
||||||
|
|
||||||
async-channel = "2.3"
|
|
||||||
114
benchmark_gpu.py
Normal file
114
benchmark_gpu.py
Normal file
@@ -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")
|
||||||
109
server.py
Normal file
109
server.py
Normal file
@@ -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)
|
||||||
68
src/lib.rs
68
src/lib.rs
@@ -1,18 +1,16 @@
|
|||||||
// --- Declare our new modules ---
|
|
||||||
mod config;
|
mod config;
|
||||||
mod parser;
|
mod parser;
|
||||||
|
|
||||||
// --- External Crates ---
|
#[cfg(feature = "gpu")]
|
||||||
|
mod llm;
|
||||||
|
|
||||||
use pyo3::prelude::*;
|
use pyo3::prelude::*;
|
||||||
use pyo3::exceptions::PyValueError;
|
use pyo3::exceptions::PyValueError;
|
||||||
use tokio;
|
use tokio;
|
||||||
use pyo3::types::{PyModule, PyAny};
|
use pyo3::types::{PyModule, PyAny};
|
||||||
use pyo3::Bound;
|
use pyo3::Bound;
|
||||||
|
|
||||||
// FIX: Use the export path you confirmed works in your IDE
|
|
||||||
use pyo3_arrow::export::Arro3RecordBatch;
|
use pyo3_arrow::export::Arro3RecordBatch;
|
||||||
|
|
||||||
// Required for global allocator (mimalloc)
|
|
||||||
#[cfg(not(target_env = "msvc"))]
|
#[cfg(not(target_env = "msvc"))]
|
||||||
use mimalloc;
|
use mimalloc;
|
||||||
|
|
||||||
@@ -20,26 +18,16 @@ use mimalloc;
|
|||||||
#[global_allocator]
|
#[global_allocator]
|
||||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||||
|
|
||||||
|
|
||||||
// --- Internal Crates ---
|
|
||||||
use crate::config::{load_and_validate_config, ConnectorConfig};
|
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};
|
use crate::parser::{run_db_logic, run_profiler_logic};
|
||||||
|
|
||||||
// NEW: Import the tracing libraries only if feature is enabled
|
|
||||||
#[cfg(feature = "profiling")]
|
#[cfg(feature = "profiling")]
|
||||||
use tracing_subscriber::layer::SubscriberExt;
|
use tracing_subscriber::layer::SubscriberExt;
|
||||||
#[cfg(feature = "profiling")]
|
#[cfg(feature = "profiling")]
|
||||||
use tracing_subscriber::util::SubscriberInitExt;
|
use tracing_subscriber::util::SubscriberInitExt;
|
||||||
|
|
||||||
|
|
||||||
// --- THE PYTHON-CALLABLE ENTRY POINT (Load Data) ---
|
|
||||||
|
|
||||||
#[pyfunction]
|
#[pyfunction]
|
||||||
#[pyo3(signature = (config_path, blast_radius=0))] // Default to 0 for auto-tuning
|
#[pyo3(signature = (config_path, blast_radius=0))]
|
||||||
#[allow(unsafe_code)]
|
|
||||||
#[allow(unsafe_op_in_unsafe_fn)]
|
|
||||||
#[allow(rust_2024_compatibility)]
|
|
||||||
fn load_data_from_config<'py>(
|
fn load_data_from_config<'py>(
|
||||||
py: Python<'py>,
|
py: Python<'py>,
|
||||||
config_path: String,
|
config_path: String,
|
||||||
@@ -51,12 +39,6 @@ fn load_data_from_config<'py>(
|
|||||||
Err(e) => return Err(PyValueError::new_err(format!("Configuration Error: {:?}", e))),
|
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::<Vec<_>>());
|
|
||||||
|
|
||||||
// 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(|| {
|
let record_batch = py.allow_threads(|| {
|
||||||
tokio::runtime::Builder::new_multi_thread()
|
tokio::runtime::Builder::new_multi_thread()
|
||||||
.enable_all()
|
.enable_all()
|
||||||
@@ -67,25 +49,19 @@ fn load_data_from_config<'py>(
|
|||||||
})
|
})
|
||||||
}).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?;
|
}).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);
|
let py_record_batch = Arro3RecordBatch::from(record_batch);
|
||||||
py_record_batch.into_pyobject(py)
|
py_record_batch.into_pyobject(py)
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- NEW PYTHON-CALLABLE FUNCTION FOR PANDAS CONVERSION (FIXED SIGNATURE) ---
|
|
||||||
#[pyfunction]
|
#[pyfunction]
|
||||||
#[pyo3(signature = (arrow_table))] // FIX: Removed 'py' from the signature macro
|
#[pyo3(signature = (arrow_table))]
|
||||||
fn to_pandas_dataframe<'py>(py: Python<'py>, arrow_table: Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
|
fn to_pandas_dataframe<'py>(_py: Python<'py>, arrow_table: Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
|
||||||
// This calls the 'to_pandas' method on the PyArrow object.
|
|
||||||
arrow_table.call_method0("to_pandas")
|
arrow_table.call_method0("to_pandas")
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- NEW PYTHON-CALLABLE ENTRY POINT FOR SCHEMA PROFILING ---
|
|
||||||
#[pyfunction]
|
#[pyfunction]
|
||||||
#[pyo3(signature = (config_path))]
|
#[pyo3(signature = (config_path))]
|
||||||
fn profile_data(config_path: String) -> PyResult<String> {
|
fn profile_data(config_path: String) -> PyResult<String> {
|
||||||
// Note: The profiler uses a small, current_thread tokio runtime since it's sequential I/O.
|
|
||||||
let output = std::thread::spawn(move || {
|
let output = std::thread::spawn(move || {
|
||||||
tokio::runtime::Builder::new_current_thread()
|
tokio::runtime::Builder::new_current_thread()
|
||||||
.enable_all()
|
.enable_all()
|
||||||
@@ -99,25 +75,47 @@ fn profile_data(config_path: String) -> PyResult<String> {
|
|||||||
Ok(output)
|
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<Self> {
|
||||||
|
Ok(TokenizerEngine {
|
||||||
|
inner: llm::PinnedBatcher::new(&model_path)?
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn encode_batch(&mut self, py: Python, texts: Vec<String>) -> PyResult<Py<PyAny>> {
|
||||||
|
self.inner.batch_encode_to_gpu(py, texts)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// --- PYTHON MODULE EXPORT ---
|
|
||||||
#[pymodule]
|
#[pymodule]
|
||||||
fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||||
|
|
||||||
// NEW: Only initialize Tracy if the feature is enabled
|
|
||||||
#[cfg(feature = "profiling")]
|
#[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()
|
let _ = tracing_subscriber::registry()
|
||||||
.with(tracing_tracy::TracyLayer::default())
|
.with(tracing_tracy::TracyLayer::default())
|
||||||
.try_init();
|
.try_init();
|
||||||
|
|
||||||
println!("UncheckedIO: Profiling Mode ENABLED 🚀");
|
println!("UncheckedIO: Profiling Mode ENABLED 🚀");
|
||||||
}
|
}
|
||||||
|
|
||||||
m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?;
|
m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?;
|
||||||
m.add_function(wrap_pyfunction!(to_pandas_dataframe, m)?)?;
|
m.add_function(wrap_pyfunction!(to_pandas_dataframe, m)?)?;
|
||||||
m.add_function(wrap_pyfunction!(profile_data, m)?)?;
|
m.add_function(wrap_pyfunction!(profile_data, m)?)?;
|
||||||
|
|
||||||
|
#[cfg(feature = "gpu")]
|
||||||
|
{
|
||||||
|
m.add_class::<TokenizerEngine>()?;
|
||||||
|
println!("UncheckedIO: GPU Acceleration ENABLED ⚡");
|
||||||
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
199
src/llm.rs
Normal file
199
src/llm.rs
Normal file
@@ -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<unsafe extern "C" fn(*mut DLManagedTensor)>,
|
||||||
|
}
|
||||||
|
|
||||||
|
const CAPSULE_NAME: &CStr = unsafe { CStr::from_bytes_with_nul_unchecked(b"dltensor\0") };
|
||||||
|
|
||||||
|
struct PinnedHostBuffer {
|
||||||
|
data: Vec<i64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PinnedHostBuffer {
|
||||||
|
fn new(capacity: usize) -> PyResult<Self> {
|
||||||
|
let data = vec![0i64; capacity];
|
||||||
|
let ptr = data.as_ptr() as *mut std::ffi::c_void;
|
||||||
|
let bytes = capacity * std::mem::size_of::<i64>();
|
||||||
|
|
||||||
|
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<CudaDevice>,
|
||||||
|
gpu_buffer: CudaSlice<i64>,
|
||||||
|
host_buffer: PinnedHostBuffer,
|
||||||
|
max_elements: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PinnedBatcher {
|
||||||
|
pub fn new(model_path: &str) -> PyResult<Self> {
|
||||||
|
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::<i64>(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<String>) -> PyResult<Py<PyAny>> {
|
||||||
|
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::<i64>();
|
||||||
|
|
||||||
|
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
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user