"""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 " in src assert "#include " in src # windows.h must sit inside a _WIN32 branch and mman.h in the #else. win_pos = src.index("#include ") mman_pos = src.index("#include ") 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& locked_page_bytes()") == 1 assert src.index("static std::atomic& 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"]