nan
5255eb4e96
- Updated ONNX Runtime to 1.24.4 and replaced OpenCV with PyAV 17.0.0 in the installation environment. - Enhanced model downloading logic to prioritize Hugging Face and fallback to OmooLab R2, with improved error handling and status reporting. - Introduced a new environment revision file for shared metadata. - Refactored video frame extraction to utilize the PyAV library instead of OpenCV. - Improved installation process to handle retries across specified PyPI indexes. - Added tests for download progress reporting and model downloading behavior.
1295 lines
44 KiB
Python
1295 lines
44 KiB
Python
import importlib.util
|
|
import io
|
|
import json
|
|
import os
|
|
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.assertNotIn("--stage", ben2_download)
|
|
self.assertNotIn("--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_allocates_contiguous_progress_slices(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)
|
|
|
|
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.ENVIRONMENT_SLICE[0])
|
|
runtime_ranges = stage_ranges[:2]
|
|
self.assertEqual(
|
|
[item[0] for item in runtime_ranges],
|
|
[1, 2],
|
|
)
|
|
self.assertTrue(
|
|
all(item[0] is None for item in stage_ranges[4:])
|
|
)
|
|
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_bundled_uv_missing(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.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,
|
|
"bundled_uv_archive",
|
|
return_value=None,
|
|
):
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"Bundled uv is missing",
|
|
):
|
|
bootstrap.install(args)
|
|
|
|
def test_install_uses_bundled_vendor_and_uv_without_network(self):
|
|
bootstrap = load_install_env()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
venv_dir = root / ".venv"
|
|
venv_dir.mkdir()
|
|
bundled = root / "bundled"
|
|
bundled.mkdir()
|
|
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
|
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
|
|
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
|
|
(bundled / "moge.tar.gz").write_bytes(b"moge")
|
|
target, archive_name = bootstrap.uv_download_target()
|
|
archive = bundled / 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 []
|
|
|
|
with patch.object(
|
|
bootstrap,
|
|
"bundled_vendor_directory",
|
|
return_value=bundled,
|
|
), patch.object(
|
|
bootstrap,
|
|
"resolve_indexes",
|
|
return_value=None,
|
|
), 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
|
|
self.assertFalse(
|
|
any(
|
|
"github.com" in str(value)
|
|
for command in all_commands
|
|
for value in command
|
|
)
|
|
)
|
|
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))
|
|
|
|
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_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_bundled_archive(self):
|
|
bootstrap = load_install_env()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
uv_path = root / "uv"
|
|
bundled = root / "bundled"
|
|
bundled.mkdir(parents=True)
|
|
target, archive_name = bootstrap.uv_download_target()
|
|
archive = bundled / 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))
|
|
with patch.object(
|
|
bootstrap,
|
|
"bundled_vendor_directory",
|
|
return_value=bundled,
|
|
):
|
|
bootstrap.install_uv(uv_path, str(root / "status.json"))
|
|
|
|
self.assertEqual(uv_path.read_bytes(), b"uv-binary")
|
|
|
|
def test_install_uv_raises_when_bundled_missing(self):
|
|
bootstrap = load_install_env()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
uv_path = root / "uv"
|
|
bundled = root / "bundled"
|
|
bundled.mkdir()
|
|
with patch.object(
|
|
bootstrap,
|
|
"bundled_vendor_directory",
|
|
return_value=bundled,
|
|
):
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"Bundled uv is missing",
|
|
):
|
|
bootstrap.install_uv(uv_path, str(root / "status.json"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|