| import os |
| import tempfile |
| from pathlib import Path |
| import unittest |
|
|
|
|
| os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") |
| os.environ["MUSE_SKIP_MODEL_LOAD"] = "1" |
|
|
| from PIL import Image |
|
|
| import app |
|
|
|
|
| class GradioContractTests(unittest.TestCase): |
| def test_click_and_enter_have_validated_serial_generators(self): |
| dependencies = app.demo.get_config_file()["dependencies"] |
| generators = [dependency for dependency in dependencies if dependency["types"]["generator"]] |
| self.assertGreaterEqual(len(generators), 2) |
|
|
| targets = set() |
| api_generator = None |
| for dependency in generators: |
| targets.update(target[1] for target in dependency["targets"]) |
| self.assertTrue(dependency["queue"]) |
| self.assertEqual(dependency.get("api_visibility", "public"), "private") |
| self.assertEqual(len(dependency["outputs"]), 6) |
| self.assertEqual(len(dependency["inputs"]), 16) |
| function = app.demo.fns[dependency["id"]] |
| self.assertEqual(function.concurrency_id, "muse-glimmer-xlarge") |
| self.assertEqual(function.concurrency_limit, 1) |
| self.assertIsNotNone(function.validator) |
| if dependency["api_name"] == "chat": |
| api_generator = dependency |
| if dependency["api_name"] == app.SUBMIT_API_NAME: |
| self.assertEqual(dependency["types"]["generator"], True) |
|
|
| self.assertEqual(targets, {"click", "submit"}) |
| self.assertIsNotNone(api_generator, "Expected one private API-visible generation path named chat.") |
|
|
| generator_functions = [app.demo.fns[dependency["id"]] for dependency in generators] |
| self.assertEqual( |
| {function.concurrency_id for function in generator_functions}, |
| {"muse-glimmer-xlarge"}, |
| ) |
| self.assertTrue(all(function.concurrency_limit == 1 for function in generator_functions)) |
| self.assertTrue(all(function.validator is not None for function in generator_functions)) |
|
|
| api_names = set() |
| for dependency in generators: |
| api_names.add(dependency["api_name"]) |
| self.assertEqual(api_names, {"chat", app.SUBMIT_API_NAME}) |
|
|
| def test_stop_and_clear_cancel_click_and_enter_generation(self): |
| dependencies = app.demo.get_config_file()["dependencies"] |
| generator_ids = { |
| dependency["id"] for dependency in dependencies if dependency["types"]["generator"] |
| } |
| cancellation_edges = [ |
| set(dependency["cancels"]) |
| for dependency in dependencies |
| if dependency["types"]["cancel"] |
| ] |
|
|
| self.assertEqual(len(cancellation_edges), 2) |
| self.assertTrue(all(edge == generator_ids for edge in cancellation_edges)) |
|
|
| def test_invalid_requests_fail_the_queue_free_validator(self): |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp_path = Path(tmpdir) |
| original_paths = { |
| model_id: spec["path"] |
| for model_id, spec in app.MODEL_REGISTRY.items() |
| } |
| for spec in app.MODEL_REGISTRY.values(): |
| spec["path"] = tmp_path |
|
|
| try: |
| verdicts = app._validate_generation_request( |
| "", |
| None, |
| app.MODEL_DEFAULT_ID, |
| [], |
| [], |
| "", |
| "high", |
| False, |
| 512, |
| 1.0, |
| 0.95, |
| 64, |
| 1.0, |
| 42, |
| False, |
| True, |
| ) |
| finally: |
| for model_id, original_path in original_paths.items(): |
| app.MODEL_REGISTRY[model_id]["path"] = original_path |
|
|
| self.assertEqual(len(verdicts), 16) |
| self.assertFalse(verdicts[0]["is_valid"]) |
| self.assertIn("prompt", verdicts[0]["message"].lower()) |
|
|
| def test_model_id_payload_formats_from_ui_are_accepted(self): |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp_path = Path(tmpdir) |
| original_paths = { |
| model_id: spec["path"] |
| for model_id, spec in app.MODEL_REGISTRY.items() |
| } |
| for spec in app.MODEL_REGISTRY.values(): |
| spec["path"] = tmp_path |
|
|
| try: |
| verdict = app._validate_generation_request( |
| "hi", |
| None, |
| ("Muse Glimmer 30B-assistant", app.ASSISTANT_MODEL_ID), |
| [], |
| [], |
| "", |
| "high", |
| False, |
| 512, |
| 1.0, |
| 0.95, |
| 64, |
| 1.0, |
| 42, |
| False, |
| True, |
| )[0] |
| finally: |
| for model_id, original_path in original_paths.items(): |
| app.MODEL_REGISTRY[model_id]["path"] = original_path |
|
|
| self.assertTrue(verdict["is_valid"]) |
|
|
| def test_legacy_model_labels_are_accepted_by_coercion(self): |
| self.assertEqual(app._coerce_model_id("/Muse-Glimmer 30B"), app.MODEL_ID) |
| self.assertEqual(app._coerce_model_id("/Muse-Glimmer-30B"), app.MODEL_ID) |
| self.assertEqual(app._coerce_model_id("Muse Glimmer 30B (full BF16)"), app.MODEL_ID) |
| self.assertEqual(app._coerce_model_id("Muse Glimmer 30B-assistant (compact)"), app.ASSISTANT_MODEL_ID) |
| self.assertEqual(app._coerce_model_id(""), app.MODEL_ID) |
|
|
| def test_validator_defaults_when_controls_are_missing(self): |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp_path = Path(tmpdir) |
| original_paths = { |
| model_id: spec["path"] |
| for model_id, spec in app.MODEL_REGISTRY.items() |
| } |
| for spec in app.MODEL_REGISTRY.values(): |
| spec["path"] = tmp_path |
|
|
| try: |
| verdict = app._validate_generation_request( |
| "hi", |
| None, |
| app.MODEL_DEFAULT_ID, |
| [], |
| [], |
| "", |
| "high", |
| False, |
| None, |
| None, |
| None, |
| None, |
| None, |
| 42, |
| False, |
| True, |
| )[0] |
| finally: |
| for model_id, original_path in original_paths.items(): |
| app.MODEL_REGISTRY[model_id]["path"] = original_path |
|
|
| self.assertTrue(verdict["is_valid"]) |
|
|
| def test_numeric_model_id_payload_formats_from_ui_are_accepted(self): |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp_path = Path(tmpdir) |
| original_paths = { |
| model_id: spec["path"] |
| for model_id, spec in app.MODEL_REGISTRY.items() |
| } |
| for spec in app.MODEL_REGISTRY.values(): |
| spec["path"] = tmp_path |
|
|
| try: |
| verdict = app._validate_generation_request( |
| "hi", |
| None, |
| "1", |
| [], |
| [], |
| "", |
| "high", |
| False, |
| 512, |
| 1.0, |
| 0.95, |
| 64, |
| 1.0, |
| 42, |
| False, |
| True, |
| )[0] |
| finally: |
| for model_id, original_path in original_paths.items(): |
| app.MODEL_REGISTRY[model_id]["path"] = original_path |
|
|
| self.assertTrue(verdict["is_valid"]) |
|
|
|
|
| class StateContractTests(unittest.TestCase): |
| def test_only_two_recent_images_remain_in_model_history(self): |
| image = Image.new("RGB", (8, 8), "purple") |
| history = [] |
| for index in range(3): |
| history.extend( |
| [ |
| { |
| "role": "user", |
| "content": [ |
| {"type": "image", "image": image.copy()}, |
| {"type": "text", "text": f"turn {index}"}, |
| ], |
| }, |
| {"role": "assistant", "content": f"answer {index}"}, |
| ] |
| ) |
|
|
| cleaned = app._clean_model_history(history) |
| image_turns = [ |
| message |
| for message in cleaned |
| if isinstance(message.get("content"), list) |
| and any(part.get("type") == "image" for part in message["content"]) |
| ] |
| self.assertEqual(len(image_turns), 2) |
| self.assertEqual(cleaned[0]["content"], "turn 0") |
|
|
| def test_stop_rolls_back_visible_chat_snapshot(self): |
| snapshot = [{"role": "user", "content": "committed"}] |
| chat, status = app._stop_conversation(snapshot) |
|
|
| self.assertEqual(chat, snapshot) |
| self.assertIsNot(chat, snapshot) |
| self.assertIn("not added", status) |
|
|
|
|
| class DurationContractTests(unittest.TestCase): |
| def setUp(self): |
| app.ACTIVE_MODEL_ID = None |
|
|
| def tearDown(self): |
| app.ACTIVE_MODEL_ID = app.MODEL_ID |
|
|
| def test_full_model_cold_start_requests_cold_start_budget(self): |
| duration = app._gpu_duration( |
| "hi", |
| None, |
| app.MODEL_ID, |
| [], |
| [], |
| "", |
| "high", |
| False, |
| 512, |
| 1.0, |
| 0.95, |
| 64, |
| 1.0, |
| 42, |
| False, |
| True, |
| ) |
| self.assertEqual(duration, 120) |
|
|
| def test_full_model_reuse_keeps_estimate_after_load(self): |
| app.ACTIVE_MODEL_ID = app.MODEL_ID |
| duration = app._gpu_duration( |
| "hi", |
| None, |
| app.MODEL_ID, |
| [], |
| [], |
| "", |
| "high", |
| False, |
| 512, |
| 1.0, |
| 0.95, |
| 64, |
| 1.0, |
| 42, |
| False, |
| True, |
| ) |
| self.assertEqual(duration, 121) |
|
|
| def test_assistant_selection_reuses_loaded_full_model(self): |
| app.ACTIVE_MODEL_ID = app.MODEL_ID |
| duration = app._gpu_duration( |
| "hi", |
| None, |
| app.ASSISTANT_MODEL_ID, |
| [], |
| [], |
| "", |
| "high", |
| False, |
| 512, |
| 1.0, |
| 0.95, |
| 64, |
| 1.0, |
| 42, |
| False, |
| True, |
| ) |
| self.assertEqual(duration, 121) |
|
|
|
|
| class RuntimeChatObjectTests(unittest.TestCase): |
| def test_coerce_chat_objects_prefers_nested_tokenizer(self): |
| class FakeResponseToken: |
| def __init__(self, is_processor=True): |
| self.is_processor = is_processor |
|
|
| def apply_chat_template(self, *_args, **_kwargs): |
| return {"input_ids": []} |
|
|
| def get_response_parser(self, *_args, **_kwargs): |
| return "parser" |
|
|
| class FakeProcessor: |
| def __init__(self): |
| self.tokenizer = FakeResponseToken() |
|
|
| original_paths = {model_id: spec["path"] for model_id, spec in app.MODEL_REGISTRY.items()} |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp_path = Path(tmpdir) |
| for spec in app.MODEL_REGISTRY.values(): |
| spec["path"] = tmp_path |
| try: |
| processor, tokenizer = app._coerce_chat_objects(FakeProcessor(), app.MODEL_ID) |
| finally: |
| for model_id, original_path in original_paths.items(): |
| app.MODEL_REGISTRY[model_id]["path"] = original_path |
|
|
| self.assertIsInstance(processor, FakeResponseToken) |
| self.assertIs(processor, tokenizer) |
|
|
|
|
| class ModelPathResolutionTests(unittest.TestCase): |
| def test_mount_path_resolution_prefers_nested_checkpoint(self): |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| nested = tmp / "Muse-Glimmer-30B" |
| nested.mkdir() |
| (nested / "config.json").write_text("{}") |
| (nested / "chat_template.jinja").write_text("") |
| resolved = app._resolve_mount_path(tmp) |
|
|
| self.assertEqual(resolved.name, "Muse-Glimmer-30B") |
|
|
|
|
| def test_mount_path_resolution_keeps_direct_checkpoint(self): |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| (tmp / "config.json").write_text("{}") |
| (tmp / "chat_template.jinja").write_text("") |
| resolved = app._resolve_mount_path(tmp) |
|
|
| self.assertEqual(resolved, tmp) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
