| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Driver-side helpers to run the inference fleet as a standard ``vllm serve`` |
| 4 | process for the sharded-RDT examples. |
| 5 | |
| 6 | The trainer reaches the server's weight-sync control plane over the RLHF HTTP |
| 7 | routes (``HTTPVLLMWeightSyncClient``); generation uses ``/v1/completions``. RDT |
| 8 | still needs the workers to be RayExecutorV2 tensor-transport actors sharing the |
| 9 | trainer's Ray cluster, so the server is launched with the Ray v2 executor and |
| 10 | inherits this process's env (``address=auto``). These helpers run only on the |
| 11 | driver — they are not shipped to or imported by the Ray actors. |
| 12 | """ |
| 13 | |
| 14 | import json |
| 15 | import os |
| 16 | import shutil |
| 17 | import subprocess |
| 18 | import sys |
| 19 | import time |
| 20 | |
| 21 | import ray |
| 22 | import requests |
| 23 | |
| 24 | |
| 25 | def launch_vllm_serve( |
| 26 | model: str, |
| 27 | *, |
| 28 | tensor_parallel_size: int = 1, |
| 29 | data_parallel_size: int = 1, |
| 30 | enable_expert_parallel: bool = False, |
| 31 | port: int = 8000, |
| 32 | gpu_memory_utilization: float = 0.7, |
| 33 | extra_args: list[str] | None = None, |
| 34 | ) -> subprocess.Popen: |
| 35 | """Start ``vllm serve`` for the RDT inference fleet. |
| 36 | |
| 37 | ``VLLM_SERVER_DEV_MODE`` exposes the RLHF weight-sync routes; the Ray v2 |
| 38 | executor makes the workers tensor-transport actors (RDT's data plane). The |
| 39 | child inherits our env, so it joins the same Ray cluster (RAY_ADDRESS below) |
| 40 | and picks up the editable vLLM install via the venv interpreter. |
| 41 | """ |
| 42 | vllm_bin = os.path.join(os.path.dirname(sys.executable), "vllm") |
| 43 | if not os.path.exists(vllm_bin): |
| 44 | vllm_bin = shutil.which("vllm") or "vllm" |
| 45 | cmd = [ |
| 46 | vllm_bin, |
| 47 | "serve", |
| 48 | model, |
| 49 | "--port", |
| 50 | str(port), |
| 51 | "--enforce-eager", |
| 52 | "--load-format", |
| 53 | "dummy", |
| 54 | "--gpu-memory-utilization", |
| 55 | str(gpu_memory_utilization), |
| 56 | "--distributed-executor-backend", |
| 57 | "ray", |
| 58 | "--weight-transfer-config", |
| 59 | json.dumps({"backend": "sharded_rdt"}), |
| 60 | ] |
| 61 | if tensor_parallel_size > 1: |
| 62 | cmd += ["--tensor-parallel-size", str(tensor_parallel_size)] |
| 63 | # Only engage the DP-ray backend when actually data-parallel: DP=1 with |
| 64 | # data_parallel_backend=ray forces vLLM's DP-placement path, which fails. |
| 65 | if data_parallel_size > 1: |
| 66 | cmd += [ |
| 67 | "--data-parallel-size", |
| 68 | str(data_parallel_size), |
| 69 | "--data-parallel-backend", |
| 70 | "ray", |
| 71 | ] |
| 72 | if enable_expert_parallel: |
| 73 | cmd.append("--enable-expert-parallel") |
| 74 | if extra_args: |
| 75 | cmd += extra_args |
| 76 | env = dict( |
| 77 | os.environ, |
| 78 | VLLM_SERVER_DEV_MODE="1", |
| 79 | VLLM_USE_RAY_V2_EXECUTOR_BACKEND="1", |
| 80 | ) |
| 81 | # Pin the child to THIS Ray cluster. vLLM's ray DP backend reaches Ray |
| 82 | # through a bare ``ray.init()``, which STARTS a cluster rather than |
| 83 | # attaching, so wherever we started Ray ourselves (a bare CI runner) the |
| 84 | # workers would come up in a second cluster and never see the producer |
| 85 | # actors -- the failure is an actor lookup that cannot succeed. A bare |
| 86 | # ``ray.init()`` does honour RAY_ADDRESS, so hand it over explicitly. |
| 87 | if ray.is_initialized(): |
| 88 | env["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address |
| 89 | # Let the server's Ray workers use the ambient interpreter/install rather |
| 90 | # than a workspace snapshot working_dir. On an Anyscale dev workspace the |
| 91 | # RAY_RUNTIME_ENV_HOOK injects a git snapshot as the workers' working_dir, |
| 92 | # which shadows an editable vLLM install (its compiled .so is gitignored, so |
| 93 | # the snapshot lacks it). No-op off Anyscale (the var is unset there). |
| 94 | env.pop("RAY_RUNTIME_ENV_HOOK", None) |
| 95 | print(f"[serve] launching: {' '.join(cmd)}", flush=True) |
| 96 | return subprocess.Popen(cmd, env=env) |
| 97 | |
| 98 | |
| 99 | def wait_for_server( |
| 100 | endpoint: str, proc: subprocess.Popen, timeout: float = 1800 |
| 101 | ) -> None: |
| 102 | """Block until ``/health`` returns 200, failing fast if the server exits.""" |
| 103 | deadline = time.time() + timeout |
| 104 | while time.time() < deadline: |
| 105 | if proc.poll() is not None: |
| 106 | raise RuntimeError(f"vllm serve exited early (code {proc.returncode})") |
| 107 | try: |
| 108 | if requests.get(f"{endpoint}/health", timeout=5).status_code == 200: |
| 109 | print("[serve] server is healthy", flush=True) |
| 110 | return |
| 111 | except requests.RequestException: |
| 112 | pass |
| 113 | time.sleep(3) |
| 114 | raise RuntimeError("vllm serve did not become healthy in time") |
| 115 | |
| 116 | |
| 117 | def http_generate( |
| 118 | endpoint: str, model: str, prompts: list[str], max_tokens: int = 16 |
| 119 | ) -> list[str]: |
| 120 | """Greedy ``/v1/completions`` for each prompt; returns the generated texts.""" |
| 121 | outs = [] |
| 122 | for p in prompts: |
| 123 | r = requests.post( |
| 124 | f"{endpoint}/v1/completions", |
| 125 | json={ |
| 126 | "model": model, |
| 127 | "prompt": p, |
| 128 | "max_tokens": max_tokens, |
| 129 | "temperature": 0, |
| 130 | }, |
| 131 | timeout=120, |
| 132 | ) |
| 133 | r.raise_for_status() |
| 134 | outs.append(r.json()["choices"][0]["text"]) |
| 135 | return outs |
| 136 | |
| 137 | |
| 138 | def pause_generation(endpoint: str, mode: str = "abort") -> None: |
| 139 | requests.post(f"{endpoint}/pause", params={"mode": mode}, timeout=60) |
| 140 | |
| 141 | |
| 142 | def resume_generation(endpoint: str) -> None: |
| 143 | requests.post(f"{endpoint}/resume", timeout=60) |
| 144 | |
| 145 | |
| 146 | def shutdown_server(proc: subprocess.Popen) -> None: |
| 147 | print("[serve] shutting down vllm serve", flush=True) |
| 148 | proc.terminate() |
| 149 | try: |
| 150 | proc.wait(timeout=30) |
| 151 | except subprocess.TimeoutExpired: |
| 152 | proc.kill() |