1244 lines
43 KiB
Python
1244 lines
43 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 patch
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_bootstrap_runtime():
|
|
specification = importlib.util.spec_from_file_location(
|
|
"snapform_bootstrap_runtime",
|
|
PROJECT_ROOT / "src" / "snapform" / "bootstrap_runtime.py",
|
|
)
|
|
module = importlib.util.module_from_spec(specification)
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class BootstrapRuntimeTest(unittest.TestCase):
|
|
def create_bundled_sources(self, root):
|
|
sources = root / "sources"
|
|
sources.mkdir()
|
|
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
|
(sources / "ben2.tar.gz").write_bytes(b"ben2")
|
|
return sources
|
|
|
|
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()
|
|
self.create_bundled_sources(root)
|
|
da3_uri = (
|
|
root / "sources" / "depth-anything-3.tar.gz"
|
|
).as_uri()
|
|
ben2_uri = (root / "sources" / "ben2.tar.gz").as_uri()
|
|
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, "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
|
|
]
|
|
da3_command = next(
|
|
command
|
|
for command in commands
|
|
if any("depth-anything-3 @" in str(value) for value in command)
|
|
)
|
|
self.assertIn(
|
|
f"addict=={bootstrap.ADDICT_VERSION}",
|
|
da3_command,
|
|
)
|
|
self.assertIn(f"depth-anything-3 @ {da3_uri}", 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
|
|
)
|
|
)
|
|
ben2_command = next(
|
|
command
|
|
for command in commands
|
|
if any("ben2 @" in str(value) for value in command)
|
|
)
|
|
self.assertIn(f"ben2 @ {ben2_uri}", ben2_command)
|
|
self.assertFalse(
|
|
any("git+" in str(value) for value in ben2_command)
|
|
)
|
|
host_command = next(
|
|
command
|
|
for command in commands
|
|
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 = bootstrap.verify_command(
|
|
bootstrap.environment_python(runtime_dir)
|
|
)
|
|
self.assertTrue(
|
|
any("import torch;" in str(value) for value in check_command)
|
|
)
|
|
self.assertTrue(
|
|
any(
|
|
"from ben2 import BEN_Base" in str(value)
|
|
for value in check_command
|
|
)
|
|
)
|
|
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)
|
|
)
|
|
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 (
|
|
"depth-anything/DA3-SMALL",
|
|
"depth-anything/DA3-BASE",
|
|
"PramaLLC/BEN2",
|
|
)
|
|
)
|
|
]
|
|
self.assertEqual(len(default_downloads), 2)
|
|
self.assertFalse(
|
|
any(
|
|
"depth-anything/DA3-BASE" in map(str, command)
|
|
for command in default_downloads
|
|
)
|
|
)
|
|
ben2_download = next(
|
|
command
|
|
for command in default_downloads
|
|
if "PramaLLC/BEN2" in command
|
|
)
|
|
patterns = [
|
|
ben2_download[index + 1]
|
|
for index, value in enumerate(ben2_download)
|
|
if value == "--allow-pattern"
|
|
]
|
|
self.assertEqual(patterns, ["config.json", "model.safetensors"])
|
|
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_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()
|
|
self.create_bundled_sources(root)
|
|
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, "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
|
|
]
|
|
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()
|
|
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_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()
|
|
self.create_bundled_sources(root)
|
|
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, "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[:4]
|
|
self.assertEqual(
|
|
[item[0] for item in runtime_ranges],
|
|
[1, 2, 3, 3],
|
|
)
|
|
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_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",
|
|
"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_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("Resolving 3 packages", messages)
|
|
self.assertIn("Downloading packages 3/3", 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 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_retries_once_with_mirror(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
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[1].args[5],
|
|
{"PIP_INDEX_URL": bootstrap.PYPI_MIRRORS[0]},
|
|
)
|
|
|
|
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_ai_component_estimate_continues_across_internal_steps(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
|
|
first = bootstrap.estimated_remaining(
|
|
"AI components",
|
|
bootstrap.XFORMERS_SLICE[0],
|
|
)
|
|
second = bootstrap.estimated_remaining(
|
|
"AI components",
|
|
bootstrap.AI_COMPONENTS_SLICE[0],
|
|
)
|
|
|
|
self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["AI components"])
|
|
self.assertGreater(second, 0)
|
|
self.assertLess(second, first)
|
|
|
|
def test_resolve_indexes_uses_configured_pypi_index(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
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_bootstrap_runtime()
|
|
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_bootstrap_runtime()
|
|
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["PIP_INDEX_URL"],
|
|
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
|
)
|
|
|
|
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()
|
|
self.create_bundled_sources(root)
|
|
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, "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(
|
|
(runtime_dir / "snapform_ready.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
self.assertEqual(
|
|
ready["packages"],
|
|
["numpy==1.26.4", "torch==2.9.1+cu128"],
|
|
)
|
|
self.assertEqual(ready["devices"], ["cuda", "cpu"])
|
|
|
|
def test_broken_runtime_is_cleaned_before_reinstall(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
self.create_bundled_sources(root)
|
|
(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, "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((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)
|
|
self.create_bundled_sources(root)
|
|
(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, "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((runtime_dir / "marker.txt").is_file())
|
|
|
|
def test_model_download_failure_keeps_runtime_ready(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
commands = []
|
|
|
|
def fake_run(command, environment=None):
|
|
commands.append(command)
|
|
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()
|
|
self.create_bundled_sources(root)
|
|
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, "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((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"])
|
|
self.assertTrue(
|
|
any(
|
|
"PramaLLC/BEN2" in map(str, command)
|
|
for command in commands
|
|
)
|
|
)
|
|
|
|
def test_install_fails_when_bundled_uv_missing(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
self.create_bundled_sources(root)
|
|
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, "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_sources_and_uv_without_network(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_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")
|
|
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),
|
|
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"),
|
|
)
|
|
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_sources_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
|
|
)
|
|
)
|
|
ai_command = next(
|
|
command
|
|
for command in install_commands
|
|
if any("depth-anything-3 @" in str(value) for value in command)
|
|
)
|
|
da3_uri = (root / "sources" / "depth-anything-3.tar.gz").as_uri()
|
|
ben2_uri = (root / "sources" / "ben2.tar.gz").as_uri()
|
|
self.assertIn(f"depth-anything-3 @ {da3_uri}", ai_command)
|
|
self.assertIn(f"ben2 @ {ben2_uri}", ai_command)
|
|
|
|
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()
|
|
self.create_bundled_sources(root)
|
|
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, "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_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")
|
|
|
|
def test_sync_bundled_sources_copies_missing_archives(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
bundled = root / "bundled"
|
|
bundled.mkdir()
|
|
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
|
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
|
|
sources = root / "sources"
|
|
|
|
bootstrap.sync_bundled_sources(bundled, sources)
|
|
|
|
self.assertEqual(
|
|
(sources / "depth-anything-3.tar.gz").read_bytes(),
|
|
b"da3",
|
|
)
|
|
self.assertEqual(
|
|
(sources / "ben2.tar.gz").read_bytes(),
|
|
b"ben2",
|
|
)
|
|
|
|
def test_install_stages_uses_local_source_archives(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
sources = root / "sources"
|
|
sources.mkdir()
|
|
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
|
(sources / "ben2.tar.gz").write_bytes(b"ben2")
|
|
uv_path = root / "uv"
|
|
uv_path.write_bytes(b"uv")
|
|
args = SimpleNamespace(
|
|
uv=str(uv_path),
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
|
|
stages = bootstrap.install_stages(
|
|
args,
|
|
runtime_dir,
|
|
root / "python",
|
|
uv_path,
|
|
)
|
|
|
|
ai_command = next(
|
|
command
|
|
for _label, command, _environment, kind in stages
|
|
if kind == "install"
|
|
and any("depth-anything-3" in str(value) for value in command)
|
|
)
|
|
da3_local = (sources / "depth-anything-3.tar.gz").as_uri()
|
|
ben2_local = (sources / "ben2.tar.gz").as_uri()
|
|
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
|
|
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
|
|
|
|
def test_install_stages_download_only_default_models(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
sources = root / "sources"
|
|
sources.mkdir()
|
|
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
|
(sources / "ben2.tar.gz").write_bytes(b"ben2")
|
|
uv_path = root / "uv"
|
|
uv_path.write_bytes(b"uv")
|
|
args = SimpleNamespace(
|
|
uv=str(uv_path),
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
|
|
stages = bootstrap.install_stages(
|
|
args,
|
|
runtime_dir,
|
|
root / "python",
|
|
uv_path,
|
|
)
|
|
|
|
downloads = {
|
|
label: command
|
|
for label, command, _environment, kind in stages
|
|
if kind == "download"
|
|
}
|
|
self.assertIn("da3_small", " ".join(map(str, downloads["DA3 Small"])))
|
|
self.assertIn("ben2_base", " ".join(map(str, downloads["BEN2 Base"])))
|
|
self.assertNotIn("DA3 Base", downloads)
|
|
|
|
def test_install_does_not_verify_runtime(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
self.create_bundled_sources(root)
|
|
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,
|
|
"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_bootstrap_runtime()
|
|
|
|
class FakeProcess:
|
|
def __init__(self, lines):
|
|
self.stdout = lines
|
|
|
|
def wait(self):
|
|
return 0
|
|
|
|
fake_process = FakeProcess(
|
|
['SNAPFORM_DEVICES={"cuda": false, "mps": true}']
|
|
)
|
|
with patch.object(
|
|
bootstrap.subprocess,
|
|
"Popen",
|
|
return_value=fake_process,
|
|
):
|
|
report = bootstrap.run_verify(
|
|
"python",
|
|
"status.json",
|
|
clock=lambda: 0.0,
|
|
)
|
|
|
|
self.assertFalse(report["cuda"])
|
|
self.assertTrue(report["mps"])
|
|
|
|
def test_run_verify_writes_heartbeat_progress(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
|
|
class FakeProcess:
|
|
def __init__(self, lines):
|
|
self.stdout = lines
|
|
|
|
def wait(self):
|
|
return 0
|
|
|
|
clock_state = {"now": 0.0}
|
|
|
|
def fake_clock():
|
|
clock_state["now"] += 6.0
|
|
return clock_state["now"]
|
|
|
|
def delayed_empty_lines():
|
|
time.sleep(0.6)
|
|
if False:
|
|
yield
|
|
|
|
fake_process = FakeProcess(delayed_empty_lines())
|
|
with patch.object(
|
|
bootstrap.subprocess,
|
|
"Popen",
|
|
return_value=fake_process,
|
|
), patch.object(bootstrap, "write_status") as mocked_write:
|
|
report = bootstrap.run_verify(
|
|
"python",
|
|
"status.json",
|
|
clock=fake_clock,
|
|
)
|
|
|
|
messages = [call.args[2] for call in mocked_write.call_args_list]
|
|
self.assertTrue(
|
|
any("Checking Runtime" in message for message in messages)
|
|
)
|
|
self.assertEqual(report, {})
|
|
|
|
def test_install_uv_extracts_bundled_archive(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
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_sources_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_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
uv_path = root / "uv"
|
|
bundled = root / "bundled"
|
|
bundled.mkdir()
|
|
with patch.object(
|
|
bootstrap,
|
|
"bundled_sources_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()
|