code/extras/vllm_rollout/rollout_worker.py

align / code /extras /vllm_rollout /rollout_worker.py

Add files using upload-large-folder tool

63f41a6 verified 15 days ago

10.1 kB

"""vLLM rollout worker: serves align/rollout/protocol.py requests from a spool dir.
OPTIONAL EXTRA -- NOT part of the `align` package. The default GRPO rollout is
the dependency-free native sampler (align/rollout/native.py); this worker is
the high-throughput alternative, selected with `backend='spool'`.
RUNS UNDER THE ISOLATED ROLLOUT VENV (.venv-rollout: vllm 0.26.0 / cu130),
never the training venv -- the 2026-08-22 sovereignty amendment confines
vLLM behind this process boundary. See extras/vllm_rollout/README.md. Launch it via
`align.rollout.client.spawn_worker`, which sets the REQUIRED environment:
    VLLM_USE_FLASHINFER_SAMPLER=0   (flashinfer's JIT ninja build fails on
                                     this box's toolchain -- measured
                                     2026-08-23, rollout pilot)
    TMPDIR / XDG_CACHE_HOME under /apps (the box's no-swap root is 13GB)
    CUDA_VISIBLE_DEVICES=<the one dedicated engine GPU>
Weight sync = filesystem: when a request's `weights_dir` differs from the
engine's current one, the engine is torn down and rebuilt from the new
checkpoint (~60s measured). The trainer amortizes that by syncing every
`sync_every` optimizer steps; the rho correction owns the resulting
off-policyness.
Failure contract: any per-request exception is written back as a RESULT
with `error` set (traceback tail in extras) so the trainer fails fast at
collect() instead of sitting at a timeout; the worker keeps serving.
"""

import argparse
import os
import sys
import time
import traceback


def _die(msg: str) -> None:
    print(f"rollout_worker refusing to start: {msg}", file=sys.stderr, flush=True)
    sys.exit(2)


# ENV hygiene BEFORE importing anything heavy: this venv has no align guard,
# and a cache or TMPDIR off /apps is the exact class of fault that wedged
# this box on 2026-08-12 (root is 13GB, /tmp is a noexec tmpfs).
for _var in ("TMPDIR", "XDG_CACHE_HOME"):
    _val = os.environ.get(_var, "")
    if not _val.startswith("/apps"):
        _die(f"{_var}={_val!r} does not resolve under /apps; set it before launch")

# <repo>/extras/vllm_rollout/rollout_worker.py -> <repo>: three levels up.
_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, _REPO_ROOT)

# align.rollout.protocol is stdlib-only BY CONTRACT (see its docstring); it is
# the single align import permitted here, taken from the repo tree directly so
# align never needs installing into the rollout venv. (`import align` and
# `import align.rollout` are both lazy/stdlib-only for exactly this reason.)
from align.rollout.protocol import (  # noqa: E402
    RolloutResult,
    read_request,
    spool_dirs,
    touch_heartbeat,
    write_result,
)


def _extract_sampled_logprobs(token_ids, sample_logprobs, request_id):
    """Per-position logprob of the SAMPLED token.
    vllm 0.26 API shape (verified against .venv-rollout's vllm/outputs.py and
    vllm/logprobs.py): CompletionOutput.logprobs is SampleLogprobs -- a
    sequence over positions where each element is dict[int, Logprob] (or the
    list-compatible FlatLogprobs view of the same); with
    SamplingParams(logprobs=0) each position carries the sampled token's own
    entry, so position t resolves as logprobs[t][token_ids[t]].logprob.
    """
    if sample_logprobs is None:
        raise RuntimeError(
            f"request {request_id!r}: engine returned no logprobs despite "
            "SamplingParams(logprobs=0); the rho correction cannot run without "
            "pi^infer_old, so this result would be unusable."
        )
    out = []
    for t, tok in enumerate(token_ids):
        pos = sample_logprobs[t]
        if pos is None or tok not in pos:
            raise RuntimeError(
                f"request {request_id!r}: position {t} lacks the sampled token "
                f"{tok} in its logprob dict (keys: "
                f"{sorted(pos) if pos else pos}); refusing to substitute."
            )
        out.append(float(pos[tok].logprob))
    return out


