"""Tests for the owner's rule in setup.py - it recommends, it never forces: a choice bigger than the recommendation is kept and its risk said (#406's 256K on 64 GB, #364/#384's two GPUs in the low-RAM mode, #403's RAM budget), and the stops that stay are the ones a user can confirm (confirm_risk: --yes with an explicit flag, or y). The defaults are tools/test_setup_golden.py's. Mocked - no GPU, no downloads, nothing written outside a temp folder. python -m unittest tools.test_setup_risk """ from __future__ import annotations import contextlib import io import json import os import sys import tempfile import unittest from pathlib import Path from unittest import mock ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) sys.path.insert(0, str(ROOT / "tools")) import setup # noqa: E402 from test_setup_golden import PROFILES, card, install # noqa: E402 def arg(cfg, flag): a = cfg["args"] return a[a.index(flag) + 1] if flag in a else None def run(fn, *args, stdin=None): """-> (exit code or None, printed text, questions asked); stdin answers every input() (None: never asked).""" out, asked, code = io.StringIO(), [], None def fake_input(prompt=""): asked.append(prompt) if stdin is None: raise AssertionError(f"asked {prompt!r} with --yes") return stdin with contextlib.redirect_stdout(out), mock.patch("builtins.input", fake_input): try: fn(*args) except SystemExit as e: code = e.code return code, out.getvalue(), asked class ConfirmRisk(unittest.TestCase): """S6: one helper for the stops a user may take on.""" def go(self, explicit, yes, stdin=None): return run(setup.confirm_risk, "it may not fit", explicit, yes, "it does not fit", "a smaller size", stdin=stdin) def test_yes_with_an_explicit_choice_is_consent(self): code, out, asked = self.go(True, True) self.assertIsNone(code) self.assertIn("it may not fit", out) self.assertEqual(asked, []) def test_yes_alone_keeps_the_stop(self): code, out, asked = self.go(False, True) self.assertEqual(code, 1) self.assertIn("it does not fit", out) self.assertIn("a smaller size", out) def test_interactive_default_is_no(self): for explicit in (True, False): code, out, asked = self.go(explicit, False, stdin="") self.assertEqual(code, 1) self.assertIn("[n]", asked[0]) self.assertIsNone(self.go(explicit, False, stdin="y")[0]) def test_paging_with_an_explicit_model(self): code, out, asked = run(setup.confirm_paging, "IQ3_XXS", 54, "auto", True, True) self.assertIsNone(code, out) self.assertIn("installing IQ3_XXS with 54 GB of RAM, as you chose (--model)", out) class Context(unittest.TestCase): """S1 (#406 #364): the RAM rule recommends, an explicit or picked context is kept with a note.""" RAM64, GPU32 = PROFILES["64GB-1x32GB"] def test_the_ram_rule(self): self.assertEqual(setup.ram_ctx("IQ3_S", 63.7), 131072) self.assertEqual(setup.ram_ctx("IQ3_XXS", 63.7), 131072) self.assertEqual(setup.ram_ctx("IQ3_XXS", 95.8), 524288) # 43 + 7.2 + 24 GB fits 96 self.assertEqual(setup.ram_ctx("IQ3_S", 76.5), 131072) self.assertEqual(setup.ram_ctx("IQ3_S", 78.5), 262144) self.assertEqual(setup.ram_ctx("Q2_0", 31.9), 524288) # other sizes: no RAM rule self.assertEqual(setup.ram_ctx("IQ3_XXS", 31.9, low_ram=True), 524288) # the KV cache is in VRAM there self.assertIsNone(setup.ctx_ram_need("IQ3_XXS", 262144, low_ram=True)) self.assertAlmostEqual(setup.ctx_ram_need("IQ3_S", 262144), 77.9, places=1) def test_an_explicit_256k_is_kept(self): code, out, cfg, _ = install(self.RAM64, self.GPU32, ["--family", "qwen", "--model", "IQ3_S", "--no-start", "--context", "262144"]) self.assertEqual(code, 0, out) self.assertEqual(arg(cfg, "--max-context"), "262144") self.assertNotIn("--rope-scaling", cfg["args"]) self.assertIn("256K with IQ3_S needs ~78 GB of RAM by setup's estimate", out) self.assertIn("this PC has 64. Kept as you chose: it may be slower or run out of RAM under load. 128K is the " "recommended size.", out) def test_a_picked_256k_is_kept_and_the_menu_says_the_risk(self): code, out, cfg, asked = install(self.RAM64, self.GPU32, ["--family", "qwen", "--model", "IQ3_S", "--no-start"], answers={"Context?": "5"}) self.assertEqual(code, 0, out) self.assertEqual(arg(cfg, "--max-context"), "262144") self.assertIn("256K tokens (needs ~78 GB RAM, this PC has 64: may run out of memory)", out) self.assertIn("128K tokens (recommended for your GPU)\n", out) self.assertIn("Kept as you chose", out) def test_low_ram_mode_keeps_262k_without_a_ram_note(self): ram, found = PROFILES["32GB-2x24GB"] # #364 code, out, cfg, _ = install(ram, found, ["--family", "qwen", "--model", "IQ3_XXS", "--no-start", "--context", "262144"]) self.assertEqual(code, 0, out) self.assertEqual(arg(cfg, "--max-context"), "262144") self.assertNotIn("Kept as you chose: it may be slower", out) def test_the_earlier_install_keeps_its_context(self): """An update (a new copy of Strata) set up like the earlier install keeps its 256K instead of capping it.""" with tempfile.TemporaryDirectory() as d: prev = Path(d) / "strata-iq3_s.json" prev.write_text(json.dumps({"args": ["--max-context", "262144", "--kv", "int8"], "port": 8080})) started = mock.Mock(return_value=0) code, out, cfg, _ = install(self.RAM64, self.GPU32, [], extra=[ mock.patch.object(setup, "previous_config", lambda *a: prev), mock.patch.object(setup, "start", started)]) self.assertEqual(code, 0, out) self.assertEqual(arg(cfg, "--max-context"), "262144") self.assertIn("Kept as you chose", out) self.assertTrue(started.called) class BrokenEarlierConfig(unittest.TestCase): """#459: a new copy of Strata set up like an earlier install skips an earlier config that does not parse (an empty strata-*.json crashed START-HERE with JSONDecodeError) and goes on as a fresh install; configs are written whole (a temporary file moved over the old one).""" RAM64, GPU32 = PROFILES["64GB-1x32GB"] def setup_with(self, files): with tempfile.TemporaryDirectory() as d: for i, (name, text) in enumerate(files): # in order, oldest first (Path(d) / name).write_text(text, encoding="utf-8") os.utime(Path(d) / name, (1_700_000_000 + i, 1_700_000_000 + i)) return install(self.RAM64, self.GPU32, [], extra=[ mock.patch.object(setup, "other_installs", lambda settings: [Path(d)]), mock.patch.object(setup, "start", mock.Mock(return_value=0))]) def test_an_empty_config_is_skipped(self): for text, why in (("", "the file is empty"), ("{\"args\": [", "not valid JSON"), ("[1]", "not a JSON object")): with self.subTest(text=text): code, out, cfg, _ = self.setup_with([("strata-iq3_s.json", text)]) self.assertEqual(code, 0, out) self.assertIn("skipped the earlier config ", out) self.assertIn(f"strata-iq3_s.json ({why}): setting this copy up without it", out) self.assertNotIn("Found your earlier install", out) self.assertIsNotNone(cfg) # the fresh install's config def test_the_newest_readable_config_is_used(self): good = json.dumps({"args": ["--max-context", "262144", "--kv", "int8"], "port": 8080}) code, out, cfg, _ = self.setup_with([("strata-iq3_s.json", good), ("strata-q2_0.json", "")]) self.assertEqual(code, 0, out) self.assertIn("strata-q2_0.json (the file is empty)", out) self.assertIn("Found your earlier install", out) self.assertIn("(iq3_s)", out) self.assertEqual(arg(cfg, "--max-context"), "262144") def test_write_config_leaves_no_partial_file(self): with tempfile.TemporaryDirectory() as d: p = Path(d) / "strata-iq3_s.json" p.write_text("old", encoding="utf-8") setup.write_config(p, {"args": ["--kv", "int8"]}) self.assertEqual(p.read_text(encoding="utf-8"), json.dumps({"args": ["--kv", "int8"]}, indent=1)) self.assertEqual([f.name for f in Path(d).iterdir()], ["strata-iq3_s.json"]) # no .tmp left with mock.patch.object(Path, "write_text", side_effect=OSError(28, "No space left on device")): with self.assertRaises(OSError): setup.write_config(p, {"args": []}) self.assertEqual(json.loads(p.read_text(encoding="utf-8")), {"args": ["--kv", "int8"]}) # kept whole class LowRamGpus(unittest.TestCase): """S2/S3 (#364 #384): one GPU recommended in the low-RAM mode, all of them when asked for.""" def test_explicit_gpus_are_kept_with_the_mapped_variant(self): for prof, model in (("32GB-2x24GB", "IQ3_XXS"), ("32GB-2x24GB", "Q2_0"), ("47GB-2x16GB", "IQ3_XXS")): with self.subTest(prof=prof, model=model): ram, found = PROFILES[prof] code, out, cfg, _ = install(ram, found, ["--family", "qwen", "--model", model, "--no-start", "--gpus", "0,1"]) self.assertEqual(code, 0, out) self.assertEqual(cfg["gpu"], [0, 1]) self.assertEqual(cfg["layer_split"], "auto") self.assertIn("--mmap-experts", cfg["args"]) self.assertNotIn("--resident-experts", cfg["args"]) self.assertIn("as you chose (--gpus)", out) self.assertIn("RAM can fill up to 0 free", out) self.assertIn("low-RAM mode on 2 GPUs", out) def test_the_share_counts_every_card(self): ram, found = PROFILES["47GB-2x16GB"] # #384: 2 x 16 GB, 47 GB, IQ3_XXS at 64K code, out, cfg, _ = install(ram, found, ["--family", "qwen", "--model", "IQ3_XXS", "--no-start", "--gpus", "0,1"]) one = setup.low_ram_gpu_gb("IQ3_XXS", 15.9, 65536, "int8") self.assertIn(f"the GPUs hold ~{100 * 2 * one / setup.MODELS['IQ3_XXS']['arena_gb']:.0f}% of them", out) def test_asked_and_answered_two(self): ram, found = PROFILES["32GB-2x24GB"] code, out, cfg, asked = install(ram, found, ["--family", "qwen", "--model", "Q2_0", "--no-start"], answers={"Low-RAM mode: which GPUs?": "2"}) self.assertEqual(code, 0, out) self.assertTrue(any("Low-RAM mode: which GPUs? [1]" in q for q in asked), asked) self.assertEqual(cfg["gpu"], [0, 1]) self.assertIn("--mmap-experts", cfg["args"]) self.assertIn("1) GPU 0 (NVIDIA GeForce RTX 3090, 24 GB) only", out) def test_yes_and_enter_keep_one_gpu(self): ram, found = PROFILES["32GB-2x24GB"] for answers in (None, ""): code, out, cfg, _ = install(ram, found, ["--family", "qwen", "--model", "Q2_0", "--no-start"], answers=answers) self.assertEqual(code, 0, out) self.assertEqual(cfg["gpu"], 0) self.assertIn("--resident-experts", cfg["args"]) def test_low_ram_resident_with_gpus(self): ram, found = PROFILES["32GB-2x24GB"] code, out, cfg, _ = install(ram, found, ["--family", "qwen", "--model", "Q2_0", "--no-start", "--gpus", "0,1", "--low-ram", "resident"]) self.assertEqual(code, 0, out) self.assertEqual(cfg["gpu"], [0, 1]) self.assertIn("--mmap-experts", cfg["args"]) self.assertIn("--low-ram resident has no layer split yet", out) def test_low_ram_resident_alone_keeps_one_gpu_without_asking(self): ram, found = PROFILES["32GB-2x24GB"] code, out, cfg, asked = install(ram, found, ["--family", "qwen", "--model", "Q2_0", "--no-start", "--low-ram", "resident"], answers="") self.assertEqual(code, 0, out) self.assertEqual(cfg["gpu"], 0) self.assertFalse(any("Low-RAM mode" in q for q in asked), asked) class StartOnSeveralGpus(unittest.TestCase): """A resident low-RAM config started on several GPUs reads the experts through the file cache (#364 #384).""" def test_split_mmap(self): cfg = {"args": ["--pack", "p", "--resident-experts", "--kv", "int8"]} out = io.StringIO() with contextlib.redirect_stdout(out): self.assertTrue(setup.split_mmap(cfg)) self.assertFalse(setup.split_mmap(cfg)) self.assertFalse(setup.split_mmap({"args": ["--mmap-experts"]})) self.assertEqual(cfg["args"], ["--pack", "p", "--mmap-experts", "--kv", "int8"]) self.assertIn("no layer split yet", out.getvalue()) def offer(self, args, stdin): found = PROFILES["32GB-2x24GB"][1] with tempfile.TemporaryDirectory() as d: p = Path(d) / "strata-q2_0.json" cfg = {"args": list(args)} with mock.patch.object(setup, "gpus", lambda: found), \ mock.patch.object(setup, "engine_runs_on", lambda g: True): code, out, asked = run(setup.offer_together, p, cfg, stdin is None, stdin=stdin) return code, out, asked, json.loads(p.read_text()) def test_offer_together_recommends_one_gpu_for_a_resident_config(self): code, out, asked, cfg = self.offer(["--resident-experts"], None) # --yes: the recommendation self.assertIsNone(code, out) self.assertNotIn("gpu", cfg) self.assertEqual(cfg["args"], ["--resident-experts"]) code, out, asked, cfg = self.offer(["--resident-experts"], "y") self.assertIn("[n]", asked[0]) self.assertEqual(cfg["gpu"], [0, 1]) self.assertEqual(cfg["args"], ["--mmap-experts"]) code, out, asked, cfg = self.offer(["--mmap-experts"], None) # other configs: as before self.assertEqual(cfg["gpu"], [0, 1]) def test_start_with_gpus(self): found = PROFILES["32GB-2x24GB"][1] with tempfile.TemporaryDirectory() as d: exe = Path(d) / "strata.exe" exe.write_bytes(b"") p = Path(d) / "strata-q2_0.json" p.write_text(json.dumps({"exe": str(exe), "args": ["--resident-experts"], "gpu": 0, "gpus_asked": True})) call = mock.Mock(return_value=0) with mock.patch.object(setup, "gpus", lambda: found), \ mock.patch.object(setup, "ensure_engine_for", lambda cards, path, cfg, yes: cfg), \ mock.patch.object(setup, "is_wsl", lambda: False), \ mock.patch.object(setup.subprocess, "call", call): code, out, _ = run(setup.start, p, None, [0, 1], False, True) cfg = json.loads(p.read_text()) self.assertIsNone(code, out) self.assertEqual(cfg["gpu"], [0, 1]) self.assertEqual(cfg["args"], ["--mmap-experts"]) self.assertIn("no layer split yet", out) self.assertTrue(call.called) class SmallCard(unittest.TestCase): """S9: a card under 8 GB named with --gpus is a risk to confirm, not a stop.""" FOUND = [card(0, "NVIDIA GeForce RTX 3090", 24.0, "86"), card(1, "NVIDIA GeForce RTX 2060", 6.0, "75")] def test_named_with_yes(self): code, out, cfg, _ = install(63.7, self.FOUND, ["--family", "qwen", "--model", "Q2_0", "--no-start", "--gpus", "0,1"]) self.assertEqual(code, 0, out) self.assertEqual(cfg["gpu"], [0, 1]) self.assertIn("GPU 1 (NVIDIA GeForce RTX 2060) has 6 GB of VRAM", out) self.assertIn("is used together with the others, as you chose", out) def test_named_and_enter_stops(self): code, out, cfg, asked = install(63.7, self.FOUND, ["--family", "qwen", "--model", "Q2_0", "--no-start", "--gpus", "0,1"], answers="") self.assertEqual(code, 1) self.assertTrue(any("Use it anyway?" in q for q in asked), asked) def test_not_named_is_not_offered(self): code, out, cfg, _ = install(63.7, self.FOUND, ["--family", "qwen", "--model", "Q2_0", "--no-start"]) self.assertEqual(code, 0, out) self.assertEqual(cfg["gpu"], 0) # one card: as before class SplitShortCard(unittest.TestCase): """#448: a second card too small to lend a split's prompt chunk: the first card alone recommended, the pair kept when named.""" FOUND = [card(0, "NVIDIA RTX PRO 4500 Blackwell", 31.8, "120"), card(1, "NVIDIA GeForce RTX 3080", 10.0, "86")] def test_rule(self): self.assertEqual([g["index"] for g in setup.split_short(self.FOUND)], [1]) self.assertEqual(setup.split_short(PROFILES["32GB-2x24GB"][1]), []) self.assertEqual(setup.split_short(PROFILES["47GB-2x16GB"][1]), []) eights = [card(i, "NVIDIA GeForce RTX 3070", 8.0, "86") for i in range(2)] self.assertEqual(setup.split_short(eights), []) # neither card alone reads more def test_yes_takes_the_first_card(self): code, out, cfg, _ = install(127.8, self.FOUND, ["--family", "qwen", "--model", "Q2_0", "--no-start"]) self.assertEqual(code, 0, out) self.assertEqual(cfg["gpu"], 0) self.assertIn("too small to lend a split its prompt buffers", out) def test_named_pair_is_kept(self): code, out, cfg, _ = install(127.8, self.FOUND, ["--family", "qwen", "--model", "Q2_0", "--no-start", "--gpus", "0,1"]) self.assertEqual(code, 0, out) self.assertEqual(cfg["gpu"], [0, 1]) def test_offer_together_defaults_to_one(self): with tempfile.TemporaryDirectory() as d: p = Path(d) / "strata-q2_0.json" with mock.patch.object(setup, "gpus", lambda: self.FOUND), \ mock.patch.object(setup, "engine_runs_on", lambda g: True): code, out, asked = run(setup.offer_together, p, {"args": ["--mmap-experts"]}, True) cfg = json.loads(p.read_text()) self.assertIsNone(code, out) self.assertNotIn("gpu", cfg) self.assertIn("#448", out) class RamFloor(unittest.TestCase): """S7: less RAM than the smallest model needs: a stop with --yes alone, a risk with --model.""" FOUND = [card(0, "NVIDIA GeForce RTX 3060", 8.0, "86")] def test_yes_alone_stops(self): code, out, cfg, _ = install(19.9, self.FOUND, ["--family", "qwen", "--no-start"]) self.assertEqual(code, 1) self.assertIn("--model NAME --yes", out) def test_an_explicit_model_goes_on(self): code, out, cfg, _ = install(19.9, self.FOUND, ["--family", "coder", "--model", "IQ1_M", "--no-start"]) self.assertEqual(code, 0, out) self.assertIn("going on with 20 GB of RAM, as you chose", out) class KvStreaming(unittest.TestCase): """8: --kv-streaming on|off|auto (auto: the RAM test, as before).""" RAM, FOUND = PROFILES["64GB-1x32GB"] def go(self, model, *flags): return install(self.RAM, self.FOUND, ["--family", "qwen", "--model", model, "--no-start", *flags]) def test_off(self): code, out, cfg, _ = self.go("IQ3_XXS", "--kv-streaming", "off") self.assertEqual(code, 0, out) self.assertNotIn("--kv-resident", cfg["args"]) self.assertIn("KV streaming off, as you chose", out) def test_on_past_the_ram_test(self): code, out, cfg, _ = self.go("IQ3_S", "--kv-streaming", "on") # auto leaves it off on 64 GB self.assertEqual(code, 0, out) self.assertEqual(arg(cfg, "--kv-resident"), "32768") self.assertIn("Kept as you chose (--kv-streaming on)", out) def test_on_where_it_cannot(self): code, out, cfg, _ = self.go("Q2_0", "--kv-streaming", "on", "--kv", "k8v4") self.assertNotIn("--kv-resident", cfg["args"]) self.assertIn("it refuses the pair", out) code, out, cfg, _ = self.go("Q2_0", "--kv-streaming", "on", "--context", "32768") self.assertNotIn("--kv-resident", cfg["args"]) self.assertIn("a context under 64K is not streamed", out) def test_budget_flag_on_another_size(self): code, out, cfg, _ = self.go("Q2_0", "--resident-budget-gib", "30") self.assertEqual(code, 0, out) self.assertNotIn("--resident-budget-gib", cfg["args"]) self.assertIn("--resident-budget-gib is for UD-Q4_K_XL", out) if __name__ == "__main__": unittest.main()