244 lines
9.4 KiB
Python
244 lines
9.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
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
|
|
_advance_architecture_history_backfill,
|
|
_complete_architecture_history_backfill,
|
|
_load_task_compatibility_contexts,
|
|
build_parser,
|
|
)
|
|
|
|
|
|
class PollPolicyTests(unittest.TestCase):
|
|
def test_architecture_history_backfill_resumes_and_keeps_only_decision_state(self) -> None:
|
|
class PagedHistoryClient:
|
|
calls: list[int] = []
|
|
|
|
def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN201
|
|
current = int(kwargs["current"])
|
|
self.calls.append(current)
|
|
records = {
|
|
1: [
|
|
{
|
|
"taskId": "history-1",
|
|
"modelId": "owner/model-1",
|
|
"gpuType": "gpu-a",
|
|
"status": "success",
|
|
"verifyResult": 1,
|
|
},
|
|
{
|
|
"taskId": "history-2",
|
|
"modelId": "owner/model-2",
|
|
"gpuType": "gpu-a",
|
|
"status": "failed",
|
|
"verifyResult": -1,
|
|
},
|
|
],
|
|
2: [
|
|
{
|
|
"taskId": "history-3",
|
|
"modelId": "owner/model-3",
|
|
"gpuType": "gpu-b",
|
|
"status": "success",
|
|
"verifyResult": 1,
|
|
}
|
|
],
|
|
}[current]
|
|
return {"data": {"records": records, "pages": 2}}
|
|
|
|
with tempfile.TemporaryDirectory() as temporary_dir, patch.dict(
|
|
"os.environ",
|
|
{"MODELHUB_ARCHITECTURE_BACKFILL_PAGES_PER_BATCH": "1"},
|
|
):
|
|
root = Path(temporary_dir)
|
|
outcomes = root / "outcomes.jsonl"
|
|
checkpoint = root / "checkpoint.json"
|
|
recent = root / "recent.jsonl"
|
|
progress_path = root / "backfill.json"
|
|
tracker = OutcomeTracker(
|
|
outcomes,
|
|
checkpoint_path=checkpoint,
|
|
recent_path=recent,
|
|
)
|
|
client = PagedHistoryClient()
|
|
first = _advance_architecture_history_backfill(
|
|
modelhub_client=client, # type: ignore[arg-type]
|
|
outcome_tracker=tracker,
|
|
ledger_path=root / "ledger.jsonl",
|
|
progress_path=progress_path,
|
|
now=datetime.now(timezone.utc),
|
|
)
|
|
self.assertFalse(first["complete"])
|
|
self.assertTrue(tracker.has_durable_checkpoint)
|
|
|
|
restored = OutcomeTracker(
|
|
outcomes,
|
|
checkpoint_path=checkpoint,
|
|
recent_path=recent,
|
|
)
|
|
second = _advance_architecture_history_backfill(
|
|
modelhub_client=client, # type: ignore[arg-type]
|
|
outcome_tracker=restored,
|
|
ledger_path=root / "ledger.jsonl",
|
|
progress_path=progress_path,
|
|
now=datetime.now(timezone.utc),
|
|
)
|
|
|
|
self.assertTrue(second["complete"])
|
|
self.assertNotIn("seenTaskIds", second)
|
|
self.assertEqual([1, 2], client.calls)
|
|
self.assertEqual(3, restored.get_stats_report()["terminalRecords"])
|
|
self.assertEqual([], json.loads(outcomes.read_text(encoding="utf-8") or "[]"))
|
|
persisted_progress = progress_path.read_text(encoding="utf-8")
|
|
self.assertNotIn("owner/model", persisted_progress)
|
|
self.assertNotIn("logs.invalid", persisted_progress)
|
|
|
|
def test_cold_start_backfill_finishes_all_pages_in_one_phase(self) -> None:
|
|
class ThreePageClient:
|
|
calls: list[int] = []
|
|
|
|
def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN201
|
|
current = int(kwargs["current"])
|
|
self.calls.append(current)
|
|
return {
|
|
"data": {
|
|
"records": [
|
|
{
|
|
"taskId": f"task-{current}",
|
|
"modelId": f"owner/model-{current}",
|
|
"gpuType": "gpu-a",
|
|
"status": "success",
|
|
"verifyResult": 1,
|
|
}
|
|
],
|
|
"pages": 3,
|
|
}
|
|
}
|
|
|
|
with tempfile.TemporaryDirectory() as temporary_dir, patch.dict(
|
|
"os.environ",
|
|
{"MODELHUB_ARCHITECTURE_BACKFILL_PAGES_PER_BATCH": "1"},
|
|
):
|
|
root = Path(temporary_dir)
|
|
tracker = OutcomeTracker(
|
|
root / "outcomes.jsonl",
|
|
checkpoint_path=root / "checkpoint.json",
|
|
recent_path=root / "recent.jsonl",
|
|
)
|
|
client = ThreePageClient()
|
|
durable_phases: list[str] = []
|
|
progress = _complete_architecture_history_backfill(
|
|
modelhub_client=client, # type: ignore[arg-type]
|
|
outcome_tracker=tracker,
|
|
ledger_path=root / "ledger.jsonl",
|
|
progress_path=root / "backfill.json",
|
|
sync_callback=lambda phase: durable_phases.append(phase) or True,
|
|
)
|
|
|
|
self.assertTrue(progress["complete"])
|
|
self.assertEqual([1, 2, 3], client.calls)
|
|
self.assertEqual(["history_backfill"], durable_phases)
|
|
self.assertEqual(3, tracker.get_stats_report()["terminalRecords"])
|
|
|
|
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_history_batch_is_classified_before_checkpoint_compaction(self) -> None:
|
|
tasks = [
|
|
{
|
|
"taskId": f"failed-{index}",
|
|
"modelId": f"owner/model-{index}",
|
|
"gpuType": "gpu-a",
|
|
"status": "failed",
|
|
"verifyResult": -1,
|
|
"logCosUrl": f"https://logs.invalid/{index}",
|
|
}
|
|
for index in range(600)
|
|
]
|
|
|
|
def classify(records, **kwargs): # noqa: ANN001, ANN003, ANN202
|
|
for record in records:
|
|
record["failureCategory"] = "model_runtime"
|
|
record["failureScope"] = "model"
|
|
record["failureDeterministic"] = True
|
|
progress = kwargs.get("progress")
|
|
if progress is not None:
|
|
progress(len(records), len(records), len(records), 0)
|
|
return len(records)
|
|
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
tracker = OutcomeTracker(
|
|
root / "outcomes.jsonl",
|
|
checkpoint_path=root / "checkpoint.json",
|
|
recent_path=root / "recent.jsonl",
|
|
)
|
|
with patch.object(tracker, "_enrich_failure_records", side_effect=classify):
|
|
summary = tracker.bootstrap_from_history_tasks(
|
|
tasks,
|
|
enrichment_limit=600,
|
|
)
|
|
|
|
report = tracker.get_stats_report()
|
|
self.assertEqual(600, summary["enrichmentAttempts"])
|
|
self.assertEqual(600, report["totals"]["attributableFailureCount"])
|
|
self.assertEqual(600, report["totals"]["failureBreakdown"]["model_runtime"])
|
|
|
|
def test_age_policy_defaults_to_admission_only_reserve_five(self) -> None:
|
|
args = build_parser().parse_args([])
|
|
|
|
self.assertEqual(5, args.recent_model_reserve_slots)
|
|
self.assertFalse(hasattr(args, "dynamic_old_model_cleanup_reserve_slots"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|