vllm_colocated.py

"""Time-share one GPU between a trainer and an isolated, sleeping vLLM engine.

The trainer environment does not import vLLM. A private authenticated Unix
socket carries control messages; CPU checkpoint files carry policy weights.
Only the generator's own subprocess and process group are managed here.
"""

from __future__ import annotations

import gc
import hashlib
import json
import os
from pathlib import Path
import secrets
import signal
import subprocess
import sys
import tempfile
import time
from typing import Any

import torch


def _release_cuda_cache(device: torch.device) -> None:
    if device.type == "cuda":
        with torch.cuda.device(device):
            torch.cuda.synchronize()
            torch.cuda.empty_cache()


class _TrainerOffload:
    """Move tensor storage, preserving Parameters, optimizer keys, and steps."""

    def __init__(self, policy, optimizer, device: torch.device):
        self.policy = policy
        self.optimizer = optimizer
        self.device = device
        self.model_tensors: list[tuple[torch.Tensor, torch.device]] = []
        self.state_tensors: list[tuple[Any, Any, torch.device]] = []

    def _offload_state(self, container):
        keys = list(container) if isinstance(container, dict) else range(len(container))
        for key in keys:
            value = container[key]
            if isinstance(value, torch.Tensor):
                self.state_tensors.append((container, key, value.device))
                container[key] = value.to("cpu")
            elif isinstance(value, (dict, list)):
                self._offload_state(value)

    def offload(self) -> None:
        self.optimizer.zero_grad(set_to_none=True)
        self.policy.zero_grad(set_to_none=True)
        try:
            # Module.to may replace Parameter objects under torch future flags.
            # Assign only storage so the optimizer keeps its original references.
            for tensor in [*self.policy.parameters(), *self.policy.buffers()]:
                self.model_tensors.append((tensor, tensor.device))
                tensor.data = tensor.detach().to("cpu")
            for state in self.optimizer.state.values():
                self._offload_state(state)
            gc.collect()
            _release_cuda_cache(self.device)
        except BaseException:
            self.restore()
            raise

    def restore(self) -> None:
        for tensor, original_device in self.model_tensors:
            tensor.data = tensor.detach().to(original_device)
        for container, key, original_device in self.state_tensors:
            container[key] = container[key].to(original_device)
        self.model_tensors.clear()
        self.state_tensors.clear()
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)