def serve(args) -> int:
    # Heartbeat BEFORE the vllm import: the import alone costs ~30s, and the
    # first smoke run proved a trainer canary can race it -- the client's
    # startup grace plus this early beat make "booting" visibly distinct from
    # "never started".
    spool_dirs(args.spool)
    touch_heartbeat(args.spool)
    from vllm import LLM, SamplingParams

    req_dir, _ = spool_dirs(args.spool)
    engine = None
    loaded_weights_dir = None
    last_work = time.time()
    print(f"rollout_worker serving spool={args.spool} tp={args.tp}", flush=True)

    while True:
        touch_heartbeat(args.spool)
        pending = sorted(
            (os.path.join(req_dir, f) for f in os.listdir(req_dir) if f.endswith(".json")),
            key=os.path.getmtime,
        )
        if not pending:
            if args.idle_exit_s and time.time() - last_work > args.idle_exit_s:
                print(f"idle {args.idle_exit_s}s; exiting", flush=True)
                return 0
            time.sleep(args.poll_s)
            continue

        for path in pending:
            request_id = os.path.basename(path)[: -len(".json")]
            reloaded = False
            try:
                req = read_request(path)
                if engine is None or req.weights_dir != loaded_weights_dir:
                    if engine is not None:
                        # vllm's own teardown first, then drop the reference;
                        # the fresh constructor re-measures free VRAM.
                        del engine
                        engine = None
                        import gc

                        gc.collect()
                        import torch

                        torch.cuda.empty_cache()
                    t0 = time.perf_counter()
                    engine = LLM(
                        model=req.weights_dir,
                        dtype="bfloat16",
                        tensor_parallel_size=args.tp,
                        gpu_memory_utilization=args.gpu_mem_util,
                    )
                    loaded_weights_dir = req.weights_dir
                    reloaded = True
                    print(
                        f"engine loaded from {req.weights_dir} in {time.perf_counter() - t0:.1f}s",
                        flush=True,
                    )

                sp = SamplingParams(
                    n=req.n_samples_per_prompt,
                    temperature=req.temperature,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    seed=req.seed,
                    logprobs=0,  # sampled token's own logprob per position
                )
                t1 = time.perf_counter()
                outs = engine.generate(req.prompts, sp)
                gen_s = time.perf_counter() - t1

                completions, logprobs, texts, prompt_token_ids = [], [], [], []
                for o in outs:
                    completions.append([list(map(int, c.token_ids)) for c in o.outputs])
                    logprobs.append(
                        [
                            _extract_sampled_logprobs(c.token_ids, c.logprobs, request_id)
                            for c in o.outputs
                        ]
                    )
                    texts.append([c.text for c in o.outputs])
                    # THE ENGINE'S OWN prompt tokenization, one list per prompt
                    # in request order. Completion ids continue from exactly
                    # this sequence; trainer-side assembly must concatenate
                    # these + completion ids and NEVER re-tokenize the prompt
                    # text (re-tokenization is not identity -- the alignment
                    # defect class the 2026-08-23 audit closed).
                    prompt_token_ids.append(list(map(int, o.prompt_token_ids)))

                write_result(
                    args.spool,
                    RolloutResult(
                        request_id=request_id,
                        completions=completions,
                        logprobs=logprobs,
                        texts=texts,
                        weights_dir=req.weights_dir,
                        engine_reloaded=reloaded,
                        generate_seconds=gen_s,
                        extras={"prompt_token_ids": prompt_token_ids},
                    ),
                )
                print(
                    f"served {request_id}: {sum(len(c) for c in completions)} "
                    f"completions in {gen_s:.1f}s (reloaded={reloaded})",
                    flush=True,
                )
            except Exception as e:  # noqa: BLE001 -- the trainer must fail fast, not time out
                tb = traceback.format_exc()
                print(f"request {request_id} FAILED: {e}\n{tb}", file=sys.stderr, flush=True)
                write_result(
                    args.spool,
                    RolloutResult(
                        request_id=request_id,
                        completions=[],
                        logprobs=[],
                        texts=[],
                        weights_dir=loaded_weights_dir or "",
                        engine_reloaded=reloaded,
                        generate_seconds=0.0,
                        error=str(e),
                        extras={"traceback_tail": tb[-2000:]},
                    ),
                )
            finally:
                # Consumed = removed; the result file is the record.
                if os.path.exists(path):
                    os.remove(path)
                last_work = time.time()


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--spool", required=True)
    p.add_argument("--gpu-mem-util", type=float, default=0.85)
    p.add_argument("--poll-s", type=float, default=2.0)
    p.add_argument("--tp", type=int, default=1)
    p.add_argument("--idle-exit-s", type=float, default=0.0, help="0 = serve forever")
    return serve(p.parse_args())


