File size: 2,449 Bytes
8129b09 7c95c69 8129b09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | #!/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()
|