#!/usr/bin/env python3
"""vLLM solver subprocess: loads the bundled AWQ model and answers everything.

Reads IOL_TEST_CSV / IOL_BUDGET_S from env; writes submission.csv incrementally.
"""
import csv
import os
import sys
import time

START = time.time()
BUDGET = float(os.environ.get("IOL_BUDGET_S", 1200))
TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")

os.environ.setdefault("VLLM_USE_V1", "0")  # V0 = in-process, Turing-safe
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
os.environ.setdefault("VLLM_NO_USAGE_STATS", "1")
os.environ.setdefault("VLLM_DO_NOT_TRACK", "1")


def _preinit_fake_dist():
    """The eval sandbox seccomp-blocks connect()/setsockopt, so neither TCPStore
    nor gloo/nccl can initialize. For world_size=1 every collective is an
    identity op, so pre-init torch.distributed with the socket-free FAKE
    process group and hand the same group back for every new_group() request
    (vLLM asks for gloo cpu groups it will never actually communicate on)."""
    try:
        import torch.distributed as dist
        from torch.testing._internal.distributed.fake_pg import FakeStore
        if not dist.is_initialized():
            dist.init_process_group(backend="fake", rank=0, world_size=1,
                                    store=FakeStore())
        dist.new_group = lambda *a, **k: dist.group.WORLD
        print("[vllm-solver] fake process group installed", flush=True)
    except Exception as e:
        print(f"[vllm-solver] fake-pg pre-init failed: {e!r}", flush=True)


def log(msg):
    print(f"[vllm-solver +{time.time()-START:5.1f}s] {msg}", flush=True)


def main():
    with open(TEST_CSV, newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    log(f"{len(rows)} rows")

    _preinit_fake_dist()
    from vllm import LLM
    model_dir = "awq14b" if os.path.isdir("awq14b") else "."
    t = time.time()
    # T4 (16GB) safety: eager mode (no cudagraph-capture OOM), shorter context,
    # prefix caching so stage1/2 reuse stage0's prefill.
    llm = LLM(model=model_dir, max_model_len=8192, gpu_memory_utilization=0.92,
              dtype="float16", swap_space=2, enforce_eager=True,
              enable_prefix_caching=True, max_num_seqs=16)
    log(f"engine up in {time.time()-t:.0f}s (model={model_dir})")

    from pipeline_v2 import solve
    solve(llm, rows, "submission.csv", START, BUDGET - 60, log)


if __name__ == "__main__":
    main()