diff options
Diffstat (limited to 'app/tests/test_backends_sglomni.py')
| -rw-r--r-- | app/tests/test_backends_sglomni.py | 84 |
1 files changed, 84 insertions, 0 deletions
diff --git a/app/tests/test_backends_sglomni.py b/app/tests/test_backends_sglomni.py index 8cb86cb..e61d51e 100644 --- a/app/tests/test_backends_sglomni.py +++ b/app/tests/test_backends_sglomni.py @@ -203,6 +203,9 @@ class CompanionPackageTests(unittest.TestCase): "audiotools") self.assertEqual( catalog.extra_import_name("descript-audio-codec==1.0.0"), "dac") + self.assertEqual( + catalog.extra_import_name("protobuf==6.33.6"), + "google.protobuf.runtime_version") def test_missing_companions_reports_absent_modules(self): entry = entry_by_key("qwen3_tts_1_7b_base") @@ -227,6 +230,28 @@ class CompanionPackageTests(unittest.TestCase): with patch.object(envs, "env_exists", return_value=False): self.assertEqual(models.missing_companions(entry), []) + def test_dac_heal_reinstalls_protobuf_after_the_downgrade(self): + # descript-audiotools carries a vestigial protobuf<3.20 pin: its + # with-deps install downgrades the protobuf 6.x the sglang-omni + # stack itself needs. The restore extra must therefore re-enter + # missing_companions whenever that downgrade happened — probed via + # google.protobuf.runtime_version, which 3.19.6 does not provide. + entry = entry_by_key("zonos2") + self.assertIn(("protobuf==6.33.6", False), entry.extras) + with patch.object(envs, "env_exists", return_value=True), \ + patch.object(envs, "module_available", + side_effect=lambda name, env_dir: + name != "google.protobuf.runtime_version"): + missing = models.missing_companions(entry) + self.assertEqual(missing, [("protobuf==6.33.6", False)]) + + def test_dac_protobuf_restore_installs_after_descript_audiotools(self): + # The restore only re-pins protobuf if it pip-installs AFTER the + # package whose dependency resolution did the downgrading. + specs = [spec for spec, _no_deps in entry_by_key("fish_s2_pro").extras] + self.assertLess(specs.index("descript-audiotools==0.7.2"), + specs.index("protobuf==6.33.6")) + def test_install_companions_installs_only_missing_with_no_deps(self): entry = entry_by_key("qwen3_tts_1_7b_base") calls = [] @@ -415,6 +440,58 @@ class HiggsConfigTests(unittest.TestCase): self.assertIsNone(entry_by_key("zonos2").chunk_words) +class EngineMemoryBudgetTests(unittest.TestCase): + """Colocated AR pipelines with an unpinned engine budget OOM on 24 GB. + + Upstream auto-sizes the engine's sglang static pool to nearly all free + VRAM when mem_fraction_static is unset (qwen3_tts, moss_tts, voxtral): + on a 24 GB card the first /v1/audio/speech request aborts with + "CUDA out of memory. Tried to allocate ~100 MiB" and every retry fails + identically. These vendored configs pin the pool at 0.70, leaving ~7 GB + for the colocated vocoder, CUDA graphs, and other GPU processes. + """ + + # (catalog key, engine stage name). Voxtral's engine stage is + # tts_generation; every other engine-bearing pipeline names it + # tts_engine. moss_tts_local is intentionally absent — its upstream + # config class already budgets its colocated stages explicitly + # (0.15 preprocessing / 0.67 AR / 0.18 vocoder), and pinning would + # fight that logic (the codec reserve derives from the fractions). + # s2-pro's OOM is model size, not budgeting (upstream issue #359). + PINNED = ( + ("qwen3_tts_0_6b_base", "tts_engine"), + ("qwen3_tts_0_6b_customvoice", "tts_engine"), + ("qwen3_tts_1_7b_base", "tts_engine"), + ("qwen3_tts_1_7b_voicedesign", "tts_engine"), + ("moss_tts", "tts_engine"), + ("voxtral_tts", "tts_generation"), + ) + UNPINNED = ("moss_tts_local", "fish_s2_pro") + + def test_engine_budget_is_pinned_under_the_engine_stage(self): + for key, stage in self.PINNED: + with self.subTest(key=key, stage=stage): + path = config_path(entry_by_key(key)) + self.assertIsNotNone(path) + self.assertTrue(path.is_file(), f"missing {path}") + text = path.read_text(encoding="utf-8") + self.assertIn( + "stages:\n" + f" {stage}:\n" + " engine:\n" + " mem_fraction_static: 0.70\n", + text) + + def test_already_budgeted_and_oversized_models_stay_unpinned(self): + for key in self.UNPINNED: + with self.subTest(key=key): + path = config_path(entry_by_key(key)) + self.assertIsNotNone(path) + self.assertTrue(path.is_file(), f"missing {path}") + text = path.read_text(encoding="utf-8") + self.assertNotIn("mem_fraction_static", text) + + class Fp8FallbackTests(unittest.TestCase): """FP8-only pipelines fall back to a vendored bf16 config on old GPUs.""" @@ -908,6 +985,13 @@ class SetupWizardTests(unittest.TestCase): self.assertEqual(settings["uninstall_keys"], []) self.assertEqual(confirms, []) + def test_setup_tree_starts_minimized(self): + """Like audio.cpp: no expand_all — every family starts collapsed.""" + from backends.sglomni import wizard + _settings, trees, _confirms = self._wizard([wizard._GO_BACK]) + self.assertEqual(trees[0][0], "Select SGLang-Omni Models to Install") + self.assertFalse(trees[0][1].get("expand_all", False)) + def test_steps_remove_before_downloading(self): from backends.sglomni import wizard gone, added = ENTRIES[0], ENTRIES[1] |