Bump Ray 2.55.1 / CUDA 12.9.1 / torch 2.10 / vllm 0.18; HAProxy ingress

[praateekmahajan]

Summary

• Bump CUDA base image to 12.9.1 and align torch / torchaudio / torchvision to the 2.10 line (cu129 wheels). Move the pytorch wheel index to cu129.
• Bump Ray to 2.55.1 ([default,data] and [serve,llm] extras).
• Cap vllm<0.19 to keep the stack on the transformers<5 / huggingface-hub<1.0 path. Resolves to vLLM 0.18.1.
• Cap torchcodec<0.11 to match torch 2.10's ABI (torchcodec 0.11 is built for torch 2.11). Resolves to torchcodec 0.10.0+cu129.
• Enable Ray Serve's HAProxy ingress:
  • Install haproxy (built from source) + socat in the Curator image and set RAY_SERVE_HAPROXY_BINARY_PATH.
  • init_cluster opportunistically opts in via RAY_SERVE_ENABLE_HA_PROXY=1 when both binaries resolve on PATH. Both are required: Ray Serve uses socat to drive HAProxy's admin socket — without it, the controller's healthcheck silently returns False and trips a 5s timeout.
  • Also pins a free RAY_SERVE_HAPROXY_METRICS_PORT so multiple Curator clusters on a single host don't collide on HAProxy's prometheus bind.
• Extend BaseModelConfig.merge_runtime_envs to handle Ray's dict-form pip / uv runtime_env entries (previously list+list only). Backwards compatible — existing list-form call sites continue to produce list-form output.
• Anchor the Dynamo backend's runtime dir under Ray's session dir (<ray_temp_dir>/<session_name>/nemo_curator_dynamo_<short_id>/) so subprocess logs / manifests sit alongside Ray's own session logs. Pin FLASHINFER_WORKSPACE_BASE per-run so cubin artifacts compiled by a prior actor venv don't leak into a fresh session.

Test plan

• Image builds end-to-end via bash benchmarking/tools/build_docker.sh.
• Existing unit tests pass (Dynamo runtime_env merging, backend start/stop, BaseModelConfig).
• HAProxy probe path: init_cluster emits the env vars when both binaries are on PATH, falls back silently otherwise.
• Full nightly benchmark on this branch (4 GPUs, all entries).

Note on vLLM 0.19

We also explored bumping vLLM to 0.19.1, which transitively requires transformers>=5 and huggingface-hub>=1.0 (and follow-on Curator-side fixes for the boundary changes those introduce). We backed off because vLLM 0.19.1 isn't compatible with Ray 2.55.1 — ray.llm._infer_supports_vision crashes during LLMServer.start() for the model configs we hit. Sticking with vLLM 0.18.1 keeps the upgrade scope minimal until the upstream Ray ↔ vLLM 0.19 path stabilizes.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[praateekmahajan]

/claude review
@greptileai

[praateekmahajan]

/ok to test 9e5b992

[claude]

Code looks good — dependency bumps are clean, HAProxy gating logic is solid (the socat check is a nice catch), and the runtime-dir anchor under Ray's session dir makes sense for the benchmarking runner's session_latest pickup.

One comment on test coverage for the new _merge_package_runtime_env dict-form merge paths.

[greptile-apps]

Greptile Summary

This PR bumps the core stack (CUDA 12.9.1, torch 2.10, Ray 2.55.1, vLLM capped at <0.19) and enables Ray Serve's HAProxy ingress mode, anchors the Dynamo backend's runtime dir under Ray's session dir for log preservation, and extends BaseModelConfig.merge_runtime_envs to handle dict-form pip/uv runtime_env entries (needed because the new DYNAMO_VLLM_RUNTIME_ENV uses a dict-form uv block with flash-attn rebuild flags and an 1800 s setup timeout).

