icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement AI Host functionality with model caching and job management

Browse Source 
- 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.
This commit is contained in:
马南
2026-08-04 07:35:40 +08:00
parent d694980da7
commit 838af0ee96
16 changed files with 2112 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files

316
tests/test_addon.py
View File

@@ -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)

223
tests/test_ai_host.py Normal file
View File

@@ -0,0 +1,223 @@
import importlib.util
import inspect
import sys
import tempfile
import unittest
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_ai_host():
specification = importlib.util.spec_from_file_location(
"snapform_ai_host",
PROJECT_ROOT / "src" / "snapform" / "ai_host.py",
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
return module
class FakeCuda:
class OutOfMemoryError(RuntimeError):
pass
@staticmethod
def is_available():
return False
@staticmethod
def empty_cache():
return None
class HostTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.host = load_ai_host()
def create_cache(self):
cache = self.host.ModelCache.__new__(self.host.ModelCache)
cache.torch = SimpleNamespace(cuda=FakeCuda())
cache.ben2_model = None
cache.ben2_key = None
cache.da3_model = None
cache.da3_key = None
return cache
def test_ben2_model_is_reused_for_matching_runtime_key(self):
loads = []
class FakeModel:
@classmethod
def from_pretrained(cls, model_dir):
loads.append(model_dir)
return cls()
def to(self, device):
self.device = device
return self
def eval(self):
return None
ben2 = ModuleType("ben2")
ben2.BEN_Base = FakeModel
cache = self.create_cache()
device = SimpleNamespace(type="cuda")
with patch.dict(sys.modules, {"ben2": ben2}):
first, _load_ms = cache.get_ben2("model", device)
second, cached_load_ms = cache.get_ben2("model", device)
self.assertIs(first, second)
self.assertEqual(loads, ["model"])
self.assertEqual(cache.status()["ben2"], "model")
self.assertEqual(cached_load_ms, 0.0)
def test_da3_model_is_reused_for_matching_runtime_key(self):
loads = []
class FakeModel:
@classmethod
def from_pretrained(cls, model_dir):
loads.append(model_dir)
return cls()
def to(self, device):
self.device = device
return self
def eval(self):
return None
package = ModuleType("depth_anything_3")
api = ModuleType("depth_anything_3.api")
api.DepthAnything3 = FakeModel
cache = self.create_cache()
device = SimpleNamespace(type="cuda")
with patch.dict(
sys.modules,
{
"depth_anything_3": package,
"depth_anything_3.api": api,
},
):
first = cache.get_da3("model", device)
second = cache.get_da3("model", device)
self.assertIs(first, second)
self.assertEqual(loads, ["model"])
self.assertEqual(cache.status()["da3"], "da3_model")
def test_http_host_has_health_jobs_auth_and_single_task_guard(self):
source = inspect.getsource(self.host.create_app)
self.assertIn('@app.get("/health")', source)
self.assertIn('@app.post("/jobs", status_code=202)', source)
self.assertIn("x_snapform_token", source)
self.assertIn("hmac.compare_digest", source)
self.assertIn("parameters[\"status\"]", source)
self.assertIn('@app.post("/clients", status_code=201)', source)
self.assertIn('@app.post("/clients/{client_id}/heartbeat")', source)
self.assertIn('@app.delete("/jobs/{job_id}", status_code=202)', source)
self.assertIn('@app.post("/models/clear")', source)
self.assertIn('@app.post("/shutdown", status_code=202)', source)
self.assertNotIn("os._exit", source)
def test_host_tracks_clients_and_rejects_shared_shutdown(self):
fake_models = SimpleNamespace(
status=lambda: {"ben2": "", "da3": ""},
close=lambda: None,
)
with patch.object(self.host, "ModelCache", return_value=fake_models):
state = self.host.HostState("token", "instance")
state.register_client("first", "Blender A")
state.register_client("second", "Blender B")
self.assertEqual(state.snapshot()["clients"], 2)
self.assertEqual(state.request_shutdown("first"), "shared")
self.assertTrue(
state.reserve("job", "first", "reconstruct", "status.json")
)
self.assertFalse(
state.reserve("other", "second", "reconstruct", "other.json")
)
def test_cancelling_job_keeps_host_and_model_cache_alive(self):
fake_models = SimpleNamespace(
status=lambda: {"ben2": "BEN2_BASE", "da3": ""},
close=lambda: None,
)
with patch.object(self.host, "ModelCache", return_value=fake_models):
state = self.host.HostState("token", "instance")
state.register_client("client", "Blender")
with tempfile.TemporaryDirectory() as directory:
status_path = Path(directory) / "status.json"
state.reserve(
"job",
"client",
"remove-background",
str(status_path),
)
active = state.cancel_job("job", "client")
self.assertTrue(active["cancel_requested"])
self.assertTrue(state.is_job_cancelled("job"))
state.run_job(
{
"job_id": "job",
"command": "remove-background",
"parameters": {"status": str(status_path)},
}
)
status = self.host.json.loads(
status_path.read_text(encoding="utf-8")
)
self.assertEqual(status["state"], "cancelled")
self.assertFalse(state.shutdown_event.is_set())
self.assertEqual(state.snapshot()["cached_models"]["ben2"], "BEN2_BASE")
def test_expired_clients_are_removed(self):
fake_models = SimpleNamespace(
status=lambda: {"ben2": "", "da3": ""},
close=lambda: None,
)
with patch.object(self.host, "ModelCache", return_value=fake_models):
state = self.host.HostState("token", "instance")
state.register_client("client", "Blender")
state.clients["client"]["last_seen"] -= (
self.host.CLIENT_TIMEOUT_SECONDS + 1.0
)
state.expire_clients()
self.assertEqual(state.snapshot()["clients"], 0)
def test_discovery_is_only_removed_by_matching_host(self):
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "host.json"
args = SimpleNamespace(
discovery=str(path),
port=51234,
token="token",
instance_id="owner",
)
with patch.object(self.host.os, "getpid", return_value=42):
self.host.write_discovery(args)
payload = self.host.json.loads(path.read_text(encoding="utf-8"))
self.assertEqual(payload["port"], 51234)
self.assertEqual(payload["pid"], 42)
self.host.remove_discovery(path, "other")
self.assertTrue(path.is_file())
self.host.remove_discovery(path, "owner")
self.assertFalse(path.exists())
if __name__ == "__main__":
unittest.main()

38
tests/test_bootstrap_runtime.py
View File

@@ -29,6 +29,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
args = SimpleNamespace(
uv="uv",
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
@@ -67,13 +68,48 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertTrue(
any("test-ben2-commit" in str(value) for value in ben2_command)
)
check_command = commands[-1]
host_command = next(
command
for command in commands
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
)
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command)
check_command = next(
command
for command in commands
if any("import torch;" in str(value) for value in command)
)
self.assertTrue(
any(
"from ben2 import BEN_Base" in str(value)
for value in check_command
)
)
self.assertTrue(
any("import fastapi, uvicorn" in str(value) for value in check_command)
)
default_downloads = [
command
for command in commands
if any(
repository in command
for repository in (
"depth-anything/DA3-SMALL",
"depth-anything/DA3-BASE",
"PramaLLC/BEN2",
)
)
]
self.assertEqual(len(default_downloads), 3)
ben2_download = next(
command
for command in default_downloads
if "PramaLLC/BEN2" in command
)
self.assertEqual(
ben2_download[-1],
'["config.json", "model.safetensors"]',
)
if __name__ == "__main__":

1
tests/test_pack.py
View File

@@ -38,5 +38,6 @@ class PackTest(unittest.TestCase):
self.assertIn("image_depth.py", names)
self.assertNotIn("snapform_depth_plane.py", names)
self.assertIn("worker.py", names)
self.assertIn("ai_host.py", names)
self.assertNotIn("importer.py", names)
self.assertNotIn("snapform/__init__.py", names)

14
tests/test_worker.py
View File

@@ -47,6 +47,20 @@ class WorkerTest(unittest.TestCase):
source.index("import torch"),
)
def test_ai_tasks_support_cooperative_cancellation(self):
with self.assertRaises(self.worker.TaskCancelled):
self.worker.check_cancelled(lambda: True)
self.worker.check_cancelled(lambda: False)
self.assertIn(
"cancel_check=cancel_check",
inspect.getsource(self.worker.reconstruct),
)
self.assertIn(
"cancel_check=cancel_check",
inspect.getsource(self.worker.remove_background),
)
def test_optional_background_removal_runs_before_da3_model_load(self):
source = inspect.getsource(self.worker.reconstruct)