Implement AI Host functionality with model caching and job management

- Added ai_host.py to handle AI tasks with FastAPI integration. - Introduced HostState and ModelCache classes for managing client connections and model loading. - Implemented job submission and cancellation endpoints in the AI host. - Enhanced tests for AI host functionality, including client tracking and job cancellation. - Updated existing tests to accommodate changes in the operator and command structure. - Refactored runtime commands to use request-based architecture instead of command-line arguments. - Improved model loading logic to handle CUDA memory management and reuse models efficiently.
2026-08-04 07:35:40 +08:00
parent d694980da7
commit 838af0ee96
16 changed files with 2112 additions and 225 deletions
--- a/tests/test_addon.py
+++ b/tests/test_addon.py
@@ -343,7 +343,7 @@ class AddonTest(unittest.TestCase):
        self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"])
        self.assertEqual(
            operator.__annotations__["snapform_process_res_factor"]["default"],
-            0.5,
+            1.0,
        )
        self.assertTrue(
            operator.__annotations__["snapform_delete_original"]["default"]
@@ -403,7 +403,7 @@ class AddonTest(unittest.TestCase):
                '"snapform_process_res_factor",'
            ),
            draw_source.index(
-                'layout.prop(self, "snapform_subdivision_level")'
+                'settings_body.prop(self, "snapform_subdivision_level")'
            ),
        )
        self.assertIn(
@@ -420,17 +420,30 @@ class AddonTest(unittest.TestCase):
            draw_source,
        )
        self.assertIn(
-            'layout.prop(self, "snapform_remove_background")',
+            'default_closed=True',
            draw_source,
        )
+        self.assertIn('settings_header.label(text="Options")', draw_source)
+        self.assertIn('settings_header.label(text="Options")', draw_source)
        self.assertNotIn("layout.separator", draw_source)
-        self.assertNotIn("layout.label", draw_source)
-        self.assertIn(
-            'layout.prop(self, "snapform_unshaded")',
-            draw_source,
+        setting_fields = (
+            "snapform_displacement_via",
+            "snapform_subdivision_level",
+            "snapform_remove_background",
+            "snapform_unshaded",
        )
-        model_items = operator.__annotations__["snapform_model"]["items"]
-        model_items = model_items(None, None)
+        field_positions = [
+            draw_source.index(f'settings_body.prop(self, "{field}")')
+            for field in setting_fields
+        ]
+        self.assertEqual(field_positions, sorted(field_positions))
+        model_items_callback = operator.__annotations__["snapform_model"]["items"]
+        with patch.object(
+            self.addon.runtime,
+            "model_ready",
+            return_value=True,
+        ):
+            model_items = model_items_callback(None, None)
        self.assertEqual(
            [item[0] for item in model_items],
            [
@@ -459,11 +472,33 @@ class AddonTest(unittest.TestCase):
        )

        self.addon.image_depth._last_depth_plane_model = "GIANT_1_1"
-        result = operator.invoke(context, None)
+        with patch.object(
+            self.addon.runtime,
+            "model_ready",
+            side_effect=lambda key: key == "GIANT_1_1",
+        ):
+            result = operator.invoke(context, None)

        self.assertIs(result, operator)
        self.assertEqual(operator.snapform_model, "GIANT_1_1")

+    def test_conversion_dialog_only_lists_installed_models(self):
+        items_callback = self.addon.ConvertImageToDepthPlane.__annotations__[
+            "snapform_model"
+        ]["items"]
+        with patch.object(
+            self.addon.runtime,
+            "model_ready",
+            side_effect=lambda key: key in {"SMALL", "BASE"},
+        ):
+            items = items_callback(None, None)
+
+        self.assertEqual([item[0] for item in items], ["SMALL", "BASE"])
+        draw_source = inspect.getsource(self.addon.ConvertImageToDepthPlane.draw)
+        self.assertIn('text="No DA3 models installed"', draw_source)
+        self.assertIn('text="Download one in Host > Models"', draw_source)
+        self.assertNotIn("_dialog_requirements_ready", inspect.getsource(self.addon.ConvertImageToDepthPlane))
+
    def test_depth_conversion_progress_is_provider_neutral(self):
        starting_source = inspect.getsource(
            self.addon.RunJob._starting_message
@@ -532,7 +567,11 @@ class AddonTest(unittest.TestCase):
                self.addon.image_depth,
                "save_packed_image",
                return_value=packed_path,
-            ) as save_packed_image:
+            ) as save_packed_image, patch.object(
+                self.addon.runtime,
+                "model_ready",
+                return_value=True,
+            ):
                result = operator.execute(context)

        self.assertEqual(result, {"FINISHED"})
@@ -640,8 +679,10 @@ class AddonTest(unittest.TestCase):
                    return_value=True,
                ) as model_ready, patch.object(
                    self.addon.runtime,
-                    "reconstruct_command",
-                    return_value=["worker"],
+                    "reconstruct_request",
+                    return_value=self.addon.runtime.HostJobRequest(
+                        {"job_id": "test"}
+                    ),
                ) as reconstruct:
                    operator.snapform_model_key = snapform_model
                    operator.snapform_use_ray_pose = snapform_ray_pose
