Update addon tests, bootstrap runtime, download progress, and worker functionality
- Modified status messages and stage labels in `test_addon.py` for clarity and consistency. - Added tests to ensure proper host termination and command execution in `test_addon.py`. - Enhanced `test_bootstrap_runtime.py` with additional tests for installation behavior, including handling of CUDA and MPS devices, and improved error handling for disk space and network checks. - Updated download progress reporting in `test_download_progress.py` to include download speed and improved message formatting. - Added new tests in `test_worker.py` for device selection, NPZ file validation, and model downloading behavior with Hugging Face endpoints. - Bumped version in `uv.lock` to 0.9.30.
This commit is contained in:
@@ -517,27 +517,27 @@ class AddonTest(unittest.TestCase):
|
||||
operator._format_status(
|
||||
{
|
||||
"progress": 0.5,
|
||||
"message": "Downloaded 3 packages",
|
||||
"message": "2 downloaded",
|
||||
"stage": 3,
|
||||
"stages": 12,
|
||||
"stage_label": "Install PyTorch",
|
||||
"stages": 8,
|
||||
"stage_label": "pytorch",
|
||||
},
|
||||
"fallback",
|
||||
),
|
||||
"Step 3/12 Install PyTorch: Downloaded 3 packages",
|
||||
"3/8 pytorch: 2 downloaded",
|
||||
)
|
||||
self.assertEqual(
|
||||
operator._format_status(
|
||||
{
|
||||
"progress": 0.0,
|
||||
"message": "",
|
||||
"stage": 10,
|
||||
"stages": 12,
|
||||
"stage_label": "Download DA3 Small",
|
||||
"stage": 6,
|
||||
"stages": 8,
|
||||
"stage_label": "da3 small",
|
||||
},
|
||||
"fallback",
|
||||
),
|
||||
"Step 10/12 Download DA3 Small",
|
||||
"6/8 da3 small",
|
||||
)
|
||||
self.assertEqual(
|
||||
operator._format_status(
|
||||
@@ -865,6 +865,31 @@ class AddonTest(unittest.TestCase):
|
||||
self.assertIn("ProxyHandler({})", inspect.getsource(runtime))
|
||||
self.assertIn('return storage_root() / "host.json"', inspect.getsource(runtime.host_discovery_path))
|
||||
|
||||
def test_closing_blender_terminates_host_it_started(self):
|
||||
runtime = importlib.import_module("snapform.runtime")
|
||||
previous_connection = runtime._host_connection
|
||||
previous_registered = runtime._registered_instance_id
|
||||
runtime._host_connection = SimpleNamespace(port=12345)
|
||||
runtime._registered_instance_id = "instance"
|
||||
try:
|
||||
with patch.object(
|
||||
runtime,
|
||||
"_host_http_request",
|
||||
) as request, patch.object(
|
||||
runtime,
|
||||
"_close_launched_host",
|
||||
) as close_host:
|
||||
runtime.detach_host_client()
|
||||
|
||||
request.assert_called_once()
|
||||
self.assertEqual(request.call_args.args[1], "DELETE")
|
||||
close_host.assert_called_once_with(terminate=True)
|
||||
finally:
|
||||
runtime._host_connection = previous_connection
|
||||
runtime._registered_instance_id = previous_registered
|
||||
self.assertIsNone(runtime._host_connection)
|
||||
self.assertIsNone(runtime._registered_instance_id)
|
||||
|
||||
def test_host_jobs_complete_from_status_state(self):
|
||||
source = inspect.getsource(self.addon.RunJob.modal)
|
||||
|
||||
@@ -1357,6 +1382,37 @@ class AddonTest(unittest.TestCase):
|
||||
if value == "--allow-pattern"
|
||||
]
|
||||
self.assertEqual(patterns, ["config.json", "model.safetensors"])
|
||||
label_index = command.index("--label")
|
||||
self.assertEqual(command[label_index + 1], "BEN2 Base")
|
||||
|
||||
def test_install_runtime_command_has_no_model_label_reference(self):
|
||||
runtime = importlib.import_module("snapform.runtime")
|
||||
with patch.object(
|
||||
runtime,
|
||||
"uv_executable",
|
||||
return_value="C:/uv/uv.exe",
|
||||
), patch.object(
|
||||
runtime,
|
||||
"preferences",
|
||||
return_value=SimpleNamespace(torch_backend="CU128"),
|
||||
), patch.object(
|
||||
runtime,
|
||||
"runtime_directory",
|
||||
return_value=Path("runtime"),
|
||||
), patch.object(
|
||||
runtime,
|
||||
"models_directory",
|
||||
return_value=Path("models"),
|
||||
):
|
||||
command = runtime.install_runtime_command(Path("status.json"))
|
||||
|
||||
self.assertNotIn("--label", command)
|
||||
self.assertEqual(command[0], sys.executable)
|
||||
self.assertIn("--uv", command)
|
||||
self.assertIn("C:/uv/uv.exe", command)
|
||||
self.assertTrue(
|
||||
any("bootstrap_runtime.py" in str(value) for value in command)
|
||||
)
|
||||
|
||||
def test_preferences_install_button_is_disabled_while_busy(self):
|
||||
draw_source = inspect.getsource(self.addon.SnapFormPreferences.draw)
|
||||
@@ -1526,7 +1582,7 @@ class AddonTest(unittest.TestCase):
|
||||
],
|
||||
)
|
||||
|
||||
def test_force_flag_triggers_window_draw_and_swap(self):
|
||||
def test_force_flag_refreshes_status_text(self):
|
||||
redrawn = []
|
||||
context = SimpleNamespace(
|
||||
window_manager=SimpleNamespace(
|
||||
@@ -1553,29 +1609,17 @@ class AddonTest(unittest.TestCase):
|
||||
),
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
self.addon.status_bar.bpy,
|
||||
"ops",
|
||||
SimpleNamespace(
|
||||
wm=SimpleNamespace(
|
||||
redraw_timer=lambda **options: redrawn.append(
|
||||
("swap", options)
|
||||
),
|
||||
),
|
||||
),
|
||||
):
|
||||
self.addon.status_bar.redraw_ui(context, force=True)
|
||||
self.addon.status_bar.redraw_ui(context, force=True)
|
||||
|
||||
self.assertEqual(
|
||||
redrawn,
|
||||
[
|
||||
"status",
|
||||
("swap", {"type": "DRAW_WIN_SWAP", "iterations": 1}),
|
||||
"status-text",
|
||||
],
|
||||
)
|
||||
|
||||
def test_force_flag_ignores_missing_redraw_helpers(self):
|
||||
def test_force_flag_ignores_missing_workspace(self):
|
||||
redrawn = []
|
||||
context = SimpleNamespace(
|
||||
window_manager=SimpleNamespace(
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
@@ -21,14 +25,16 @@ def load_bootstrap_runtime():
|
||||
|
||||
|
||||
class BootstrapRuntimeTest(unittest.TestCase):
|
||||
def test_installs_da3_and_non_compiled_gsplat(self):
|
||||
def test_installs_da3_from_archives_without_gsplat(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv"
|
||||
uv_path.write_bytes(b"uv")
|
||||
args = SimpleNamespace(
|
||||
uv="uv",
|
||||
uv=str(uv_path),
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
@@ -36,7 +42,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "run") as mocked_run, patch.object(
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
) as mocked_run, patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
) as mocked_run_install:
|
||||
@@ -55,25 +64,44 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
f"addict=={bootstrap.ADDICT_VERSION}",
|
||||
da3_command,
|
||||
)
|
||||
gsplat_call = next(
|
||||
call
|
||||
for call in mocked_run_install.call_args_list
|
||||
if any("gsplat @" in str(value) for value in call.args[4])
|
||||
)
|
||||
gsplat_command = gsplat_call.args[4]
|
||||
self.assertIn("--no-build-isolation", gsplat_command)
|
||||
self.assertTrue(
|
||||
any(bootstrap.GSPLAT_COMMIT in str(value) for value in gsplat_command)
|
||||
any(
|
||||
"https://github.com/ByteDance-Seed/Depth-Anything-3/"
|
||||
"archive/test-commit.tar.gz" in str(value)
|
||||
for value in da3_command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("git+" in str(value) for value in da3_command)
|
||||
)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"gsplat" in str(value)
|
||||
for command in commands
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"--no-build-isolation" in str(value)
|
||||
for command in commands
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
self.assertEqual(gsplat_call.args[5], {"BUILD_NO_CUDA": "1"})
|
||||
|
||||
ben2_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if any("ben2 @" in str(value) for value in command)
|
||||
)
|
||||
self.assertTrue(
|
||||
any("test-ben2-commit" in str(value) for value in ben2_command)
|
||||
any(
|
||||
"https://github.com/PramaLLC/BEN2/archive/"
|
||||
"test-ben2-commit.tar.gz" in str(value)
|
||||
for value in ben2_command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("git+" in str(value) for value in ben2_command)
|
||||
)
|
||||
host_command = next(
|
||||
command
|
||||
@@ -81,6 +109,8 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
|
||||
)
|
||||
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
|
||||
@@ -95,6 +125,9 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
self.assertTrue(
|
||||
any("import fastapi, uvicorn" in str(value) for value in check_command)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("import gsplat" in str(value) for value in check_command)
|
||||
)
|
||||
default_downloads = [
|
||||
command
|
||||
for command in commands
|
||||
@@ -121,7 +154,86 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
self.assertEqual(patterns, ["config.json", "model.safetensors"])
|
||||
self.assertIn("--stage", ben2_download)
|
||||
self.assertIn("--stage-label", ben2_download)
|
||||
self.assertEqual(ben2_download[-1], "Download BEN2 Base")
|
||||
self.assertEqual(ben2_download[-1], "ben2 base")
|
||||
|
||||
def test_darwin_install_skips_cuda_index_and_xformers(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
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"),
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.sys,
|
||||
"platform",
|
||||
"darwin",
|
||||
), patch.object(
|
||||
bootstrap.platform,
|
||||
"machine",
|
||||
return_value="arm64",
|
||||
), patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
) as mocked_run_install:
|
||||
bootstrap.install(args)
|
||||
|
||||
commands = [
|
||||
call.args[4] for call in mocked_run_install.call_args_list
|
||||
]
|
||||
torch_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if any(
|
||||
f"torch=={bootstrap.TORCH_VERSION}" in str(value)
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
self.assertNotIn("--index-url", torch_command)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"xformers" in str(value)
|
||||
for command in commands
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
|
||||
def test_intel_mac_is_rejected_with_clear_error(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
uv="uv",
|
||||
runtime_dir="runtime",
|
||||
models_dir="models",
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status="status.json",
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.sys,
|
||||
"platform",
|
||||
"darwin",
|
||||
), patch.object(
|
||||
bootstrap.platform,
|
||||
"machine",
|
||||
return_value="x86_64",
|
||||
), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree:
|
||||
with self.assertRaisesRegex(RuntimeError, "Apple Silicon"):
|
||||
bootstrap.install(args)
|
||||
|
||||
mocked_rmtree.assert_not_called()
|
||||
|
||||
def test_run_stage_writes_stage_progress(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
@@ -131,8 +243,8 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
bootstrap.run_stage(
|
||||
str(status_path),
|
||||
3,
|
||||
12,
|
||||
"Install PyTorch",
|
||||
8,
|
||||
"pytorch",
|
||||
["uv", "pip", "install", "torch"],
|
||||
None,
|
||||
"plain",
|
||||
@@ -145,8 +257,131 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
self.assertEqual(status["stage"], 3)
|
||||
self.assertEqual(status["stages"], 12)
|
||||
self.assertEqual(status["stage_label"], "Install PyTorch")
|
||||
self.assertEqual(status["stages"], 8)
|
||||
self.assertEqual(status["stage_label"], "pytorch")
|
||||
|
||||
def test_run_stage_writes_progress_range(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
with patch.object(bootstrap, "run"), patch.object(
|
||||
bootstrap,
|
||||
"write_status",
|
||||
) as mocked_write:
|
||||
bootstrap.run_stage(
|
||||
str(status_path),
|
||||
3,
|
||||
8,
|
||||
"pytorch",
|
||||
["uv", "pip", "install", "torch"],
|
||||
None,
|
||||
"plain",
|
||||
start=0.2,
|
||||
end=0.4,
|
||||
)
|
||||
|
||||
progresses = [
|
||||
call.args[1] for call in mocked_write.call_args_list
|
||||
]
|
||||
self.assertEqual(progresses[0], 0.2)
|
||||
self.assertEqual(progresses[-1], 0.4)
|
||||
|
||||
def test_run_stage_passes_progress_range_to_download_command(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
command = [
|
||||
"python",
|
||||
"worker.py",
|
||||
"download",
|
||||
"--repository",
|
||||
"depth-anything/DA3-SMALL",
|
||||
]
|
||||
with patch.object(bootstrap, "run") as mocked_run:
|
||||
bootstrap.run_stage(
|
||||
str(status_path),
|
||||
6,
|
||||
8,
|
||||
"da3 small",
|
||||
command,
|
||||
None,
|
||||
"download",
|
||||
start=0.9,
|
||||
end=0.9333,
|
||||
)
|
||||
|
||||
called_command = mocked_run.call_args.args[0]
|
||||
self.assertIn("--progress-start", called_command)
|
||||
self.assertIn("--progress-end", called_command)
|
||||
self.assertEqual(
|
||||
called_command[called_command.index("--progress-start") + 1],
|
||||
"0.9",
|
||||
)
|
||||
self.assertEqual(
|
||||
called_command[called_command.index("--progress-end") + 1],
|
||||
"0.9333",
|
||||
)
|
||||
self.assertEqual(called_command[-1], "da3 small")
|
||||
|
||||
def test_install_allocates_contiguous_progress_slices(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
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"),
|
||||
)
|
||||
stage_ranges = []
|
||||
|
||||
def fake_run_stage(
|
||||
status_path,
|
||||
stage,
|
||||
stages,
|
||||
label,
|
||||
command,
|
||||
environment,
|
||||
kind,
|
||||
start=0.0,
|
||||
end=1.0,
|
||||
):
|
||||
stage_ranges.append((stage, start, end))
|
||||
bootstrap.write_status(
|
||||
status_path,
|
||||
end,
|
||||
"",
|
||||
stage=stage,
|
||||
stages=stages,
|
||||
stage_label=label,
|
||||
)
|
||||
return []
|
||||
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run_stage",
|
||||
side_effect=fake_run_stage,
|
||||
), patch.object(bootstrap, "run"):
|
||||
bootstrap.install(args)
|
||||
|
||||
starts = [item[1] for item in stage_ranges]
|
||||
ends = [item[2] for item in stage_ranges]
|
||||
self.assertTrue(
|
||||
all(starts[i] < starts[i + 1] for i in range(len(starts) - 1))
|
||||
)
|
||||
self.assertEqual(ends[-1], 1.0)
|
||||
self.assertEqual(starts[0], bootstrap.RUNTIME_SLICE_START)
|
||||
status = json.loads(
|
||||
(root / "status.json").read_text(encoding="utf-8")
|
||||
)
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
|
||||
def test_run_install_tracks_package_download_events(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
@@ -161,33 +396,515 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
fake_process = FakeProcess(
|
||||
[
|
||||
"Resolved 3 packages in 1.0s",
|
||||
"Downloading torch (10MiB)",
|
||||
" Downloaded torch",
|
||||
"DEBUG Adding transitive dependency for torch==2.9.1: filelock*",
|
||||
"Prepared 3 packages in 2.0s",
|
||||
"Installed 3 packages in 0.5s",
|
||||
"+ torch==2.9.1",
|
||||
"+ numpy==1.26.4",
|
||||
]
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.subprocess,
|
||||
"Popen",
|
||||
return_value=fake_process,
|
||||
), patch.object(bootstrap, "write_status") as mocked_write:
|
||||
installed = bootstrap.run_install(
|
||||
"status.json",
|
||||
2,
|
||||
8,
|
||||
"pytorch deps",
|
||||
["uv", "pip", "install", "torch"],
|
||||
)
|
||||
|
||||
messages = [call.args[2] for call in mocked_write.call_args_list]
|
||||
self.assertEqual(
|
||||
messages,
|
||||
[
|
||||
"3 packages",
|
||||
"3/3 packages",
|
||||
"3/3 packages",
|
||||
],
|
||||
)
|
||||
self.assertEqual(mocked_write.call_args_list[-1].args[1], 0.02)
|
||||
first = mocked_write.call_args_list[0]
|
||||
self.assertEqual(first.kwargs["stage"], 2)
|
||||
self.assertEqual(first.kwargs["stages"], 8)
|
||||
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
|
||||
self.assertEqual(installed, ["torch==2.9.1", "numpy==1.26.4"])
|
||||
|
||||
def test_run_install_smooths_progress_between_events(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
|
||||
class SlowLines:
|
||||
def __init__(self, lines, delay):
|
||||
self.lines = list(lines)
|
||||
self.delay = delay
|
||||
|
||||
def __iter__(self):
|
||||
for line in self.lines:
|
||||
time.sleep(self.delay)
|
||||
yield line
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, lines):
|
||||
self.stdout = lines
|
||||
|
||||
def wait(self):
|
||||
return 0
|
||||
|
||||
fake_process = FakeProcess(
|
||||
SlowLines(
|
||||
[
|
||||
"Resolved 3 packages in 1.0s",
|
||||
"Prepared 3 packages in 2.0s",
|
||||
"Installed 3 packages in 0.5s",
|
||||
],
|
||||
0.3,
|
||||
)
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.subprocess,
|
||||
"Popen",
|
||||
return_value=fake_process,
|
||||
), patch.object(bootstrap, "write_status") as mocked_write:
|
||||
bootstrap.run_install(
|
||||
"status.json",
|
||||
2,
|
||||
8,
|
||||
"pytorch deps",
|
||||
["uv", "pip", "install", "torch"],
|
||||
)
|
||||
|
||||
progresses = [call.args[1] for call in mocked_write.call_args_list]
|
||||
self.assertGreater(progresses[-1], 0.02)
|
||||
self.assertTrue(
|
||||
all(
|
||||
progresses[i] <= progresses[i + 1] + 1e-9
|
||||
for i in range(len(progresses) - 1)
|
||||
)
|
||||
)
|
||||
messages = [call.args[2] for call in mocked_write.call_args_list]
|
||||
self.assertIn("3 packages", messages)
|
||||
self.assertIn("3/3 packages", messages)
|
||||
|
||||
def test_run_install_maps_progress_into_stage_range(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, lines):
|
||||
self.stdout = lines
|
||||
|
||||
def wait(self):
|
||||
return 0
|
||||
|
||||
fake_process = FakeProcess(
|
||||
[
|
||||
"Resolved 3 packages in 1.0s",
|
||||
"Prepared 3 packages in 2.0s",
|
||||
"Installed 3 packages in 0.5s",
|
||||
]
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
with patch.object(
|
||||
bootstrap.subprocess,
|
||||
"Popen",
|
||||
return_value=fake_process,
|
||||
):
|
||||
bootstrap.run_install(
|
||||
str(status_path),
|
||||
2,
|
||||
12,
|
||||
"Install PyTorch",
|
||||
["uv", "pip", "install", "torch"],
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.subprocess,
|
||||
"Popen",
|
||||
return_value=fake_process,
|
||||
), patch.object(bootstrap, "write_status") as mocked_write:
|
||||
bootstrap.run_install(
|
||||
"status.json",
|
||||
2,
|
||||
8,
|
||||
"pytorch deps",
|
||||
["uv", "pip", "install", "torch"],
|
||||
start=0.2,
|
||||
end=0.4,
|
||||
)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(status["progress"], 0.95)
|
||||
self.assertEqual(status["message"], "Installing packages")
|
||||
self.assertEqual(status["stage"], 2)
|
||||
self.assertEqual(status["stages"], 12)
|
||||
self.assertEqual(status["stage_label"], "Install PyTorch")
|
||||
first = mocked_write.call_args_list[0]
|
||||
self.assertAlmostEqual(first.args[1], 0.204)
|
||||
self.assertEqual(first.kwargs["stage"], 2)
|
||||
self.assertEqual(first.kwargs["stages"], 8)
|
||||
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
|
||||
|
||||
def test_formats_elapsed_time(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
|
||||
self.assertEqual(bootstrap.format_elapsed(12), "12s")
|
||||
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
|
||||
|
||||
def test_preflight_rejects_low_disk_space(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=1024**3),
|
||||
), patch.object(bootstrap, "reachable", return_value=True):
|
||||
with self.assertRaisesRegex(RuntimeError, "disk space"):
|
||||
bootstrap.preflight(args)
|
||||
|
||||
def test_preflight_rejects_unreachable_network(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
||||
), patch.object(bootstrap, "reachable", return_value=False):
|
||||
with self.assertRaisesRegex(RuntimeError, "Cannot reach"):
|
||||
bootstrap.preflight(args)
|
||||
|
||||
def test_preflight_passes_when_checks_ok(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
||||
), patch.object(bootstrap, "reachable", return_value=True):
|
||||
resolved, used_mirror = bootstrap.preflight(args)
|
||||
|
||||
self.assertFalse(used_mirror)
|
||||
self.assertEqual(
|
||||
resolved["pypi"],
|
||||
"https://pypi.org/simple/",
|
||||
)
|
||||
|
||||
def test_preflight_falls_back_to_mirrors(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
|
||||
def fake_reachable(url, timeout=bootstrap.PREFLIGHT_TIMEOUT):
|
||||
if "github.com" in url:
|
||||
return True
|
||||
if "tuna" in url or "hf-mirror" in url or "aliyun" in url:
|
||||
return True
|
||||
return False
|
||||
|
||||
with patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"reachable",
|
||||
side_effect=fake_reachable,
|
||||
):
|
||||
resolved, used_mirror = bootstrap.preflight(args)
|
||||
|
||||
self.assertTrue(used_mirror)
|
||||
self.assertEqual(
|
||||
resolved["pypi"],
|
||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
||||
)
|
||||
self.assertEqual(resolved["hf"], "https://hf-mirror.com")
|
||||
self.assertEqual(
|
||||
resolved["torch"],
|
||||
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
|
||||
)
|
||||
self.assertEqual(
|
||||
bootstrap._resolved_torch_index,
|
||||
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
|
||||
)
|
||||
|
||||
def test_preflight_uses_environment_overrides(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"PIP_INDEX_URL": "https://pypi.example/simple",
|
||||
"HF_ENDPOINT": "https://hf.example",
|
||||
"SNAPFORM_TORCH_INDEX": "https://torch.example/whl/cu128",
|
||||
},
|
||||
), patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
||||
), patch.object(bootstrap, "reachable", return_value=True):
|
||||
resolved, used_mirror = bootstrap.preflight(args)
|
||||
|
||||
self.assertFalse(used_mirror)
|
||||
self.assertEqual(resolved["pypi"], "https://pypi.example/simple")
|
||||
self.assertEqual(resolved["hf"], "https://hf.example")
|
||||
self.assertEqual(
|
||||
resolved["torch"],
|
||||
"https://torch.example/whl/cu128",
|
||||
)
|
||||
|
||||
def test_child_environment_injects_resolved_mirrors(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
previous_pypi = bootstrap._resolved_pypi_index
|
||||
previous_hf = bootstrap._resolved_hf_endpoint
|
||||
bootstrap._resolved_pypi_index = (
|
||||
"https://pypi.tuna.tsinghua.edu.cn/simple"
|
||||
)
|
||||
bootstrap._resolved_hf_endpoint = "https://hf-mirror.com"
|
||||
try:
|
||||
env = bootstrap.child_environment()
|
||||
finally:
|
||||
bootstrap._resolved_pypi_index = previous_pypi
|
||||
bootstrap._resolved_hf_endpoint = previous_hf
|
||||
|
||||
self.assertEqual(
|
||||
env["PIP_INDEX_URL"],
|
||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
||||
)
|
||||
self.assertEqual(env["HF_ENDPOINT"], "https://hf-mirror.com")
|
||||
|
||||
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:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
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"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_stage",
|
||||
return_value=["torch==2.9.1+cu128", "numpy==1.26.4"],
|
||||
):
|
||||
bootstrap.install(args)
|
||||
|
||||
ready = json.loads(
|
||||
(runtime_dir / "snapform_ready.json").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
)
|
||||
self.assertEqual(
|
||||
ready["packages"],
|
||||
["numpy==1.26.4", "torch==2.9.1+cu128"],
|
||||
)
|
||||
|
||||
def test_broken_runtime_is_cleaned_before_reinstall(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
(runtime_dir / "partial-python.exe").write_text("junk")
|
||||
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",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
):
|
||||
bootstrap.install(args)
|
||||
|
||||
self.assertFalse((runtime_dir / "partial-python.exe").exists())
|
||||
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
|
||||
|
||||
def test_ready_runtime_is_not_cleaned(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir(parents=True)
|
||||
(runtime_dir / "Scripts").mkdir()
|
||||
(runtime_dir / "Scripts" / "python.exe").write_text("python")
|
||||
(runtime_dir / "snapform_ready.json").write_text(
|
||||
json.dumps({"runtime_schema": bootstrap.RUNTIME_SCHEMA}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(runtime_dir / "marker.txt").write_text("keep")
|
||||
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",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
):
|
||||
bootstrap.install(args)
|
||||
|
||||
self.assertTrue((runtime_dir / "marker.txt").is_file())
|
||||
|
||||
def test_model_download_failure_keeps_runtime_ready(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
|
||||
def fake_run(command, environment=None):
|
||||
if any("DA3-SMALL" in str(value) for value in command):
|
||||
raise subprocess.CalledProcessError(1, command)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
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",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
side_effect=fake_run,
|
||||
), patch.object(bootstrap, "run_install"):
|
||||
bootstrap.install(args)
|
||||
|
||||
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
|
||||
status = json.loads(
|
||||
(root / "status.json").read_text(encoding="utf-8")
|
||||
)
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
self.assertIn("Runtime is ready", status["message"])
|
||||
|
||||
def test_install_downloads_missing_uv(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv" / "uv.exe"
|
||||
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"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
), patch.object(
|
||||
bootstrap.shutil,
|
||||
"which",
|
||||
return_value=None,
|
||||
), patch.object(
|
||||
bootstrap.shutil,
|
||||
"rmtree",
|
||||
), patch.object(bootstrap, "download_uv") as mocked_download:
|
||||
bootstrap.install(args)
|
||||
|
||||
mocked_download.assert_called_once()
|
||||
call_args = mocked_download.call_args
|
||||
self.assertEqual(call_args.args[0], uv_path)
|
||||
self.assertEqual(call_args.args[1], str(root / "status.json"))
|
||||
self.assertEqual(call_args.kwargs["stage"], 1)
|
||||
self.assertEqual(call_args.kwargs["stage_label"], "uv")
|
||||
|
||||
def test_install_reuses_existing_uv(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv.exe"
|
||||
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"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
), patch.object(bootstrap, "download_uv") as mocked_download:
|
||||
bootstrap.install(args)
|
||||
|
||||
mocked_download.assert_not_called()
|
||||
|
||||
def test_extract_uv_unpacks_standalone_binary(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
archive_path = root / "uv.zip"
|
||||
with zipfile.ZipFile(archive_path, "w") as bundle:
|
||||
bundle.writestr("uv-x86_64-pc-windows-msvc/uv.exe", "binary")
|
||||
|
||||
uv_path = root / "uv.exe"
|
||||
bootstrap.extract_uv(archive_path, uv_path)
|
||||
|
||||
self.assertEqual(uv_path.read_bytes(), b"binary")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
@@ -56,7 +58,70 @@ class DownloadProgressTest(unittest.TestCase):
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertAlmostEqual(status["progress"], 0.275)
|
||||
self.assertEqual(status["message"], "Downloading owner/model")
|
||||
self.assertEqual(status["message"], "downloading owner/model")
|
||||
|
||||
def test_reports_download_speed(self):
|
||||
class FakeTqdmWithRate(FakeTqdm):
|
||||
@property
|
||||
def format_dict(self):
|
||||
return {"rate": 14_000_000}
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
progress_class = self.worker.create_download_progress_class(
|
||||
status_path,
|
||||
"owner/model",
|
||||
tqdm_type=FakeTqdmWithRate,
|
||||
clock=lambda: 1.0,
|
||||
label="da3 small",
|
||||
)
|
||||
progress = progress_class(
|
||||
desc="Downloading bytes",
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
|
||||
progress.update(50)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(
|
||||
status["message"],
|
||||
"downloading da3 small 13.4 MB/s",
|
||||
)
|
||||
|
||||
def test_staged_download_prefixes_speed_with_downloading(self):
|
||||
class FakeTqdmWithRate(FakeTqdm):
|
||||
@property
|
||||
def format_dict(self):
|
||||
return {"rate": 14_000_000}
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
progress_class = self.worker.create_download_progress_class(
|
||||
status_path,
|
||||
"owner/model",
|
||||
tqdm_type=FakeTqdmWithRate,
|
||||
clock=lambda: 1.0,
|
||||
label="da3 small",
|
||||
stage=6,
|
||||
stages=8,
|
||||
stage_label="da3 small",
|
||||
)
|
||||
progress = progress_class(
|
||||
desc="Downloading bytes",
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
|
||||
progress.update(50)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(status["message"], "downloading 13.4 MB/s")
|
||||
self.assertEqual(status["stage"], 6)
|
||||
self.assertEqual(status["stages"], 8)
|
||||
self.assertEqual(status["stage_label"], "da3 small")
|
||||
|
||||
def test_ignores_unrelated_progress_bars(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
@@ -78,6 +143,50 @@ class DownloadProgressTest(unittest.TestCase):
|
||||
|
||||
self.assertFalse(status_path.exists())
|
||||
|
||||
def test_resolve_hf_endpoint_keeps_direct_endpoint(self):
|
||||
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||
self.worker,
|
||||
"reachable",
|
||||
return_value=True,
|
||||
):
|
||||
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
|
||||
|
||||
self.assertEqual(endpoint, "https://huggingface.co/")
|
||||
self.assertFalse(used_mirror)
|
||||
|
||||
def test_resolve_hf_endpoint_falls_back_to_mirror(self):
|
||||
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||
self.worker,
|
||||
"reachable",
|
||||
side_effect=lambda url, timeout=None: "hf-mirror" in url,
|
||||
):
|
||||
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
|
||||
|
||||
self.assertEqual(endpoint, "https://hf-mirror.com")
|
||||
self.assertTrue(used_mirror)
|
||||
|
||||
def test_resolve_hf_endpoint_respects_environment_override(self):
|
||||
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}), patch.object(
|
||||
self.worker,
|
||||
"reachable",
|
||||
) as mocked_reachable:
|
||||
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
|
||||
|
||||
self.assertEqual(endpoint, "https://hf.example")
|
||||
self.assertFalse(used_mirror)
|
||||
mocked_reachable.assert_not_called()
|
||||
|
||||
def test_resolve_hf_endpoint_keeps_direct_when_mirrors_unreachable(self):
|
||||
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||
self.worker,
|
||||
"reachable",
|
||||
return_value=False,
|
||||
):
|
||||
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
|
||||
|
||||
self.assertEqual(endpoint, "https://huggingface.co/")
|
||||
self.assertFalse(used_mirror)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -29,7 +29,7 @@ class PackTest(unittest.TestCase):
|
||||
finally:
|
||||
pack.DIST_DIR = original_dist
|
||||
|
||||
self.assertEqual(output_path.name, "SnapForm.v0.9.29.zip")
|
||||
self.assertEqual(output_path.name, "SnapForm.v0.9.32.zip")
|
||||
with ZipFile(output_path) as archive:
|
||||
names = set(archive.namelist())
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import importlib.util
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
@@ -9,6 +10,8 @@ from pathlib import Path
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
@@ -61,6 +64,57 @@ class WorkerTest(unittest.TestCase):
|
||||
inspect.getsource(self.worker.remove_background),
|
||||
)
|
||||
|
||||
def test_select_device_supports_mps_and_auto_fallback(self):
|
||||
mps_torch = SimpleNamespace(
|
||||
cuda=SimpleNamespace(is_available=lambda: False),
|
||||
backends=SimpleNamespace(
|
||||
mps=SimpleNamespace(is_available=lambda: True),
|
||||
),
|
||||
)
|
||||
mps_torch.device = lambda device_type: SimpleNamespace(
|
||||
type=device_type
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
self.worker.select_device("mps", mps_torch).type,
|
||||
"mps",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.worker.select_device("auto", mps_torch).type,
|
||||
"mps",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.worker.select_device("cpu", mps_torch).type,
|
||||
"cpu",
|
||||
)
|
||||
|
||||
cuda_torch = SimpleNamespace(
|
||||
cuda=SimpleNamespace(is_available=lambda: True),
|
||||
)
|
||||
cuda_torch.device = lambda device_type: SimpleNamespace(
|
||||
type=device_type
|
||||
)
|
||||
self.assertEqual(
|
||||
self.worker.select_device("auto", cuda_torch).type,
|
||||
"cuda",
|
||||
)
|
||||
|
||||
cpu_only_torch = SimpleNamespace(
|
||||
cuda=SimpleNamespace(is_available=lambda: False),
|
||||
backends=SimpleNamespace(
|
||||
mps=SimpleNamespace(is_available=lambda: False),
|
||||
),
|
||||
)
|
||||
cpu_only_torch.device = lambda device_type: SimpleNamespace(
|
||||
type=device_type
|
||||
)
|
||||
self.assertEqual(
|
||||
self.worker.select_device("auto", cpu_only_torch).type,
|
||||
"cpu",
|
||||
)
|
||||
with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
|
||||
self.worker.select_device("mps", cpu_only_torch)
|
||||
|
||||
def test_optional_background_removal_runs_before_da3_model_load(self):
|
||||
source = inspect.getsource(self.worker.reconstruct)
|
||||
|
||||
@@ -121,6 +175,38 @@ class WorkerTest(unittest.TestCase):
|
||||
self.assertEqual(images, [image_path])
|
||||
self.assertIsNone(temporary_directory)
|
||||
|
||||
def test_npz_is_complete_requires_readable_archive(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
valid_path = root / "valid.npz"
|
||||
np.savez_compressed(
|
||||
valid_path,
|
||||
depth=np.zeros((4, 4)),
|
||||
)
|
||||
self.assertTrue(self.worker.npz_is_complete(valid_path))
|
||||
|
||||
empty_path = root / "empty.npz"
|
||||
empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18)
|
||||
self.assertFalse(self.worker.npz_is_complete(empty_path))
|
||||
|
||||
def test_wait_for_exported_paths_waits_for_complete_npz(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
npz_path = root / "exports" / "mini_npz" / "results.npz"
|
||||
npz_path.parent.mkdir(parents=True)
|
||||
npz_path.touch()
|
||||
with patch.object(
|
||||
self.worker,
|
||||
"npz_is_complete",
|
||||
side_effect=[False, True],
|
||||
):
|
||||
result = self.worker.wait_for_exported_paths(
|
||||
root,
|
||||
"mini_npz",
|
||||
timeout=2.0,
|
||||
)
|
||||
self.assertEqual(result, [npz_path])
|
||||
|
||||
def test_rejects_unsupported_input_file(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
input_path = Path(directory) / "input.txt"
|
||||
@@ -149,6 +235,10 @@ class WorkerTest(unittest.TestCase):
|
||||
self.worker,
|
||||
"create_download_progress_class",
|
||||
return_value=object,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"resolve_hf_endpoint",
|
||||
return_value=(self.worker.HF_ENDPOINT, False),
|
||||
):
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
@@ -159,6 +249,44 @@ class WorkerTest(unittest.TestCase):
|
||||
["config.json", "model.safetensors"],
|
||||
)
|
||||
|
||||
def test_download_model_switches_to_mirror_before_hf_import(self):
|
||||
calls = []
|
||||
fake_hugging_face = ModuleType("huggingface_hub")
|
||||
fake_hugging_face.snapshot_download = lambda **options: calls.append(
|
||||
options
|
||||
)
|
||||
args = SimpleNamespace(
|
||||
repository="PramaLLC/BEN2",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
allow_pattern=[],
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||
sys.modules,
|
||||
{"huggingface_hub": fake_hugging_face},
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"create_download_progress_class",
|
||||
return_value=object,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"resolve_hf_endpoint",
|
||||
return_value=("https://hf-mirror.com", True),
|
||||
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||
self.worker,
|
||||
"write_status",
|
||||
) as mocked_write:
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
self.worker.download_model(args)
|
||||
self.assertEqual(
|
||||
os.environ["HF_ENDPOINT"],
|
||||
"https://hf-mirror.com",
|
||||
)
|
||||
messages = [call.args[2] for call in mocked_write.call_args_list]
|
||||
self.assertTrue(any("(mirror)" in message for message in messages))
|
||||
|
||||
def test_ben2_background_removal_writes_fast_transparent_png(self):
|
||||
inference_options = []
|
||||
save_options = []
|
||||
@@ -279,7 +407,10 @@ class WorkerTest(unittest.TestCase):
|
||||
else:
|
||||
mini_directory = output_directory / "exports" / "mini_npz"
|
||||
mini_directory.mkdir(parents=True)
|
||||
(mini_directory / "results.npz").touch()
|
||||
np.savez_compressed(
|
||||
mini_directory / "results.npz",
|
||||
depth=np.zeros((14, 14)),
|
||||
)
|
||||
|
||||
fake_torch = ModuleType("torch")
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(
|
||||
|
||||
Reference in New Issue
Block a user