def _tensor_signature(tensor: torch.Tensor) -> dict[str, Any]:
    """Hash every weight byte using at most 2 MiB of CPU tensor workspace."""
    detached = tensor.detach()
    flat = detached.view(-1) if detached.is_contiguous() else None
    chunk_elements = max(1, (2 * 1024 * 1024) // tensor.element_size())
    digest = hashlib.sha256()
    for start in range(0, tensor.numel(), chunk_elements):
        end = min(start + chunk_elements, tensor.numel())
        if flat is not None:
            chunk = flat[start:end]
        else:
            # Avoid materializing a whole contiguous copy of a strided tensor.
            offsets = torch.arange(start, end, device=tensor.device)
            chunk = torch.take(detached, offsets)
        cpu_bytes = chunk.cpu().contiguous().view(torch.uint8).numpy()
        digest.update(memoryview(cpu_bytes))
    return {"shape": list(tensor.shape), "dtype": str(tensor.dtype), "sha256": digest.hexdigest()}


def _weight_manifest(state: dict[str, torch.Tensor]) -> tuple[dict, str]:
    manifest = {name: _tensor_signature(value) for name, value in sorted(state.items())}
    fingerprint = hashlib.sha256(
        json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode(),
    ).hexdigest()
    return manifest, fingerprint


def _stage_policy_weights(policy: torch.nn.Module, path: Path) -> dict[str, Any]:
    state = {name: parameter.detach() for name, parameter in policy.named_parameters()}
    if not state or any(value.device.type != "cpu" for value in state.values()):
        raise RuntimeError("Policy parameters must be on CPU before weight staging")
    manifest, fingerprint = _weight_manifest(state)
    # The unique path is never published to the worker until save/fsync finish.
    with path.open("xb") as output:
        torch.save({"state_dict": state, "manifest": manifest, "fingerprint": fingerprint}, output)
        output.flush()
        os.fsync(output.fileno())
    return {"fingerprint": fingerprint, "source_parameter_count": len(state)}


def _loaded_parameter_slice(name, state, parameters):
    if name in parameters:
        return parameters[name]
    # vLLM packs OLMo2's Q/K/V and gate/up projections along the first axis.
    for pieces, packed in [(("q", "k", "v"), "qkv"), (("gate", "up"), "gate_up")]:
        for index, piece in enumerate(pieces):
            marker = f"{piece}_proj."
            if marker not in name:
                continue
            target_name = name.replace(marker, f"{packed}_proj.")
            if target_name not in parameters:
                continue
            preceding = [name.replace(marker, f"{prior}_proj.") for prior in pieces[:index]]
            offset = sum(state[prior].shape[0] for prior in preceding)
            return parameters[target_name].narrow(0, offset, state[name].shape[0])
    raise RuntimeError(f"Weight coverage cannot map source parameter {name!r}")


def _load_staged_weights(model, path, expected_fingerprint: str) -> dict[str, Any]:
    package = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
    state = package["state_dict"]
    manifest, fingerprint = _weight_manifest(state)
    if (
        fingerprint != expected_fingerprint
        or fingerprint != package["fingerprint"]
        or manifest != package["manifest"]
    ):
        raise RuntimeError("Staged weight fingerprint mismatch")
    loaded = set(model.load_weights(state.items()))
    # Alias-aware coverage accounts for tied embedding/output weights.
    parameters = dict(model.named_parameters(remove_duplicate=False))
    loaded_ids = {id(parameters[name]) for name in loaded if name in parameters}
    missing = [name for name, value in model.named_parameters() if id(value) not in loaded_ids]
    unexpected = loaded - parameters.keys()
    if missing or unexpected:
        raise RuntimeError(f"Weight coverage mismatch: missing={missing}, unexpected={sorted(unexpected)}")
    for name, expected in manifest.items():
        actual = _loaded_parameter_slice(name, state, parameters)
        if _tensor_signature(actual) != expected:
            raise RuntimeError(f"Loaded weight fingerprint mismatch for {name}")
    return {
        "fingerprint": fingerprint,
        "source_parameter_count": len(state),
        "loaded_parameter_count": len(loaded),
        "loaded_parameter_names": sorted(loaded),
    }


class ColocatedWeightLoader:
    """Methods mixed into vLLM's GPU worker through worker_extension_cls."""

    def load_colocated_policy_weights(self, path: str, fingerprint: str):
        result = _load_staged_weights(self.model_runner.model, path, fingerprint)
        torch.cuda.synchronize()
        return result


def _worker_environment(device: torch.device) -> dict[str, str]:
    env = os.environ.copy()
    env.pop("PYTORCH_ALLOC_CONF", None)
    env.pop("PYTORCH_CUDA_ALLOC_CONF", None)
    visible = env.get("CUDA_VISIBLE_DEVICES")
    index = device.index or 0
    if device.type == "cuda":
        env["CUDA_VISIBLE_DEVICES"] = visible.split(",")[index].strip() if visible else str(index)
    else:
        env["CUDA_VISIBLE_DEVICES"] = ""
    env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
    env["TOKENIZERS_PARALLELISM"] = "false"
    env["HF_HUB_OFFLINE"] = "1"
    root = str(Path(__file__).resolve().parents[1])
    env["PYTHONPATH"] = os.pathsep.join(filter(None, [root, env.get("PYTHONPATH")]))
    return env


def _worker_sigterm(signum, _frame) -> None:
    # Raising on the controller's main thread reaches _worker_main's shutdown
    # finally even when the trainer died during a long generation request.
    raise SystemExit(128 + signum)


def _arm_parent_death_signal(expected_parent_pid: int) -> None:
    """Have Linux notify this worker if its original trainer process dies."""
    if sys.platform != "linux":
        raise RuntimeError("The vLLM parent-death guard requires Linux")
    if os.getppid() != expected_parent_pid:
        raise RuntimeError("The trainer parent exited before the worker guard was installed")

    import ctypes

    signal.signal(signal.SIGTERM, _worker_sigterm)
    libc = ctypes.CDLL(None, use_errno=True)
    prctl = libc.prctl
    prctl.argtypes = [ctypes.c_int, ctypes.c_ulong, ctypes.c_ulong, ctypes.c_ulong, ctypes.c_ulong]
    prctl.restype = ctypes.c_int
    if prctl(1, signal.SIGTERM, 0, 0, 0) != 0:  # PR_SET_PDEATHSIG
        error_number = ctypes.get_errno()
        raise OSError(error_number, "Could not install the vLLM parent-death signal")
    # The parent may have exited between the first check and prctl, before
    # Linux started delivering its death signal to this worker.
    if os.getppid() != expected_parent_pid:
        _worker_sigterm(signal.SIGTERM, None)


class _VLLMProcess:
    def __init__(self, *, python, work_dir, device, engine_kwargs, timeout=1800):
        self.python = str(python)
        self.work_dir = Path(work_dir)
        self.device = device
        self.engine_kwargs = engine_kwargs
        self.timeout = timeout
        self.process = None
        self.connection = None
        self.runtime = None
        self.log_file = None
        self.log_path = None
        self._request_id = 0

    def _receive(self, timeout, expected_id=0):
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            if self.connection.poll(min(0.5, max(0.0, deadline - time.monotonic()))):
                response = self.connection.recv()
                if not response.get("ok"):
                    raise RuntimeError(f"vLLM worker failed: {response.get('error')}; log: {self.log_path}")
                if response.get("id") != expected_id:
                    # An interrupted request can leave its reply in the socket.
                    # In particular, generation output is never a sleep ACK.
                    continue
                return response.get("result")
            if self.process.poll() is not None:
                raise RuntimeError(f"vLLM worker exited {self.process.returncode}; log: {self.log_path}")
        raise TimeoutError(f"Timed out waiting for vLLM worker; log: {self.log_path}")

    def start(self):
        if self.process is not None:
            return
        from multiprocessing.connection import Client

        self.work_dir.mkdir(parents=True, exist_ok=True)
        self.runtime = tempfile.TemporaryDirectory(prefix="cs336-vllm-")
        runtime_path = Path(self.runtime.name)
        socket_path = runtime_path / "engine.sock"
        authkey = secrets.token_bytes(32)
        config_path = runtime_path / "config.json"
        config_path.write_text(json.dumps({
            "socket": str(socket_path), "authkey": authkey.hex(),
            "engine_kwargs": self.engine_kwargs, "parent_pid": os.getpid(),
        }), encoding="utf-8")
        config_path.chmod(0o600)
        self.log_path = self.work_dir / f"vllm-server-{runtime_path.name}.log"
        self.log_file = self.log_path.open("ab", buffering=0)
        self.process = subprocess.Popen(
            [self.python, "-u", "-m", "cs336_alignment.vllm_colocated", "--worker", str(config_path)],
            env=_worker_environment(self.device),
            cwd=Path(__file__).resolve().parents[1],
            stdin=subprocess.DEVNULL, stdout=self.log_file, stderr=subprocess.STDOUT,
            start_new_session=True,
        )
        deadline = time.monotonic() + self.timeout
        while time.monotonic() < deadline:
            if self.process.poll() is not None:
                raise RuntimeError(f"vLLM worker startup exited {self.process.returncode}; log: {self.log_path}")
            try:
                self.connection = Client(str(socket_path), family="AF_UNIX", authkey=authkey)
                self._receive(max(1.0, deadline - time.monotonic()))
                return
            except (FileNotFoundError, ConnectionRefusedError):
                time.sleep(0.1)
        raise TimeoutError(f"vLLM worker startup timed out; log: {self.log_path}")

    def request(self, operation, *, timeout=None, **payload):
        if self.connection is None or self.process.poll() is not None:
            raise RuntimeError("vLLM worker is not connected")
        self._request_id += 1
        self.connection.send({"id": self._request_id, "operation": operation, **payload})
        return self._receive(
            self.timeout if timeout is None else timeout, expected_id=self._request_id,
        )

    def close(self):
        process = self.process
        if process is not None and process.poll() is None:
            try:
                if self.connection is not None:
                    self.request("close", timeout=30)
                process.wait(timeout=30)
            except (OSError, RuntimeError, TimeoutError, EOFError, subprocess.TimeoutExpired):
                # The group was created by this Popen call: never match names
                # or terminate an unrelated trainer/server already on the GPU.
                if process.poll() is None:
                    os.killpg(process.pid, signal.SIGTERM)
                    try:
                        process.wait(timeout=10)
                    except subprocess.TimeoutExpired:
                        os.killpg(process.pid, signal.SIGKILL)
                        process.wait(timeout=10)
        exit_code = process.poll() if process is not None else None
        if exit_code not in (None, 0):
            # A crashed controller may leave its own spawned workers behind.
            # Its dedicated process group cannot contain the trainer.
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
        if self.connection is not None:
            self.connection.close()
            self.connection = None
        if self.log_file is not None:
            self.log_file.close()
            self.log_file = None
        if self.runtime is not None:
            self.runtime.cleanup()
            self.runtime = None
        if exit_code not in (None, 0):
            raise RuntimeError(f"vLLM worker exited with code {exit_code}; log: {self.log_path}")


def _generate_with_engine(engine, sampling_params_class, prompts, sampling):
    outputs = engine.generate(prompts, sampling_params_class(**sampling), use_tqdm=False)
    if len(outputs) != len(prompts):
        raise RuntimeError("vLLM returned the wrong number of prompt outputs")
    texts, lengths = [], []
    for output in outputs:
        completions = sorted(output.outputs, key=lambda completion: completion.index)
        if [completion.index for completion in completions] != list(range(sampling["n"])):
            raise RuntimeError("vLLM returned incomplete or duplicate completion indices")
        texts.extend(completion.text for completion in completions)
        lengths.extend(len(completion.token_ids) for completion in completions)
    return {"texts": texts, "token_lengths": lengths}


def _worker_main(config_path: str) -> None:
    from multiprocessing.connection import Listener

    config = json.loads(Path(config_path).read_text(encoding="utf-8"))
    _arm_parent_death_signal(config["parent_pid"])
    engine = None
    request_id = 0
    with Listener(config["socket"], family="AF_UNIX", authkey=bytes.fromhex(config["authkey"])) as listener:
        with listener.accept() as connection:
            try:
                from vllm import LLM, SamplingParams

                engine = LLM(**config["engine_kwargs"])
                engine.sleep(level=2)
                connection.send({"ok": True, "id": 0, "result": {"pid": os.getpid()}})
                while True:
                    request = connection.recv()
                    request_id = request["id"]
                    operation = request["operation"]
                    if operation == "close":
                        engine.llm_engine.engine_core.shutdown(timeout=30)
                        engine = None
                        connection.send({"ok": True, "id": request_id, "result": {}})
                        break
                    if operation == "wake_weights":
                        engine.wake_up(tags=["weights"])
                        result = {}
                    elif operation == "load_weights":
                        result = engine.collective_rpc(
                            "load_colocated_policy_weights",
                            args=(request["path"], request["fingerprint"]),
                        )
                    elif operation == "reset_prefix_cache":
                        engine.reset_prefix_cache()
                        result = {}
                    elif operation == "wake_kv":
                        engine.wake_up(tags=["kv_cache"])
                        result = {}
                    elif operation == "generate":
                        sampling = dict(request["sampling"], n=request["n"], seed=request["seed"])
                        result = _generate_with_engine(engine, SamplingParams, request["prompts"], sampling)
                    elif operation == "sleep":
                        engine.sleep(level=2)
                        result = {}
                    else:
                        raise ValueError(f"Unknown worker operation {operation!r}")
                    connection.send({"ok": True, "id": request_id, "result": result})
            except BaseException as error:
                import traceback

                traceback.print_exc()
                try:
                    connection.send({"ok": False, "id": request_id, "error": repr(error)})
                except (OSError, EOFError):
                    pass
                raise
            finally:
                if engine is not None:
                    engine.llm_engine.engine_core.shutdown(timeout=30)


class SingleGPUVLLMGenerator:
    def __init__(
        self, policy, optimizer, model_id, vllm_python, work_dir, device="cuda:0", *,
        max_tokens, temperature, top_p, stop, seed, max_num_seqs=32,
        max_model_len=1024, gpu_memory_utilization=0.65, staging_dir=None,
    ):
        self.policy = policy
        self.optimizer = optimizer
        self.device = torch.device(device)
        self.vllm_python = str(vllm_python)
        self.work_dir = Path(work_dir)
        self.staging_dir = Path(staging_dir) if staging_dir else (
            Path("/dev/shm") if Path("/dev/shm").is_dir() else Path(tempfile.gettempdir())
        )
        self.sampling = {
            "max_tokens": max_tokens, "temperature": temperature, "top_p": top_p,
            "stop": [stop] if isinstance(stop, str) else stop,
            "include_stop_str_in_output": stop is not None,
        }
        self.engine_kwargs = {
            "model": str(model_id), "dtype": str(next(policy.parameters()).dtype).removeprefix("torch."),
            "seed": seed, "tensor_parallel_size": 1, "enable_sleep_mode": True,
            "worker_extension_cls": "cs336_alignment.vllm_colocated.ColocatedWeightLoader",
            "gpu_memory_utilization": gpu_memory_utilization,
            "max_num_seqs": max_num_seqs, "max_model_len": max_model_len,
            "max_num_batched_tokens": max(max_model_len, max_num_seqs),
            "enable_prefix_caching": True, "load_format": "dummy", "enforce_eager": True,
        }
        self._server = None
        self._server_ready = False
        self._closed = False
        self._dirty = True
        self.last_generation_metadata: dict[str, Any] = {}

    def sync(self) -> None:
        self._dirty = True

    def generate(self, prompts: list[str], n: int, seed: int) -> tuple[list[str], list[int]]:
        if self._closed:
            raise RuntimeError("SingleGPUVLLMGenerator is closed")
        if n <= 0:
            raise ValueError("n must be positive")
        if not prompts:
            return [], []
        started = time.perf_counter()
        metadata: dict[str, Any] = {"sync_requested": self._dirty}
        offload = _TrainerOffload(self.policy, self.optimizer, self.device)
        result, error = None, None
        try:
            before = time.perf_counter()
            offload.offload()
            metadata["offload_seconds"] = time.perf_counter() - before
            with tempfile.TemporaryDirectory(prefix="grpo-vllm-weights-", dir=self.staging_dir) as stage:
                before = time.perf_counter()
                path = Path(stage) / "policy.pt"
                manifest = _stage_policy_weights(self.policy, path)
                metadata.update(manifest)
                metadata["stage_seconds"] = time.perf_counter() - before
                before = time.perf_counter()
                if self._server is None:
                    self._server = _VLLMProcess(
                        python=self.vllm_python, work_dir=self.work_dir, device=self.device,
                        engine_kwargs=self.engine_kwargs,
                    )
                    self._server.start()
                    self._server_ready = True
                metadata["server_start_seconds"] = time.perf_counter() - before
                metadata["worker_pid"] = self._server.process.pid
                before = time.perf_counter()
                self._server.request("wake_weights")
                reports = self._server.request("load_weights", path=str(path), fingerprint=manifest["fingerprint"])
                if len(reports) != 1 or reports[0]["fingerprint"] != manifest["fingerprint"]:
                    raise RuntimeError("vLLM worker did not confirm the current policy fingerprint")
                metadata["loaded_parameter_count"] = reports[0]["loaded_parameter_count"]
                path.unlink()
                self._server.request("reset_prefix_cache")
                self._server.request("wake_kv")
                metadata["weight_load_seconds"] = time.perf_counter() - before
                self._dirty = False
                before = time.perf_counter()
                result = self._server.request(
                    "generate", prompts=prompts, n=n, seed=seed, sampling=self.sampling,
                )
                metadata["generation_seconds"] = time.perf_counter() - before
                if len(result["texts"]) != len(prompts) * n or len(result["token_lengths"]) != len(result["texts"]):
                    raise RuntimeError("vLLM returned the wrong completion count")
        except BaseException as caught:
            error = caught
        finally:
            before = time.perf_counter()
            try:
                if self._server_ready:
                    self._server.request("sleep")
            except BaseException as caught:
                error = error or caught
            metadata["sleep_seconds"] = time.perf_counter() - before
            if error is not None and self._server is not None:
                # Confirm the private engine has released its GPU allocation
                # before attempting to restore trainer tensors to that GPU.
                try:
                    self.close()
                except BaseException as close_error:
                    process = self._server.process
                    if process is not None and process.poll() is None:
                        raise RuntimeError(
                            "vLLM worker could not be stopped; trainer remains on CPU"
                        ) from close_error
                    # The child is dead, so trainer restoration is safe even
                    # though its nonzero exit must remain visible to the caller.
                    error.add_note(f"Worker cleanup also reported: {close_error}")
                    self._closed = True
                    self._server_ready = False
            before = time.perf_counter()
            offload.restore()
            metadata["restore_seconds"] = time.perf_counter() - before
            metadata["total_seconds"] = time.perf_counter() - started
            self.last_generation_metadata = metadata
        if error is not None:
            raise error
        return result["texts"], result["token_lengths"]

    def close(self) -> None:
        if not self._closed:
            if self._server is not None:
                self._server.close()
            self._closed = True
            self._server_ready = False

    def __enter__(self):
        return self

    def __exit__(self, *_exc):
        self.close()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--worker", required=True)
    _worker_main(parser.parse_args().worker)