• The config key in merge_runtime_envs is handled by the initial shallow {**base, **override} — unlike env_vars, pip, and uv it is not deep-merged. A user passing any config dict (e.g. {\"config\": {\"pip_check\": True}}) will silently replace the Curator-owned {\"setup_timeout_seconds\": 1800}, causing the flash-attn from-source rebuild to hit Ray's 600 s default and fail with RuntimeEnvSetupError.

Confidence Score: 4/5

Safe to merge with attention to the config merge gap; all other changes are well-structured and tested

One P1 finding: the config key in merge_runtime_envs is shallow-merged, which silently drops the critical 1800s setup timeout whenever a user includes any config dict in their runtime_env. Everything else — HAProxy integration, runtime_dir anchoring, version bumps, vLLM API adaptation — is clean and well-tested.

nemo_curator/core/serve/base.py — the config merge gap affects the new 1800s flash-attn timeout

Important Files Changed

Filename	Overview
nemo_curator/core/serve/base.py	Extends merge_runtime_envs to handle dict-form pip/uv entries; config key is shallow-merged, silently overwriting the new 1800s timeout if a user passes any config dict
docker/common/install_haproxy.sh	New script that builds HAProxy from source with SHA256 verification; fetches from GitHub mirror due to haproxy.org cert expiry
nemo_curator/core/serve/dynamo/backend.py	Moves runtime_dir from tempfile to Ray session dir; removes cleanup on stop (intentional — logs are preserved for benchmarking runner)
nemo_curator/core/serve/dynamo/vllm.py	Upgrades DYNAMO_VLLM_RUNTIME_ENV to dict-form uv with flash-attn rebuild flags and 1800s timeout; adds FLASHINFER_WORKSPACE_BASE env per worker
nemo_curator/core/utils.py	Adds opportunistic HAProxy ingress opt-in via shutil.which checks before spawning Ray; correctly sets env vars before Popen
nemo_curator/stages/text/embedders/vllm.py	Adapts LLM.embed() call to vLLM 0.18 API — moves truncate_prompt_tokens into tokenization_kwargs
pyproject.toml	Bumps ray to >=2.55.1, pins vllm<0.19, upgrades torch/torchaudio/torchvision to 2.10.0/0.25.0, pins torchcodec~=0.10.0 for ABI compatibility
benchmarking/runner/ray_cluster.py	Extends session copy exclusion set to also skip runtime_resources/ (venvs) to avoid bloating benchmark archives
tests/core/serve/test_runtime_env.py	New tests covering dict-form uv merging, flash-attn rebuild flags, and config timeout regression guards

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["init_cluster()"] --> B{haproxy + socat on PATH?}
    B -- yes --> C[Set RAY_SERVE_ENABLE_HA_PROXY=1
RAY_SERVE_HAPROXY_METRICS_PORT=free_port]
    B -- no --> D[Log debug: fallback to Python proxy]
    C --> E[Popen ray start]
    D --> E
    E --> F["DynamoBackend.start()"]
    F --> G["ray.init(namespace=nemo_curator_dynamo)"]
    G --> H[Anchor runtime_dir under Ray session dir]
    H --> I[_deploy_and_healthcheck]
    I -- success --> J[Backend running]
    I -- failure --> K[_teardown_actors_and_pgs
runtime_dir retained for logs]
    L["merge_runtime_envs(base, override)"] --> M["{**base, **override}"]
    M --> N[deep-merge env_vars]
    M --> O[deep-merge pip/uv via _merge_package_runtime_env]
    M --> P["config: shallow only ⚠️"]

Loading

Comments Outside Diff (1)

1. nemo_curator/core/serve/base.py, line 53-64 (link)

   config key shallow-merged — user dict silently drops the 1800 s timeout

   merge_runtime_envs deep-merges env_vars, pip, and uv, but the config key falls through the initial {**base, **override} shallow merge. A user who includes {"config": {"setup_timeout_seconds": X}} in their model's runtime_env will completely replace the Curator-owned {"setup_timeout_seconds": 1800} introduced in this PR, causing the flash-attn from-source rebuild — which the PR explicitly notes takes ~15 min — to hit Ray's 600 s default and fail with RuntimeEnvSetupError.

   Consider deep-merging config the same way env_vars is handled:

   base_config = base.get("config", {})
   override_config = override.get("config", {})
   if base_config or override_config:
       merged["config"] = {**base_config, **override_config}

   A matching test case (user passes {"config": {"pip_check": True}}) would guard against future regressions.

_{Reviews (6): Last reviewed commit: "Bump Dynamo runtime_env setup_timeout to..." | Re-trigger Greptile}

[praateekmahajan]

/ok to test 972245d

[praateekmahajan]

/ok to test 1133e3d

[praateekmahajan]

/ok to test 78c883d

[praateekmahajan]

/claude review
@greptileai

[praateekmahajan]

/ok to test 632f341

[praateekmahajan]

/ok to test bb58caf

[praateekmahajan]

/ok to test 7dc9341

[ayushdg]

I guess going forward the override here always has to strictly match what vllm supports?

[praateekmahajan]

Yup that's my understanding, unless we start doing sub venvs

[praateekmahajan]

/ok to test d8535ff