优化 Runtime 安装进度体验
- 状态栏与窗口强制刷新,进度更新无需移动鼠标即可实时显示 - 安装过程拆分为 12 个阶段,每段进度从 0 到 100% - 安装阶段解析 uv 下载事件推进进度,模型下载阶段实时反映字节进度
This commit is contained in:
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user