Merge pull request #39523 from BerriAI/litellm_fix_ai_gateway_image_build

fix(ai-gateway): build the release image again and cover it in CI

.github/ci-coverage-allowlist.yml

@@ -106,11 +106,6 @@ dockerfiles:
       and lint workflows already exercise that output, so building the image adds no signal about it
     paths:
       - ui/Dockerfile
-  - reason: >-
-      The Rust gateway ships as its own chart and package with a separate release pipeline, so its
-      image is not part of this repo's Python image set
-    paths:
-      - litellm-rust/crates/ai-gateway/Dockerfile
   - reason: >-
       An example image under cookbook/ that is documentation rather than a shipped artifact
     paths:

.github/workflows/ai-gateway-image.yml

@@ -0,0 +1,73 @@
+name: ai-gateway image
+
+on:
+  push:
+    paths:
+      - "litellm-rust/**"
+      - "litellm/**"
+      - "enterprise/**"
+      - "litellm-proxy-extras/**"
+      - "pyproject.toml"
+      - "rust-toolchain.toml"
+      - ".github/workflows/ai-gateway-image.yml"
+  pull_request:
+    branches:
+      - main
+      - litellm_internal_staging
+      - litellm_oss_staging
+      - "litellm_**"
+    paths:
+      - "litellm-rust/**"
+      - "litellm/**"
+      - "enterprise/**"
+      - "litellm-proxy-extras/**"
+      - "pyproject.toml"
+      - "rust-toolchain.toml"
+      - ".github/workflows/ai-gateway-image.yml"
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  ai-gateway-image:
+    name: ai-gateway release image
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    permissions:
+      contents: read
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+      - name: Build the release image
+        run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} .
+      - name: Start the gateway and wait for readiness
+        env:
+          IMAGE: litellm-ai-gateway:${{ github.sha }}
+        run: |
+          docker run -d --name ai-gateway -p 4001:4001 \
+            -e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \
+            -e OPENAI_API_KEY=sk-ci-not-a-real-key \
+            "$IMAGE"
+          for _ in $(seq 1 60); do
+            if curl -fsS http://127.0.0.1:4001/health/readiness; then
+              echo "gateway is serving readiness"
+              exit 0
+            fi
+            sleep 2
+          done
+          echo "gateway never became ready" >&2
+          docker logs ai-gateway >&2
+          exit 1
+      - name: Assert the gateway loaded the baked config
+        run: |
+          docker logs ai-gateway 2>&1 | tee gateway.log
+          grep 'via python config reader' gateway.log
+      - name: Stop the gateway
+        if: always()
+        run: docker rm -f ai-gateway || true

litellm-rust/crates/ai-gateway/Dockerfile

@@ -14,15 +14,20 @@
 # ---- Chef -------------------------------------------------------------------
 # cargo-chef caches the dependency build so only the gateway crate recompiles on
 # a source-only change. python3-dev is present in every rust stage because the
-# `python-config` feature links libpython via pyo3 (even in the cook step).
-FROM rust:1.90-slim-bookworm AS chef
+# `python-config` feature links libpython via pyo3 (even in the cook step), and
+# python3-pip builds the litellm wheel in the builder stage.
+FROM rust:1.98-slim-bookworm AS chef
 ENV PYO3_PYTHON=python3.11
+# rustup reads rust-toolchain.toml from any parent of the working directory, so
+# copying it in is what keeps every cargo call below on the repo's pinned
+# channel rather than on whatever the base image happens to ship.
+COPY rust-toolchain.toml /build/rust-toolchain.toml
+WORKDIR /build/litellm-rust
 RUN apt-get update \
     && apt-get install -y --no-install-recommends \
-        python3 python3-dev pkg-config libssl-dev clang \
+        python3 python3-dev python3-pip pkg-config libssl-dev clang \
     && rm -rf /var/lib/apt/lists/* \
     && cargo install cargo-chef --locked --version 0.1.77
-WORKDIR /build/litellm-rust
 
 # ---- Planner ----------------------------------------------------------------
 # Produce the dependency recipe from the rust workspace manifests + Cargo.lock.
@@ -43,6 +48,19 @@ RUN cargo chef cook --locked --release \
 COPY litellm-rust/ .
 RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config
 
+# The root pyproject builds with maturin against litellm-rust/crates/python-bridge,
+# so the wheel is built here, next to the crate sources and the cargo toolchain,
+# and the runtime stage installs the artifact instead of compiling anything.
+# litellm[proxy] pins litellm-enterprise and litellm-proxy-extras to the versions
+# in this repo, and those hit PyPI hours after every version bump merges, so both
+# wheels are built from the repo too instead of being resolved from PyPI.
+COPY pyproject.toml README.md LICENSE /build/
+COPY litellm/ /build/litellm/
+COPY enterprise/ /build/enterprise/
+COPY litellm-proxy-extras/ /build/litellm-proxy-extras/
+RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist \
+        /build /build/enterprise /build/litellm-proxy-extras
+
 # ---- Runtime ----------------------------------------------------------------
 # python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3
 # 3.11 ABI so the embedded interpreter links and imports cleanly.
@@ -56,11 +74,16 @@ RUN apt-get update \
 WORKDIR /app
 
 # Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so
-# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the
-# package + packaging metadata, then pip install the proxy extra.
-COPY pyproject.toml README.md LICENSE ./
-COPY litellm/ ./litellm/
-RUN pip install --no-cache-dir ".[proxy]"
+# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. The two
+# sibling wheels come from the builder as well, so the pins in litellm[proxy]
+# resolve against them and never wait on a PyPI publish.
+COPY --from=builder /build/dist/*.whl /tmp/wheels/
+RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \
+    && pip install --no-cache-dir \
+        /tmp/wheels/litellm_enterprise-*.whl \
+        /tmp/wheels/litellm_proxy_extras-*.whl \
+        "${wheel}[proxy]" \
+    && rm -rf /tmp/wheels
 
 # The compiled gateway binary (pure-Rust realtime hot path; Python is load-time
 # only).

litellm-rust/crates/ai-gateway/Dockerfile.dockerignore

@@ -9,19 +9,28 @@
 # Strategy: ignore everything, then re-include only what the build needs:
 #   - litellm/            (pip install . needs the full package + proxy reader)
 #   - litellm-rust/       (the rust workspace; Cargo.lock + crate sources)
-#   - pyproject.toml / README.md / LICENSE  (packaging metadata for pip install)
+#   - enterprise/         (litellm/proxy/enterprise symlinks into it; maturin walks it)
+#   - litellm-proxy-extras/ (built into a wheel alongside enterprise/ for litellm[proxy])
+#   - pyproject.toml / README.md / LICENSE  (packaging metadata for the wheel build)
+#   - rust-toolchain.toml (the pinned channel every cargo call in the build uses)
 *
 
 # --- re-include the build inputs ---
 !litellm/
 !litellm-rust/
+!enterprise/
+!litellm-proxy-extras/
 !pyproject.toml
+!rust-toolchain.toml
 !README.md
 !LICENSE
 
 # --- prune heavy / irrelevant subpaths back out of the re-included trees ---
 # Rust build artifacts (huge; regenerated in the builder).
 **/target/
+# Committed python distribution artifacts; the wheel build does not read them.
+enterprise/dist/
+litellm-proxy-extras/dist/
 # Python caches and compiled bytecode.
 **/__pycache__/
 **/*.pyc