if __name__ == "__main__":
    sys.exit(main())

	"""vLLM rollout worker: serves align/rollout/protocol.py requests from a spool dir.

	OPTIONAL EXTRA -- NOT part of the `align` package. The default GRPO rollout is
	the dependency-free native sampler (align/rollout/native.py); this worker is
	the high-throughput alternative, selected with `backend='spool'`.

	RUNS UNDER THE ISOLATED ROLLOUT VENV (.venv-rollout: vllm 0.26.0 / cu130),
	never the training venv -- the 2026-08-22 sovereignty amendment confines
	vLLM behind this process boundary. See extras/vllm_rollout/README.md. Launch it via
	`align.rollout.client.spawn_worker`, which sets the REQUIRED environment:

	VLLM_USE_FLASHINFER_SAMPLER=0 (flashinfer's JIT ninja build fails on
	this box's toolchain -- measured
	2026-08-23, rollout pilot)
	TMPDIR / XDG_CACHE_HOME under /apps (the box's no-swap root is 13GB)
	CUDA_VISIBLE_DEVICES=<the one dedicated engine GPU>

	Weight sync = filesystem: when a request's `weights_dir` differs from the
	engine's current one, the engine is torn down and rebuilt from the new
	checkpoint (~60s measured). The trainer amortizes that by syncing every
	`sync_every` optimizer steps; the rho correction owns the resulting
	off-policyness.

	Failure contract: any per-request exception is written back as a RESULT
	with `error` set (traceback tail in extras) so the trainer fails fast at
	collect() instead of sitting at a timeout; the worker keeps serving.
	"""

	import argparse
	import os
	import sys
	import time
	import traceback


	def _die(msg: str) -> None:
	print(f"rollout_worker refusing to start: {msg}", file=sys.stderr, flush=True)
	sys.exit(2)


	# ENV hygiene BEFORE importing anything heavy: this venv has no align guard,
	# and a cache or TMPDIR off /apps is the exact class of fault that wedged
	# this box on 2026-08-12 (root is 13GB, /tmp is a noexec tmpfs).
	for _var in ("TMPDIR", "XDG_CACHE_HOME"):
	_val = os.environ.get(_var, "")
	if not _val.startswith("/apps"):
	_die(f"{_var}={_val!r} does not resolve under /apps; set it before launch")

	# <repo>/extras/vllm_rollout/rollout_worker.py -> <repo>: three levels up.
	_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
	sys.path.insert(0, _REPO_ROOT)

	# align.rollout.protocol is stdlib-only BY CONTRACT (see its docstring); it is
	# the single align import permitted here, taken from the repo tree directly so
	# align never needs installing into the rollout venv. (`import align` and
	# `import align.rollout` are both lazy/stdlib-only for exactly this reason.)
	from align.rollout.protocol import ( # noqa: E402
	RolloutResult,
	read_request,
	spool_dirs,
	touch_heartbeat,
	write_result,
	)


	def _extract_sampled_logprobs(token_ids, sample_logprobs, request_id):
	"""Per-position logprob of the SAMPLED token.

	vllm 0.26 API shape (verified against .venv-rollout's vllm/outputs.py and
	vllm/logprobs.py): CompletionOutput.logprobs is SampleLogprobs -- a
	sequence over positions where each element is dict[int, Logprob] (or the
	list-compatible FlatLogprobs view of the same); with
	SamplingParams(logprobs=0) each position carries the sampled token's own
	entry, so position t resolves as logprobs[t][token_ids[t]].logprob.
	"""
	if sample_logprobs is None:
	raise RuntimeError(
	f"request {request_id!r}: engine returned no logprobs despite "
	"SamplingParams(logprobs=0); the rho correction cannot run without "
	"pi^infer_old, so this result would be unusable."
	)
	out = []
	for t, tok in enumerate(token_ids):
	pos = sample_logprobs[t]
	if pos is None or tok not in pos:
	raise RuntimeError(
	f"request {request_id!r}: position {t} lacks the sampled token "
	f"{tok} in its logprob dict (keys: "
	f"{sorted(pos) if pos else pos}); refusing to substitute."
	)
	out.append(float(pos[tok].logprob))
	return out


	def serve(args) -> int:
	# Heartbeat BEFORE the vllm import: the import alone costs ~30s, and the
	# first smoke run proved a trainer canary can race it -- the client's
	# startup grace plus this early beat make "booting" visibly distinct from
	# "never started".
	spool_dirs(args.spool)
	touch_heartbeat(args.spool)
	from vllm import LLM, SamplingParams

	req_dir, _ = spool_dirs(args.spool)
	engine = None
	loaded_weights_dir = None
	last_work = time.time()
	print(f"rollout_worker serving spool={args.spool} tp={args.tp}", flush=True)

	while True:
	touch_heartbeat(args.spool)
	pending = sorted(
	(os.path.join(req_dir, f) for f in os.listdir(req_dir) if f.endswith(".json")),
	key=os.path.getmtime,
	)
	if not pending:
	if args.idle_exit_s and time.time() - last_work > args.idle_exit_s:
	print(f"idle {args.idle_exit_s}s; exiting", flush=True)
	return 0
	time.sleep(args.poll_s)
	continue

	for path in pending:
	request_id = os.path.basename(path)[: -len(".json")]
	reloaded = False
	try:
	req = read_request(path)
	if engine is None or req.weights_dir != loaded_weights_dir:
	if engine is not None:
	# vllm's own teardown first, then drop the reference;
	# the fresh constructor re-measures free VRAM.
	del engine
	engine = None
	import gc

	gc.collect()
	import torch

	torch.cuda.empty_cache()
	t0 = time.perf_counter()
	engine = LLM(
	model=req.weights_dir,
	dtype="bfloat16",
	tensor_parallel_size=args.tp,
	gpu_memory_utilization=args.gpu_mem_util,
	)
	loaded_weights_dir = req.weights_dir
	reloaded = True
	print(
	f"engine loaded from {req.weights_dir} in {time.perf_counter() - t0:.1f}s",
	flush=True,
	)

	sp = SamplingParams(
	n=req.n_samples_per_prompt,
	temperature=req.temperature,
	top_p=req.top_p,
	max_tokens=req.max_tokens,
	seed=req.seed,
	logprobs=0, # sampled token's own logprob per position
	)
	t1 = time.perf_counter()
	outs = engine.generate(req.prompts, sp)
	gen_s = time.perf_counter() - t1

	completions, logprobs, texts, prompt_token_ids = [], [], [], []
	for o in outs:
	completions.append([list(map(int, c.token_ids)) for c in o.outputs])
	logprobs.append(
	[
	_extract_sampled_logprobs(c.token_ids, c.logprobs, request_id)
	for c in o.outputs
	]
	)
	texts.append([c.text for c in o.outputs])
	# THE ENGINE'S OWN prompt tokenization, one list per prompt
	# in request order. Completion ids continue from exactly
	# this sequence; trainer-side assembly must concatenate
	# these + completion ids and NEVER re-tokenize the prompt
	# text (re-tokenization is not identity -- the alignment
	# defect class the 2026-08-23 audit closed).
	prompt_token_ids.append(list(map(int, o.prompt_token_ids)))

	write_result(
	args.spool,
	RolloutResult(
	request_id=request_id,
	completions=completions,
	logprobs=logprobs,
	texts=texts,
	weights_dir=req.weights_dir,
	engine_reloaded=reloaded,
	generate_seconds=gen_s,
	extras={"prompt_token_ids": prompt_token_ids},
	),
	)
	print(
	f"served {request_id}: {sum(len(c) for c in completions)} "
	f"completions in {gen_s:.1f}s (reloaded={reloaded})",
	flush=True,
	)
	except Exception as e: # noqa: BLE001 -- the trainer must fail fast, not time out
	tb = traceback.format_exc()
	print(f"request {request_id} FAILED: {e}\n{tb}", file=sys.stderr, flush=True)
	write_result(
	args.spool,
	RolloutResult(
	request_id=request_id,
	completions=[],
	logprobs=[],
	texts=[],
	weights_dir=loaded_weights_dir or "",
	engine_reloaded=reloaded,
	generate_seconds=0.0,
	error=str(e),
	extras={"traceback_tail": tb[-2000:]},
	),
	)
	finally:
	# Consumed = removed; the result file is the record.
	if os.path.exists(path):
	os.remove(path)
	last_work = time.time()


	def main() -> int:
	p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
	p.add_argument("--spool", required=True)
	p.add_argument("--gpu-mem-util", type=float, default=0.85)
	p.add_argument("--poll-s", type=float, default=2.0)
	p.add_argument("--tp", type=int, default=1)
	p.add_argument("--idle-exit-s", type=float, default=0.0, help="0 = serve forever")
	return serve(p.parse_args())


	if __name__ == "__main__":
	sys.exit(main())