code/tests/conftest.py · viai957/align at main

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

7.5 kB

"""Shared fixtures + house test conventions.
Conventions (inherited from the predecessor's 65-test suite):
  * regression tests are NAMED AFTER THE INCIDENT they prevent, incident in
    the docstring;
  * every guard is tested in BOTH directions (fires when it must, silent when
    it must);
  * undefined statistics assert NaN, never 0;
  * every test file's docstring carries its run command;
  * end-to-end tests use a tiny random-init model on CPU with >=2 layers and
    >=2 KV groups (1-layer/1-group models cannot surface interior/GQA bugs).
Run all CPU tests: uv run pytest -m "not gpu"
"""

from __future__ import annotations

import os

import pytest

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def heavy_job_blocking_gpu_test(lock_path: str | None = None) -> str | None:
    """The label of ANOTHER process's heavy-job lock, or None if the lock is free (or ours).
    THE UNGUARDED-GPU-TEST DEFECT (P2 R2 review). The `gpu` marker is a label, not a guard: `pytest -k <pattern>`
    without `-m "not gpu"` ran CUDA tests on GPU 0 while a 4-rank training run held the cards (21.7 GB in use of 23);
    the test OOM'd, and the only reason the run survived is that its allocator had already reached steady state. A gpu
    test is a heavy job like any other and must not start beside one. Probes the flock and releases it at once, so tests
    that take the lock themselves (`apply_resource_guards(heavy=True)`) behave exactly as before."""
    import fcntl

    from align.utils import guard

    path = lock_path or guard.LOCK_PATH
    if not os.path.exists(path):
        return None
    with open(path, "a+") as fh:
        fh.seek(0)
        label = fh.read().strip()
        try:
            fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            return None if f"pid={os.getpid()}" in label.split() else (label or "unknown holder")
        fcntl.flock(fh.fileno(), fcntl.LOCK_UN)
    return None


@pytest.fixture(autouse=True)
def _gpu_tests_do_not_start_beside_a_heavy_job(request):
    if request.node.get_closest_marker("gpu") is None:
        return
    holder = heavy_job_blocking_gpu_test()
    if holder is not None:
        pytest.skip(f"gpu test not started: the heavy-job lock is held by {holder!r} (never two heavy jobs at once)")


@pytest.fixture
def tmp_workdir(tmp_path_factory):
    """A scratch dir that honors the path policy: under the repo on /apps,
    never /tmp (which is noexec AND on a different device on this box)."""
    base = os.path.join(REPO_ROOT, ".tmp", "tests")
    os.makedirs(base, exist_ok=True)
    import tempfile

    d = tempfile.mkdtemp(dir=base)
    yield d
    # best-effort cleanup; stale dirs under .tmp/tests are harmless
    import shutil

    shutil.rmtree(d, ignore_errors=True)


@pytest.fixture
def tiny_model_constants():
    """Fixture constants for the CPU tiny-model harness, consumed by
    section-04's model tests. >=2 layers and >=2 KV groups is a hard house
    rule: a 1-layer or 1-group model cannot surface interior-layer or
    GQA-indexing bugs."""
    return {
        "n_layers": 2,
        "hidden_size": 64,
        "n_heads": 4,
        "n_kv_heads": 2,
        "head_dim": 16,
        "intermediate_size": 128,
        "vocab_size": 256,
        "max_seq_len": 128,
        "rope_theta": 1e6,
        "seed": 0,
    }


#: HF-cache snapshot of Mistral-7B-v0.3 (config + 3 shards + tokenizer). NOT present
#: on this box as of 2026-10-12 (hub/ has no models--mistralai--* entry); anything that
#: needs the WEIGHTS or the snapshot config.json must still skip when it is absent.
V03_TOKENIZER_SNAPSHOT = (
    "/apps/voxtral/.cache/huggingface/hub/models--mistralai--Mistral-7B-v0.3/"
    "snapshots/caa1feb0e54d415e2df31207e5f4e273e33509b1"
)

