248 lines
10 KiB
Python
248 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import sys
|
|
|
|
from dulwich import porcelain
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
MODULE_ROOT = ROOT / "modelhub_submmit_api"
|
|
if str(MODULE_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(MODULE_ROOT))
|
|
|
|
from common import read_jsonl, write_json, write_jsonl # noqa: E402
|
|
from config_optimizer import SafeConfigOptimizer # noqa: E402
|
|
from hf_discovery import HuggingFaceDiscovery, parse_model_card_front_matter # noqa: E402
|
|
from official_capabilities import OfficialCapabilityRegistry # noqa: E402
|
|
from routing_engine import SuccessFirstRoutingEngine # noqa: E402
|
|
from state_sync import StateGitSync # noqa: E402
|
|
|
|
|
|
class OfficialClient:
|
|
def __init__(self, *, fail_catalog: bool = False) -> None:
|
|
self.fail_catalog = fail_catalog
|
|
|
|
def list_machine_info(self): # noqa: ANN201
|
|
if self.fail_catalog:
|
|
raise RuntimeError("offline")
|
|
return [
|
|
{"gpuType": "gpu-fast", "canVerify": True, "maxConcurrentTasks": 2},
|
|
{"gpuType": "gpu-disabled", "canVerify": False, "maxConcurrentTasks": 8},
|
|
]
|
|
|
|
def list_task_levels(self): # noqa: ANN201
|
|
return {"data": [{"taskType": "text-generation"}, {"taskType": "new-task"}]}
|
|
|
|
def list_model_task_types(self, target_gpu, model_address): # noqa: ANN001, ANN201
|
|
del target_gpu, model_address
|
|
return {"data": [{"taskType": "text-generation"}]}
|
|
|
|
|
|
class SuperAgentTests(unittest.TestCase):
|
|
def test_official_registry_discovers_catalog_and_exact_model_routes(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
registry = OfficialCapabilityRegistry(Path(temporary_dir) / "official.json", log_fn=lambda _: None)
|
|
client = OfficialClient()
|
|
registry.prepare(
|
|
client,
|
|
fallback_gpus=["legacy"],
|
|
task_types=["text-generation"],
|
|
now=datetime(2026, 8, 15, tzinfo=timezone.utc),
|
|
)
|
|
self.assertTrue(registry.ready)
|
|
self.assertEqual(["gpu-fast"], registry.eligible_gpus())
|
|
self.assertEqual(
|
|
["text-generation"],
|
|
registry.task_types_for(
|
|
client,
|
|
model_address="https://modelscope.cn/models/owner/model",
|
|
model_last_modified="2026-08-15T00:00:00+00:00",
|
|
gpu="gpu-fast",
|
|
),
|
|
)
|
|
|
|
def test_official_registry_fails_closed_without_catalog_cache(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
registry = OfficialCapabilityRegistry(Path(temporary_dir) / "official.json", log_fn=lambda _: None)
|
|
registry.prepare(
|
|
OfficialClient(fail_catalog=True),
|
|
fallback_gpus=["legacy"],
|
|
task_types=["text-generation"],
|
|
now=datetime(2026, 8, 15, tzinfo=timezone.utc),
|
|
)
|
|
self.assertFalse(registry.ready)
|
|
self.assertEqual("critical_official_signal_unavailable", registry.pause_reason)
|
|
|
|
def test_success_band_beats_shorter_queue(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
engine = SuccessFirstRoutingEngine(Path(temporary_dir) / "routing.json", log_fn=lambda _: None)
|
|
candidates = [
|
|
{
|
|
"repoId": "owner/model-a",
|
|
"targetGpu": "reliable",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"frameworkMarketSamples": 1000,
|
|
"frameworkMarketSuccessRate": 0.9,
|
|
"queueBacklogHours": 12,
|
|
},
|
|
{
|
|
"repoId": "owner/model-b",
|
|
"targetGpu": "fast",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"frameworkMarketSamples": 1000,
|
|
"frameworkMarketSuccessRate": 0.4,
|
|
"queueBacklogHours": 0.25,
|
|
},
|
|
]
|
|
ordered = engine.order_candidates(candidates)
|
|
self.assertEqual("reliable", ordered[0]["targetGpu"])
|
|
self.assertGreater(ordered[0]["routingSuccessBand"], ordered[1]["routingSuccessBand"])
|
|
|
|
def test_modelscope_metadata_and_model_card_lineage_are_structured(self) -> None:
|
|
item = {
|
|
"id": "owner/model",
|
|
"downloads": 123,
|
|
"params": 7_000_000_000,
|
|
"file_size": 14_000_000_000,
|
|
"tags": ["qwen", "chat"],
|
|
"tasks": ["text-generation"],
|
|
"license": "apache-2.0",
|
|
"likes": 9,
|
|
}
|
|
model = HuggingFaceDiscovery._parse_model(
|
|
item,
|
|
fallback_pipeline_tag="text-generation",
|
|
min_downloads=0,
|
|
)
|
|
self.assertIsNotNone(model)
|
|
assert model is not None
|
|
self.assertEqual(7_000_000_000, model.params)
|
|
self.assertEqual(("qwen", "chat"), model.tags)
|
|
metadata = parse_model_card_front_matter(
|
|
"---\nbase_model: Qwen/base\nframeworks:\n - transformers\ntasks:\n - text-generation\n---\nbody"
|
|
)
|
|
self.assertEqual("Qwen/base", metadata["base_model"])
|
|
self.assertEqual(["transformers"], metadata["frameworks"])
|
|
|
|
def test_state_branch_round_trip_persists_intent_without_config_or_secret(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
remote = root / "remote.git"
|
|
project = root / "project"
|
|
restored_project = root / "restored"
|
|
project.mkdir()
|
|
restored_project.mkdir()
|
|
porcelain.init(remote, bare=True)
|
|
write_json(project / ".modelhub_state" / "account_capacity.json", {"version": 1})
|
|
credentials = {"username": "tester", "email": "tester@example.com", "password": "secret-value"}
|
|
manager = StateGitSync(
|
|
project_root=project,
|
|
credentials=credentials,
|
|
remote=str(remote),
|
|
log_fn=lambda _: None,
|
|
)
|
|
manager.acquire_process_lock()
|
|
self.assertTrue(manager.restore())
|
|
batch_id = manager.begin_batch(
|
|
[
|
|
{
|
|
"repoId": "owner/model",
|
|
"modelAddress": "https://modelscope.cn/models/owner/model",
|
|
"targetGpu": "gpu-a",
|
|
"taskType": "text-generation",
|
|
"framework": "vllm",
|
|
"configParams": "password: must-not-be-copied",
|
|
}
|
|
]
|
|
)
|
|
self.assertIsNotNone(batch_id)
|
|
manager.close()
|
|
|
|
restored = StateGitSync(
|
|
project_root=restored_project,
|
|
credentials=credentials,
|
|
remote=str(remote),
|
|
log_fn=lambda _: None,
|
|
)
|
|
restored.acquire_process_lock()
|
|
self.assertTrue(restored.restore())
|
|
intents = read_jsonl(restored_project / ".modelhub_state" / "recovery_intents.jsonl")
|
|
self.assertEqual("owner/model", intents[0]["repoId"])
|
|
state_text = "\n".join(
|
|
path.read_text(encoding="utf-8")
|
|
for path in restored._workspace.rglob("*")
|
|
if path.is_file() and ".git" not in path.parts
|
|
)
|
|
self.assertNotIn("secret-value", state_text)
|
|
self.assertNotIn("must-not-be-copied", state_text)
|
|
restored.close()
|
|
|
|
def test_failed_intent_push_returns_no_batch_id(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
manager = StateGitSync(
|
|
project_root=Path(temporary_dir),
|
|
credentials={"username": "u", "email": "e@example.com", "password": "p"},
|
|
remote="unused",
|
|
log_fn=lambda _: None,
|
|
)
|
|
manager.healthy = True
|
|
with patch.object(manager, "sync", return_value=False):
|
|
self.assertIsNone(
|
|
manager.begin_batch(
|
|
[
|
|
{
|
|
"repoId": "owner/model",
|
|
"targetGpu": "gpu",
|
|
"taskType": "text-generation",
|
|
"framework": "vllm",
|
|
"configParams": "safe",
|
|
}
|
|
]
|
|
)
|
|
)
|
|
|
|
def test_config_patch_requires_repeated_cross_model_success(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
intents = []
|
|
outcomes = []
|
|
for index in range(5):
|
|
intents.append(
|
|
{
|
|
"taskId": str(index),
|
|
"taskType": "text-generation",
|
|
"targetGpu": "gpu-a",
|
|
"framework": "vllm",
|
|
"repoId": f"owner/model-{index % 2}",
|
|
"configFingerprint": "proven",
|
|
"safeConfigVector": {"gpuNum": 1, "tensorParallel": 1},
|
|
}
|
|
)
|
|
outcomes.append({"taskId": str(index), "outcome": "success"})
|
|
write_jsonl(root / "intents.jsonl", intents)
|
|
write_jsonl(root / "outcomes.jsonl", outcomes)
|
|
optimizer = SafeConfigOptimizer(
|
|
intents_path=root / "intents.jsonl",
|
|
outcomes_path=root / "outcomes.jsonl",
|
|
)
|
|
config, metadata = optimizer.optimize(
|
|
task_type="text-generation",
|
|
target_gpu="gpu-a",
|
|
framework="vllm",
|
|
official_config="framework: vllm\nsut_config:\n gpu_num: 2\nref_config:\n gpu_num: 2\n",
|
|
official_lower_bound=0.40,
|
|
)
|
|
self.assertTrue(metadata["applied"])
|
|
self.assertNotIn("gpu_num: 2", config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|