diff --git a/api/runner.py b/api/runner.py index 12f2bf23..9afccdd9 100644 --- a/api/runner.py +++ b/api/runner.py @@ -98,6 +98,46 @@ def _resolve_ready_schema(GenClass, node: dict, manifest: dict) -> list: return node.get("params_schema") or manifest.get("params_schema", []) +def _generator_is_loaded(gen) -> bool: + """ + Best-effort read of the generator's loaded state. + + Never raises: this is used on the error path, where a generator left in a + half-initialised state can make is_loaded() itself blow up. An unreadable + state is reported as "not loaded" so the caller reloads rather than reusing + a worker that may be broken. + """ + try: + return bool(gen.is_loaded()) + except Exception: + return False + + +def _ensure_model_loaded(gen) -> bool: + """ + Guarantees the model is in memory before inference. Returns True if a + reload was needed. + + Texture pipelines are built lazily on first use, and extensions typically + free the shape pipeline first to make room for them. When that setup fails + (a missing `xatlas` being the common case) the generator is left with + _model = None, yet the worker process stays alive and both ExtensionProcess + and GeneratorRegistry still consider it loaded — so load() is never called + again and every later generate() dies with + + TypeError: 'NoneType' object is not callable + + until the process is killed by hand. Re-loading here mirrors what + GeneratorRegistry.get_active() already does host-side, which makes a failed + setup recoverable on the next attempt instead of permanently poisoning the + worker. + """ + if _generator_is_loaded(gen): + return False + gen.load() + return True + + def _apply_manifest_metadata(gen, manifest: dict, node: dict) -> None: gen.hf_repo = node.get("hf_repo") or manifest.get("hf_repo", "") gen.hf_skip_prefixes = node.get("hf_skip_prefixes") or manifest.get("hf_skip_prefixes", []) @@ -167,6 +207,10 @@ def progress_cb(pct: int, step: str = "") -> None: send({"type": "progress", "id": rid, "pct": pct, "step": step}) try: + if _ensure_model_loaded(gen): + send({"type": "log", "level": "warning", + "message": ("Model was not loaded (earlier setup failure?); " + "reloaded before generating.")}) output_path = gen.generate(image_bytes, params, progress_cb, cancel_evt) send({"type": "done", "id": rid, "output_path": str(output_path)}) except Exception as exc: @@ -174,9 +218,13 @@ def progress_cb(pct: int, step: str = "") -> None: if type(exc).__name__ == "GenerationCancelled": send({"type": "cancelled", "id": rid}) else: + # Report whether the model survived the failure so the + # host can drop its cached "loaded" flag and reload + # instead of reusing a worker whose _model is None. send({"type": "error", "id": rid, "message": str(exc), - "traceback": traceback.format_exc()}) + "traceback": traceback.format_exc(), + "loaded": _generator_is_loaded(gen)}) finally: _cancel.pop(rid, None) diff --git a/api/services/extension_process.py b/api/services/extension_process.py index 341df6d6..dff48886 100644 --- a/api/services/extension_process.py +++ b/api/services/extension_process.py @@ -351,6 +351,13 @@ def generate( return Path(msg["output_path"]) elif t == "error": + # A failed generation can leave the worker without a model: a + # lazy texture-setup failure frees the shape pipeline and then + # raises. The worker reports its post-failure state, so drop + # our cached flag and let GeneratorRegistry.get_active() reload + # before the next run rather than reusing a broken worker. + if msg.get("loaded") is False: + self._loaded = False raise RuntimeError(msg.get("traceback") or msg.get("message", "Unknown error")) elif t == "cancelled": diff --git a/api/tests/test_extension_process.py b/api/tests/test_extension_process.py index 9cd5c5d6..e82e170f 100644 --- a/api/tests/test_extension_process.py +++ b/api/tests/test_extension_process.py @@ -142,5 +142,50 @@ def test_empty_queue_raises_timeout_error(self) -> None: proc._recv(timeout=0.05) +class GenerateErrorLoadedFlagTests(unittest.TestCase): + """ + Issue #239: a generation that fails during lazy texture setup leaves the + worker without a model. If _loaded stays True, GeneratorRegistry.get_active() + skips load() forever and every later run reuses the broken worker. + """ + + def _failing_generate(self, error_msg: dict) -> ExtensionProcess: + proc = _make_proc() + proc._loaded = True + proc._send = lambda msg: None # type: ignore[assignment] + proc._queue.put(error_msg) + with self.assertRaises(RuntimeError): + proc.generate(b"", {}) + return proc + + def test_clears_loaded_when_worker_reports_model_lost(self) -> None: + proc = self._failing_generate( + {"type": "error", "message": "No module named 'xatlas'", "loaded": False} + ) + self.assertFalse(proc._loaded) + + def test_keeps_loaded_when_worker_still_has_its_model(self) -> None: + proc = self._failing_generate( + {"type": "error", "message": "bad input image", "loaded": True} + ) + self.assertTrue(proc._loaded) + + def test_keeps_loaded_when_worker_reports_no_state(self) -> None: + proc = self._failing_generate({"type": "error", "message": "boom"}) + self.assertTrue(proc._loaded) + + def test_error_still_propagates_the_original_cause(self) -> None: + proc = _make_proc() + proc._loaded = True + proc._send = lambda msg: None # type: ignore[assignment] + proc._queue.put( + {"type": "error", "message": "short", "traceback": "full traceback here", + "loaded": False} + ) + with self.assertRaises(RuntimeError) as ctx: + proc.generate(b"", {}) + self.assertIn("full traceback here", str(ctx.exception)) + + if __name__ == "__main__": unittest.main() diff --git a/api/tests/test_runner.py b/api/tests/test_runner.py index f11a9c22..8fce3d31 100644 --- a/api/tests/test_runner.py +++ b/api/tests/test_runner.py @@ -154,5 +154,209 @@ def test_send_writes_single_json_line(self) -> None: self.assertEqual(json.loads(written), {"type": "ready", "params_schema": []}) +_FAKE_TEXGEN_GENERATOR = ''' +from pathlib import Path + +INSTANCES = [] + + +class FakeTexGen: + """ + Mimics an extension whose texture pipeline is built lazily on first use and + frees the shape pipeline first to make room for it. + + The first texture setup fails (missing `xatlas`), which is what issue #239 + reports; a later attempt succeeds once the dependency is available. + """ + + def __init__(self, model_dir, outputs_dir): + self.model_dir = model_dir + self.outputs_dir = outputs_dir + self._model = None + self.load_calls = 0 + self.texgen_attempts = 0 + INSTANCES.append(self) + + def is_loaded(self): + return self._model is not None + + def load(self): + self.load_calls += 1 + self._model = lambda image: "mesh" + + def unload(self): + self._model = None + + def _setup_texgen(self): + # Free the shape pipeline before building the texture pipeline. + self._model = None + self.texgen_attempts += 1 + if self.texgen_attempts == 1: + raise RuntimeError("No module named 'xatlas'") + # Texture setup succeeded: shape pipeline comes back. + self.load() + + def generate(self, image_bytes, params, progress_cb=None, cancel_event=None): + # Stands in for `self._model(image)` on a generator whose model is gone. + if self._model is None: + raise TypeError("'NoneType' object is not callable") + self._model(image_bytes) + if params.get("enable_texture"): + self._setup_texgen() + return Path("out.glb") +''' + + +class _RunnerDriver: + """Runs runner.main() against a throwaway extension dir.""" + + def __init__(self, generator_src: str, generator_class: str) -> None: + self.ext_dir = Path(tempfile.mkdtemp(prefix="modly-texgen-test-")) + (self.ext_dir / "generator.py").write_text(generator_src, encoding="utf-8") + (self.ext_dir / "manifest.json").write_text( + json.dumps({"id": "demo-ext", "generator_class": generator_class}), + encoding="utf-8", + ) + + def run(self, actions: list) -> list: + """Feeds actions on stdin, returns the parsed messages runner emitted.""" + original_ext_dir = runner.EXT_DIR + original_stdin = sys.stdin + original_module = sys.modules.pop("generator", None) + runner.EXT_DIR = self.ext_dir + sys.stdin = io.StringIO("".join(json.dumps(a) + "\n" for a in actions)) + out = io.StringIO() + try: + with redirect_stdout(out): + runner.main() + self.generator_module = sys.modules["generator"] + finally: + runner.EXT_DIR = original_ext_dir + sys.stdin = original_stdin + sys.modules.pop("generator", None) + if original_module is not None: + sys.modules["generator"] = original_module + return [json.loads(line) for line in out.getvalue().splitlines() if line.strip()] + + +class GeneratorLoadedStateTests(unittest.TestCase): + def test_reports_loaded_state(self) -> None: + gen = type("Gen", (), {"is_loaded": lambda self: True})() + self.assertTrue(runner._generator_is_loaded(gen)) + + def test_treats_raising_is_loaded_as_not_loaded(self) -> None: + class Gen: + def is_loaded(self): + raise RuntimeError("half-initialised") + + self.assertFalse(runner._generator_is_loaded(Gen())) + + def test_ensure_model_loaded_is_a_noop_when_already_loaded(self) -> None: + class Gen: + load_calls = 0 + + def is_loaded(self): + return True + + def load(self): + self.load_calls += 1 + + gen = Gen() + self.assertFalse(runner._ensure_model_loaded(gen)) + self.assertEqual(gen.load_calls, 0) + + def test_ensure_model_loaded_reloads_when_model_is_gone(self) -> None: + class Gen: + def __init__(self): + self._model = None + self.load_calls = 0 + + def is_loaded(self): + return self._model is not None + + def load(self): + self.load_calls += 1 + self._model = object() + + gen = Gen() + self.assertTrue(runner._ensure_model_loaded(gen)) + self.assertEqual(gen.load_calls, 1) + self.assertIsNotNone(gen._model) + + +class TextureSetupRecoveryTests(unittest.TestCase): + """ + Regression tests for issue #239: a failed lazy texture setup left the worker + with _model = None and no recovery path, so every later generation raised + "TypeError: 'NoneType' object is not callable" until the process was killed. + """ + + def setUp(self) -> None: + self.driver = _RunnerDriver(_FAKE_TEXGEN_GENERATOR, "FakeTexGen") + self.texture_run = { + "action": "generate", + "image_b64": "", + "params": {"enable_texture": True}, + } + + def test_worker_recovers_after_failed_texture_setup(self) -> None: + messages = self.driver.run([ + {"action": "load"}, + dict(self.texture_run, id="run-1"), + dict(self.texture_run, id="run-2"), + ]) + + by_id = {m.get("id"): m for m in messages if m.get("type") in ("done", "error")} + + # First run fails during texture setup and surfaces the real cause. + self.assertEqual(by_id["run-1"]["type"], "error") + self.assertIn("xatlas", by_id["run-1"]["message"]) + + # The retry must succeed instead of dying on a None model. + self.assertEqual( + by_id["run-2"]["type"], "done", + msg=f"retry did not recover: {by_id['run-2']}", + ) + self.assertNotIn("NoneType", json.dumps(by_id["run-2"])) + + # …and the worker's model is genuinely back. + gen = self.driver.generator_module.INSTANCES[0] + self.assertIsNotNone(gen._model) + self.assertTrue(gen.is_loaded()) + + def test_failed_run_reports_that_the_model_was_lost(self) -> None: + messages = self.driver.run([ + {"action": "load"}, + dict(self.texture_run, id="run-1"), + ]) + + error = next(m for m in messages if m.get("type") == "error") + self.assertIs(error["loaded"], False) + + def test_reload_before_generate_is_logged(self) -> None: + messages = self.driver.run([ + {"action": "load"}, + dict(self.texture_run, id="run-1"), + dict(self.texture_run, id="run-2"), + ]) + + logs = [m for m in messages if m.get("type") == "log"] + self.assertTrue( + any("reloaded before generating" in m.get("message", "") for m in logs), + msg=f"expected a reload log, got {logs}", + ) + + def test_successful_run_does_not_reload_the_model(self) -> None: + messages = self.driver.run([ + {"action": "load"}, + {"action": "generate", "id": "run-1", "image_b64": "", "params": {}}, + ]) + + self.assertEqual( + [m["type"] for m in messages if m.get("id") == "run-1"], ["done"] + ) + self.assertEqual(self.driver.generator_module.INSTANCES[0].load_calls, 1) + + if __name__ == "__main__": unittest.main()