[Rust Frontend] Replace external protoc with pure Rust lib protox (#52892)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Bugen Zhao
2026-08-20 18:43:51 -07:00
committed by GitHub
co-authored by OpenAI Codex Nick Hill
parent 91a893de64
commit 83c5d59209
14 changed files with 181 additions and 125 deletions
-1
View File
@@ -25,7 +25,6 @@ run_all_patterns:
- "rust/"
- "rust-toolchain.toml"
- "tools/build_rust.py"
- "tools/install_protoc.sh"
- "tools/install_torchcodec_rocm.sh"
- "tests/vllm_test_utils/"
- "vllm/envs.py"
-5
View File
@@ -10,11 +10,6 @@ set -euo pipefail
# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
git fetch --tags --force origin
# The Rust frontend build needs protoc.
if ! command -v protoc >/dev/null 2>&1; then
brew install protobuf
fi
# upload-nightly-wheels.sh expects exactly one wheel.
rm -rf artifacts/dist
mkdir -p artifacts/dist
+2 -2
View File
@@ -15,7 +15,7 @@ set -euo pipefail
DEFAULT_REPO_SLUG="vllm-project/vllm"
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils"
DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh rust-toolchain.toml tests/vllm_test_utils"
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_METADATA_VERSION="3"
@@ -23,7 +23,7 @@ DEFAULT_CI_BASE_METADATA_VERSION="3"
# local-source stages rather than unreachable remote-fetch alternatives.
DEFAULT_ROCM_CSRC_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt pyproject.toml setup.py CMakeLists.txt cmake csrc vllm/envs.py vllm/__init__.py tools/build_rust.py"
DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm build_vllm_dependencies rocm-triton-kernels csrc-build"
DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore .git_archival.txt pyproject.toml requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py tools/install_protoc.sh build_rust.sh"
DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore .git_archival.txt pyproject.toml requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py build_rust.sh"
DEFAULT_ROCM_RUST_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm vllm-version rust_toolchain_input_0 rust-toolchain-input rust_input_0 rust-input rust-toolchain rust-build"
# Docker's 128-character tag limit minus the longest cache prefix
# ("csrc-rocm-branch-" and "rust-rocm-branch-", both 17 characters).
@@ -21,7 +21,6 @@ export CARGO_HOME="${CARGO_HOME:-$HOME/.cargo}"
export RUSTUP_HOME="${RUSTUP_HOME:-$HOME/.rustup}"
export PATH="$CARGO_HOME/bin:$PATH"
PROTOC_VERSION="${PROTOC_VERSION:-31.1}"
CARGO_BINSTALL_VERSION="${CARGO_BINSTALL_VERSION:-1.20.1}"
UV_VERSION="${UV_VERSION:-0.11.28}"
PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
@@ -34,34 +33,6 @@ log_section() {
echo "--- $*"
}
install_protoc() {
local arch
case "$(uname -m)" in
x86_64)
arch="x86_64"
;;
aarch64|arm64)
arch="aarch_64"
;;
*)
echo "Unsupported protoc architecture: $(uname -m)" >&2
return 1
;;
esac
local url="https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-${arch}.zip"
local tmp_dir
tmp_dir="$(mktemp -d)"
log_section "Installing protoc ${PROTOC_VERSION}"
curl -L --proto '=https' --tlsv1.2 -sSf "$url" -o "$tmp_dir/protoc.zip"
mkdir -p "$CARGO_HOME/bin"
unzip -q "$tmp_dir/protoc.zip" bin/protoc 'include/*' -d "$CARGO_HOME"
chmod +x "$CARGO_HOME/bin/protoc"
rm -rf "$tmp_dir"
protoc --version
}
rust_toolchain() {
awk -F '"' '/channel[[:space:]]*=/ { print $2; exit }' rust-toolchain.toml
}
@@ -186,7 +157,6 @@ run_tests() {
--no-fail-fast
}
install_protoc
install_rust_toolchain
case "$MODE" in
+3 -6
View File
@@ -243,21 +243,18 @@ ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
#################### RUST BUILD IMAGE ####################
# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel
# build stage doesn't need the rust toolchain, protoc, or the rust source.
# build stage doesn't need the rust toolchain or the rust source.
# This stage reuses the Python environment from base and runs in parallel with
# csrc-build/extensions-build.
FROM base AS rust-build
ARG USE_SCCACHE
ARG SCCACHE_ENDPOINT
# Install native tools needed only for Rust/protoc builds.
RUN dnf install -y --setopt=install_weak_deps=False make unzip \
# Install native tools needed only for Rust builds.
RUN dnf install -y --setopt=install_weak_deps=False make \
&& dnf clean all \
&& rm -rf /var/cache/dnf
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
WORKDIR /workspace
COPY requirements/build/rust.txt requirements/build/rust.txt
+2 -5
View File
@@ -127,19 +127,16 @@ RUN echo 'ulimit -c 0' >> ~/.bashrc
######################### RUST BUILD IMAGE #########################
# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the wheel build
# stage doesn't need the rust toolchain or protoc. This stage runs in parallel
# stage doesn't need the rust toolchain. This stage runs in parallel
# with the main vllm-build stage.
FROM base AS rust-build
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
build-essential unzip python3 python3-pip \
build-essential python3 python3-pip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
WORKDIR /workspace
COPY requirements/build/rust.txt requirements/build/rust.txt
+5 -13
View File
@@ -138,13 +138,12 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
> /vllm-version.txt
# -----------------------
# Rust/protoc toolchain inputs
# Rust toolchain inputs
#
# Keep these separate from Rust source inputs so Rust source-only changes do not
# invalidate the long-lived ci_base image that carries the toolchain.
FROM scratch AS rust_toolchain_input_0
COPY rust-toolchain.toml /rust-toolchain-input/vllm/rust-toolchain.toml
COPY tools/install_protoc.sh /rust-toolchain-input/vllm/tools/install_protoc.sh
FROM base AS rust_toolchain_input_1
ARG VLLM_REPO
@@ -155,7 +154,6 @@ RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-toolchain-inpu
&& git sparse-checkout init --no-cone \
&& git sparse-checkout set \
rust-toolchain.toml \
tools/install_protoc.sh \
&& git checkout FETCH_HEAD
FROM rust_toolchain_input_${REMOTE_VLLM} AS rust-toolchain-input
@@ -195,21 +193,17 @@ RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-input/vllm \
FROM rust_input_${REMOTE_VLLM} AS rust-input
# -----------------------
# Rust/protoc toolchain
# Rust toolchain
FROM base AS rust-toolchain
ENV CARGO_HOME=/root/.cargo
ENV RUSTUP_HOME=/root/.rustup
ENV PATH=${CARGO_HOME}/bin:${PATH}
# protoc is used by tonic-build/prost-build.
RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
ca-certificates curl unzip \
ca-certificates curl \
&& rm -rf /var/lib/apt/lists/*
COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/rust-toolchain.toml /tmp/rust-toolchain.toml
RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \
&& if ! command -v rustup >/dev/null 2>&1; then \
@@ -224,7 +218,7 @@ RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"
# -----------------------
# Rust build stage
# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
# don't need the rust toolchain or protoc.
# don't need the rust toolchain.
FROM rust-toolchain AS rust-build
ARG COMMON_WORKDIR
ARG USE_SCCACHE
@@ -780,12 +774,10 @@ ENV CARGO_HOME=/root/.cargo
ENV RUSTUP_HOME=/root/.rustup
ENV PATH=${CARGO_HOME}/bin:${PATH}
# Bake the Rust/protoc toolchain into ci_base so Rust build/test steps do not
# Bake the Rust toolchain into ci_base so Rust build/test steps do not
# need to fetch it on each per-commit image build.
COPY --from=rust-toolchain /root/.cargo /root/.cargo
COPY --from=rust-toolchain /root/.rustup /root/.rustup
COPY --from=rust-toolchain /usr/local/bin/protoc /usr/local/bin/protoc
COPY --from=rust-toolchain /usr/local/include/google /usr/local/include/google
# Update rdma-core to support latest rocshmem.
ARG DEEPEP_NIC
+1 -9
View File
@@ -141,18 +141,10 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
# -----------------------
# Rust build stage
# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
# don't need the rust toolchain or protoc.
# don't need the rust toolchain.
FROM fetch_vllm AS rust-build
ARG COMMON_WORKDIR
# protoc is used by tonic-build/prost-build.
RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
ca-certificates curl unzip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
ENV CARGO_BUILD_JOBS=4
+2 -5
View File
@@ -1,17 +1,14 @@
######################### RUST BUILD IMAGE #########################
# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main image
# doesn't need the rust toolchain or protoc. Runs in parallel with vllm-base.
# doesn't need the rust toolchain. Runs in parallel with vllm-base.
FROM ubuntu:22.04 AS rust-build
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip python3 python3-pip \
ca-certificates curl git build-essential python3 python3-pip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
WORKDIR /workspace
COPY requirements/build/rust.txt requirements/build/rust.txt
+152 -3
View File
@@ -385,6 +385,12 @@ version = "1.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
[[package]]
name = "beef"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3a8241f3ebb85c056b509d4327ad0358fbbba6ffb340bf388f26350aeda225b1"
[[package]]
name = "bit-set"
version = "0.5.3"
@@ -661,7 +667,7 @@ dependencies = [
"encode_unicode",
"libc",
"once_cell",
"unicode-width",
"unicode-width 0.2.2",
"windows-sys 0.61.2",
]
@@ -1499,7 +1505,7 @@ version = "0.2.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df"
dependencies = [
"unicode-width",
"unicode-width 0.2.2",
]
[[package]]
@@ -2005,7 +2011,7 @@ checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb"
dependencies = [
"console",
"portable-atomic",
"unicode-width",
"unicode-width 0.2.2",
"unit-prefix",
"web-time",
]
@@ -2264,6 +2270,72 @@ version = "0.4.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
[[package]]
name = "logos"
version = "0.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ff472f899b4ec2d99161c51f60ff7075eeb3097069a36050d8037a6325eb8154"
dependencies = [
"logos-derive 0.15.1",
]
[[package]]
name = "logos"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eb2c55a318a87600ea870ff8c2012148b44bf18b74fad48d0f835c38c7d07c5f"
dependencies = [
"logos-derive 0.16.1",
]
[[package]]
name = "logos-codegen"
version = "0.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "192a3a2b90b0c05b27a0b2c43eecdb7c415e29243acc3f89cc8247a5b693045c"
dependencies = [
"beef",
"fnv",
"lazy_static",
"proc-macro2",
"quote",
"regex-syntax",
"rustc_version",
"syn 2.0.117",
]
[[package]]
name = "logos-codegen"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "58b3ffaa284e1350d017a57d04ada118c4583cf260c8fb01e0fe28a2e9cf8970"
dependencies = [
"fnv",
"proc-macro2",
"quote",
"regex-automata",
"regex-syntax",
"syn 2.0.117",
]
[[package]]
name = "logos-derive"
version = "0.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "605d9697bcd5ef3a42d38efc51541aa3d6a4a25f7ab6d1ed0da5ac632a26b470"
dependencies = [
"logos-codegen 0.15.1",
]
[[package]]
name = "logos-derive"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "52d3a9855747c17eaf4383823f135220716ab49bea5fbea7dd42cc9a92f8aa31"
dependencies = [
"logos-codegen 0.16.1",
]
[[package]]
name = "macro_rules_attribute"
version = "0.2.2"
@@ -2375,6 +2447,28 @@ version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b"
[[package]]
name = "miette"
version = "7.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5f98efec8807c63c752b5bd61f862c165c115b0a35685bdcfd9238c7aeb592b7"
dependencies = [
"cfg-if",
"miette-derive",
"unicode-width 0.1.14",
]
[[package]]
name = "miette-derive"
version = "7.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "db5b29714e950dbb20d5e6f74f9dcec4edbcc1067bb7f8ed198c097b8c1a818b"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "mimalloc"
version = "0.1.52"
@@ -3116,6 +3210,18 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "prost-reflect"
version = "0.16.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "01b80ea363c31af2de2b92e3c07ed1156628f7838c4afb4df75ee78a37fedbd1"
dependencies = [
"logos 0.16.1",
"miette",
"prost",
"prost-types",
]
[[package]]
name = "prost-types"
version = "0.14.3"
@@ -3125,6 +3231,33 @@ dependencies = [
"prost",
]
[[package]]
name = "protox"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4f25a07a73c6717f0b9bbbd685918f5df9815f7efba450b83d9c9dea41f0e3a1"
dependencies = [
"bytes",
"miette",
"prost",
"prost-reflect",
"prost-types",
"protox-parse",
"thiserror 2.0.18",
]
[[package]]
name = "protox-parse"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "072eee358134396a4643dff81cfff1c255c9fbd3fb296be14bdb6a26f9156366"
dependencies = [
"logos 0.15.1",
"miette",
"prost-types",
"thiserror 2.0.18",
]
[[package]]
name = "pulldown-cmark"
version = "0.13.3"
@@ -3623,6 +3756,15 @@ version = "2.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d"
[[package]]
name = "rustc_version"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92"
dependencies = [
"semver",
]
[[package]]
name = "rustfft"
version = "6.4.1"
@@ -5284,6 +5426,12 @@ version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "da36089a805484bcccfffe0739803392c8298778a2d2f09febf76fac5ad9025b"
[[package]]
name = "unicode-width"
version = "0.1.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af"
[[package]]
name = "unicode-width"
version = "0.2.2"
@@ -5733,6 +5881,7 @@ dependencies = [
"openssl",
"prost",
"prost-types",
"protox",
"rmp-serde",
"rmpv",
"serde",
+1
View File
@@ -75,6 +75,7 @@ prometheus-client = "0.24.0"
prometheus-client-derive-encode = "0.5.0"
prost = "0.14.3"
prost-types = "0.14.3"
protox = "0.9.1"
pyo3 = "0.28.3"
pythonize = "0.28.0"
rand = "0.9.2"
+1
View File
@@ -53,6 +53,7 @@ vllm-metrics.workspace = true
vllm-text.workspace = true
[build-dependencies]
protox.workspace = true
tonic-prost-build.workspace = true
[dev-dependencies]
+12 -9
View File
@@ -3,19 +3,22 @@
fn main() -> Result<(), Box<dyn std::error::Error>> {
let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
let proto_dir = format!("{manifest_dir}/../../proto");
let proto_dir = std::path::PathBuf::from(format!("{manifest_dir}/../../proto"));
let protos = [
proto_dir.join("control.proto"),
proto_dir.join("inference.proto"),
];
for proto in &protos {
println!("cargo:rerun-if-changed={}", proto.display());
}
let file_descriptor_set = protox::compile(&protos, [&proto_dir])?;
tonic_prost_build::configure()
.build_server(true)
.build_client(true)
.protoc_arg("--experimental_allow_proto3_optional") // be compatible with old compilers
.compile_protos(
&[
format!("{proto_dir}/control.proto"),
format!("{proto_dir}/inference.proto"),
],
&[proto_dir],
)?;
.compile_fds(file_descriptor_set)?;
Ok(())
}
-37
View File
@@ -1,37 +0,0 @@
#!/usr/bin/env bash
set -euo pipefail
# Install a pinned protoc binary from upstream GitHub releases.
#
# Distro protobuf-compiler packages vary widely in version (e.g.
# AlmaLinux/RHEL 8 ships protoc 3.5, predating the
# --experimental_allow_proto3_optional flag the rust frontend's build.rs
# passes), so we pin the protoc version here instead.
#
# Override the version via the PROTOC_VERSION env var.
# Requires: curl, unzip, root privileges.
if [[ $(id -u) -ne 0 ]]; then
echo "Must be run as root" >&2
exit 1
fi
VERSION="${PROTOC_VERSION:-34.2}"
ARCH="$(uname -m)"
case "${ARCH}" in
# protoc release archives use "aarch_64" (with an underscore), not
# "aarch64". Don't "fix" this.
aarch64|arm64) URL_ARCH="aarch_64" ;;
x86_64|amd64) URL_ARCH="x86_64" ;;
*) echo "Unsupported arch for protoc binary: ${ARCH}" >&2; exit 1 ;;
esac
URL="https://github.com/protocolbuffers/protobuf/releases/download/v${VERSION}/protoc-${VERSION}-linux-${URL_ARCH}.zip"
TMPDIR="$(mktemp -d)"
trap 'rm -rf "${TMPDIR}"' EXIT
echo "Downloading: ${URL}"
curl -fsSL -o "${TMPDIR}/protoc.zip" "${URL}"
unzip -q -o "${TMPDIR}/protoc.zip" -d /usr/local
echo "Installed $(protoc --version)"