align / code /tests /test_audit_g6_eval_scripts_grpo_worker.py
viai957's picture
Add files using upload-large-folder tool
201dc77 verified
Raw History Blame Contribute Delete
1.36 kB
"""U11-7: the vLLM rollout worker (own session, serves forever) must not outlive rank 0."""
import importlib.util
import os
import sys
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def _load():
spec = importlib.util.spec_from_file_location("grpo_under_test", os.path.join(REPO, "scripts", "grpo.py"))
m = importlib.util.module_from_spec(spec)
spec.loader.exec_module(m)
return m
class _Proc:
pid = 4242
def test_spawn_bounds_worker_lifetime_and_registers_teardown(monkeypatch, tmp_path):
grpo = _load()
import align.rollout.client as client
seen = {}
monkeypatch.setattr(client, "spawn_worker", lambda *a, **kw: seen.update(kw) or _Proc())
killed = []
monkeypatch.setattr(client, "terminate_worker", lambda p, *a, **kw: killed.append(p))
registered = []
monkeypatch.setattr(grpo.atexit, "register", lambda fn, *a, **kw: registered.append((fn, a)))
proc = grpo._spawn_rollout_worker(str(tmp_path / "spool"), 3, str(tmp_path / "w.log"))
assert proc.pid == 4242
# not 'serve forever': an orphaned worker (rank 0 SIGKILLed) must exit by itself
assert seen["idle_exit_s"] > 0
# an exception anywhere between spawn and the train try/finally still tears it down
assert len(registered) == 1
fn, args = registered[0]
fn(*args)
assert killed == [proc]