#: Tokenizer files vendored in-repo (byte-identical to the snapshot's; see
#: assets/tokenizer/mistral-7b-v0.3/PROVENANCE.md). Durable: the legacy SFT
#: checkpoint dirs these were copied from rotate.
V03_TOKENIZER_VENDORED = os.path.join(REPO_ROOT, "assets", "tokenizer", "mistral-7b-v0.3")


def resolve_v03_tokenizer_dir() -> str | None:
    """First directory holding a v0.3 tokenizer.json: ALIGN_TOKENIZER_SRC override,
    the HF snapshot, then the vendored copy. None if none exists."""
    env = os.environ.get("ALIGN_TOKENIZER_SRC")
    for c in ([env] if env else []) + [V03_TOKENIZER_SNAPSHOT, V03_TOKENIZER_VENDORED]:
        if c and os.path.isfile(os.path.join(c, "tokenizer.json")):
            return c
    return None


#: Resolved at import so tests can `isdir`-style gate on it.
V03_TOKENIZER_DIR = resolve_v03_tokenizer_dir()

#: Predecessor eval sets used as REAL gold sources for the ceiling tests
#: (MATH golds are LaTeX, GSM8K golds are bare integers -- they break
#: extraction in different ways; synthetic '42' golds prove nothing).
#: Section-03 regenerates these under nmo with hash receipts; these paths are
#: parameterized into the ceiling functions, never hardcoded inside nmo/.
GOLD_SOURCES = (
    "/apps/development/rlhf/mistral-prm-ppo/artifacts/data/eval/math500.json",
    "/apps/development/rlhf/mistral-prm-ppo/artifacts/data/eval/gsm8k_test.json",
)


@pytest.fixture(scope="session")
def v03_tokenizer():
    """The real Mistral-7B-v0.3 tokenizer, loaded by the in-repo align.tokenizer
    (stdlib-only; its parity with transformers is certified by the oracle test in
    tests/test_tokenizer.py). CPU-only, fast. Skips LOUDLY if no copy exists --
    never downloads."""
    d = resolve_v03_tokenizer_dir()
    if d is None:
        pytest.skip(
            f"no v0.3 tokenizer.json at {V03_TOKENIZER_SNAPSHOT} or {V03_TOKENIZER_VENDORED}; "
            "boundary certification cannot run against a fake tokenizer"
        )
    from align.tokenizer import load_tokenizer

    return load_tokenizer(d)


@pytest.fixture
def clean_script_env():
    """Environment for tests that SPAWN nmo scripts in a subprocess.
    The launching harness exports TMPDIR pointing at a scratchpad on the
    HOME device; the section-01 path guard correctly REFUSES that env at
    script startup (it did so in production on this guard's first day).
    Tests that spawn scripts must therefore pass an env whose TMPDIR is on
    /apps — via this fixture, so no lane rediscovers the leak. The guard's
    own refusal test poisons TMPDIR deliberately and must NOT use this."""
    env = dict(os.environ)
    env["TMPDIR"] = os.path.join(REPO_ROOT, ".tmp")
    os.makedirs(env["TMPDIR"], exist_ok=True)
    return env


