nan
9c896213a5
- Updated frame extraction logic to allow for single frame extraction. - Integrated model readiness checks for BEN2 and DA3 models in relevant operators. - Added directory name attribute to model specifications for better organization. - Improved device selection logic for production and debug tasks. - Refactored preferences to expose device settings and ensure they reflect available hardware. - Enhanced runtime device detection and management, including support for CUDA and MPS. - Updated tests to cover new device management features and model directory handling. - Removed deprecated download model tests to streamline codebase.
253 lines
8.2 KiB
Python
253 lines
8.2 KiB
Python
import importlib.util
|
|
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"], "model")
|
|
|
|
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())
|
|
|
|
def test_host_dispatches_frame_extraction_without_models(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")
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
input_directory = root / "frames"
|
|
input_directory.mkdir()
|
|
(input_directory / "frame-0001.jpg").touch()
|
|
(input_directory / "frame-0002.jpg").touch()
|
|
result_directory = root / "result"
|
|
status_path = root / "status.json"
|
|
state.reserve(
|
|
"job",
|
|
"client",
|
|
"extract-frames",
|
|
str(status_path),
|
|
)
|
|
|
|
state.run_job(
|
|
{
|
|
"job_id": "job",
|
|
"command": "extract-frames",
|
|
"parameters": {
|
|
"input": str(input_directory),
|
|
"output_dir": str(result_directory),
|
|
"status": str(status_path),
|
|
"video_frames": 16,
|
|
"start_frame": 1,
|
|
},
|
|
}
|
|
)
|
|
status = self.host.json.loads(
|
|
status_path.read_text(encoding="utf-8")
|
|
)
|
|
self.assertTrue(
|
|
(result_directory / "frame.0001.jpg").is_file()
|
|
)
|
|
|
|
self.assertEqual(status["state"], "succeeded")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|