Triton on Ray: Ray Serve's gRPC proxy in front of Triton

[kondratyevd]

What

Triton on Ray, with Ray Serve's gRPC proxy carrying Triton's protocol: every worker pod runs NVIDIA's Triton Inference Server, pointed at the models CMSSW ships on CVMFS (four --model-repository directories inside CMSSW_17_0_0_pre2, an explicit list of ten models, the cluster's read-only cvmfs claim), beside a Ray container; Serve's gRPC proxy is handed Triton's own generated servicer, so it accepts exactly Triton's RPCs, counts each one, load-balances it, and forwards it to the Triton in the chosen pod. Ray Serve sizes the deployment from those requests and the Ray autoscaler adds a GPU pod for each replica with nowhere to go. Triton does all of the inference.

clients ──▶ sonic-ray-serve (LoadBalancer, geddes-private-pool) :8001  Triton gRPC
              ▼  Serve gRPC proxy — counted, balanced, autoscaled
   worker pod × 1…4:  ray-worker (raylet `triton: 1`, proxy, forwarder replica)
                      triton     (1 GPU, 4 CPU, 16G, /cvmfs ro)    ◀ localhost:8001 (HTTP moved to 8100)
   head pod:          Serve controller + Ray autoscaler (no GPU, 0 CPUs for work)

• All of Triton, none of ours. The ten CMSSW models (deepmet, deeptau_2018v2p5, six ParticleNet variants, particlenet_PT, unifiedparticletransformer_AK4_V01) straight from CVMFS — nothing uploaded, no model manager; a new CMSSW release is a path change in the values. Every backend, config.pbtxt semantics, dynamic batching, repository index — it is Triton. CMSSW's TritonClient and tritonclient.grpc point at sonic-ray-serve:8001 like at any Triton endpoint. The one RPC not forwarded is ModelStreamInfer (bidirectional; Serve's proxy carries unary and server-streaming only; CMSSW uses unary ModelInfer).
• The glue is ~100 lines, apps/ray/sonic-ray/chart/files/sonic_ray/serve_app.py: a Serve deployment that, for every unary RPC of GRPCInferenceService (derived from the generated servicer at import), forwards the protobuf message to Triton on localhost and returns Triton's answer. It blocks on ServerReady before becoming ready and polls ServerLive as its health check. Nothing in it parses a request.
• One replica per Triton pod. Each worker advertises a triton: 1 resource; each replica claims one. That pins a replica next to its Triton, leaves a pod without a replica idle (reclaimable), and makes a replica without a pod the pending request that grows the group. One setting, replicas: {min: 1, max: 4}, bounds Serve and the worker group alike (target_ongoing_requests: 16; the group's own floor is 0 since Serve's minimum keeps pods alive); the chart refuses to render if a pod has other than one GPU or if the grace period would cut Triton's drain short.
• No custom image. Official Ray (2.52.0-py312-cpu, no Ray process touches a GPU) via the Docker Hub proxy cache, official Triton (26.04-py3, the tag supersonic runs). Triton's Python stubs (tritonclient==2.48.0, the last release whose protobuf-4 stubs match the Ray image, --no-deps) are pip-installed into an emptyDir on PYTHONPATH by an init container on every pod — Serve's proxies import the servicer outside any runtime_env, so replicas-only installation would not do.
• Ports. Both containers share the pod's network namespace, so Triton's HTTP moves to 8100: Ray Serve's HTTP proxy holds 8000 and KubeRay probes it there, and Triton otherwise exits with Socket '0.0.0.0:8000' already in use. Triton's gRPC keeps 8001 (what the forwarder dials). Both are chart values, and the chart refuses to render if they collide with Ray's ports or disagree with triton.args.
• No KEDA, no Prometheus in the loop. nv_* and ray_serve_* metrics are exposed for the AF Prometheus under release="sonic-ray".

Layout

path	what
apps/ray/operator/	kuberay-operator 1.7.0, namespaced to cms
apps/ray/sonic-ray/chart/	RayService (Triton sidecar per worker, Serve gRPC proxy with Triton's servicer, pip init container), the ConfigMap carrying the forwarder, two metrics Services; refusals for values that cannot work
apps/ray/sonic-ray/values.yaml	the AF release: Triton image/args/resources, the CVMFS model repositories and load list, the CMS GPU nodes, the private address pool
tests/sonic_ray/	source-level checks of the forwarder (every RPC but the stream, readiness gated on Triton, no request parsing)
tests/manifests/test_ray.py	rendered chart: the Triton container is the one in the values, models from the read-only cvmfs claim with host-to-container propagation, one replica per pod, Serve ≤ worker bounds, stubs installed where the proxies run, chart refusals

Also: validate-manifests.sh now helm templates charts sourced from this repository (the RayService kind has no kubeconform schema), README badges for the two new components, mypy.ini scope.

Verified

• pytest tests green; validate-manifests.sh green; ruff/mypy/prettier clean.
• Not run on the cluster (no VPN from here). The first deploy has to confirm: the init container's pip from the GPU nodes, the CVMFS mount inside the Triton container (first load pulls the models over the network; the startup probe allows four minutes), Serve's gRPC proxy loading tritonclient's servicer against the image's grpcio 1.74 / protobuf 4.25, a CMSSW client end to end, autoscaling, and the LoadBalancer allocation. Nothing needs to be built first — merging is enough for Flux to reconcile it.

History of this PR

Three earlier shapes were tried and dropped in review: a custom Ray image with ONNX Runtime and a KServe v2 HTTP server; the same on the stock image; then plain-JSON endpoints. Each served fewer models or fewer clients than Triton already does. This shape keeps Triton and reduces our part to the forwarder.

Follow-ups (deliberately out of scope)

1. HTTP through Serve, if HTTP clients turn up: a second forwarder for Triton's HTTP, or Serve's HTTP proxy passing bytes to Triton's port.
2. A Grafana panel on ray_serve_* next to the SuperSONIC nv_* dashboards.
3. Retuning target_ongoing_requests against real CMSSW load once it is measured.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!