@pytest.fixture
def fake_host_state(monkeypatch):
    """Factory: monkeypatch align.utils.guard.host_state to a fixed reading.
    Guard tests must never depend on real host load (CI would flake exactly
    when the shared box is busy -- the moment the guards matter most)."""

    def _set(
        mem_available_gb: float,
        load_per_core: float = 0.5,
        disk_in_flight: int = 0,
        d_state: int = 0,
    ):
        from align.utils import guard

        state = {
            "ram_total_gb": 186.0,
            "ram_available_gb": mem_available_gb,
            "swap_total_gb": 0.0,
            "load1": load_per_core * 48,
            "load5": 0.0,
            "n_cores": 48,
            "load_per_core": load_per_core,
            "disk_in_flight": disk_in_flight,
            "d_state_procs": d_state,
        }
        monkeypatch.setattr(guard, "host_state", lambda: dict(state))
        return state

    return _set

	"""Shared fixtures + house test conventions.

	Conventions (inherited from the predecessor's 65-test suite):
	* regression tests are NAMED AFTER THE INCIDENT they prevent, incident in
	the docstring;
	* every guard is tested in BOTH directions (fires when it must, silent when
	it must);
	* undefined statistics assert NaN, never 0;
	* every test file's docstring carries its run command;
	* end-to-end tests use a tiny random-init model on CPU with >=2 layers and
	>=2 KV groups (1-layer/1-group models cannot surface interior/GQA bugs).

	Run all CPU tests: uv run pytest -m "not gpu"
	"""

	from __future__ import annotations

	import os

	import pytest

	REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


	def heavy_job_blocking_gpu_test(lock_path: str \| None = None) -> str \| None:
	"""The label of ANOTHER process's heavy-job lock, or None if the lock is free (or ours).

	THE UNGUARDED-GPU-TEST DEFECT (P2 R2 review). The `gpu` marker is a label, not a guard: `pytest -k <pattern>`
	without `-m "not gpu"` ran CUDA tests on GPU 0 while a 4-rank training run held the cards (21.7 GB in use of 23);
	the test OOM'd, and the only reason the run survived is that its allocator had already reached steady state. A gpu
	test is a heavy job like any other and must not start beside one. Probes the flock and releases it at once, so tests
	that take the lock themselves (`apply_resource_guards(heavy=True)`) behave exactly as before."""
	import fcntl

	from align.utils import guard

	path = lock_path or guard.LOCK_PATH
	if not os.path.exists(path):
	return None
	with open(path, "a+") as fh:
	fh.seek(0)
	label = fh.read().strip()
	try:
	fcntl.flock(fh.fileno(), fcntl.LOCK_EX \| fcntl.LOCK_NB)
	except BlockingIOError:
	return None if f"pid={os.getpid()}" in label.split() else (label or "unknown holder")
	fcntl.flock(fh.fileno(), fcntl.LOCK_UN)
	return None


	@pytest.fixture(autouse=True)
	def _gpu_tests_do_not_start_beside_a_heavy_job(request):
	if request.node.get_closest_marker("gpu") is None:
	return
	holder = heavy_job_blocking_gpu_test()
	if holder is not None:
	pytest.skip(f"gpu test not started: the heavy-job lock is held by {holder!r} (never two heavy jobs at once)")


	@pytest.fixture
	def tmp_workdir(tmp_path_factory):
	"""A scratch dir that honors the path policy: under the repo on /apps,
	never /tmp (which is noexec AND on a different device on this box)."""
	base = os.path.join(REPO_ROOT, ".tmp", "tests")
	os.makedirs(base, exist_ok=True)
	import tempfile

	d = tempfile.mkdtemp(dir=base)
	yield d
	# best-effort cleanup; stale dirs under .tmp/tests are harmless
	import shutil

	shutil.rmtree(d, ignore_errors=True)


	@pytest.fixture
	def tiny_model_constants():
	"""Fixture constants for the CPU tiny-model harness, consumed by
	section-04's model tests. >=2 layers and >=2 KV groups is a hard house
	rule: a 1-layer or 1-group model cannot surface interior-layer or
	GQA-indexing bugs."""
	return {
	"n_layers": 2,
	"hidden_size": 64,
	"n_heads": 4,
	"n_kv_heads": 2,
	"head_dim": 16,
	"intermediate_size": 128,
	"vocab_size": 256,
	"max_seq_len": 128,
	"rope_theta": 1e6,
	"seed": 0,
	}


	#: HF-cache snapshot of Mistral-7B-v0.3 (config + 3 shards + tokenizer). NOT present
	#: on this box as of 2026-10-12 (hub/ has no models--mistralai--* entry); anything that
	#: needs the WEIGHTS or the snapshot config.json must still skip when it is absent.
	V03_TOKENIZER_SNAPSHOT = (
	"/apps/voxtral/.cache/huggingface/hub/models--mistralai--Mistral-7B-v0.3/"
	"snapshots/caa1feb0e54d415e2df31207e5f4e273e33509b1"
	)

	#: Tokenizer files vendored in-repo (byte-identical to the snapshot's; see
	#: assets/tokenizer/mistral-7b-v0.3/PROVENANCE.md). Durable: the legacy SFT
	#: checkpoint dirs these were copied from rotate.
	V03_TOKENIZER_VENDORED = os.path.join(REPO_ROOT, "assets", "tokenizer", "mistral-7b-v0.3")


	def resolve_v03_tokenizer_dir() -> str \| None:
	"""First directory holding a v0.3 tokenizer.json: ALIGN_TOKENIZER_SRC override,
	the HF snapshot, then the vendored copy. None if none exists."""
	env = os.environ.get("ALIGN_TOKENIZER_SRC")
	for c in ([env] if env else []) + [V03_TOKENIZER_SNAPSHOT, V03_TOKENIZER_VENDORED]:
	if c and os.path.isfile(os.path.join(c, "tokenizer.json")):
	return c
	return None


	#: Resolved at import so tests can `isdir`-style gate on it.
	V03_TOKENIZER_DIR = resolve_v03_tokenizer_dir()

	#: Predecessor eval sets used as REAL gold sources for the ceiling tests
	#: (MATH golds are LaTeX, GSM8K golds are bare integers -- they break
	#: extraction in different ways; synthetic '42' golds prove nothing).
	#: Section-03 regenerates these under nmo with hash receipts; these paths are
	#: parameterized into the ceiling functions, never hardcoded inside nmo/.
	GOLD_SOURCES = (
	"/apps/development/rlhf/mistral-prm-ppo/artifacts/data/eval/math500.json",
	"/apps/development/rlhf/mistral-prm-ppo/artifacts/data/eval/gsm8k_test.json",
	)


	@pytest.fixture(scope="session")
	def v03_tokenizer():
	"""The real Mistral-7B-v0.3 tokenizer, loaded by the in-repo align.tokenizer
	(stdlib-only; its parity with transformers is certified by the oracle test in
	tests/test_tokenizer.py). CPU-only, fast. Skips LOUDLY if no copy exists --
	never downloads."""
	d = resolve_v03_tokenizer_dir()
	if d is None:
	pytest.skip(
	f"no v0.3 tokenizer.json at {V03_TOKENIZER_SNAPSHOT} or {V03_TOKENIZER_VENDORED}; "
	"boundary certification cannot run against a fake tokenizer"
	)
	from align.tokenizer import load_tokenizer

	return load_tokenizer(d)


	@pytest.fixture
	def clean_script_env():
	"""Environment for tests that SPAWN nmo scripts in a subprocess.

	The launching harness exports TMPDIR pointing at a scratchpad on the
	HOME device; the section-01 path guard correctly REFUSES that env at
	script startup (it did so in production on this guard's first day).
	Tests that spawn scripts must therefore pass an env whose TMPDIR is on
	/apps — via this fixture, so no lane rediscovers the leak. The guard's
	own refusal test poisons TMPDIR deliberately and must NOT use this."""
	env = dict(os.environ)
	env["TMPDIR"] = os.path.join(REPO_ROOT, ".tmp")
	os.makedirs(env["TMPDIR"], exist_ok=True)
	return env


	@pytest.fixture
	def fake_host_state(monkeypatch):
	"""Factory: monkeypatch align.utils.guard.host_state to a fixed reading.
	Guard tests must never depend on real host load (CI would flake exactly
	when the shared box is busy -- the moment the guards matter most)."""

	def _set(
	mem_available_gb: float,
	load_per_core: float = 0.5,
	disk_in_flight: int = 0,
	d_state: int = 0,
	):
	from align.utils import guard

	state = {
	"ram_total_gb": 186.0,
	"ram_available_gb": mem_available_gb,
	"swap_total_gb": 0.0,
	"load1": load_per_core * 48,
	"load5": 0.0,
	"n_cores": 48,
	"load_per_core": load_per_core,
	"disk_in_flight": disk_in_flight,
	"d_state_procs": d_state,
	}
	monkeypatch.setattr(guard, "host_state", lambda: dict(state))
	return state

	return _set