submmit/tests/test_takeover_accuracy.py

import copy
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "modelhub_submmit_api"))
from models import ModelInspection
from candidate_preflight import CandidatePreflightAdvisor
from profile_admission import filter_profile_batch, profile_key
from outcome_tracker import OutcomeTracker
from common import write_json
from state_sync import StateGitSync
from dulwich import porcelain


class TakeoverAccuracyTests(unittest.TestCase):
    def tracker(self, root):
        return OutcomeTracker(root / "outcomes.jsonl", checkpoint_path=root / "checkpoint.json",
                              recent_path=root / "recent.jsonl")

    def test_mlx_and_unknown_quantization_not_plain_weights(self):
        for config, expected in [({"quantization": {"bits": 4, "group_size": 64}}, "mlx"),
                                 ({"quantization_config": {"bits": 4}}, "unknown")]:
            inspection = ModelInspection(repo_id="owner/model", model_config=config)
            self.assertEqual(expected, inspection.quantization_method)
            advisor = CandidatePreflightAdvisor()
            result = advisor.assess(inspection=inspection, task_type="text-generation",
                                    target_gpu="MetaX_c-500", framework="vllm", config_params="x")
            self.assertFalse(result.allowed)
            self.assertEqual("preflight_unverified_weight_format", result.reason)
        self.assertEqual("mlx", ModelInspection(repo_id="m", model_card_metadata={
            "library_name": "mlx"}).quantization_method)
        self.assertIsNone(ModelInspection(repo_id="m", model_config={}).quantization_method)

    def test_encoding_fingerprints_distinguish_compressed_formats(self):
        a = ModelInspection(repo_id="a", model_config={"quantization_config": {
            "quant_method": "compressed-tensors", "format": "float-quantized"}})
        b = ModelInspection(repo_id="b", model_config={"quantization_config": {
            "quant_method": "compressed-tensors", "format": "pack-quantized"}})
        self.assertNotEqual(a.quantization_signature, b.quantization_signature)

    def candidate(self, model="a"):
        return {"repoId": model, "targetGpu": "gpu", "framework": "vllm", "taskType": "text-generation",
                "preflightMetadata": {"modelType": "family", "quantizationSignature": "awq:abc"}}

    def test_batch_guard_counts_pending_and_same_batch_and_is_retryable(self):
        candidate = self.candidate()
        pending = {**candidate, "outcome": "pending"}
        selected, deferred = filter_profile_batch([self.candidate(str(i)) for i in range(4)], [pending], {})
        self.assertEqual(1, len(selected))
        self.assertEqual(3, len(deferred))
        selected, deferred = filter_profile_batch(deferred, [], {})
        self.assertEqual(2, len(selected))

    def test_proven_encoding_releases_cap_but_failures_do_not(self):
        candidates = [self.candidate(str(i)) for i in range(4)]
        key = "|".join(profile_key(candidates[0]))
        stats = {"encodingCombinationStats": {key: {"successCount": 3, "failureCount": 0}}}
        self.assertEqual(4, len(filter_profile_batch(candidates, [], stats)[0]))
        stats["encodingCombinationStats"][key]["failureCount"] = 10
        self.assertEqual(2, len(filter_profile_batch(candidates, [], stats)[0]))

    def test_takeover_is_idempotent_and_does_not_invent_version(self):
        with tempfile.TemporaryDirectory() as tmp:
            tracker = self.tracker(Path(tmp))
            tasks = [{"taskId": "old", "modelId": "owner/model", "gpuType": "gpu", "status": "waiting"}]
            self.assertEqual(1, tracker.adopt_active_tasks(tasks))
            self.assertEqual(0, tracker.adopt_active_tasks(tasks))
            self.assertEqual(1, tracker.pending_route_count_for_model("owner/model"))
            self.assertNotIn("agentVersion", tracker.pending_records()[0])

    def test_encoding_and_version_statistics_survive_compaction(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            tracker = self.tracker(root)
            for i in range(3):
                tracker.record_submission("m"+str(i), "gpu", "vllm", "text-generation", str(i),
                    "2026-10-26T00:00:00+00:00", self.candidate()["preflightMetadata"])
                tracker._update_record_from_task(tracker._by_task_id[str(i)], {"status": "success", "verifyResult": 1})
            before = tracker.get_stats_report()
            tracker._compact_if_needed(force=True)
            restored = self.tracker(root).get_stats_report()
            self.assertEqual(before["encodingCombinationStats"], restored["encodingCombinationStats"])
            self.assertEqual(before["agentSummaries"], restored["agentSummaries"])

    def test_state_rotation_preserves_remote_and_bounds_local_objects(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            remote = root / "remote.git"
            with porcelain.init(remote, bare=True):
                pass
            project = root / "project"
            project.mkdir()
            manager = StateGitSync(project_root=project, credentials={"username": "u", "email": "u@x", "password": "p"},
                                   remote=str(remote), history_depth=2, log_fn=lambda _: None)
            try:
                self.assertTrue(manager.restore())
                for i in range(5):
                    write_json(project / ".modelhub_state/account_capacity.json", {"generation": i})
                    self.assertTrue(manager.sync("result"), manager.last_error)
                self.assertTrue(manager.healthy)
                self.assertEqual(manager._expected_remote_oid, manager._remote_oid())
                with patch.object(manager, "_copy_snapshot", side_effect=AssertionError("idle writes should throttle")):
                    self.assertTrue(manager.sync("cycle_no_slots"))
                write_json(project / ".modelhub_state/account_capacity.json", {"generation": 99})
                self.assertTrue(manager.sync("intent"))
            finally:
                manager.close()


if __name__ == "__main__":
    unittest.main()