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

Enhance device management and model handling in SnapForm

Browse Source 
- 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.
This commit is contained in:
马南
2026-08-08 20:16:11 +08:00
parent aee460fe7f
commit 9c896213a5
18 changed files with 699 additions and 196 deletions
Download Patch File Download Diff File Expand all files Collapse all files

192
tests/test_bootstrap_runtime.py
View File

@@ -62,7 +62,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
) as mocked_run, patch.object(
bootstrap,
"run_install",
) as mocked_run_install:
) as mocked_run_install, patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
commands = [call.args[0] for call in mocked_run.call_args_list]
@@ -113,10 +117,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command)
self.assertIn("setuptools", host_command)
self.assertNotIn("ninja", host_command)
check_command = next(
command
for command in commands
if any("import torch;" in str(value) for value in command)
check_command = bootstrap.verify_command(
bootstrap.environment_python(runtime_dir)
)
self.assertTrue(
any("import torch;" in str(value) for value in check_command)
)
self.assertTrue(
any(
@@ -130,6 +135,12 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertFalse(
any("import gsplat" in str(value) for value in check_command)
)
self.assertTrue(
any(
"SNAPFORM_DEVICES=" in str(value)
for value in check_command
)
)
default_downloads = [
command
for command in commands
@@ -190,7 +201,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object(
bootstrap,
"run_install",
) as mocked_run_install:
) as mocked_run_install, patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
commands = [
@@ -372,7 +387,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap,
"run_stage",
side_effect=fake_run_stage,
), patch.object(bootstrap, "run"):
), patch.object(bootstrap, "run"), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
starts = [item[1] for item in stage_ranges]
@@ -604,34 +623,6 @@ class BootstrapRuntimeTest(unittest.TestCase):
"https://pypi.tuna.tsinghua.edu.cn/simple",
)
def test_child_environment_injects_registry_proxy(self):
bootstrap = load_bootstrap_runtime()
with patch.object(
bootstrap.urllib.request,
"getproxies",
return_value={
"http": "http://proxy:8080",
"https": "http://proxy:8080",
},
):
env = bootstrap.child_environment()
self.assertEqual(env["HTTP_PROXY"], "http://proxy:8080")
self.assertEqual(env["HTTPS_PROXY"], "http://proxy:8080")
def test_child_environment_keeps_explicit_proxy(self):
bootstrap = load_bootstrap_runtime()
with patch.object(
bootstrap.urllib.request,
"getproxies",
return_value={"https": "http://proxy:8080"},
):
env = bootstrap.child_environment(
{"HTTPS_PROXY": "http://explicit:3128"}
)
self.assertEqual(env["HTTPS_PROXY"], "http://explicit:3128")
def test_ready_json_records_installed_packages(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
@@ -657,6 +648,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap,
"run_stage",
return_value=["torch==2.9.1+cu128", "numpy==1.26.4"],
), patch.object(
bootstrap,
"run_verify",
return_value={"cuda": True, "mps": False},
):
bootstrap.install(args)
@@ -669,6 +664,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
ready["packages"],
["numpy==1.26.4", "torch==2.9.1+cu128"],
)
self.assertEqual(ready["devices"], ["cuda", "cpu"])
def test_broken_runtime_is_cleaned_before_reinstall(self):
bootstrap = load_bootstrap_runtime()
@@ -693,6 +689,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
@@ -728,6 +728,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
@@ -758,7 +762,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap,
"run",
side_effect=fake_run,
), patch.object(bootstrap, "run_install"):
), patch.object(bootstrap, "run_install"), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
@@ -872,6 +880,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap,
"run_install",
side_effect=fake_run_install,
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
@@ -918,6 +930,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
@@ -995,6 +1011,110 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
def test_install_stages_download_models_into_prefixed_directories(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
sources = root / "sources"
sources.mkdir()
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
(sources / "ben2.tar.gz").write_bytes(b"ben2")
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
stages = bootstrap.install_stages(
args,
runtime_dir,
root / "python",
uv_path,
)
downloads = {
label: command
for label, command, _environment, kind in stages
if kind == "download"
}
self.assertIn("da3_small", " ".join(map(str, downloads["da3 small"])))
self.assertIn("da3_base", " ".join(map(str, downloads["da3 base"])))
self.assertIn("ben2_base", " ".join(map(str, downloads["ben2 base"])))
def test_run_verify_returns_detected_devices(self):
bootstrap = load_bootstrap_runtime()
class FakeProcess:
def __init__(self, lines):
self.stdout = lines
def wait(self):
return 0
fake_process = FakeProcess(
['SNAPFORM_DEVICES={"cuda": false, "mps": true}']
)
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
):
report = bootstrap.run_verify(
"python",
"status.json",
clock=lambda: 0.0,
)
self.assertFalse(report["cuda"])
self.assertTrue(report["mps"])
def test_run_verify_writes_heartbeat_progress(self):
bootstrap = load_bootstrap_runtime()
class FakeProcess:
def __init__(self, lines):
self.stdout = lines
def wait(self):
return 0
clock_state = {"now": 0.0}
def fake_clock():
clock_state["now"] += 6.0
return clock_state["now"]
def delayed_empty_lines():
time.sleep(0.6)
if False:
yield
fake_process = FakeProcess(delayed_empty_lines())
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
), patch.object(bootstrap, "write_status") as mocked_write:
report = bootstrap.run_verify(
"python",
"status.json",
clock=fake_clock,
)
messages = [call.args[2] for call in mocked_write.call_args_list]
self.assertTrue(
any("verifying runtime" in message for message in messages)
)
self.assertEqual(report, {})
def test_install_uv_extracts_bundled_archive(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory: