nan
d08d0e9c1a
- Modified status messages and stage labels in `test_addon.py` for clarity and consistency. - Added tests to ensure proper host termination and command execution in `test_addon.py`. - Enhanced `test_bootstrap_runtime.py` with additional tests for installation behavior, including handling of CUDA and MPS devices, and improved error handling for disk space and network checks. - Updated download progress reporting in `test_download_progress.py` to include download speed and improved message formatting. - Added new tests in `test_worker.py` for device selection, NPZ file validation, and model downloading behavior with Hugging Face endpoints. - Bumped version in `uv.lock` to 0.9.30.
912 lines
32 KiB
Python
912 lines
32 KiB
Python
import importlib.util
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
import zipfile
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
|
|
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 test_installs_da3_from_archives_without_gsplat(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
uv_path = root / "uv"
|
|
uv_path.write_bytes(b"uv")
|
|
args = SimpleNamespace(
|
|
uv=str(uv_path),
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run",
|
|
) as mocked_run, patch.object(
|
|
bootstrap,
|
|
"run_install",
|
|
) as mocked_run_install:
|
|
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.assertTrue(
|
|
any(
|
|
"https://github.com/ByteDance-Seed/Depth-Anything-3/"
|
|
"archive/test-commit.tar.gz" in str(value)
|
|
for value in da3_command
|
|
)
|
|
)
|
|
self.assertFalse(
|
|
any("git+" in str(value) for value in da3_command)
|
|
)
|
|
self.assertFalse(
|
|
any(
|
|
"gsplat" in str(value)
|
|
for command in commands
|
|
for value in command
|
|
)
|
|
)
|
|
self.assertFalse(
|
|
any(
|
|
"--no-build-isolation" in str(value)
|
|
for command in commands
|
|
for value in command
|
|
)
|
|
)
|
|
ben2_command = next(
|
|
command
|
|
for command in commands
|
|
if any("ben2 @" in str(value) for value in command)
|
|
)
|
|
self.assertTrue(
|
|
any(
|
|
"https://github.com/PramaLLC/BEN2/archive/"
|
|
"test-ben2-commit.tar.gz" in str(value)
|
|
for value in ben2_command
|
|
)
|
|
)
|
|
self.assertFalse(
|
|
any("git+" in str(value) for value in ben2_command)
|
|
)
|
|
host_command = next(
|
|
command
|
|
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 = next(
|
|
command
|
|
for command in commands
|
|
if any("import torch;" in str(value) for value in 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)
|
|
)
|
|
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), 3)
|
|
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.assertIn("--stage", ben2_download)
|
|
self.assertIn("--stage-label", ben2_download)
|
|
self.assertEqual(ben2_download[-1], "ben2 base")
|
|
|
|
def test_darwin_install_skips_cuda_index_and_xformers(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
uv_path = root / "uv"
|
|
uv_path.write_bytes(b"uv")
|
|
args = SimpleNamespace(
|
|
uv=str(uv_path),
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
with patch.object(
|
|
bootstrap.sys,
|
|
"platform",
|
|
"darwin",
|
|
), patch.object(
|
|
bootstrap.platform,
|
|
"machine",
|
|
return_value="arm64",
|
|
), patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run",
|
|
), patch.object(
|
|
bootstrap,
|
|
"run_install",
|
|
) as mocked_run_install:
|
|
bootstrap.install(args)
|
|
|
|
commands = [
|
|
call.args[4] for call in mocked_run_install.call_args_list
|
|
]
|
|
torch_command = next(
|
|
command
|
|
for command in commands
|
|
if any(
|
|
f"torch=={bootstrap.TORCH_VERSION}" in str(value)
|
|
for value in command
|
|
)
|
|
)
|
|
self.assertNotIn("--index-url", torch_command)
|
|
self.assertFalse(
|
|
any(
|
|
"xformers" in str(value)
|
|
for command in commands
|
|
for value in command
|
|
)
|
|
)
|
|
|
|
def test_intel_mac_is_rejected_with_clear_error(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
args = SimpleNamespace(
|
|
uv="uv",
|
|
runtime_dir="runtime",
|
|
models_dir="models",
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status="status.json",
|
|
)
|
|
with patch.object(
|
|
bootstrap.sys,
|
|
"platform",
|
|
"darwin",
|
|
), patch.object(
|
|
bootstrap.platform,
|
|
"machine",
|
|
return_value="x86_64",
|
|
), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree:
|
|
with self.assertRaisesRegex(RuntimeError, "Apple Silicon"):
|
|
bootstrap.install(args)
|
|
|
|
mocked_rmtree.assert_not_called()
|
|
|
|
def test_run_stage_writes_stage_progress(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
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()
|
|
uv_path = root / "uv"
|
|
uv_path.write_bytes(b"uv")
|
|
args = SimpleNamespace(
|
|
uv=str(uv_path),
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
stage_ranges = []
|
|
|
|
def fake_run_stage(
|
|
status_path,
|
|
stage,
|
|
stages,
|
|
label,
|
|
command,
|
|
environment,
|
|
kind,
|
|
start=0.0,
|
|
end=1.0,
|
|
):
|
|
stage_ranges.append((stage, start, end))
|
|
bootstrap.write_status(
|
|
status_path,
|
|
end,
|
|
"",
|
|
stage=stage,
|
|
stages=stages,
|
|
stage_label=label,
|
|
)
|
|
return []
|
|
|
|
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run_stage",
|
|
side_effect=fake_run_stage,
|
|
), patch.object(bootstrap, "run"):
|
|
bootstrap.install(args)
|
|
|
|
starts = [item[1] for item in stage_ranges]
|
|
ends = [item[2] for item in stage_ranges]
|
|
self.assertTrue(
|
|
all(starts[i] < starts[i + 1] for i in range(len(starts) - 1))
|
|
)
|
|
self.assertEqual(ends[-1], 1.0)
|
|
self.assertEqual(starts[0], bootstrap.RUNTIME_SLICE_START)
|
|
status = json.loads(
|
|
(root / "status.json").read_text(encoding="utf-8")
|
|
)
|
|
self.assertEqual(status["progress"], 1.0)
|
|
|
|
def test_run_install_tracks_package_download_events(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
|
|
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,
|
|
[
|
|
"3 packages",
|
|
"3/3 packages",
|
|
"3/3 packages",
|
|
],
|
|
)
|
|
self.assertEqual(mocked_write.call_args_list[-1].args[1], 0.02)
|
|
first = mocked_write.call_args_list[0]
|
|
self.assertEqual(first.kwargs["stage"], 2)
|
|
self.assertEqual(first.kwargs["stages"], 8)
|
|
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
|
|
self.assertEqual(installed, ["torch==2.9.1", "numpy==1.26.4"])
|
|
|
|
def test_run_install_smooths_progress_between_events(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
|
|
class SlowLines:
|
|
def __init__(self, lines, delay):
|
|
self.lines = list(lines)
|
|
self.delay = delay
|
|
|
|
def __iter__(self):
|
|
for line in self.lines:
|
|
time.sleep(self.delay)
|
|
yield line
|
|
|
|
class FakeProcess:
|
|
def __init__(self, lines):
|
|
self.stdout = lines
|
|
|
|
def wait(self):
|
|
return 0
|
|
|
|
fake_process = FakeProcess(
|
|
SlowLines(
|
|
[
|
|
"Resolved 3 packages in 1.0s",
|
|
"Prepared 3 packages in 2.0s",
|
|
"Installed 3 packages in 0.5s",
|
|
],
|
|
0.3,
|
|
)
|
|
)
|
|
with patch.object(
|
|
bootstrap.subprocess,
|
|
"Popen",
|
|
return_value=fake_process,
|
|
), patch.object(bootstrap, "write_status") as mocked_write:
|
|
bootstrap.run_install(
|
|
"status.json",
|
|
2,
|
|
8,
|
|
"pytorch deps",
|
|
["uv", "pip", "install", "torch"],
|
|
)
|
|
|
|
progresses = [call.args[1] for call in mocked_write.call_args_list]
|
|
self.assertGreater(progresses[-1], 0.02)
|
|
self.assertTrue(
|
|
all(
|
|
progresses[i] <= progresses[i + 1] + 1e-9
|
|
for i in range(len(progresses) - 1)
|
|
)
|
|
)
|
|
messages = [call.args[2] for call in mocked_write.call_args_list]
|
|
self.assertIn("3 packages", messages)
|
|
self.assertIn("3/3 packages", messages)
|
|
|
|
def test_run_install_maps_progress_into_stage_range(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
|
|
class FakeProcess:
|
|
def __init__(self, lines):
|
|
self.stdout = lines
|
|
|
|
def wait(self):
|
|
return 0
|
|
|
|
fake_process = FakeProcess(
|
|
[
|
|
"Resolved 3 packages in 1.0s",
|
|
"Prepared 3 packages in 2.0s",
|
|
"Installed 3 packages in 0.5s",
|
|
]
|
|
)
|
|
with 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_formats_elapsed_time(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
|
|
self.assertEqual(bootstrap.format_elapsed(12), "12s")
|
|
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
|
|
|
|
def test_preflight_rejects_low_disk_space(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
args = SimpleNamespace(
|
|
backend="cu128",
|
|
runtime_dir=str(Path.cwd()),
|
|
)
|
|
with patch.object(
|
|
bootstrap.shutil,
|
|
"disk_usage",
|
|
return_value=SimpleNamespace(free=1024**3),
|
|
), patch.object(bootstrap, "reachable", return_value=True):
|
|
with self.assertRaisesRegex(RuntimeError, "disk space"):
|
|
bootstrap.preflight(args)
|
|
|
|
def test_preflight_rejects_unreachable_network(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
args = SimpleNamespace(
|
|
backend="cu128",
|
|
runtime_dir=str(Path.cwd()),
|
|
)
|
|
with patch.object(
|
|
bootstrap.shutil,
|
|
"disk_usage",
|
|
return_value=SimpleNamespace(free=100 * 1024**3),
|
|
), patch.object(bootstrap, "reachable", return_value=False):
|
|
with self.assertRaisesRegex(RuntimeError, "Cannot reach"):
|
|
bootstrap.preflight(args)
|
|
|
|
def test_preflight_passes_when_checks_ok(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
args = SimpleNamespace(
|
|
backend="cu128",
|
|
runtime_dir=str(Path.cwd()),
|
|
)
|
|
with patch.object(
|
|
bootstrap.shutil,
|
|
"disk_usage",
|
|
return_value=SimpleNamespace(free=100 * 1024**3),
|
|
), patch.object(bootstrap, "reachable", return_value=True):
|
|
resolved, used_mirror = bootstrap.preflight(args)
|
|
|
|
self.assertFalse(used_mirror)
|
|
self.assertEqual(
|
|
resolved["pypi"],
|
|
"https://pypi.org/simple/",
|
|
)
|
|
|
|
def test_preflight_falls_back_to_mirrors(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
args = SimpleNamespace(
|
|
backend="cu128",
|
|
runtime_dir=str(Path.cwd()),
|
|
)
|
|
|
|
def fake_reachable(url, timeout=bootstrap.PREFLIGHT_TIMEOUT):
|
|
if "github.com" in url:
|
|
return True
|
|
if "tuna" in url or "hf-mirror" in url or "aliyun" in url:
|
|
return True
|
|
return False
|
|
|
|
with patch.object(
|
|
bootstrap.shutil,
|
|
"disk_usage",
|
|
return_value=SimpleNamespace(free=100 * 1024**3),
|
|
), patch.object(
|
|
bootstrap,
|
|
"reachable",
|
|
side_effect=fake_reachable,
|
|
):
|
|
resolved, used_mirror = bootstrap.preflight(args)
|
|
|
|
self.assertTrue(used_mirror)
|
|
self.assertEqual(
|
|
resolved["pypi"],
|
|
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
|
)
|
|
self.assertEqual(resolved["hf"], "https://hf-mirror.com")
|
|
self.assertEqual(
|
|
resolved["torch"],
|
|
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
|
|
)
|
|
self.assertEqual(
|
|
bootstrap._resolved_torch_index,
|
|
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
|
|
)
|
|
|
|
def test_preflight_uses_environment_overrides(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
args = SimpleNamespace(
|
|
backend="cu128",
|
|
runtime_dir=str(Path.cwd()),
|
|
)
|
|
with patch.dict(
|
|
os.environ,
|
|
{
|
|
"PIP_INDEX_URL": "https://pypi.example/simple",
|
|
"HF_ENDPOINT": "https://hf.example",
|
|
"SNAPFORM_TORCH_INDEX": "https://torch.example/whl/cu128",
|
|
},
|
|
), patch.object(
|
|
bootstrap.shutil,
|
|
"disk_usage",
|
|
return_value=SimpleNamespace(free=100 * 1024**3),
|
|
), patch.object(bootstrap, "reachable", return_value=True):
|
|
resolved, used_mirror = bootstrap.preflight(args)
|
|
|
|
self.assertFalse(used_mirror)
|
|
self.assertEqual(resolved["pypi"], "https://pypi.example/simple")
|
|
self.assertEqual(resolved["hf"], "https://hf.example")
|
|
self.assertEqual(
|
|
resolved["torch"],
|
|
"https://torch.example/whl/cu128",
|
|
)
|
|
|
|
def test_child_environment_injects_resolved_mirrors(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
previous_pypi = bootstrap._resolved_pypi_index
|
|
previous_hf = bootstrap._resolved_hf_endpoint
|
|
bootstrap._resolved_pypi_index = (
|
|
"https://pypi.tuna.tsinghua.edu.cn/simple"
|
|
)
|
|
bootstrap._resolved_hf_endpoint = "https://hf-mirror.com"
|
|
try:
|
|
env = bootstrap.child_environment()
|
|
finally:
|
|
bootstrap._resolved_pypi_index = previous_pypi
|
|
bootstrap._resolved_hf_endpoint = previous_hf
|
|
|
|
self.assertEqual(
|
|
env["PIP_INDEX_URL"],
|
|
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
|
)
|
|
self.assertEqual(env["HF_ENDPOINT"], "https://hf-mirror.com")
|
|
|
|
def test_child_environment_injects_registry_proxy(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with patch.object(
|
|
bootstrap.urllib.request,
|
|
"getproxies",
|
|
return_value={
|
|
"http": "http://proxy:8080",
|
|
"https": "http://proxy:8080",
|
|
},
|
|
):
|
|
env = bootstrap.child_environment()
|
|
|
|
self.assertEqual(env["HTTP_PROXY"], "http://proxy:8080")
|
|
self.assertEqual(env["HTTPS_PROXY"], "http://proxy:8080")
|
|
|
|
def test_child_environment_keeps_explicit_proxy(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with patch.object(
|
|
bootstrap.urllib.request,
|
|
"getproxies",
|
|
return_value={"https": "http://proxy:8080"},
|
|
):
|
|
env = bootstrap.child_environment(
|
|
{"HTTPS_PROXY": "http://explicit:3128"}
|
|
)
|
|
|
|
self.assertEqual(env["HTTPS_PROXY"], "http://explicit:3128")
|
|
|
|
def test_ready_json_records_installed_packages(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
uv_path = root / "uv"
|
|
uv_path.write_bytes(b"uv")
|
|
args = SimpleNamespace(
|
|
uv=str(uv_path),
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run",
|
|
), patch.object(
|
|
bootstrap,
|
|
"run_stage",
|
|
return_value=["torch==2.9.1+cu128", "numpy==1.26.4"],
|
|
):
|
|
bootstrap.install(args)
|
|
|
|
ready = json.loads(
|
|
(runtime_dir / "snapform_ready.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
self.assertEqual(
|
|
ready["packages"],
|
|
["numpy==1.26.4", "torch==2.9.1+cu128"],
|
|
)
|
|
|
|
def test_broken_runtime_is_cleaned_before_reinstall(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
(runtime_dir / "partial-python.exe").write_text("junk")
|
|
args = SimpleNamespace(
|
|
uv="uv",
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run",
|
|
), patch.object(
|
|
bootstrap,
|
|
"run_install",
|
|
):
|
|
bootstrap.install(args)
|
|
|
|
self.assertFalse((runtime_dir / "partial-python.exe").exists())
|
|
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
|
|
|
|
def test_ready_runtime_is_not_cleaned(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir(parents=True)
|
|
(runtime_dir / "Scripts").mkdir()
|
|
(runtime_dir / "Scripts" / "python.exe").write_text("python")
|
|
(runtime_dir / "snapform_ready.json").write_text(
|
|
json.dumps({"runtime_schema": bootstrap.RUNTIME_SCHEMA}),
|
|
encoding="utf-8",
|
|
)
|
|
(runtime_dir / "marker.txt").write_text("keep")
|
|
args = SimpleNamespace(
|
|
uv="uv",
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run",
|
|
), patch.object(
|
|
bootstrap,
|
|
"run_install",
|
|
):
|
|
bootstrap.install(args)
|
|
|
|
self.assertTrue((runtime_dir / "marker.txt").is_file())
|
|
|
|
def test_model_download_failure_keeps_runtime_ready(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
|
|
def fake_run(command, environment=None):
|
|
if any("DA3-SMALL" in str(value) for value in command):
|
|
raise subprocess.CalledProcessError(1, command)
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
args = SimpleNamespace(
|
|
uv="uv",
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run",
|
|
side_effect=fake_run,
|
|
), patch.object(bootstrap, "run_install"):
|
|
bootstrap.install(args)
|
|
|
|
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
|
|
status = json.loads(
|
|
(root / "status.json").read_text(encoding="utf-8")
|
|
)
|
|
self.assertEqual(status["progress"], 1.0)
|
|
self.assertIn("Runtime is ready", status["message"])
|
|
|
|
def test_install_downloads_missing_uv(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
uv_path = root / "uv" / "uv.exe"
|
|
args = SimpleNamespace(
|
|
uv=str(uv_path),
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run",
|
|
), patch.object(
|
|
bootstrap,
|
|
"run_install",
|
|
), patch.object(
|
|
bootstrap.shutil,
|
|
"which",
|
|
return_value=None,
|
|
), patch.object(
|
|
bootstrap.shutil,
|
|
"rmtree",
|
|
), patch.object(bootstrap, "download_uv") as mocked_download:
|
|
bootstrap.install(args)
|
|
|
|
mocked_download.assert_called_once()
|
|
call_args = mocked_download.call_args
|
|
self.assertEqual(call_args.args[0], uv_path)
|
|
self.assertEqual(call_args.args[1], str(root / "status.json"))
|
|
self.assertEqual(call_args.kwargs["stage"], 1)
|
|
self.assertEqual(call_args.kwargs["stage_label"], "uv")
|
|
|
|
def test_install_reuses_existing_uv(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
runtime_dir = root / "runtime"
|
|
runtime_dir.mkdir()
|
|
uv_path = root / "uv.exe"
|
|
uv_path.write_bytes(b"uv")
|
|
args = SimpleNamespace(
|
|
uv=str(uv_path),
|
|
runtime_dir=str(runtime_dir),
|
|
models_dir=str(root / "models"),
|
|
backend="cu128",
|
|
da3_commit="test-commit",
|
|
ben2_commit="test-ben2-commit",
|
|
status=str(root / "status.json"),
|
|
)
|
|
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
|
bootstrap,
|
|
"run",
|
|
), patch.object(
|
|
bootstrap,
|
|
"run_install",
|
|
), patch.object(bootstrap, "download_uv") as mocked_download:
|
|
bootstrap.install(args)
|
|
|
|
mocked_download.assert_not_called()
|
|
|
|
def test_extract_uv_unpacks_standalone_binary(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
archive_path = root / "uv.zip"
|
|
with zipfile.ZipFile(archive_path, "w") as bundle:
|
|
bundle.writestr("uv-x86_64-pc-windows-msvc/uv.exe", "binary")
|
|
|
|
uv_path = root / "uv.exe"
|
|
bootstrap.extract_uv(archive_path, uv_path)
|
|
|
|
self.assertEqual(uv_path.read_bytes(), b"binary")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|