ci(rust): shard unit tests and disable wheel LTO

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

.github/workflows/test-rust.yml

@@ -113,7 +113,7 @@ jobs:
 
       - run: cargo clippy --workspace --all-targets --locked -- -D warnings
 
-  rust-test:
+  rust-test-build:
     runs-on: ubuntu-latest
     timeout-minutes: 30
     # Trim debuginfo to keep the job's target dir within the runner's disk
@@ -162,27 +162,161 @@ jobs:
         working-directory: .
         run: uv run scripts/generate_trace_types.py --check
 
-      - run: cargo nextest run --workspace --locked
+      - run: cargo nextest archive --workspace --locked --archive-file "$RUNNER_TEMP/rust-tests.tar.zst"
+
+      - uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1
+        with:
+          name: rust-tests
+          path: ${{ runner.temp }}/rust-tests.tar.zst
+          compression-level: 0
+          retention-days: 1
+          if-no-files-found: error
 
       - run: cargo test --workspace --doc --locked
 
-      - name: Test token counter feature combinations
+  rust-test-shard:
+    needs: rust-test-build
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    strategy:
+      fail-fast: false
+      matrix:
+        shard: [1, 2, 3, 4]
+    env:
+      RUSTC_WRAPPER: ""
+    defaults:
+      run:
+        working-directory: litellm-rust
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955
+        with:
+          persist-credentials: false
+
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065
+        with:
+          python-version: "3.12"
+
+      - uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: "0.10.9"
+
+      - name: Install Python dependencies for the bridge tests
+        working-directory: .
         run: |
-          for features in '' fast huggingface tiktoken fast,huggingface fast,tiktoken huggingface,tiktoken fast,huggingface,tiktoken; do
-            cargo test -p litellm-token-counter --locked --no-default-features --features "$features"
+          uv sync --frozen --no-install-project
+          echo "PYTHONPATH=$PWD/.venv/lib/$(ls .venv/lib)/site-packages" >> "$GITHUB_ENV"
+
+      - run: rustup toolchain install --no-self-update
+
+      - uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff
+        with:
+          tool: cargo-nextest@0.9.143
+
+      - uses: actions/download-artifact@95815c38cf2ff2164869cbab79da8d1f422bc89e
+        with:
+          name: rust-tests
+          path: ${{ runner.temp }}
+
+      - name: Run unit-test shard
+        env:
+          SHARD: ${{ matrix.shard }}
+        run: >-
+          cargo nextest run --archive-file "$RUNNER_TEMP/rust-tests.tar.zst"
+          --workspace-remap "$PWD" --partition "hash:$SHARD/4"
+
+  rust-test-features:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - suite: token-counter-single
+            package: litellm-token-counter
+            features: "none fast huggingface tiktoken"
+          - suite: token-counter-combined
+            package: litellm-token-counter
+            features: "fast,huggingface fast,tiktoken huggingface,tiktoken fast,huggingface,tiktoken"
+          - suite: secrets-single
+            package: litellm-secrets
+            features: "none aws google hashicorp azure cyberark"
+          - suite: secrets-combined
+            package: litellm-secrets
+            features: "aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark aws,google,hashicorp,azure,cyberark"
+    name: rust-test-features (${{ matrix.suite }})
+    env:
+      CARGO_PROFILE_DEV_DEBUG: line-tables-only
+      PACKAGE: ${{ matrix.package }}
+      FEATURE_SETS: ${{ matrix.features }}
+    defaults:
+      run:
+        working-directory: litellm-rust
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955
+        with:
+          persist-credentials: false
+
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065
+        with:
+          python-version: "3.12"
+
+      - run: rustup toolchain install --no-self-update
+
+      - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6
+        with:
+          workspaces: litellm-rust
+          key: ${{ matrix.suite }}
+          cache-on-failure: true
+          save-if: ${{ github.ref == 'refs/heads/main' }}
+
+      - name: Set up sccache
+        uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a
+        with:
+          version: "v0.18.0"
+
+      - name: Test GCP authentication without default features
+        if: matrix.suite == 'secrets-single'
+        run: cargo test -p litellm-auth-gcp --locked --no-default-features
+
+      - name: Test token counter feature combinations
+        if: matrix.package == 'litellm-token-counter'
+        run: |
+          for feature_set in $FEATURE_SETS; do
+            features="${feature_set/none/}"
+            cargo test -p "$PACKAGE" --locked --no-default-features --features "$features"
             cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}"
           done
 
       - name: Test secret manager feature combinations
+        if: matrix.package == 'litellm-secrets'
         run: |
-          cargo test -p litellm-auth-gcp --locked --no-default-features
-          for features in '' aws google hashicorp azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark aws,google,hashicorp,azure,cyberark; do
-            cargo test -p litellm-secrets --locked --no-default-features --features "$features"
+          for feature_set in $FEATURE_SETS; do
+            features="${feature_set/none/}"
+            cargo test -p "$PACKAGE" --locked --no-default-features --features "$features"
           done
 
+  rust-test:
+    if: always()
+    needs: [rust-test-build, rust-test-shard, rust-test-features]
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Require every Rust test job to pass
+        env:
+          BUILD_RESULT: ${{ needs.rust-test-build.result }}
+          SHARD_RESULT: ${{ needs.rust-test-shard.result }}
+          FEATURE_RESULT: ${{ needs.rust-test-features.result }}
+        run: |
+          test "$BUILD_RESULT" = success
+          test "$SHARD_RESULT" = success
+          test "$FEATURE_RESULT" = success
+
   rust-wheel:
     runs-on: ubuntu-latest
     timeout-minutes: 30
+    env:
+      CARGO_PROFILE_RELEASE_LTO: "off"
+      CARGO_PROFILE_RELEASE_CODEGEN_UNITS: "16"
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
         with: