Merge pull request #41480 from BerriAI/litellm_rust_ci_nextest_rust_cache

ci(rust): split rust jobs, use nextest and Swatinem/rust-cache

.github/workflows/test-rust.yml

@@ -70,7 +70,7 @@ env:
 jobs:
   rust-lint:
     runs-on: ubuntu-latest
-    timeout-minutes: 10
+    timeout-minutes: 15
     defaults:
       run:
         working-directory: litellm-rust
@@ -81,24 +81,48 @@ jobs:
 
       - run: rustup toolchain install --no-self-update
 
-      - run: cargo fmt --check
+      - run: cargo fmt --all --check
 
-      - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+      - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
         with:
-          path: |
-            ~/.cargo/registry
-            ~/.cargo/git
-            litellm-rust/target
-          key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
-          restore-keys: |
-            ${{ runner.os }}-cargo-${{ github.job }}-
+          workspaces: litellm-rust
+          cache-on-failure: true
 
       - run: cargo clippy --workspace --all-targets --locked -- -D warnings
 
   rust-test:
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 20
+    defaults:
+      run:
+        working-directory: litellm-rust
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
 
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - run: rustup toolchain install --no-self-update
+
+      - uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8
+        with:
+          tool: cargo-nextest@0.9.143
+
+      - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
+        with:
+          workspaces: litellm-rust
+          cache-on-failure: true
+
+      - run: cargo nextest run --workspace --locked
+
+      - run: cargo test --workspace --doc --locked
+
+  rust-wheel:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
         with:
@@ -114,18 +138,10 @@ jobs:
 
       - run: rustup toolchain install --no-self-update
 
-      - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+      - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
         with:
-          path: |
-            ~/.cargo/registry
-            ~/.cargo/git
-            litellm-rust/target
-          key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
-          restore-keys: |
-            ${{ runner.os }}-cargo-${{ github.job }}-
-
-      - run: cargo test --workspace --locked
-        working-directory: litellm-rust
+          workspaces: litellm-rust
+          cache-on-failure: true
 
       - run: uv build --wheel --out-dir dist
 

litellm-rust/crates/python-bridge/src/lifecycle/mod.rs

@@ -599,6 +599,34 @@ mod tests {
 
     static PYTHON_GLOBALS: Mutex<()> = Mutex::new(());
 
+    fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
+        py.run(
+            pyo3::ffi::c_str!(
+                r#"
+import sys
+import types
+
+sys.modules.setdefault('litellm', types.ModuleType('litellm'))
+sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge'))
+"#
+            ),
+            None,
+            None,
+        )
+        .unwrap();
+        let source = std::ffi::CString::new(include_str!(
+            "../../../../../litellm/rust_bridge/lifecycle.py"
+        ))
+        .unwrap();
+        PyModule::from_code(
+            py,
+            &source,
+            pyo3::ffi::c_str!("lifecycle.py"),
+            pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"),
+        )
+        .unwrap()
+    }
+
     fn install_logging_worker(py: Python<'_>, worker: &Bound<'_, PyAny>) -> PyResult<()> {
         py.import("litellm.litellm_core_utils.logging_worker")?
             .setattr("GLOBAL_LOGGING_WORKER", worker)
@@ -773,17 +801,7 @@ mod tests {
             .unwrap_or_else(|error| error.into_inner());
         Python::initialize();
         Python::attach(|py| {
-            let source = std::ffi::CString::new(include_str!(
-                "../../../../../litellm/rust_bridge/lifecycle.py"
-            ))
-            .unwrap();
-            PyModule::from_code(
-                py,
-                &source,
-                pyo3::ffi::c_str!("lifecycle.py"),
-                pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"),
-            )
-            .unwrap();
+            install_lifecycle_module(py);
             let route = SyntheticRoute(
                 PythonCallState::new(
                     py,
@@ -819,17 +837,7 @@ mod tests {
         Python::initialize();
         Python::attach(|py| {
             py.import("asyncio").unwrap();
-            let source = std::ffi::CString::new(include_str!(
-                "../../../../../litellm/rust_bridge/lifecycle.py"
-            ))
-            .unwrap();
-            let module = PyModule::from_code(
-                py,
-                &source,
-                pyo3::ffi::c_str!("lifecycle.py"),
-                pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"),
-            )
-            .unwrap();
+            let module = install_lifecycle_module(py);
             let locals = PyDict::new(py);
             locals
                 .set_item("drive", module.getattr("drive").unwrap())

litellm-rust/crates/python-bridge/src/marshal.rs

@@ -190,6 +190,7 @@ mod tests {
 
     #[test]
     fn required_shapes_preserve_nested_values_and_existing_errors() {
+        Python::initialize();
         let nested = json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]);
         assert_eq!(
             Value::Array(required_array("messages", nested.clone()).unwrap()),

litellm-rust/crates/token-counter/src/tiktoken.rs

@@ -195,14 +195,21 @@ mod tests {
     }
 
     #[test]
-    fn long_repeated_runs_stay_cheap() {
+    fn long_repeated_runs_cost_close_to_linear() {
         let ranks = ranks();
         let mut scratch = MergeScratch::default();
-        let piece = vec![b' '; 1 << 20];
-        let started = std::time::Instant::now();
-        let count = ranks.count_piece(&piece, &mut scratch);
-        assert!(count > 0);
-        assert!(started.elapsed().as_secs() < 5, "{:?}", started.elapsed());
+        let mut time = |len: usize| {
+            let piece = vec![b' '; len];
+            let started = std::time::Instant::now();
+            assert!(ranks.count_piece(&piece, &mut scratch) > 0);
+            started.elapsed()
+        };
+        let small = (0..3).map(|_| time(1 << 14)).min().unwrap();
+        let large = time(1 << 18);
+        assert!(
+            large < small * 64,
+            "{small:?} for 2^14 bytes, {large:?} for 2^18"
+        );
     }
 
     #[test]