@@ -650,7 +691,10 @@ class AddonTest(unittest.TestCase):
                        settings,
                    )

-                self.assertEqual(command, ["worker"])
+                self.assertIsInstance(
+                    command,
+                    self.addon.runtime.HostJobRequest,
+                )
                self.assertEqual(prepared_paths, paths)
                model_ready.assert_called_once_with(snapform_model)
                self.assertEqual(reconstruct.call_args.args[1], "mini_npz")
@@ -713,7 +757,7 @@ class AddonTest(unittest.TestCase):
                with self.assertRaisesRegex(RuntimeError, "BEN2 Base"):
                    operator._prepare_command(None, settings)

-    def test_image_depth_background_option_reaches_worker_command(self):
+    def test_image_depth_background_option_reaches_host_request(self):
        runtime = importlib.import_module("snapform.runtime")
        settings = SimpleNamespace(
            input_path="input.png",
@@ -739,7 +783,7 @@ class AddonTest(unittest.TestCase):
            "model_directory",
            side_effect=lambda key: Path(key.lower()),
        ):
-            command = runtime.reconstruct_command(
+            request = runtime.reconstruct_request(
                settings,
                "mini_npz",
                Path("status.json"),
@@ -747,8 +791,10 @@ class AddonTest(unittest.TestCase):
                remove_background=True,
            )

-        option_index = command.index("--remove-background-model-dir")
-        self.assertEqual(command[option_index + 1], str(Path("ben2_base")))
+        self.assertEqual(
+            request.payload["parameters"]["remove_background_model_dir"],
+            str(Path("ben2_base")),
+        )

    def test_preferences_use_package_identifier(self):
        self.assertEqual(
@@ -767,6 +813,25 @@ class AddonTest(unittest.TestCase):
            str(runtime.DEFAULT_STORAGE_ROOT),
        )

+    def test_ai_host_uses_authenticated_loopback_http(self):
+        runtime = importlib.import_module("snapform.runtime")
+        request_source = inspect.getsource(runtime._host_http_request)
+        start_source = inspect.getsource(runtime.start_host)
+
+        self.assertEqual(runtime.HOST_ADDRESS, "127.0.0.1")
+        self.assertIn('"X-SnapForm-Token"', request_source)
+        self.assertIn('"/health"', inspect.getsource(runtime._check_host))
+        self.assertIn('with_name("ai_host.py")', start_source)
+        self.assertIn("ProxyHandler({})", inspect.getsource(runtime))
+        self.assertIn('return storage_root() / "host.json"', inspect.getsource(runtime.host_discovery_path))
+
+    def test_host_jobs_complete_from_status_state(self):
+        source = inspect.getsource(self.addon.RunJob.modal)
+
+        self.assertIn("if job.uses_host:", source)
+        self.assertIn('state == "succeeded"', source)
+        self.assertIn('{"failed", "cancelled"}', source)
+
    def test_process_res_method_defaults_to_upper_bound_resize(self):
        process_res = self.addon.SnapFormSettings.__annotations__[
            "process_res"
@@ -780,7 +845,7 @@ class AddonTest(unittest.TestCase):
        self.assertEqual(process_res["default"], 504)
        self.assertEqual(process_res_method["default"], "upper_bound_resize")

-    def test_reconstruct_command_preserves_da3_parameter_names(self):
+    def test_reconstruct_request_preserves_da3_parameter_names(self):
        runtime = importlib.import_module("snapform.runtime")
        settings = SimpleNamespace(
            input_path="input.png",
@@ -806,29 +871,30 @@ class AddonTest(unittest.TestCase):
            "model_directory",
            return_value=Path("model"),
        ):
-            command = runtime.reconstruct_command(
+            request = runtime.reconstruct_request(
                settings,
                "gs_ply",
                Path("status.json"),
                Path("result"),
            )

-        self.assertIn("--export-format", command)
-        self.assertIn("gs_ply", command)
+        parameters = request.payload["parameters"]
+        self.assertEqual(request.payload["command"], "reconstruct")
+        self.assertEqual(parameters["export_format"], "gs_ply")
        for parameter in (
-            "--process-res",
-            "--process-res-method",
-            "--conf-thresh-percentile",
-            "--num-max-points",
-            "--show-cameras",
-            "--gs-views-interval",
-            "--use-ray-pose",
-            "--ref-view-strategy",
+            "process_res",
+            "process_res_method",
+            "conf_thresh_percentile",
+            "num_max_points",
+            "show_cameras",
+            "gs_views_interval",
+            "use_ray_pose",
+            "ref_view_strategy",
        ):
            with self.subTest(parameter=parameter):
-                self.assertIn(parameter, command)
+                self.assertIn(parameter, parameters)

-    def test_reconstruct_command_accepts_image_depth_overrides(self):
+    def test_reconstruct_request_accepts_image_depth_overrides(self):
        runtime = importlib.import_module("snapform.runtime")
        settings = SimpleNamespace(
            input_path="input.png",
@@ -854,7 +920,7 @@ class AddonTest(unittest.TestCase):
            "model_directory",
            return_value=Path("large-model"),
        ) as snapform_model_directory:
-            snapform_command = runtime.reconstruct_command(
+            request = runtime.reconstruct_request(
                settings,
                "mini_npz",
                Path("status.json"),
@@ -866,15 +932,13 @@ class AddonTest(unittest.TestCase):
            )

        snapform_model_directory.assert_called_once_with("BASE")
-        snapform_process_res_index = snapform_command.index("--process-res")
-        self.assertEqual(snapform_command[snapform_process_res_index + 1], "252")
-        snapform_process_res_method_index = snapform_command.index("--process-res-method")
+        parameters = request.payload["parameters"]
+        self.assertEqual(parameters["process_res"], 252)
        self.assertEqual(
-            snapform_command[snapform_process_res_method_index + 1],
+            parameters["process_res_method"],
            "lower_bound_resize",
        )
-        snapform_ray_pose_index = snapform_command.index("--use-ray-pose")
-        self.assertEqual(snapform_command[snapform_ray_pose_index + 1], "0")
+        self.assertEqual(parameters["use_ray_pose"], 0)

    def test_export_format_property_includes_all_da3_formats(self):
        export_format = self.addon.SnapFormSettings.__annotations__[
@@ -904,37 +968,139 @@ class AddonTest(unittest.TestCase):
            "DEFAULT_CLOSED",
            self.addon.DebugPanel.bl_options,
        )
+        for panel in (self.addon.Da3DebugPanel, self.addon.Ben2DebugPanel):
+            self.assertEqual(panel.bl_parent_id, self.addon.DebugPanel.bl_idname)
+            self.assertIn("DEFAULT_CLOSED", panel.bl_options)

    def test_debug_panel_uses_one_export_button_and_dynamic_settings(self):
-        draw_source = inspect.getsource(
-            self.addon.DebugPanel.draw
-        )
+        debug_source = inspect.getsource(self.addon.DebugPanel.draw)
+        da3_source = inspect.getsource(self.addon.Da3DebugPanel.draw)
+        ben2_source = inspect.getsource(self.addon.Ben2DebugPanel.draw)

-        self.assertEqual(draw_source.count("export_row.operator("), 1)
-        self.assertIn('settings.export_format == "glb"', draw_source)
-        self.assertIn('settings.export_format == "gs_ply"', draw_source)
-        self.assertIn('"gs_model" if settings.export_format', draw_source)
-        self.assertIn("layout.prop(settings, model_property)", draw_source)
-        self.assertNotIn("model_row", draw_source)
-        self.assertIn("export_row.enabled = active_job() is None", draw_source)
-        self.assertNotIn('text="Ready"', draw_source)
-        self.assertNotIn('text="Missing"', draw_source)
-        self.assertNotIn("License:", draw_source)
-        self.assertIn('operator.job_type = "RECONSTRUCT"', draw_source)
-        self.assertIn('text="BEN2 Background Removal"', draw_source)
+        self.assertIn('layout.prop(settings, "device")', debug_source)
+        self.assertEqual(da3_source.count("export_row.operator("), 1)
+        self.assertIn('settings.export_format == "glb"', da3_source)
+        self.assertIn('settings.export_format == "gs_ply"', da3_source)
+        self.assertIn('"gs_model" if settings.export_format', da3_source)
+        self.assertIn("layout.prop(settings, model_property)", da3_source)
+        self.assertNotIn("model_row", da3_source)
+        self.assertIn("and model_ready(selected_model_key(settings))", da3_source)
+        self.assertNotIn('text="Ready"', da3_source)
+        self.assertNotIn('text="Missing"', da3_source)
+        self.assertNotIn("License:", da3_source)
+        self.assertIn('operator.job_type = "RECONSTRUCT"', da3_source)
        self.assertIn(
-            'background_box.prop(settings, "background_input_path")',
-            draw_source,
+            'layout.prop(settings, "background_input_path")',
+            ben2_source,
        )
-        self.assertIn('text="Remove Background"', draw_source)
-        self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', draw_source)
+        self.assertIn('text="Remove Background"', ben2_source)
+        self.assertIn("and model_ready(BEN2_MODEL.key)", ben2_source)
+        self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', ben2_source)

-    def test_runtime_models_panel_assigns_each_download_target(self):
+    def test_host_panel_exposes_shared_host_state_and_controls(self):
+        self.assertIn(self.addon.HostPanel, self.addon.CLASSES)
+        self.assertIn(self.addon.ControlHost, self.addon.CLASSES)
+        source = inspect.getsource(self.addon.HostPanel.draw)
+
+        self.assertEqual(self.addon.HostPanel.bl_label, "Host")
+        self.assertIn('text="Runtime is not installed"', source)
+        self.assertIn('text="Install Runtime"', source)
+        self.assertIn("runtime.host_status()", source)
+        self.assertIn('f"Running (Clients:', source)
+        self.assertIn('else "Running (Busy)"', source)
+        self.assertIn('"READY": "RADIOBUT_ON"', source)
+        self.assertIn('"BUSY": "TIME"', source)
+        self.assertNotIn('"READY": "PLAY"', source)
+        self.assertNotIn('text=f"Task:', source)
+        self.assertIn('elif state == "STOPPED":', source)
+        self.assertNotIn("Blender Clients:", source)
+        self.assertNotIn("status['port']", source)
+        self.assertIn('icon="TRASH"', source)
+        self.assertIn('icon="FILE_REFRESH"', source)
+        self.assertIn('icon="PAUSE"', source)
+        self.assertNotIn('text="Open Log"', source)
+        preferences_source = inspect.getsource(self.addon.SnapFormPreferences.draw)
+        self.assertIn('text="Open Host Log"', preferences_source)
+
+    def test_host_start_control_does_not_wait_for_model_runtime(self):
+        source = inspect.getsource(self.addon.runtime.enable_host)
+
+        self.assertIn("start_host()", source)
+        self.assertNotIn("ensure_host()", source)
+
+    def test_ai_operation_restarts_a_manually_stopped_host(self):
+        runtime = self.addon.runtime
+        connection = runtime.HostConnection(51234, "token", 42, "host")
+        health = {
+            "host_schema": runtime.HOST_SCHEMA,
+            "instance_id": "host",
+            "busy": False,
+            "clients": 1,
+            "cached_models": {},
+        }
+        previous_auto_start = runtime._host_auto_start
+        previous_connection = runtime._host_connection
+        previous_snapshot = runtime._host_snapshot
+        runtime._host_auto_start = False
+        runtime._host_connection = None
+        runtime._host_snapshot = {"state": "STOPPED", "message": "Stopped"}
+        try:
+            with patch.object(
+                runtime,
+                "_discover_live_host",
+                return_value=(connection, health),
+            ), patch.object(
+                runtime,
+                "_register_host_client",
+            ), patch.object(
+                runtime,
+                "_check_host",
+                return_value=health,
+            ), patch.object(runtime, "_redraw_host_ui"):
+                result = runtime.ensure_host(timeout=0.1)
+
+            self.assertIs(result, connection)
+            self.assertTrue(runtime._host_auto_start)
+            self.assertEqual(runtime.host_status()["state"], "READY")
+        finally:
+            runtime._host_auto_start = previous_auto_start
+            runtime._host_connection = previous_connection
+            runtime._host_snapshot = previous_snapshot
+
+    def test_blender_startup_does_not_start_ai_host(self):
+        runtime = self.addon.runtime
+        previous_auto_start = runtime._host_auto_start
+        previous_snapshot = runtime._host_snapshot
+        runtime._host_auto_start = False
+        runtime._host_snapshot = {"state": "STOPPED", "message": "Stopped"}
+        try:
+            with patch.object(
+                runtime,
+                "runtime_ready",
+                return_value=True,
+            ), patch.object(
+                runtime,
+                "_discover_live_host",
+                return_value=(None, None),
+            ), patch.object(
+                runtime,
+                "_close_finished_launched_host",
+            ), patch.object(runtime, "start_host") as start_host:
+                interval = runtime._service_host_timer()
+
+            start_host.assert_not_called()
+            self.assertEqual(interval, 2.0)
+            self.assertEqual(runtime.host_status()["state"], "STOPPED")
+        finally:
+            runtime._host_auto_start = previous_auto_start
+            runtime._host_snapshot = previous_snapshot
+
+    def test_models_panel_is_collapsed_and_assigns_each_download_target(self):
        draw_source = inspect.getsource(
-            self.addon.RuntimeModelsPanel.draw
+            self.addon.ModelsPanel.draw
        )
        model_source = inspect.getsource(
-            self.addon.RuntimeModelsPanel._draw_model
+            self.addon.ModelsPanel._draw_model
        )

        self.assertIn("for model in MODELS.values()", draw_source)
@@ -951,8 +1117,16 @@ class AddonTest(unittest.TestCase):
        self.assertNotIn("Isolated Runtime ready", draw_source)
        self.assertNotIn("License:", model_source)
        self.assertEqual(
-            self.addon.RuntimeModelsPanel.bl_label,
-            "Setup",
+            self.addon.ModelsPanel.bl_label,
+            "Models",
+        )
+        self.assertEqual(
+            self.addon.ModelsPanel.bl_parent_id,
+            self.addon.HostPanel.bl_idname,
+        )
+        self.assertEqual(
+            self.addon.ModelsPanel.bl_options,
+            {"DEFAULT_CLOSED"},
        )

    def test_ben2_debug_job_uses_its_image_and_model(self):
@@ -987,15 +1161,17 @@ class AddonTest(unittest.TestCase):
                return_value=True,
            ) as model_ready, patch.object(
                self.addon.runtime,
-                "remove_background_command",
-                return_value=["worker"],
+                "remove_background_request",
+                return_value=self.addon.runtime.HostJobRequest(
+                    {"job_id": "test"}
+                ),
            ) as remove_background:
                command, prepared_paths = operator._prepare_command(
                    None,
                    settings,
                )

-        self.assertEqual(command, ["worker"])
+        self.assertIsInstance(command, self.addon.runtime.HostJobRequest)
        self.assertEqual(prepared_paths, paths)
        model_ready.assert_called_once_with("BEN2_BASE")
        self.assertEqual(remove_background.call_args.args[1], input_path)
@@ -1030,15 +1206,17 @@ class AddonTest(unittest.TestCase):
                return_value=True,
            ), patch.object(
                self.addon.runtime,
-                "remove_background_command",
-                return_value=["worker"],
+                "remove_background_request",
+                return_value=self.addon.runtime.HostJobRequest(
+                    {"job_id": "test"}
+                ),
            ) as remove_background:
                command, prepared_paths = operator._prepare_command(
                    None,
                    settings,
                )

-        self.assertEqual(command, ["worker"])
+        self.assertIsInstance(command, self.addon.runtime.HostJobRequest)
        self.assertEqual(prepared_paths, paths)
        self.assertEqual(remove_background.call_args.args[1], input_path)