From e0cb9954736b75d3a32e11758d505e58770cfb02 Mon Sep 17 00:00:00 2001 From: nan Date: Wed, 5 Aug 2026 22:45:58 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=20Runtime=20=E5=AE=89?= =?UTF-8?q?=E8=A3=85=E8=BF=9B=E5=BA=A6=E4=BD=93=E9=AA=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 状态栏与窗口强制刷新，进度更新无需移动鼠标即可实时显示 - 安装过程拆分为 12 个阶段，每段进度从 0 到 100% - 安装阶段解析 uv 下载事件推进进度，模型下载阶段实时反映字节进度 --- README.md | 5 +- pyproject.toml | 2 +- src/snapform/blender_manifest.toml | 2 +- src/snapform/bootstrap_runtime.py | 421 ++++++++++++++++++----------- src/snapform/operators.py | 16 +- src/snapform/status_bar.py | 14 + src/snapform/worker.py | 83 +++++- tests/test_addon.py | 94 ++++++- tests/test_bootstrap_runtime.py | 94 ++++++- tests/test_pack.py | 2 +- uv.lock | 2 +- 11 files changed, 552 insertions(+), 183 deletions(-) diff --git a/README.md b/README.md index d05c296..5a0743d 100644 --- a/README.md +++ b/README.md @@ -43,7 +43,8 @@ 2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。 4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、DA3、BEN2、FastAPI、Uvicorn， - 并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。 + 并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 12 个阶段依次推进，每段进度 + 从 0 到 100%，模型下载阶段会实时反映已下载字节。 5. 在 3D View 按 `N`，打开 `SnapForm` 页签。 6. 如需其他模型，展开 `Host > Models`，找到模型并点击该行的 `Download`。 7. 展开默认收起的 `Debug`，选择 `Export Format` 和本次执行使用的模型。 @@ -99,7 +100,7 @@ 所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的紧凑进度条内部，取消按钮 位于进度条右侧。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存 -写出和面片构建。状态更新仅标记相关区域重绘，不强制执行窗口 Draw/Swap。 +写出和面片构建。进度更新会强制刷新状态栏与窗口，无需移动鼠标即可实时显示。 导入过程精确保留 NPZ 内深度值之间的差值。例如深度 `2.0` 和 `5.5` 在 `Displacement Strength = -1` 时，沿局部 Z 轴相差 `3.5 Blender Unit`。它不会把数值压缩到 diff --git a/pyproject.toml b/pyproject.toml index 30d93a6..aa5fb44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "snapform" -version = "0.9.28" +version = "0.9.29" description = "Fast image-guided structure building tools for Blender" readme = "README.md" requires-python = ">=3.11" diff --git a/src/snapform/blender_manifest.toml b/src/snapform/blender_manifest.toml index 52b6448..af9ca74 100644 --- a/src/snapform/blender_manifest.toml +++ b/src/snapform/blender_manifest.toml @@ -1,7 +1,7 @@ schema_version = "1.0.0" id = "snapform" -version = "0.9.28" +version = "0.9.29" name = "SnapForm" tagline = "Build spatial structures from images, fast" maintainer = "OmooLab" diff --git a/src/snapform/bootstrap_runtime.py b/src/snapform/bootstrap_runtime.py index cd3e6fe..1241c50 100644 --- a/src/snapform/bootstrap_runtime.py +++ b/src/snapform/bootstrap_runtime.py @@ -28,10 +28,22 @@ TORCH_INDEXES = { DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key) -def write_status(path, progress, message, error=None): +def write_status( + path, + progress, + message, + error=None, + stage=None, + stages=None, + stage_label=None, +): payload = {"progress": progress, "message": message} if error: payload["error"] = error + if stage is not None: + payload["stage"] = stage + payload["stages"] = stages + payload["stage_label"] = stage_label Path(path).write_text( json.dumps(payload, ensure_ascii=False), encoding="utf-8", @@ -50,185 +62,282 @@ def run(command, environment=None): ) +def run_install(status_path, stage, stages, label, command, environment=None): + """Run a uv install command and translate its verbose package events into progress.""" + print("Running:", subprocess.list2cmdline([str(value) for value in command])) + process_environment = os.environ.copy() + process_environment["UV_LOG"] = "debug" + if environment: + process_environment.update(environment) + process = subprocess.Popen( + [str(value) for value in command], + env=process_environment, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + ) + progress = 0.0 + downloaded = 0 + for raw_line in process.stdout: + line = raw_line.rstrip() + print(line) + lowered = line.lower() + if "resolved" in lowered and "package" in lowered: + new_progress, message = 0.05, "Resolving dependencies" + elif "downloaded" in lowered and "package" not in lowered: + downloaded += 1 + new_progress = min(0.85, 0.10 + downloaded * 0.08) + message = f"Downloaded {downloaded} package(s)" + elif "prepared" in lowered: + new_progress, message = 0.70, "Preparing packages" + elif "installed" in lowered or "built" in lowered: + new_progress, message = 0.95, "Installing packages" + else: + continue + if new_progress <= progress: + continue + progress = new_progress + write_status( + status_path, + progress, + message, + stage=stage, + stages=stages, + stage_label=label, + ) + return_code = process.wait() + if return_code != 0: + raise subprocess.CalledProcessError(return_code, command) + + def environment_python(runtime_dir): if sys.platform == "win32": return Path(runtime_dir) / "Scripts" / "python.exe" return Path(runtime_dir) / "bin" / "python" -def download_default_models(args, python): - models = {**MODELS, BEN2_MODEL.key: BEN2_MODEL} - download_script = ( - "import json, sys; " - "from huggingface_hub import snapshot_download; " - "patterns = json.loads(sys.argv[3]); " - "snapshot_download(repo_id=sys.argv[1], local_dir=sys.argv[2], " - "allow_patterns=patterns or None)" - ) - models_dir = Path(args.models_dir) - models_dir.mkdir(parents=True, exist_ok=True) - for index, model_key in enumerate(DEFAULT_MODEL_KEYS): - model = models[model_key] - progress = 0.90 + index * 0.03 - write_status(args.status, progress, f"Downloading {model.label}") - run( - [ - python, - "-c", - download_script, - model.repository, - models_dir / model.key.lower(), - json.dumps(model.download_patterns), - ] - ) - - -def install(args): - runtime_dir = Path(args.runtime_dir) - runtime_dir.parent.mkdir(parents=True, exist_ok=True) - - write_status(args.status, 0.03, "Creating isolated Python Runtime") - run( - [ - args.uv, - "venv", - "--clear", - "--python", - "3.12", - runtime_dir, - ] - ) - python = environment_python(runtime_dir) - - write_status(args.status, 0.12, "Installing PyTorch") - run( - [ - args.uv, - "pip", - "install", - "--python", - python, - f"torch=={TORCH_VERSION}", - f"torchvision=={TORCHVISION_VERSION}", - "--index-url", - TORCH_INDEXES[args.backend], - ] - ) - - write_status(args.status, 0.38, "Installing matching xFormers") - xformers_command = [ - args.uv, - "pip", - "install", - "--python", - python, - f"xformers=={XFORMERS_VERSION}", - ] - if args.backend == "cu128": - xformers_command.extend( - ["--index-url", TORCH_INDEXES[args.backend]] - ) - else: - xformers_command.append("--no-deps") - run(xformers_command) - - write_status(args.status, 0.50, "Installing Depth Anything 3") +def install_stages(args, runtime_dir, python): + uv = args.uv + torch_index = TORCH_INDEXES[args.backend] da3_requirement = ( "depth-anything-3 @ " "git+https://github.com/ByteDance-Seed/Depth-Anything-3.git" f"@{args.da3_commit}" ) - run( - [ - args.uv, - "pip", - "install", - "--python", - python, - da3_requirement, - f"addict=={ADDICT_VERSION}", - ] - ) - - write_status(args.status, 0.68, "Installing BEN2") ben2_requirement = ( "ben2 @ git+https://github.com/PramaLLC/BEN2.git" f"@{args.ben2_commit}" ) - run( - [ - args.uv, - "pip", - "install", - "--python", - python, - ben2_requirement, - ] + xformers_command = [ + uv, + "pip", + "install", + "-v", + "--python", + python, + f"xformers=={XFORMERS_VERSION}", + ] + if args.backend == "cu128": + xformers_command.extend(["--index-url", torch_index]) + else: + xformers_command.append("--no-deps") + gsplat_requirement = ( + "gsplat @ git+https://github.com/nerfstudio-project/" + f"gsplat.git@{GSPLAT_COMMIT}" ) - - write_status(args.status, 0.75, "Installing AI Host") - run( - [ - args.uv, - "pip", - "install", - "--python", - python, - f"fastapi=={FASTAPI_VERSION}", - f"uvicorn=={UVICORN_VERSION}", - ] + check_code = ( + "import torch; " + "from depth_anything_3.api import DepthAnything3; " + "from ben2 import BEN_Base; " + "import fastapi, uvicorn; " + "import gsplat; " + "print(torch.__version__, torch.cuda.is_available())" ) - - write_status(args.status, 0.78, "Installing gsplat build tools") - run( - [ - args.uv, - "pip", + stages = [ + ( + "Create Python Runtime", + [uv, "venv", "--clear", "--python", "3.12", runtime_dir], + None, + "plain", + ), + ( + "Install PyTorch", + [ + uv, + "pip", + "install", + "-v", + "--python", + python, + f"torch=={TORCH_VERSION}", + f"torchvision=={TORCHVISION_VERSION}", + "--index-url", + torch_index, + ], + None, "install", - "--python", + ), + ("Install xFormers", xformers_command, None, "install"), + ( + "Install Depth Anything 3", + [ + uv, + "pip", + "install", + "-v", + "--python", + python, + da3_requirement, + f"addict=={ADDICT_VERSION}", + ], + None, + "install", + ), + ( + "Install BEN2", + [ + uv, + "pip", + "install", + "-v", + "--python", + python, + ben2_requirement, + ], + None, + "install", + ), + ( + "Install AI Host", + [ + uv, + "pip", + "install", + "-v", + "--python", + python, + f"fastapi=={FASTAPI_VERSION}", + f"uvicorn=={UVICORN_VERSION}", + ], + None, + "install", + ), + ( + "Install gsplat build tools", + [ + uv, + "pip", + "install", + "-v", + "--python", + python, + "setuptools", + "wheel", + "ninja", + ], + None, + "install", + ), + ( + "Install gsplat", + [ + uv, + "pip", + "install", + "-v", + "--python", + python, + "--no-build-isolation", + gsplat_requirement, + ], + # PLY 导出不调用 rasterizer，禁用本地扩展可避免 CUDA 与 MSVC 工具链耦合 + {"BUILD_NO_CUDA": "1"}, + "install", + ), + ("Check Runtime", [python, "-c", check_code], None, "plain"), + ] + models = {**MODELS, BEN2_MODEL.key: BEN2_MODEL} + worker_script = Path(__file__).with_name("worker.py") + models_dir = Path(args.models_dir) + for model_key in DEFAULT_MODEL_KEYS: + model = models[model_key] + download_command = [ python, - "setuptools", - "wheel", - "ninja", + worker_script, + "download", + "--repository", + model.repository, + "--model-dir", + models_dir / model.key.lower(), + "--status", + args.status, ] - ) - - write_status(args.status, 0.84, "Installing gsplat") - # PLY 导出不调用 rasterizer，禁用本地扩展可避免 CUDA 与 MSVC 工具链耦合。 - gsplat_environment = {"BUILD_NO_CUDA": "1"} - run( - [ - args.uv, - "pip", - "install", - "--python", - python, - "--no-build-isolation", + for pattern in model.download_patterns: + download_command.extend(["--allow-pattern", pattern]) + stages.append( ( - "gsplat @ git+https://github.com/nerfstudio-project/" - f"gsplat.git@{GSPLAT_COMMIT}" - ), - ], - gsplat_environment, + f"Download {model.label}", + download_command, + None, + "download", + ) + ) + return stages + + +def run_stage(status_path, stage, stages, label, command, environment, kind): + write_status( + status_path, + 0.0, + "", + stage=stage, + stages=stages, + stage_label=label, + ) + if kind == "install": + run_install(status_path, stage, stages, label, command, environment) + elif kind == "download": + run( + [ + *command, + "--stage", + str(stage), + "--stages", + str(stages), + "--stage-label", + label, + ], + environment, + ) + else: + run(command, environment) + write_status( + status_path, + 1.0, + "", + stage=stage, + stages=stages, + stage_label=label, ) - write_status(args.status, 0.88, "Checking Runtime") - run( - [ - python, - "-c", - ( - "import torch; " - "from depth_anything_3.api import DepthAnything3; " - "from ben2 import BEN_Base; " - "import fastapi, uvicorn; " - "import gsplat; " - "print(torch.__version__, torch.cuda.is_available())" - ), - ] - ) - download_default_models(args, python) +def install(args): + runtime_dir = Path(args.runtime_dir) + runtime_dir.parent.mkdir(parents=True, exist_ok=True) + python = environment_python(runtime_dir) + stages = install_stages(args, runtime_dir, python) + for index, (label, command, environment, kind) in enumerate(stages, start=1): + run_stage( + args.status, + index, + len(stages), + label, + command, + environment, + kind, + ) ready_path = runtime_dir / "snapform_ready.json" ready_path.write_text( diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 91ae12c..3882ed6 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -412,8 +412,20 @@ class RunJob(bpy.types.Operator): return settings = context.scene.snapform_settings settings.progress = float(status.get("progress", settings.progress)) - settings.status = status.get("message", settings.status) - self._redraw_ui(context) + settings.status = self._format_status(status, settings.status) + self._redraw_ui(context, force=True) + + def _format_status(self, status, fallback): + message = status.get("message", fallback) + stage = status.get("stage") + stages = status.get("stages") + stage_label = status.get("stage_label") + if not stage or not stages: + return message + stage_text = f"Step {stage}/{stages}" + if stage_label: + stage_text = f"{stage_text} {stage_label}" + return f"{stage_text}: {message}" if message else stage_text def _finish(self, context, job, return_code): settings = context.scene.snapform_settings diff --git a/src/snapform/status_bar.py b/src/snapform/status_bar.py index 6082921..7d6712b 100644 --- a/src/snapform/status_bar.py +++ b/src/snapform/status_bar.py @@ -1,3 +1,5 @@ +import bpy + from .operators import active_job @@ -45,3 +47,15 @@ def tag_ui_regions(context): def redraw_ui(context, force=False): tag_ui_regions(context) + if not force: + return + try: + bpy.ops.wm.redraw_timer(type="DRAW_WIN_SWAP", iterations=1) + except (AttributeError, RuntimeError, TypeError): + pass + workspace = getattr(context, "workspace", None) + if workspace is not None: + try: + workspace.status_text_set_internal(None) + except (AttributeError, RuntimeError, TypeError): + pass diff --git a/src/snapform/worker.py b/src/snapform/worker.py index f18748f..0c492cd 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -32,10 +32,23 @@ def check_cancelled(cancel_check): raise TaskCancelled("Task cancelled") -def write_status(path, progress, message, error=None, state=None): +def write_status( + path, + progress, + message, + error=None, + state=None, + stage=None, + stages=None, + stage_label=None, +): payload = {"progress": progress, "message": message} if error: payload["error"] = error + if stage is not None: + payload["stage"] = stage + payload["stages"] = stages + payload["stage_label"] = stage_label if state is None: if error: state = "failed" @@ -136,12 +149,27 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None): class DownloadProgressReporter: - def __init__(self, status_path, repository, clock=time.monotonic): + def __init__( + self, + status_path, + repository, + clock=time.monotonic, + start=0.05, + span=0.90, + stage=None, + stages=None, + stage_label=None, + ): self.status_path = status_path self.repository = repository self.clock = clock self.lock = threading.Lock() - self.progress = 0.05 + self.start = start + self.span = span + self.stage = stage + self.stages = stages + self.stage_label = stage_label + self.progress = start self.last_write_time = 0.0 def report(self, current, total): @@ -149,7 +177,7 @@ class DownloadProgressReporter: return factor = min(max(current / total, 0.0), 1.0) - progress = 0.05 + factor * 0.90 + progress = self.start + factor * self.span now = self.clock() with self.lock: if progress <= self.progress: @@ -162,6 +190,9 @@ class DownloadProgressReporter: self.status_path, progress, f"Downloading {self.repository}", + stage=self.stage, + stages=self.stages, + stage_label=self.stage_label, ) @@ -170,11 +201,25 @@ def create_download_progress_class( repository, tqdm_type=None, clock=time.monotonic, + start=0.05, + span=0.90, + stage=None, + stages=None, + stage_label=None, ): if tqdm_type is None: from tqdm.auto import tqdm as tqdm_type - reporter = DownloadProgressReporter(status_path, repository, clock) + reporter = DownloadProgressReporter( + status_path, + repository, + clock, + start, + span, + stage, + stages, + stage_label, + ) class DownloadProgress(tqdm_type): def __init__(self, *args, **kwargs): @@ -205,10 +250,26 @@ def download_model(args): model_dir = Path(args.model_dir) model_dir.mkdir(parents=True, exist_ok=True) - write_status(args.status, 0.05, f"Downloading {args.repository}") + staged = getattr(args, "stage", None) is not None + start = 0.0 if staged else 0.05 + span = 1.0 if staged else 0.90 + stage_kwargs = { + "stage": getattr(args, "stage", None), + "stages": getattr(args, "stages", None), + "stage_label": getattr(args, "stage_label", None), + } + write_status( + args.status, + start, + f"Downloading {args.repository}", + **stage_kwargs, + ) progress_class = create_download_progress_class( args.status, args.repository, + start=start, + span=span, + **stage_kwargs, ) snapshot_download( repo_id=args.repository, @@ -216,7 +277,12 @@ def download_model(args): tqdm_class=progress_class, allow_patterns=args.allow_pattern or None, ) - write_status(args.status, 1.0, "Model is ready") + write_status( + args.status, + 1.0, + "Model is ready", + **stage_kwargs, + ) def select_device(requested, torch): @@ -480,6 +546,9 @@ def parser(): download.add_argument("--model-dir", required=True) download.add_argument("--status", required=True) download.add_argument("--allow-pattern", action="append", default=[]) + download.add_argument("--stage", type=int, default=None) + download.add_argument("--stages", type=int, default=None) + download.add_argument("--stage-label", default=None) background = commands.add_parser("remove-background") background.add_argument("--input", required=True) diff --git a/tests/test_addon.py b/tests/test_addon.py index 69f805d..70b7c42 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -510,6 +510,46 @@ class AddonTest(unittest.TestCase): self.assertNotIn("DA3 image depth", starting_source) self.assertNotIn("DA3 depth plane", finish_source) + def test_update_status_formats_stage_progress(self): + operator = self.addon.RunJob() + + self.assertEqual( + operator._format_status( + { + "progress": 0.5, + "message": "Downloaded 3 packages", + "stage": 3, + "stages": 12, + "stage_label": "Install PyTorch", + }, + "fallback", + ), + "Step 3/12 Install PyTorch: Downloaded 3 packages", + ) + self.assertEqual( + operator._format_status( + { + "progress": 0.0, + "message": "", + "stage": 10, + "stages": 12, + "stage_label": "Download DA3 Small", + }, + "fallback", + ), + "Step 10/12 Download DA3 Small", + ) + self.assertEqual( + operator._format_status( + { + "progress": 0.5, + "message": "Estimating depth", + }, + "fallback", + ), + "Estimating depth", + ) + def test_depth_background_is_loaded_without_changing_source_empty(self): source = inspect.getsource(self.addon.RunJob._build_depth_plane) @@ -1486,7 +1526,56 @@ class AddonTest(unittest.TestCase): ], ) - def test_force_flag_does_not_trigger_window_draw_and_swap(self): + def test_force_flag_triggers_window_draw_and_swap(self): + redrawn = [] + context = SimpleNamespace( + window_manager=SimpleNamespace( + windows=( + SimpleNamespace( + screen=SimpleNamespace( + areas=( + SimpleNamespace( + type="STATUSBAR", + tag_redraw=lambda: redrawn.append( + "status" + ), + regions=(), + ), + ), + ), + ), + ), + ), + workspace=SimpleNamespace( + status_text_set_internal=lambda _value: redrawn.append( + "status-text" + ), + ), + ) + + with patch.object( + self.addon.status_bar.bpy, + "ops", + SimpleNamespace( + wm=SimpleNamespace( + redraw_timer=lambda **options: redrawn.append( + ("swap", options) + ), + ), + ), + ): + self.addon.status_bar.redraw_ui(context, force=True) + + self.assertEqual( + redrawn, + [ + "status", + ("swap", {"type": "DRAW_WIN_SWAP", "iterations": 1}), + "status-text", + ], + ) + + def test_force_flag_ignores_missing_redraw_helpers(self): redrawn = [] context = SimpleNamespace( window_manager=SimpleNamespace( @@ -1511,6 +1600,3 @@ class AddonTest(unittest.TestCase): self.addon.status_bar.redraw_ui(context, force=True) self.assertEqual(redrawn, ["status"]) - source = inspect.getsource(self.addon.status_bar) - self.assertNotIn("redraw_timer", source) - self.assertNotIn("DRAW_WIN_SWAP", source) diff --git a/tests/test_bootstrap_runtime.py b/tests/test_bootstrap_runtime.py index 1949d78..215aa72 100644 --- a/tests/test_bootstrap_runtime.py +++ b/tests/test_bootstrap_runtime.py @@ -1,4 +1,5 @@ import importlib.util +import json import tempfile import unittest from pathlib import Path @@ -35,10 +36,16 @@ class BootstrapRuntimeTest(unittest.TestCase): ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) - with patch.object(bootstrap, "run") as mocked_run: + with 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 @@ -50,15 +57,15 @@ class BootstrapRuntimeTest(unittest.TestCase): ) gsplat_call = next( call - for call in mocked_run.call_args_list - if any("gsplat @" in str(value) for value in call.args[0]) + 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[0] + 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) ) - self.assertEqual(gsplat_call.args[1], {"BUILD_NO_CUDA": "1"}) + self.assertEqual(gsplat_call.args[5], {"BUILD_NO_CUDA": "1"}) ben2_command = next( command @@ -106,10 +113,81 @@ class BootstrapRuntimeTest(unittest.TestCase): for command in default_downloads if "PramaLLC/BEN2" in command ) - self.assertEqual( - ben2_download[-1], - '["config.json", "model.safetensors"]', + 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], "Download BEN2 Base") + + 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, + 12, + "Install 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"], 12) + self.assertEqual(status["stage_label"], "Install PyTorch") + + 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", + "Downloading torch (10MiB)", + " Downloaded torch", + "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"], + ) + + 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") if __name__ == "__main__": diff --git a/tests/test_pack.py b/tests/test_pack.py index 69dca31..9386ee4 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -29,7 +29,7 @@ class PackTest(unittest.TestCase): finally: pack.DIST_DIR = original_dist - self.assertEqual(output_path.name, "SnapForm.v0.9.28.zip") + self.assertEqual(output_path.name, "SnapForm.v0.9.29.zip") with ZipFile(output_path) as archive: names = set(archive.namelist()) diff --git a/uv.lock b/uv.lock index 1f1d316..2e8c96e 100644 --- a/uv.lock +++ b/uv.lock @@ -4,5 +4,5 @@ requires-python = ">=3.11" [[package]] name = "snapform" -version = "0.9.28" +version = "0.9.29" source = { virtual = "." }