Spaces:
Running on Zero
Running on Zero
File size: 4,934 Bytes
5c49122 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """Test repository migration without downloading any model or reading tokens."""
import unittest
from unittest.mock import Mock, patch
from fastapi.testclient import TestClient
import app as demo
from runtime import (
DEFAULT_MODEL_REPO,
LEGACY_MODEL_REPO,
InferenceEngine,
configured_model_source,
)
class ModelSourceTests(unittest.TestCase):
"""Check new defaults, old Space variables, explicit pins and local models."""
def test_default_repository_and_revision(self) -> None:
"""Unset and blank configuration selects the owner's current repository."""
for environment in [{}, {"LLMON_MODEL_REPO": " ", "LLMON_MODEL_REVISION": ""}]:
with self.subTest(environment=environment), patch.dict("os.environ", environment, clear=True):
self.assertEqual(configured_model_source(), (DEFAULT_MODEL_REPO, "main"))
def test_old_repository_does_not_reuse_its_revision(self) -> None:
"""An old pinned commit must never be applied to the new repository."""
environment = {"LLMON_MODEL_REPO": LEGACY_MODEL_REPO, "LLMON_MODEL_REVISION": "old-repo-only-sha"}
with patch.dict("os.environ", environment, clear=True), self.assertLogs("uvicorn.error", level="WARNING"):
self.assertEqual(configured_model_source(), (DEFAULT_MODEL_REPO, "main"))
def test_explicit_new_repository_pin_is_respected(self) -> None:
"""Operators can pin the new checkpoint after verifying its revision."""
for repo in [DEFAULT_MODEL_REPO, "test-org/explicit-checkpoint"]:
environment = {"LLMON_MODEL_REPO": repo, "LLMON_MODEL_REVISION": "a" * 40}
with self.subTest(repo=repo), patch.dict("os.environ", environment, clear=True):
self.assertEqual(configured_model_source(), (repo, "a" * 40))
def test_prepare_downloads_migrated_source_and_reports_commit(self) -> None:
"""Verify the actual downloader call, not just the settings resolver."""
commit = "b" * 40
snapshot = f"/test-only/cache/snapshots/{commit}"
model = Mock()
environment = {
"LLMON_MODEL_REPO": LEGACY_MODEL_REPO,
"LLMON_MODEL_REVISION": "old-repo-only-sha",
"LLMON_DEVICE": "cpu",
}
with (
patch.dict("os.environ", environment, clear=True),
patch("huggingface_hub.snapshot_download", return_value=snapshot) as download,
patch("vendor.attention_extractor.load_model", return_value=(model, Mock(), {})) as load,
self.assertLogs("uvicorn.error", level="WARNING"),
):
engine = InferenceEngine()
engine.prepare()
self.assertEqual(download.call_args.kwargs["repo_id"], DEFAULT_MODEL_REPO)
self.assertEqual(download.call_args.kwargs["revision"], "main")
self.assertEqual(download.call_count, 1)
load.assert_called_once_with(snapshot)
model.to.assert_called_once_with("cpu")
self.assertEqual(engine.status, "ready")
with patch.object(demo, "engine", engine):
health = TestClient(demo.app).get("/api/health").json()
self.assertEqual(health["model_repo"], DEFAULT_MODEL_REPO)
self.assertEqual(health["model_revision"], commit)
self.assertEqual(health["model_requested_revision"], "main")
self.assertNotIn(snapshot, str(health))
def test_local_checkpoint_takes_precedence_without_download(self) -> None:
"""Keep explicit CPU/MPS development paths separate from Hub models."""
environment = {"LLMON_MODEL_PATH": "/test-only/local", "LLMON_DEVICE": "cpu"}
with (
patch.dict("os.environ", environment, clear=True),
patch("huggingface_hub.snapshot_download") as download,
patch("vendor.attention_extractor.load_model", return_value=(Mock(), Mock(), {})) as load,
):
engine = InferenceEngine()
engine.prepare()
download.assert_not_called()
load.assert_called_once_with("/test-only/local")
self.assertIsNone(engine.model_info["model_repo"])
self.assertIsNone(engine.model_info["model_revision"])
def test_failed_download_never_falls_back_to_old_model(self) -> None:
"""An unavailable release must surface an error, not a different model."""
with (
patch.dict("os.environ", {"LLMON_DEVICE": "cpu"}, clear=True),
patch("huggingface_hub.snapshot_download", side_effect=OSError("unavailable")) as download,
):
engine = InferenceEngine()
with self.assertRaises(OSError):
engine.prepare()
self.assertEqual(download.call_count, 1)
self.assertEqual(download.call_args.kwargs["repo_id"], DEFAULT_MODEL_REPO)
self.assertIsNone(engine.model)
self.assertNotEqual(engine.status, "ready")
if __name__ == "__main__":
unittest.main()
|