"""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())