diff --git a/Dockerfile.cuda b/Dockerfile.cuda index 5c3e9c1e5a..35e8d7f283 100644 --- a/Dockerfile.cuda +++ b/Dockerfile.cuda @@ -18,12 +18,42 @@ ENV DEBIAN_FRONTEND=noninteractive ENV TZ=Etc/UTC RUN apt-get update && apt-get install -y --no-install-recommends --force-yes \ build-essential \ + autoconf \ + automake \ + libtool \ + pkg-config \ + ca-certificates \ curl \ sudo \ git \ ninja-build \ + libnuma-dev \ + libnl-3-dev \ + libnl-route-3-dev \ + libibverbs-dev \ + librdmacm-dev \ && apt-get clean autoclean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* +ARG UCX_VERSION=1.19.1 +RUN git clone --depth 1 --branch v${UCX_VERSION} https://github.com/openucx/ucx.git /tmp/ucx \ + && cd /tmp/ucx \ + && ./autogen.sh \ + && ./configure \ + --prefix=/opt/ucx \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-cma \ + --enable-devel-headers \ + --enable-mt \ + --with-verbs \ + --with-cuda=/usr/local/cuda \ + --with-ze=no \ + && make -j"$(nproc)" \ + && make install \ + && rm -rf /tmp/ucx + # Download the latest installer ADD https://astral.sh/uv/install.sh /uv-installer.sh @@ -49,7 +79,7 @@ COPY examples /app/examples COPY benchmarks/scripts /app/benchmarks/scripts RUN --mount=type=cache,target=/app/.cache/uv \ - uv sync --extra flash-attn --extra flash-attn-3 --extra flash-attn-cute --extra envs --extra gpt-oss --group mamba-ssm --locked --no-dev + uv sync --extra flash-attn --extra flash-attn-3 --extra flash-attn-cute --extra envs --extra gpt-oss --extra modelexpress --group mamba-ssm --locked --no-dev # arm64: build flash-attn from source, fix namespace conflicts, apply workarounds ARG TARGETARCH @@ -74,8 +104,12 @@ RUN apt-get update && apt-get install -y \ net-tools \ curl \ vim \ + libnuma1 \ + libnl-3-200 \ + libnl-route-3-200 \ libibverbs1 \ ibverbs-providers \ + librdmacm1 \ && apt-get clean autoclean \ && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* @@ -96,6 +130,7 @@ ENV PATH="/usr/local/bin:$PATH" WORKDIR /app # Copy the application from the builder COPY --from=builder --chown=appuser:appuser /app /app +COPY --from=builder /opt/ucx /opt/ucx # Copy and set up entrypoint script COPY --chown=appuser:appuser scripts/docker-entrypoint.sh /app/docker-entrypoint.sh @@ -107,6 +142,8 @@ RUN rm /app/.venv/bin/python3.12 && ln -s /usr/local/bin/python /app/.venv/bin/p # Place executables in the environment at the front of the path ENV PATH="/app/.venv/bin:$PATH" +ENV UCX_HOME=/opt/ucx +ENV LD_LIBRARY_PATH="/opt/ucx/lib:/opt/ucx/lib/ucx${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" # HuggingFace Hub timeouts (defaults are 10s which causes issues on slow networks) ENV HF_HUB_ETAG_TIMEOUT=500 diff --git a/pyproject.toml b/pyproject.toml index fdf35fb747..d6c860f524 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,6 +99,9 @@ disagg = [ gpt-oss = [ "kernels", ] +modelexpress = [ + "modelexpress==0.3.0", +] quack = [ "quack-kernels>=0.4.1", ] @@ -178,6 +181,22 @@ override-dependencies = [ "openenv-core", ] +# ModelExpress 0.3.0 publishes protobuf<6 metadata, but its generated proto is +# compatible with protobuf 6. prime-sandboxes requires protobuf>=6.31.1; keep +# this capped to the validated protobuf major. +[[tool.uv.dependency-metadata]] +name = "modelexpress" +version = "0.3.0" +requires-dist = [ + "grpcio>=1.66.2", + "huggingface_hub>=0.20.0", + "nixl[cu12]", + "numpy>=1.24.0", + "protobuf>=5.27.0,<7.0.0", + "pydantic>=2.0.0", + "torch>=2.6.0", +] + [tool.uv.exclude-newer-package] # we want latest vllm, remove next patch vllm = false diff --git a/uv.lock b/uv.lock index 26b1eb3b29..e9937a56f0 100644 --- a/uv.lock +++ b/uv.lock @@ -77,6 +77,11 @@ overrides = [ { name = "transformers", git = "https://github.com/huggingface/transformers.git?rev=c1c3424" }, ] +[[manifest.dependency-metadata]] +name = "modelexpress" +version = "0.3.0" +requires-dist = ["grpcio>=1.66.2", "huggingface-hub>=0.20.0", "nixl[cu12]", "numpy>=1.24.0", "protobuf>=5.27.0,<7.0.0", "pydantic>=2.0.0", "torch>=2.6.0"] + [[package]] name = "absl-py" version = "2.4.0" @@ -2832,6 +2837,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/61/26/c7aea197f1719f31d0dd686eb4475982fe9efd7668ce259cb52b62c676b6/model_hosting_container_standards-0.1.15-py3-none-any.whl", hash = "sha256:849e08c4732203ee861c8c24966b4e916ea4420fa324b430f7f74a1e1fe8811a", size = 125418, upload-time = "2026-05-05T18:22:27.819Z" }, ] +[[package]] +name = "modelexpress" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "grpcio", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, + { name = "huggingface-hub", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, + { name = "nixl", extra = ["cu12"], marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, + { name = "numpy", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, + { name = "protobuf", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, + { name = "pydantic", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, + { name = "torch", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/b5/12f940a41940fd83b9e50ce46124695b68a80feab02612779f8fb584db09/modelexpress-0.3.0-py3-none-any.whl", hash = "sha256:6f93d8de74903f6c11ebf9b4621fa02e3c493a5e05207c38fb2c98395a774618", size = 44333, upload-time = "2026-04-17T20:17:29.478Z" }, +] + [[package]] name = "more-itertools" version = "11.0.2" @@ -2954,6 +2976,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bb/a4/10e80e623790d3fb070e966b36bced009419d483c92ae0df6645230f606f/nixl-0.10.1-py3-none-any.whl", hash = "sha256:616465673dae5180d296525a03237af4cd5f2c00c3228d185bc06dbe621509b7", size = 6680, upload-time = "2026-03-03T19:55:44.848Z" }, ] +[package.optional-dependencies] +cu12 = [ + { name = "nixl-cu12", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, + { name = "nixl-cu12", version = "0.10.1", source = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl" }, marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, +] + [[package]] name = "nixl-cu12" version = "0.10.1" @@ -4019,6 +4047,9 @@ flash-attn-cute = [ gpt-oss = [ { name = "kernels", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, ] +modelexpress = [ + { name = "modelexpress", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, +] quack = [ { name = "quack-kernels", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-9-verifiers-openenv' and extra == 'group-9-verifiers-policy')" }, ] @@ -4070,6 +4101,7 @@ requires-dist = [ { name = "mini-swe-agent-plus", marker = "extra == 'envs'", editable = "deps/research-environments/environments/mini_swe_agent_plus" }, { name = "mini-swe-agent-plus-rlm", marker = "extra == 'envs'", editable = "deps/research-environments/environments/mini_swe_agent_plus_rlm" }, { name = "mmlu-pro", marker = "extra == 'envs'", editable = "deps/research-environments/environments/mmlu_pro" }, + { name = "modelexpress", marker = "extra == 'modelexpress'", specifier = "==0.3.0" }, { name = "nixl", marker = "extra == 'disagg'" }, { name = "nixl-cu12", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl" }, { name = "numpy", specifier = ">=2.2.6" }, @@ -4116,7 +4148,7 @@ requires-dist = [ { name = "wandb", specifier = ">=0.26.1" }, { name = "wiki-search", marker = "extra == 'envs'", editable = "deps/verifiers/environments/wiki_search" }, ] -provides-extras = ["flash-attn", "flash-attn-3", "flash-attn-cute", "envs", "disagg", "gpt-oss", "quack", "all"] +provides-extras = ["flash-attn", "flash-attn-3", "flash-attn-cute", "envs", "disagg", "gpt-oss", "modelexpress", "quack", "all"] [package.metadata.requires-dev] dev = [