tests/test_linux_platform.py

"""Platform-conditional page-locking tests (Windows-runnable, no Linux needed).

Windows (this box) exercises the VirtualLock path; the mlock path is
source-audited here and runtime-exercised on Linux CI where these tests
also run (Linux-only tests skip on win32).
"""

import importlib
import os
import pathlib
import sys

import numpy as np
import pytest

native = importlib.import_module(os.environ.get("CISM_NATIVE_TEST_MODULE", "cism._native"))

REPO_ROOT = pathlib.Path(__file__).resolve().parents[1]
RUNTIME_CPP = REPO_ROOT / "native" / "runtime.cpp"

IS_WINDOWS = sys.platform == "win32"
IS_LINUX = sys.platform.startswith("linux")


def tiny_model(hidden=24, intermediate=37):
    rng = np.random.default_rng(312)
    config = dict(
        model_type="llama",
        hidden_size=hidden,
        intermediate_size=intermediate,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        head_dim=8,
        vocab_size=43,
        max_position_embeddings=80,
        rms_norm_eps=1e-5,
        rope_theta=10000.0,
        tie_word_embeddings=True,
        hidden_act="silu",
    )
    weights = {}

    def matrix(name, shape):
        weights[name] = rng.normal(0, 0.14, shape).astype(np.float32)

    def norm(name, width):
        weights[name] = rng.uniform(0.75, 1.25, width).astype(np.float32)

    matrix("model.embed_tokens.weight", (43, hidden))
    norm("model.norm.weight", hidden)
    for i in range(2):
        prefix = f"model.layers.{i}."
        norm(prefix + "input_layernorm.weight", hidden)
        norm(prefix + "post_attention_layernorm.weight", hidden)
        for name, shape in (
            ("q_proj", (32, hidden)),
            ("k_proj", (16, hidden)),
            ("v_proj", (16, hidden)),
            ("o_proj", (hidden, 32)),
        ):
            matrix(prefix + f"self_attn.{name}.weight", shape)
        for name, shape in (
            ("gate_proj", (intermediate, hidden)),
            ("up_proj", (intermediate, hidden)),
            ("down_proj", (hidden, intermediate)),
        ):
            matrix(prefix + f"mlp.{name}.weight", shape)
    return config, weights


def _read_runtime():
    return RUNTIME_CPP.read_text(encoding="utf-8")


# 1. Source audit: includes are platform-guarded.
def test_includes_are_platform_guarded():
    src = _read_runtime()
    assert "#include <windows.h>" in src
    assert "#include <sys/mman.h>" in src
    # windows.h must sit inside a _WIN32 branch and mman.h in the #else.
    win_pos = src.index("#include <windows.h>")
    mman_pos = src.index("#include <sys/mman.h>")
    guard_pos = src.index("#if defined(_WIN32)")
    else_pos = src.index("#else", guard_pos)
    assert guard_pos < win_pos < else_pos < mman_pos


# 2. Source audit: locked_page_bytes counter is platform-independent.
def test_locked_page_bytes_defined_outside_win32_guard():
    src = _read_runtime()
    assert "locked_page_bytes" in src
    # The counter definition must not live only inside the _WIN32 block:
    # there is exactly one definition and it precedes the _WIN32 branch.
    assert src.count("static std::atomic<std::size_t>& locked_page_bytes()") == 1
    assert src.index("static std::atomic<std::size_t>& locked_page_bytes()") < src.index(
        "#if defined(_WIN32)\nstatic bool lock_range"
    )
    # Both backends still update the shared counter.
    assert "locked_page_bytes().fetch_add(locked" in src


# 3. Source audit: cpu_relax has MSVC + GCC paths and a fallback.
def test_cpu_relax_has_gcc_path():
    src = _read_runtime()
    assert "static inline void cpu_relax()" in src
    assert "_mm_pause()" in src  # MSVC path
    assert "__builtin_ia32_pause()" in src  # GCC/Clang x86 path
    assert "std::this_thread::yield()" in src  # portable fallback


# 4. Source audit: Linux mlock limit documented in code.
def test_mlock_rlimit_documented():
    src = _read_runtime()
    assert "mlock(address, bytes)" in src
    assert "munlock(address, bytes)" in src
    assert "RLIMIT_MEMLOCK" in src
    assert "ulimit -l" in src


# 5. Runtime (Windows VirtualLock path + Linux smoke): touch/scan stable.
@pytest.mark.skipif(native is None, reason="native extension missing")
def test_touch_and_scan_stable():
    config, weights = tiny_model()
    model = native.Model(config, weights, "fp32")
    info_bytes = model.info["weight_bytes"]
    assert not model.info["pages_locked"]
    first = model.touch()
    rotary_bytes = (config["head_dim"] // 2) * 4
    assert info_bytes < first <= info_bytes + rotary_bytes
    assert model.touch() == first
    assert model.scan() == first


# 6. Runtime: lock_pages is honest (returns bytes, never throws for limits).
@pytest.mark.skipif(native is None, reason="native extension missing")
def test_lock_pages_returns_bytes_without_throwing():
    config, weights = tiny_model()
    model = native.Model(config, weights, "fp32")
    total = model.touch()
    locked = model.lock_pages()  # must not raise even under RLIMIT_MEMLOCK
    assert isinstance(locked, int)
    assert 0 <= locked <= total
    if IS_WINDOWS:
        # Tiny model comfortably fits the working-set path: expect real bytes.
        assert locked > 0
        assert model.info["pages_locked"]
    unlocked = model.unlock_pages()
    assert isinstance(unlocked, int)
    assert 0 < unlocked <= max(locked, 1) or locked == 0
    assert not model.info["pages_locked"]


# 7. Runtime: double lock / double unlock are errors; inference survives cycle.
@pytest.mark.skipif(native is None, reason="native extension missing")
def test_lock_cycle_errors_and_inference_survives():
    config, weights = tiny_model()
    model = native.Model(config, weights, "fp32")
    with pytest.raises(ValueError, match="not locked"):
        model.unlock_pages()
    model.lock_pages()
    with pytest.raises(ValueError, match="already locked"):
        model.lock_pages()
    model.unlock_pages()
    with pytest.raises(ValueError, match="not locked"):
        model.unlock_pages()
    assert not model.info["pages_locked"]
    logits = model.logits([1, 2, 3])
    assert logits.shape == (43,)
    assert np.isfinite(logits).all()
    session = model.create_session([1, 2], max_new_tokens=2)
    assert len(session.next_tokens(2)) == 2


# 8. Linux-only runtime: mlock path honors RLIMIT_MEMLOCK (skipped on Windows).
@pytest.mark.skipif(not IS_LINUX, reason="mlock runtime path needs Linux")
def test_linux_mlock_path_and_rlimit():
    resource = pytest.importorskip("resource")
    soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK)
    # Document the cap; tiny-model lock may legitimately be partial.
    print(f"\nRLIMIT_MEMLOCK soft={soft} hard={hard} (see `ulimit -l`)")
    config, weights = tiny_model()
    model = native.Model(config, weights, "fp32")
    total = model.touch()
    locked = model.lock_pages()
    assert 0 <= locked <= total
    assert model.info["pages_locked"]
    unlocked = model.unlock_pages()
    assert 0 <= unlocked <= max(locked, 1) or locked == 0
    assert not model.info["pages_locked"]