nan
838af0ee96
- Added ai_host.py to handle AI tasks with FastAPI integration. - Introduced HostState and ModelCache classes for managing client connections and model loading. - Implemented job submission and cancellation endpoints in the AI host. - Enhanced tests for AI host functionality, including client tracking and job cancellation. - Updated existing tests to accommodate changes in the operator and command structure. - Refactored runtime commands to use request-based architecture instead of command-line arguments. - Improved model loading logic to handle CUDA memory management and reuse models efficiently.
117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
import importlib.util
|
|
import tempfile
|
|
import unittest
|
|
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_and_non_compiled_gsplat(self):
|
|
bootstrap = load_bootstrap_runtime()
|
|
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, "run") as mocked_run:
|
|
bootstrap.install(args)
|
|
|
|
commands = [call.args[0] for call in mocked_run.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,
|
|
)
|
|
gsplat_call = next(
|
|
call
|
|
for call in mocked_run.call_args_list
|
|
if any("gsplat @" in str(value) for value in call.args[0])
|
|
)
|
|
gsplat_command = gsplat_call.args[0]
|
|
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"})
|
|
|
|
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)
|
|
)
|
|
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)
|
|
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)
|
|
)
|
|
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
|
|
)
|
|
self.assertEqual(
|
|
ben2_download[-1],
|
|
'["config.json", "model.safetensors"]',
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|