Update addon tests, bootstrap runtime, download progress, and worker functionality
- Modified status messages and stage labels in `test_addon.py` for clarity and consistency. - Added tests to ensure proper host termination and command execution in `test_addon.py`. - Enhanced `test_bootstrap_runtime.py` with additional tests for installation behavior, including handling of CUDA and MPS devices, and improved error handling for disk space and network checks. - Updated download progress reporting in `test_download_progress.py` to include download speed and improved message formatting. - Added new tests in `test_worker.py` for device selection, NPZ file validation, and model downloading behavior with Hugging Face endpoints. - Bumped version in `uv.lock` to 0.9.30.
This commit is contained in:
@@ -1,7 +1,11 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
@@ -21,14 +25,16 @@ def load_bootstrap_runtime():
|
||||
|
||||
|
||||
class BootstrapRuntimeTest(unittest.TestCase):
|
||||
def test_installs_da3_and_non_compiled_gsplat(self):
|
||||
def test_installs_da3_from_archives_without_gsplat(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv"
|
||||
uv_path.write_bytes(b"uv")
|
||||
args = SimpleNamespace(
|
||||
uv="uv",
|
||||
uv=str(uv_path),
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
@@ -36,7 +42,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "run") as mocked_run, patch.object(
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
) as mocked_run, patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
) as mocked_run_install:
|
||||
@@ -55,25 +64,44 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
f"addict=={bootstrap.ADDICT_VERSION}",
|
||||
da3_command,
|
||||
)
|
||||
gsplat_call = next(
|
||||
call
|
||||
for call in mocked_run_install.call_args_list
|
||||
if any("gsplat @" in str(value) for value in call.args[4])
|
||||
)
|
||||
gsplat_command = gsplat_call.args[4]
|
||||
self.assertIn("--no-build-isolation", gsplat_command)
|
||||
self.assertTrue(
|
||||
any(bootstrap.GSPLAT_COMMIT in str(value) for value in gsplat_command)
|
||||
any(
|
||||
"https://github.com/ByteDance-Seed/Depth-Anything-3/"
|
||||
"archive/test-commit.tar.gz" in str(value)
|
||||
for value in da3_command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("git+" in str(value) for value in da3_command)
|
||||
)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"gsplat" in str(value)
|
||||
for command in commands
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"--no-build-isolation" in str(value)
|
||||
for command in commands
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
self.assertEqual(gsplat_call.args[5], {"BUILD_NO_CUDA": "1"})
|
||||
|
||||
ben2_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if any("ben2 @" in str(value) for value in command)
|
||||
)
|
||||
self.assertTrue(
|
||||
any("test-ben2-commit" in str(value) for value in ben2_command)
|
||||
any(
|
||||
"https://github.com/PramaLLC/BEN2/archive/"
|
||||
"test-ben2-commit.tar.gz" in str(value)
|
||||
for value in ben2_command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("git+" in str(value) for value in ben2_command)
|
||||
)
|
||||
host_command = next(
|
||||
command
|
||||
@@ -81,6 +109,8 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
|
||||
)
|
||||
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command)
|
||||
self.assertIn("setuptools", host_command)
|
||||
self.assertNotIn("ninja", host_command)
|
||||
check_command = next(
|
||||
command
|
||||
for command in commands
|
||||
@@ -95,6 +125,9 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
self.assertTrue(
|
||||
any("import fastapi, uvicorn" in str(value) for value in check_command)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("import gsplat" in str(value) for value in check_command)
|
||||
)
|
||||
default_downloads = [
|
||||
command
|
||||
for command in commands
|
||||
@@ -121,7 +154,86 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
self.assertEqual(patterns, ["config.json", "model.safetensors"])
|
||||
self.assertIn("--stage", ben2_download)
|
||||
self.assertIn("--stage-label", ben2_download)
|
||||
self.assertEqual(ben2_download[-1], "Download BEN2 Base")
|
||||
self.assertEqual(ben2_download[-1], "ben2 base")
|
||||
|
||||
def test_darwin_install_skips_cuda_index_and_xformers(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv"
|
||||
uv_path.write_bytes(b"uv")
|
||||
args = SimpleNamespace(
|
||||
uv=str(uv_path),
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.sys,
|
||||
"platform",
|
||||
"darwin",
|
||||
), patch.object(
|
||||
bootstrap.platform,
|
||||
"machine",
|
||||
return_value="arm64",
|
||||
), patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
) as mocked_run_install:
|
||||
bootstrap.install(args)
|
||||
|
||||
commands = [
|
||||
call.args[4] for call in mocked_run_install.call_args_list
|
||||
]
|
||||
torch_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if any(
|
||||
f"torch=={bootstrap.TORCH_VERSION}" in str(value)
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
self.assertNotIn("--index-url", torch_command)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"xformers" in str(value)
|
||||
for command in commands
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
|
||||
def test_intel_mac_is_rejected_with_clear_error(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
uv="uv",
|
||||
runtime_dir="runtime",
|
||||
models_dir="models",
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status="status.json",
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.sys,
|
||||
"platform",
|
||||
"darwin",
|
||||
), patch.object(
|
||||
bootstrap.platform,
|
||||
"machine",
|
||||
return_value="x86_64",
|
||||
), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree:
|
||||
with self.assertRaisesRegex(RuntimeError, "Apple Silicon"):
|
||||
bootstrap.install(args)
|
||||
|
||||
mocked_rmtree.assert_not_called()
|
||||
|
||||
def test_run_stage_writes_stage_progress(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
@@ -131,8 +243,8 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
bootstrap.run_stage(
|
||||
str(status_path),
|
||||
3,
|
||||
12,
|
||||
"Install PyTorch",
|
||||
8,
|
||||
"pytorch",
|
||||
["uv", "pip", "install", "torch"],
|
||||
None,
|
||||
"plain",
|
||||
@@ -145,8 +257,131 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
self.assertEqual(status["stage"], 3)
|
||||
self.assertEqual(status["stages"], 12)
|
||||
self.assertEqual(status["stage_label"], "Install PyTorch")
|
||||
self.assertEqual(status["stages"], 8)
|
||||
self.assertEqual(status["stage_label"], "pytorch")
|
||||
|
||||
def test_run_stage_writes_progress_range(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
with patch.object(bootstrap, "run"), patch.object(
|
||||
bootstrap,
|
||||
"write_status",
|
||||
) as mocked_write:
|
||||
bootstrap.run_stage(
|
||||
str(status_path),
|
||||
3,
|
||||
8,
|
||||
"pytorch",
|
||||
["uv", "pip", "install", "torch"],
|
||||
None,
|
||||
"plain",
|
||||
start=0.2,
|
||||
end=0.4,
|
||||
)
|
||||
|
||||
progresses = [
|
||||
call.args[1] for call in mocked_write.call_args_list
|
||||
]
|
||||
self.assertEqual(progresses[0], 0.2)
|
||||
self.assertEqual(progresses[-1], 0.4)
|
||||
|
||||
def test_run_stage_passes_progress_range_to_download_command(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
command = [
|
||||
"python",
|
||||
"worker.py",
|
||||
"download",
|
||||
"--repository",
|
||||
"depth-anything/DA3-SMALL",
|
||||
]
|
||||
with patch.object(bootstrap, "run") as mocked_run:
|
||||
bootstrap.run_stage(
|
||||
str(status_path),
|
||||
6,
|
||||
8,
|
||||
"da3 small",
|
||||
command,
|
||||
None,
|
||||
"download",
|
||||
start=0.9,
|
||||
end=0.9333,
|
||||
)
|
||||
|
||||
called_command = mocked_run.call_args.args[0]
|
||||
self.assertIn("--progress-start", called_command)
|
||||
self.assertIn("--progress-end", called_command)
|
||||
self.assertEqual(
|
||||
called_command[called_command.index("--progress-start") + 1],
|
||||
"0.9",
|
||||
)
|
||||
self.assertEqual(
|
||||
called_command[called_command.index("--progress-end") + 1],
|
||||
"0.9333",
|
||||
)
|
||||
self.assertEqual(called_command[-1], "da3 small")
|
||||
|
||||
def test_install_allocates_contiguous_progress_slices(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv"
|
||||
uv_path.write_bytes(b"uv")
|
||||
args = SimpleNamespace(
|
||||
uv=str(uv_path),
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
stage_ranges = []
|
||||
|
||||
def fake_run_stage(
|
||||
status_path,
|
||||
stage,
|
||||
stages,
|
||||
label,
|
||||
command,
|
||||
environment,
|
||||
kind,
|
||||
start=0.0,
|
||||
end=1.0,
|
||||
):
|
||||
stage_ranges.append((stage, start, end))
|
||||
bootstrap.write_status(
|
||||
status_path,
|
||||
end,
|
||||
"",
|
||||
stage=stage,
|
||||
stages=stages,
|
||||
stage_label=label,
|
||||
)
|
||||
return []
|
||||
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run_stage",
|
||||
side_effect=fake_run_stage,
|
||||
), patch.object(bootstrap, "run"):
|
||||
bootstrap.install(args)
|
||||
|
||||
starts = [item[1] for item in stage_ranges]
|
||||
ends = [item[2] for item in stage_ranges]
|
||||
self.assertTrue(
|
||||
all(starts[i] < starts[i + 1] for i in range(len(starts) - 1))
|
||||
)
|
||||
self.assertEqual(ends[-1], 1.0)
|
||||
self.assertEqual(starts[0], bootstrap.RUNTIME_SLICE_START)
|
||||
status = json.loads(
|
||||
(root / "status.json").read_text(encoding="utf-8")
|
||||
)
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
|
||||
def test_run_install_tracks_package_download_events(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
@@ -161,33 +396,515 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
fake_process = FakeProcess(
|
||||
[
|
||||
"Resolved 3 packages in 1.0s",
|
||||
"Downloading torch (10MiB)",
|
||||
" Downloaded torch",
|
||||
"DEBUG Adding transitive dependency for torch==2.9.1: filelock*",
|
||||
"Prepared 3 packages in 2.0s",
|
||||
"Installed 3 packages in 0.5s",
|
||||
"+ torch==2.9.1",
|
||||
"+ numpy==1.26.4",
|
||||
]
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.subprocess,
|
||||
"Popen",
|
||||
return_value=fake_process,
|
||||
), patch.object(bootstrap, "write_status") as mocked_write:
|
||||
installed = bootstrap.run_install(
|
||||
"status.json",
|
||||
2,
|
||||
8,
|
||||
"pytorch deps",
|
||||
["uv", "pip", "install", "torch"],
|
||||
)
|
||||
|
||||
messages = [call.args[2] for call in mocked_write.call_args_list]
|
||||
self.assertEqual(
|
||||
messages,
|
||||
[
|
||||
"3 packages",
|
||||
"3/3 packages",
|
||||
"3/3 packages",
|
||||
],
|
||||
)
|
||||
self.assertEqual(mocked_write.call_args_list[-1].args[1], 0.02)
|
||||
first = mocked_write.call_args_list[0]
|
||||
self.assertEqual(first.kwargs["stage"], 2)
|
||||
self.assertEqual(first.kwargs["stages"], 8)
|
||||
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
|
||||
self.assertEqual(installed, ["torch==2.9.1", "numpy==1.26.4"])
|
||||
|
||||
def test_run_install_smooths_progress_between_events(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
|
||||
class SlowLines:
|
||||
def __init__(self, lines, delay):
|
||||
self.lines = list(lines)
|
||||
self.delay = delay
|
||||
|
||||
def __iter__(self):
|
||||
for line in self.lines:
|
||||
time.sleep(self.delay)
|
||||
yield line
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, lines):
|
||||
self.stdout = lines
|
||||
|
||||
def wait(self):
|
||||
return 0
|
||||
|
||||
fake_process = FakeProcess(
|
||||
SlowLines(
|
||||
[
|
||||
"Resolved 3 packages in 1.0s",
|
||||
"Prepared 3 packages in 2.0s",
|
||||
"Installed 3 packages in 0.5s",
|
||||
],
|
||||
0.3,
|
||||
)
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.subprocess,
|
||||
"Popen",
|
||||
return_value=fake_process,
|
||||
), patch.object(bootstrap, "write_status") as mocked_write:
|
||||
bootstrap.run_install(
|
||||
"status.json",
|
||||
2,
|
||||
8,
|
||||
"pytorch deps",
|
||||
["uv", "pip", "install", "torch"],
|
||||
)
|
||||
|
||||
progresses = [call.args[1] for call in mocked_write.call_args_list]
|
||||
self.assertGreater(progresses[-1], 0.02)
|
||||
self.assertTrue(
|
||||
all(
|
||||
progresses[i] <= progresses[i + 1] + 1e-9
|
||||
for i in range(len(progresses) - 1)
|
||||
)
|
||||
)
|
||||
messages = [call.args[2] for call in mocked_write.call_args_list]
|
||||
self.assertIn("3 packages", messages)
|
||||
self.assertIn("3/3 packages", messages)
|
||||
|
||||
def test_run_install_maps_progress_into_stage_range(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, lines):
|
||||
self.stdout = lines
|
||||
|
||||
def wait(self):
|
||||
return 0
|
||||
|
||||
fake_process = FakeProcess(
|
||||
[
|
||||
"Resolved 3 packages in 1.0s",
|
||||
"Prepared 3 packages in 2.0s",
|
||||
"Installed 3 packages in 0.5s",
|
||||
]
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
with patch.object(
|
||||
bootstrap.subprocess,
|
||||
"Popen",
|
||||
return_value=fake_process,
|
||||
):
|
||||
bootstrap.run_install(
|
||||
str(status_path),
|
||||
2,
|
||||
12,
|
||||
"Install PyTorch",
|
||||
["uv", "pip", "install", "torch"],
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.subprocess,
|
||||
"Popen",
|
||||
return_value=fake_process,
|
||||
), patch.object(bootstrap, "write_status") as mocked_write:
|
||||
bootstrap.run_install(
|
||||
"status.json",
|
||||
2,
|
||||
8,
|
||||
"pytorch deps",
|
||||
["uv", "pip", "install", "torch"],
|
||||
start=0.2,
|
||||
end=0.4,
|
||||
)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(status["progress"], 0.95)
|
||||
self.assertEqual(status["message"], "Installing packages")
|
||||
self.assertEqual(status["stage"], 2)
|
||||
self.assertEqual(status["stages"], 12)
|
||||
self.assertEqual(status["stage_label"], "Install PyTorch")
|
||||
first = mocked_write.call_args_list[0]
|
||||
self.assertAlmostEqual(first.args[1], 0.204)
|
||||
self.assertEqual(first.kwargs["stage"], 2)
|
||||
self.assertEqual(first.kwargs["stages"], 8)
|
||||
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
|
||||
|
||||
def test_formats_elapsed_time(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
|
||||
self.assertEqual(bootstrap.format_elapsed(12), "12s")
|
||||
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
|
||||
|
||||
def test_preflight_rejects_low_disk_space(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=1024**3),
|
||||
), patch.object(bootstrap, "reachable", return_value=True):
|
||||
with self.assertRaisesRegex(RuntimeError, "disk space"):
|
||||
bootstrap.preflight(args)
|
||||
|
||||
def test_preflight_rejects_unreachable_network(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
||||
), patch.object(bootstrap, "reachable", return_value=False):
|
||||
with self.assertRaisesRegex(RuntimeError, "Cannot reach"):
|
||||
bootstrap.preflight(args)
|
||||
|
||||
def test_preflight_passes_when_checks_ok(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
||||
), patch.object(bootstrap, "reachable", return_value=True):
|
||||
resolved, used_mirror = bootstrap.preflight(args)
|
||||
|
||||
self.assertFalse(used_mirror)
|
||||
self.assertEqual(
|
||||
resolved["pypi"],
|
||||
"https://pypi.org/simple/",
|
||||
)
|
||||
|
||||
def test_preflight_falls_back_to_mirrors(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
|
||||
def fake_reachable(url, timeout=bootstrap.PREFLIGHT_TIMEOUT):
|
||||
if "github.com" in url:
|
||||
return True
|
||||
if "tuna" in url or "hf-mirror" in url or "aliyun" in url:
|
||||
return True
|
||||
return False
|
||||
|
||||
with patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"reachable",
|
||||
side_effect=fake_reachable,
|
||||
):
|
||||
resolved, used_mirror = bootstrap.preflight(args)
|
||||
|
||||
self.assertTrue(used_mirror)
|
||||
self.assertEqual(
|
||||
resolved["pypi"],
|
||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
||||
)
|
||||
self.assertEqual(resolved["hf"], "https://hf-mirror.com")
|
||||
self.assertEqual(
|
||||
resolved["torch"],
|
||||
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
|
||||
)
|
||||
self.assertEqual(
|
||||
bootstrap._resolved_torch_index,
|
||||
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
|
||||
)
|
||||
|
||||
def test_preflight_uses_environment_overrides(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
args = SimpleNamespace(
|
||||
backend="cu128",
|
||||
runtime_dir=str(Path.cwd()),
|
||||
)
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"PIP_INDEX_URL": "https://pypi.example/simple",
|
||||
"HF_ENDPOINT": "https://hf.example",
|
||||
"SNAPFORM_TORCH_INDEX": "https://torch.example/whl/cu128",
|
||||
},
|
||||
), patch.object(
|
||||
bootstrap.shutil,
|
||||
"disk_usage",
|
||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
||||
), patch.object(bootstrap, "reachable", return_value=True):
|
||||
resolved, used_mirror = bootstrap.preflight(args)
|
||||
|
||||
self.assertFalse(used_mirror)
|
||||
self.assertEqual(resolved["pypi"], "https://pypi.example/simple")
|
||||
self.assertEqual(resolved["hf"], "https://hf.example")
|
||||
self.assertEqual(
|
||||
resolved["torch"],
|
||||
"https://torch.example/whl/cu128",
|
||||
)
|
||||
|
||||
def test_child_environment_injects_resolved_mirrors(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
previous_pypi = bootstrap._resolved_pypi_index
|
||||
previous_hf = bootstrap._resolved_hf_endpoint
|
||||
bootstrap._resolved_pypi_index = (
|
||||
"https://pypi.tuna.tsinghua.edu.cn/simple"
|
||||
)
|
||||
bootstrap._resolved_hf_endpoint = "https://hf-mirror.com"
|
||||
try:
|
||||
env = bootstrap.child_environment()
|
||||
finally:
|
||||
bootstrap._resolved_pypi_index = previous_pypi
|
||||
bootstrap._resolved_hf_endpoint = previous_hf
|
||||
|
||||
self.assertEqual(
|
||||
env["PIP_INDEX_URL"],
|
||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
||||
)
|
||||
self.assertEqual(env["HF_ENDPOINT"], "https://hf-mirror.com")
|
||||
|
||||
def test_child_environment_injects_registry_proxy(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with patch.object(
|
||||
bootstrap.urllib.request,
|
||||
"getproxies",
|
||||
return_value={
|
||||
"http": "http://proxy:8080",
|
||||
"https": "http://proxy:8080",
|
||||
},
|
||||
):
|
||||
env = bootstrap.child_environment()
|
||||
|
||||
self.assertEqual(env["HTTP_PROXY"], "http://proxy:8080")
|
||||
self.assertEqual(env["HTTPS_PROXY"], "http://proxy:8080")
|
||||
|
||||
def test_child_environment_keeps_explicit_proxy(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with patch.object(
|
||||
bootstrap.urllib.request,
|
||||
"getproxies",
|
||||
return_value={"https": "http://proxy:8080"},
|
||||
):
|
||||
env = bootstrap.child_environment(
|
||||
{"HTTPS_PROXY": "http://explicit:3128"}
|
||||
)
|
||||
|
||||
self.assertEqual(env["HTTPS_PROXY"], "http://explicit:3128")
|
||||
|
||||
def test_ready_json_records_installed_packages(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv"
|
||||
uv_path.write_bytes(b"uv")
|
||||
args = SimpleNamespace(
|
||||
uv=str(uv_path),
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_stage",
|
||||
return_value=["torch==2.9.1+cu128", "numpy==1.26.4"],
|
||||
):
|
||||
bootstrap.install(args)
|
||||
|
||||
ready = json.loads(
|
||||
(runtime_dir / "snapform_ready.json").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
)
|
||||
self.assertEqual(
|
||||
ready["packages"],
|
||||
["numpy==1.26.4", "torch==2.9.1+cu128"],
|
||||
)
|
||||
|
||||
def test_broken_runtime_is_cleaned_before_reinstall(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
(runtime_dir / "partial-python.exe").write_text("junk")
|
||||
args = SimpleNamespace(
|
||||
uv="uv",
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
):
|
||||
bootstrap.install(args)
|
||||
|
||||
self.assertFalse((runtime_dir / "partial-python.exe").exists())
|
||||
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
|
||||
|
||||
def test_ready_runtime_is_not_cleaned(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir(parents=True)
|
||||
(runtime_dir / "Scripts").mkdir()
|
||||
(runtime_dir / "Scripts" / "python.exe").write_text("python")
|
||||
(runtime_dir / "snapform_ready.json").write_text(
|
||||
json.dumps({"runtime_schema": bootstrap.RUNTIME_SCHEMA}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(runtime_dir / "marker.txt").write_text("keep")
|
||||
args = SimpleNamespace(
|
||||
uv="uv",
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
):
|
||||
bootstrap.install(args)
|
||||
|
||||
self.assertTrue((runtime_dir / "marker.txt").is_file())
|
||||
|
||||
def test_model_download_failure_keeps_runtime_ready(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
|
||||
def fake_run(command, environment=None):
|
||||
if any("DA3-SMALL" in str(value) for value in command):
|
||||
raise subprocess.CalledProcessError(1, command)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
args = SimpleNamespace(
|
||||
uv="uv",
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
side_effect=fake_run,
|
||||
), patch.object(bootstrap, "run_install"):
|
||||
bootstrap.install(args)
|
||||
|
||||
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
|
||||
status = json.loads(
|
||||
(root / "status.json").read_text(encoding="utf-8")
|
||||
)
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
self.assertIn("Runtime is ready", status["message"])
|
||||
|
||||
def test_install_downloads_missing_uv(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv" / "uv.exe"
|
||||
args = SimpleNamespace(
|
||||
uv=str(uv_path),
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
), patch.object(
|
||||
bootstrap.shutil,
|
||||
"which",
|
||||
return_value=None,
|
||||
), patch.object(
|
||||
bootstrap.shutil,
|
||||
"rmtree",
|
||||
), patch.object(bootstrap, "download_uv") as mocked_download:
|
||||
bootstrap.install(args)
|
||||
|
||||
mocked_download.assert_called_once()
|
||||
call_args = mocked_download.call_args
|
||||
self.assertEqual(call_args.args[0], uv_path)
|
||||
self.assertEqual(call_args.args[1], str(root / "status.json"))
|
||||
self.assertEqual(call_args.kwargs["stage"], 1)
|
||||
self.assertEqual(call_args.kwargs["stage_label"], "uv")
|
||||
|
||||
def test_install_reuses_existing_uv(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
runtime_dir = root / "runtime"
|
||||
runtime_dir.mkdir()
|
||||
uv_path = root / "uv.exe"
|
||||
uv_path.write_bytes(b"uv")
|
||||
args = SimpleNamespace(
|
||||
uv=str(uv_path),
|
||||
runtime_dir=str(runtime_dir),
|
||||
models_dir=str(root / "models"),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
||||
bootstrap,
|
||||
"run",
|
||||
), patch.object(
|
||||
bootstrap,
|
||||
"run_install",
|
||||
), patch.object(bootstrap, "download_uv") as mocked_download:
|
||||
bootstrap.install(args)
|
||||
|
||||
mocked_download.assert_not_called()
|
||||
|
||||
def test_extract_uv_unpacks_standalone_binary(self):
|
||||
bootstrap = load_bootstrap_runtime()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
archive_path = root / "uv.zip"
|
||||
with zipfile.ZipFile(archive_path, "w") as bundle:
|
||||
bundle.writestr("uv-x86_64-pc-windows-msvc/uv.exe", "binary")
|
||||
|
||||
uv_path = root / "uv.exe"
|
||||
bootstrap.extract_uv(archive_path, uv_path)
|
||||
|
||||
self.assertEqual(uv_path.read_bytes(), b"binary")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user