File size: 1,937 Bytes
201dc77 | 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 | """U11-16: shards 1..N-1 of scripts/generate.py must not run beside a live heavy job."""
import importlib.util
import os
import pytest
from align.exit_codes import PreflightRefused
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
INODE = 12345
def _load():
spec = importlib.util.spec_from_file_location(
"generate_under_test", os.path.join(REPO, "scripts", "generate.py")
)
m = importlib.util.module_from_spec(spec)
spec.loader.exec_module(m)
return m
def _locks(pid):
return f"1: FLOCK ADVISORY WRITE {pid} 103:06:{INODE} 0 EOF\n2: POSIX ADVISORY READ 99 fd:00:777 0 EOF\n"
def _cmd(table):
return lambda pid: table[pid]
def test_non_zero_shard_refuses_when_a_trainer_holds_the_lock():
g = _load()
with pytest.raises(PreflightRefused, match="heavy"):
g.verify_shard_lock(
rank=1, lock_inode=INODE, locks_text=_locks(500),
cmdline_of=_cmd({500: "python scripts/grpo.py configs/x.toml"}),
)
def test_non_zero_shard_allows_its_own_shard0_sibling_and_a_free_lock():
g = _load()
g.verify_shard_lock(
rank=2, lock_inode=INODE, locks_text=_locks(500),
cmdline_of=_cmd({500: "python scripts/generate.py --shard-index 0 --num-shards 4"}),
)
g.verify_shard_lock(rank=3, lock_inode=INODE, locks_text="", cmdline_of=_cmd({}))
def test_rank0_and_no_heavy_lock_are_untouched():
g = _load()
# rank 0 takes the real lock in apply_resource_guards; nothing to verify here
g.verify_shard_lock(rank=0, lock_inode=INODE, locks_text=_locks(500), cmdline_of=_cmd({500: "x"}))
@pytest.mark.parametrize("idx,n", [(4, 4), (-1, 4), (0, 0), (5, 2)])
def test_shard_bounds_validated(idx, n):
g = _load()
with pytest.raises(ValueError, match="shard"):
g.validate_shards(idx, n)
def test_valid_shard_bounds():
g = _load()
g.validate_shards(0, 1)
g.validate_shards(3, 4)
|