139 lines
4.8 KiB
Python
139 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
|
|
PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api"
|
|
if str(PACKAGE_DIR) in sys.path:
|
|
sys.path.remove(str(PACKAGE_DIR))
|
|
sys.path.insert(0, str(PACKAGE_DIR))
|
|
|
|
from outcome_tracker import OutcomeTracker # noqa: E402
|
|
from poll_runner import ( # noqa: E402
|
|
_bootstrap_architecture_history,
|
|
_load_task_compatibility_contexts,
|
|
resolve_age_cleanup_policy,
|
|
)
|
|
|
|
|
|
class PollPolicyTests(unittest.TestCase):
|
|
def test_architecture_bootstrap_falls_back_when_public_logs_are_hidden(self) -> None:
|
|
class HistoryClient:
|
|
@staticmethod
|
|
def list_tasks_page(**_kwargs): # noqa: ANN003
|
|
return {
|
|
"data": {
|
|
"records": [
|
|
{
|
|
"taskId": "public-failure",
|
|
"modelId": "public/model",
|
|
"gpuType": "gpu",
|
|
"status": "success",
|
|
"verifyResult": -1,
|
|
"logCosUrl": None,
|
|
}
|
|
]
|
|
}
|
|
}
|
|
|
|
@staticmethod
|
|
def list_tasks(**_kwargs): # noqa: ANN003
|
|
return [
|
|
{
|
|
"taskId": "owned-success",
|
|
"modelId": "owner/model",
|
|
"gpuType": "gpu",
|
|
"modelTaskLevelId": 23,
|
|
"status": "success",
|
|
"verifyResult": 1,
|
|
"updateTime": datetime.now(timezone.utc).isoformat(),
|
|
}
|
|
]
|
|
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
tracker = OutcomeTracker(root / "outcomes.jsonl")
|
|
summary = _bootstrap_architecture_history(
|
|
modelhub_client=HistoryClient(), # type: ignore[arg-type]
|
|
outcome_tracker=tracker,
|
|
ledger_path=root / "ledger.jsonl",
|
|
now=datetime.now(timezone.utc),
|
|
)
|
|
|
|
self.assertEqual("owned_full_history", summary["source"])
|
|
self.assertEqual(1, summary["terminalRecords"])
|
|
self.assertEqual(0, summary["communityUsableFailureDetails"])
|
|
|
|
def test_cleanup_contexts_merge_outcomes_with_older_ledger_entries(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
tracker = OutcomeTracker(root / "outcomes.jsonl")
|
|
tracker.record_submission(
|
|
"owner/new",
|
|
"gpu-a",
|
|
"vllm",
|
|
"text-generation",
|
|
"task-new",
|
|
datetime.now(timezone.utc).isoformat(),
|
|
model_profile={"architectures": ["Qwen2ForCausalLM"]},
|
|
)
|
|
ledger_path = root / "ledger.jsonl"
|
|
ledger_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"taskId": "task-old",
|
|
"modelId": "owner/old",
|
|
"targetGpu": "gpu-b",
|
|
"framework": "mindie",
|
|
"taskType": "text-generation",
|
|
}
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
contexts = _load_task_compatibility_contexts(
|
|
tracker,
|
|
ledger_path=ledger_path,
|
|
)
|
|
|
|
self.assertEqual(["Qwen2ForCausalLM"], contexts["task-new"]["modelProfile"]["architectures"])
|
|
self.assertEqual("mindie", contexts["task-old"]["framework"])
|
|
self.assertEqual({}, contexts["task-old"]["modelProfile"])
|
|
|
|
def test_age_cleanup_uses_minus_ten_once_then_minus_five(self) -> None:
|
|
args = argparse.Namespace(
|
|
recent_model_reserve_slots=10,
|
|
dynamic_old_model_cleanup_reserve_slots=5,
|
|
)
|
|
|
|
self.assertEqual(
|
|
("initial", 10),
|
|
resolve_age_cleanup_policy(args, initial_cleanup_pending=True),
|
|
)
|
|
self.assertEqual(
|
|
("dynamic", 5),
|
|
resolve_age_cleanup_policy(args, initial_cleanup_pending=False),
|
|
)
|
|
|
|
def test_dynamic_cleanup_cannot_be_stricter_than_initial_cleanup(self) -> None:
|
|
args = argparse.Namespace(
|
|
recent_model_reserve_slots=10,
|
|
dynamic_old_model_cleanup_reserve_slots=20,
|
|
)
|
|
|
|
self.assertEqual(
|
|
("dynamic", 10),
|
|
resolve_age_cleanup_policy(args, initial_cleanup_pending=False),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|