Spaces:
Sleeping
Sleeping
File size: 7,055 Bytes
421b8c2 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | """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"]
|