Files
AnyImage/tests/test_install_env.py
nan 0c8d3a8988 Bump version to 0.11.2 and update uv version to 0.11.3
- Updated version in pyproject.toml and blender_manifest.toml to 0.11.2.
- Enhanced uv executable path documentation in client.py.
- Added inference_mode parameter to reconstruct_request function in client.py.
- Implemented uv download functionality in install_env.py with improved error handling and checksum verification.
- Modified install_uv function to download and extract uv binary instead of using bundled version.
- Updated job progress handling in operators.py to track progress stages.
- Improved download progress reporting in worker.py and adjusted related tests.
- Refactored tests to accommodate changes in status reporting and model downloading.
- Removed vendor.lock.toml as part of the cleanup process.
2026-08-14 23:43:01 +08:00

1349 lines
46 KiB
Python

import importlib.util
import io
import hashlib
import json
import os
import shutil
import subprocess
import tarfile
import tempfile
import time
import unittest
import zipfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_install_env():
specification = importlib.util.spec_from_file_location(
"snapform_install_env",
PROJECT_ROOT / "src" / "snapform" / "install_env.py",
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
return module
class InstallEnvTest(unittest.TestCase):
def create_bundled_vendor(self, root):
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(vendor / "ben2.tar.gz").write_bytes(b"ben2")
(vendor / "utils3d.tar.gz").write_bytes(b"utils3d")
(vendor / "moge.tar.gz").write_bytes(b"moge")
return vendor
def test_installs_onnx_runtime_without_training_frameworks(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
self.create_bundled_vendor(root)
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap,
"run",
) as mocked_run, patch.object(
bootstrap,
"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]
commands += [
call.args[4] for call in mocked_run_install.call_args_list
]
service_command = next(
command
for command in commands
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
)
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command)
self.assertIn(
f"onnxruntime-directml=={bootstrap.ONNXRUNTIME_VERSION}",
service_command,
)
for forbidden in ("torch", "torchvision", "xformers", "ben2 @", "moge @"):
self.assertFalse(
any(forbidden in str(value) for value in service_command),
forbidden,
)
check_command = bootstrap.verify_command(
bootstrap.venv_python(venv_dir)
)
self.assertFalse(any("torch" in str(value) for value in check_command))
self.assertTrue(
any("import fastapi, uvicorn" in str(value) for value in check_command)
)
self.assertTrue(
any("onnxruntime" in str(value) for value in check_command)
)
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
if any(
repository in command
for repository in (
"Ruicheng/moge-2-vits-normal-onnx",
"onnx-community/BEN2-ONNX",
)
)
]
self.assertEqual(len(default_downloads), 2)
self.assertTrue(any("Ruicheng/moge-2-vits-normal-onnx" in map(str, command) for command in default_downloads))
ben2_download = next(
command
for command in default_downloads
if "onnx-community/BEN2-ONNX" in command
)
patterns = [
ben2_download[index + 1]
for index, value in enumerate(ben2_download)
if value == "--allow-pattern"
]
self.assertEqual(patterns, ["onnx/model_fp16.onnx"])
self.assertIn("--stage", ben2_download)
self.assertIn("--stage-label", ben2_download)
self.assertEqual(
ben2_download[ben2_download.index("--label") + 1],
"BEN2 Base",
)
def test_darwin_install_uses_default_onnx_runtime(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
self.create_bundled_vendor(root)
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="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, "resolve_indexes", return_value=None), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
) as mocked_run_install, patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
commands = [
call.args[4] for call in mocked_run_install.call_args_list
]
runtime_command = next(
command
for command in commands
if any(
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}" in str(value)
for value in command
)
)
self.assertNotIn("--index-url", runtime_command)
self.assertFalse(
any(
"xformers" in str(value)
for command in commands
for value in command
)
)
def test_intel_mac_uses_portable_onnx_runtime(self):
bootstrap = load_install_env()
with patch.object(
bootstrap.sys,
"platform",
"darwin",
):
stages = bootstrap.install_stages(
SimpleNamespace(
models_dir="models",
status="status.json",
),
".venv",
"python",
"uv",
)
runtime = stages[1][1]
self.assertIn(
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}",
runtime,
)
def test_run_stage_writes_stage_progress(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
status_path = Path(directory) / "status.json"
with patch.object(bootstrap, "run") as mocked_run:
bootstrap.run_stage(
str(status_path),
3,
8,
"pytorch",
["uv", "pip", "install", "torch"],
None,
"plain",
)
self.assertEqual(
mocked_run.call_args.args[0],
["uv", "pip", "install", "torch"],
)
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"], 8)
self.assertEqual(status["stage_label"], "pytorch")
def test_run_stage_writes_progress_range(self):
bootstrap = load_install_env()
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_install_env()
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_resets_progress_for_each_stage(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
self.create_bundled_vendor(root)
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="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, "resolve_indexes", return_value=None), patch.object(
bootstrap,
"run_stage",
side_effect=fake_run_stage,
), patch.object(bootstrap, "run"), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
self.assertEqual(stage_ranges[:2], [(1, 0.0, 0.05), (1, 0.05, 1.0)])
self.assertEqual(stage_ranges[2:], [(2, 0.0, 1.0), (3, 0.0, 1.0)])
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_install_env()
class FakeProcess:
def __init__(self, lines):
self.stdout = lines
def wait(self):
return 0
fake_process = FakeProcess(
[
"Resolved 3 packages in 1.0s",
"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,
[
"Resolving 3 packages",
"Downloading packages 3/3",
"Installing packages 3/3",
],
)
self.assertEqual(mocked_write.call_args_list[0].args[1], 0.02)
self.assertLess(mocked_write.call_args_list[-1].args[1], 0.03)
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_install_env()
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("Resolving 3 packages", messages)
self.assertIn("Downloading packages 3/3", messages)
def test_run_install_maps_progress_into_stage_range(self):
bootstrap = load_install_env()
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 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,
)
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_run_install_heartbeat_reports_package_count(self):
bootstrap = load_install_env()
class SlowLines:
def __iter__(self):
yield "Resolved 3 packages in 1.0s\n"
time.sleep(0.6)
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"]
fake_process = FakeProcess(SlowLines())
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
), patch.object(bootstrap, "write_status") as mocked_write:
bootstrap._run_install_once(
"status.json",
2,
2,
"ONNX Runtime",
["uv", "pip", "install", "onnxruntime"],
None,
0.10,
0.85,
fake_clock,
)
heartbeats = [
call
for call in mocked_write.call_args_list
if call.args[2].startswith("Downloading")
]
self.assertTrue(heartbeats)
self.assertTrue(
any("3 packages" in call.args[2] for call in heartbeats)
)
for call in heartbeats:
self.assertNotIn("elapsed", call.args[2])
self.assertNotIn("eta_seconds", call.kwargs)
def test_run_install_retries_across_indexes(self):
bootstrap = load_install_env()
previous = bootstrap._resolved_pypi_index
bootstrap._resolved_pypi_index = None
try:
with patch.object(
bootstrap,
"_run_install_once",
side_effect=[
subprocess.CalledProcessError(1, ["uv", "pip"]),
["numpy==1.26.4"],
],
) as mocked_install:
installed = bootstrap.run_install(
"status.json",
3,
8,
"ai deps",
["uv", "pip", "install", "x"],
)
finally:
bootstrap._resolved_pypi_index = previous
self.assertEqual(installed, ["numpy==1.26.4"])
self.assertEqual(mocked_install.call_count, 2)
# 未指定源时默认国内镜像优先,第一个镜像失败后切到第二个镜像。
self.assertEqual(
mocked_install.call_args_list[0].args[5],
{"UV_INDEX_URL": bootstrap.PYPI_MIRRORS[0]},
)
self.assertEqual(
mocked_install.call_args_list[1].args[5],
{"UV_INDEX_URL": bootstrap.PYPI_MIRRORS[1]},
)
def test_run_install_uses_only_configured_index(self):
bootstrap = load_install_env()
previous = bootstrap._resolved_pypi_index
bootstrap._resolved_pypi_index = (
"https://pypi.tuna.tsinghua.edu.cn/simple"
)
try:
with patch.object(
bootstrap,
"_run_install_once",
return_value=["numpy==1.26.4"],
) as mocked_install:
installed = bootstrap.run_install(
"status.json",
3,
8,
"ai deps",
["uv", "pip", "install", "x"],
)
finally:
bootstrap._resolved_pypi_index = previous
self.assertEqual(installed, ["numpy==1.26.4"])
self.assertEqual(mocked_install.call_count, 1)
self.assertEqual(
mocked_install.call_args.args[5],
{"UV_INDEX_URL": "https://pypi.tuna.tsinghua.edu.cn/simple"},
)
def test_resolve_indexes_uses_configured_pypi_index(self):
bootstrap = load_install_env()
bootstrap.resolve_indexes(
SimpleNamespace(pypi_index="https://pypi.tuna.tsinghua.edu.cn/simple")
)
self.assertEqual(
bootstrap._resolved_pypi_index,
"https://pypi.tuna.tsinghua.edu.cn/simple",
)
def test_resolve_indexes_uses_environment_override(self):
bootstrap = load_install_env()
with patch.dict(
os.environ,
{"UV_INDEX_URL": "https://pypi.example/simple"},
):
bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None))
self.assertEqual(
bootstrap._resolved_pypi_index,
"https://pypi.example/simple",
)
def test_resolve_indexes_falls_back_to_pip_index_url(self):
bootstrap = load_install_env()
with patch.dict(
os.environ,
{"PIP_INDEX_URL": "https://pypi.example/simple"},
):
bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None))
self.assertEqual(
bootstrap._resolved_pypi_index,
"https://pypi.example/simple",
)
def test_child_environment_injects_resolved_pypi_index(self):
bootstrap = load_install_env()
previous_pypi = bootstrap._resolved_pypi_index
bootstrap._resolved_pypi_index = (
"https://pypi.tuna.tsinghua.edu.cn/simple"
)
try:
env = bootstrap.child_environment()
finally:
bootstrap._resolved_pypi_index = previous_pypi
self.assertEqual(
env["UV_INDEX_URL"],
"https://pypi.tuna.tsinghua.edu.cn/simple",
)
self.assertEqual(
env["PYTHONIOENCODING"],
"utf-8:backslashreplace",
)
self.assertEqual(env["PYTHONUTF8"], "1")
def test_standard_streams_are_reconfigured_for_unicode_logs(self):
bootstrap = load_install_env()
stdout = Mock()
stderr = Mock()
with patch.object(bootstrap.sys, "stdout", stdout), patch.object(
bootstrap.sys,
"stderr",
stderr,
):
bootstrap.configure_standard_streams()
stdout.reconfigure.assert_called_once_with(
encoding="utf-8",
errors="backslashreplace",
)
stderr.reconfigure.assert_called_once_with(
encoding="utf-8",
errors="backslashreplace",
)
def test_ready_json_records_installed_packages(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
self.create_bundled_vendor(root)
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap,
"run",
), patch.object(
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)
ready = json.loads(
(root / "manifest.json").read_text(
encoding="utf-8"
)
)
self.assertEqual(
ready["packages"],
["numpy==1.26.4", "torch==2.9.1+cu128"],
)
self.assertEqual(ready["devices"], ["directml", "cpu"])
self.assertEqual(
ready["manifest_revision"],
bootstrap.MANIFEST_REVISION,
)
def test_broken_env_is_cleaned_before_reinstall(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
self.create_bundled_vendor(root)
(venv_dir / "partial-python.exe").write_text("junk")
args = SimpleNamespace(
uv="uv",
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
self.assertFalse((venv_dir / "partial-python.exe").exists())
self.assertTrue(
(root / "manifest.json").is_file()
)
def test_installed_env_is_not_cleaned(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir(parents=True)
self.create_bundled_vendor(root)
(venv_dir / "Scripts").mkdir()
(venv_dir / "Scripts" / "python.exe").write_text("python")
(root / "manifest.json").write_text(
json.dumps(
{"manifest_revision": bootstrap.MANIFEST_REVISION}
),
encoding="utf-8",
)
(venv_dir / "marker.txt").write_text("keep")
args = SimpleNamespace(
uv="uv",
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
self.assertTrue((venv_dir / "marker.txt").is_file())
def test_model_download_failure_keeps_env_installed(self):
bootstrap = load_install_env()
commands = []
def fake_run(command, environment=None):
commands.append(command)
if any("moge-2-vits-normal" in str(value) for value in command):
raise subprocess.CalledProcessError(1, command)
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
self.create_bundled_vendor(root)
args = SimpleNamespace(
uv="uv",
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap,
"run",
side_effect=fake_run,
), patch.object(bootstrap, "run_install"), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
self.assertTrue(
(root / "manifest.json").is_file()
)
status = json.loads(
(root / "status.json").read_text(encoding="utf-8")
)
self.assertEqual(status["progress"], 1.0)
self.assertIn("Environment is installed", status["message"])
self.assertTrue(
any(
"onnx-community/BEN2-ONNX" in map(str, command)
for command in commands
)
)
def test_install_fails_when_all_uv_sources_are_unavailable(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
uv_path = root / "uv" / "uv.exe"
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "resolve_indexes", return_value=None), 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_archive",
side_effect=RuntimeError("all sources unavailable"),
):
with self.assertRaisesRegex(
RuntimeError,
"all sources unavailable",
):
bootstrap.install(args)
def test_install_downloads_and_extracts_uv(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
target, archive_name = bootstrap.uv_download_target()
archive = root / f"uv-{target}.{archive_name}"
if archive_name.endswith("zip"):
with zipfile.ZipFile(archive, "w") as bundle:
bundle.writestr("uv/uv.exe", b"uv-binary")
else:
with tarfile.open(archive, "w:gz") as tar:
content = b"uv-binary"
info = tarfile.TarInfo("uv")
info.size = len(content)
tar.addfile(info, io.BytesIO(content))
uv_path = root / "uv" / ("uv.exe" if os.name == "nt" else "uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
run_commands = []
install_commands = []
def fake_run(command, environment=None):
run_commands.append(command)
def fake_run_install(
status_path,
stage,
stages,
label,
command,
environment=None,
start=0.0,
end=1.0,
clock=None,
):
install_commands.append(command)
return []
def fake_download(destination, *_args):
shutil.copy2(archive, destination)
return destination
with patch.object(
bootstrap,
"resolve_indexes",
return_value=None,
), patch.object(
bootstrap,
"download_uv_archive",
side_effect=fake_download,
), patch.object(
bootstrap,
"run",
side_effect=fake_run,
), patch.object(
bootstrap,
"run_install",
side_effect=fake_run_install,
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
self.assertEqual(uv_path.read_bytes(), b"uv-binary")
all_commands = run_commands + install_commands
runtime_command = next(
command
for command in install_commands
if any("onnxruntime" in str(value) for value in command)
)
self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
self.assertFalse(any("torch" in str(value) for value in runtime_command))
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
self.assertEqual(status["stage"], 4)
self.assertEqual(status["stages"], 4)
def test_install_reuses_existing_uv(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
self.create_bundled_vendor(root)
uv_path = root / "uv.exe"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
def test_extract_uv_unpacks_standalone_binary(self):
bootstrap = load_install_env()
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")
def test_uv_download_falls_back_to_omoolab_and_verifies_checksum(self):
bootstrap = load_install_env()
content = io.BytesIO()
with zipfile.ZipFile(content, "w") as archive:
archive.writestr("uv/uv.exe", b"uv-binary")
payload = content.getvalue()
class Response(io.BytesIO):
headers = {"Content-Length": str(len(payload))}
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
destination = root / "uv.tmp"
status = root / "status.json"
urls = []
clock_state = {"now": 0.0}
def clock():
clock_state["now"] += 0.25
return clock_state["now"]
def open_url(request):
urls.append(request.full_url)
if len(urls) == 1:
raise OSError("official unavailable")
return Response(payload)
with patch.object(
bootstrap,
"uv_download_target",
return_value=("test-target", "zip"),
), patch.dict(
bootstrap.UV_ARTIFACTS,
{
"test-target": (
"uv-test.zip",
hashlib.sha256(payload).hexdigest(),
)
},
), patch.dict(
bootstrap._uv_download_openers,
{
bootstrap.UV_OFFICIAL_BASE_URL: SimpleNamespace(
open=open_url
),
bootstrap.UV_MIRROR_BASE_URL: SimpleNamespace(
open=open_url
),
},
), patch.object(
bootstrap,
"write_status",
) as write_status:
result = bootstrap.download_uv_archive(
destination,
status,
0.0,
1.0,
{},
clock=clock,
)
self.assertEqual(result.read_bytes(), payload)
self.assertTrue(urls[0].startswith(bootstrap.UV_OFFICIAL_BASE_URL))
self.assertTrue(urls[1].startswith(bootstrap.UV_MIRROR_BASE_URL))
self.assertEqual(bootstrap.UV_OFFICIAL_CONNECT_TIMEOUT, 3.0)
self.assertEqual(bootstrap.UV_MIRROR_CONNECT_TIMEOUT, 15.0)
messages = [call.args[2] for call in write_status.call_args_list]
self.assertTrue(any("uv (mirror)" in value for value in messages))
self.assertTrue(any("B/s" in value for value in messages))
self.assertFalse(any("left" in value for value in messages))
def test_install_stages_use_only_runtime_packages(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
stages = bootstrap.install_stages(
args,
venv_dir,
root / "python",
uv_path,
)
runtime_command = next(
command
for _label, command, _environment, kind in stages
if kind == "install"
and any("onnxruntime" in str(value) for value in command)
)
self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
self.assertFalse(any("torch" in str(value) for value in runtime_command))
def test_install_stages_download_only_default_models(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(vendor / "ben2.tar.gz").write_bytes(b"ben2")
(vendor / "utils3d.tar.gz").write_bytes(b"utils3d")
(vendor / "moge.tar.gz").write_bytes(b"moge")
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
stages = bootstrap.install_stages(
args,
venv_dir,
root / "python",
uv_path,
)
downloads = {
label: command
for label, command, _environment, kind in stages
if kind == "download"
}
self.assertIn(
"moge-2-vits-normal-onnx",
" ".join(map(str, downloads["MoGe-2 ViT-S Normal"])),
)
self.assertIn("BEN2-ONNX", " ".join(map(str, downloads["BEN2 Base"])))
self.assertEqual(
set(downloads),
{"MoGe-2 ViT-S Normal", "BEN2 Base"},
)
def test_install_does_not_verify_env(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
self.create_bundled_vendor(root)
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
venv_dir=str(venv_dir),
models_dir=str(root / "models"),
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(
bootstrap,
"resolve_indexes",
return_value=None,
), patch.object(bootstrap, "run"), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
) as mocked_verify:
bootstrap.install(args)
mocked_verify.assert_not_called()
def test_run_verify_returns_detected_devices(self):
bootstrap = load_install_env()
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_install_env()
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(
"Checking Environment" in message
for message in messages
)
)
self.assertEqual(report, {})
def test_install_uv_extracts_downloaded_archive(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
uv_path = root / "uv"
target, archive_name = bootstrap.uv_download_target()
archive = root / f"uv-{target}.{archive_name}"
if archive_name.endswith("zip"):
with zipfile.ZipFile(archive, "w") as bundle:
bundle.writestr("uv/uv.exe", b"uv-binary")
else:
with tarfile.open(archive, "w:gz") as tar:
content = b"uv-binary"
info = tarfile.TarInfo("uv")
info.size = len(content)
tar.addfile(info, io.BytesIO(content))
def fake_download(destination, *_args):
shutil.copy2(archive, destination)
return destination
with patch.object(
bootstrap,
"download_uv_archive",
side_effect=fake_download,
):
bootstrap.install_uv(uv_path, str(root / "status.json"))
self.assertEqual(uv_path.read_bytes(), b"uv-binary")
def test_install_uv_raises_when_download_fails(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
uv_path = root / "uv"
with patch.object(
bootstrap,
"download_uv_archive",
side_effect=RuntimeError("all uv sources failed"),
):
with self.assertRaisesRegex(
RuntimeError,
"all uv sources failed",
):
bootstrap.install_uv(uv_path, str(root / "status.json"))
if __name__ == "__main__":
unittest.main()