feat: calibrate account capacity weekly

tests/test_concurrency.py

@@ -20,6 +20,7 @@ from common import read_json, read_jsonl  # noqa: E402
 from main import build_adaptive_scan_stages, build_parser, make_run_dir, run_submission, submit_candidate  # noqa: E402
 from modelhub_client import (  # noqa: E402
     ModelHubAPIError,
+    ModelHubClient,
     ModelHubClientPool,
     OldModelQueuePolicyError,
     is_duplicate_submission_error,
@@ -120,8 +121,8 @@ class DynamicCapacityClient:
         self.limit = limit
         self.token = token
 
-    def count_active_tasks(self, *, max_count: int, **_kwargs) -> int:  # noqa: ANN003
-        return min(self.active, max_count)
+    def count_active_tasks(self, *, max_count: int | None, **_kwargs) -> int:  # noqa: ANN003
+        return self.active if max_count is None else min(self.active, max_count)
 
     def add_task(self, _payload: dict) -> dict:
         if self.active >= self.limit:
@@ -215,6 +216,195 @@ def make_candidate(index: int) -> dict:
 
 
 class ClientPoolConcurrencyTests(unittest.TestCase):
+    def test_exact_active_count_uses_filtered_platform_totals(self) -> None:
+        client = ModelHubClient(token="test-token")
+        calls: list[str] = []
+
+        def page(*, status: str, **_kwargs) -> dict:  # noqa: ANN003
+            calls.append(status)
+            return {"data": {"total": {"waiting": 175, "running": 25}[status]}}
+
+        client.list_tasks_page = page  # type: ignore[method-assign]
+        self.assertEqual(200, client.count_active_tasks_exact())
+        self.assertEqual(["waiting", "running"], calls)
+
+    @staticmethod
+    def monday_calibration_time(hour: int = 10) -> datetime:
+        # 2026-09-21 10:00 Beijing time.
+        return datetime(2026, 9, 21, hour - 8, tzinfo=timezone.utc)
+
+    def test_weekly_calibration_borrows_five_slots_only_during_monday_window(self) -> None:
+        client = DynamicCapacityClient(active=95, limit=100)
+        pool = ModelHubClientPool(
+            [client],  # type: ignore[list-item]
+            active_task_cap=100,
+            recent_model_reserve_slots=5,
+        )
+        self.assertEqual([95], pool.active_task_counts())
+        old_time = datetime(2026, 9, 1, tzinfo=timezone.utc)
+        pool.configure_weekly_calibration(self.monday_calibration_time(9))
+        with self.assertRaises(OldModelQueuePolicyError):
+            pool.add_task_for_model({"model": "before"}, model_last_modified=old_time)
+
+        pool.configure_weekly_calibration(self.monday_calibration_time())
+        self.assertEqual(1, pool.old_model_submit_slots())
+        for index in range(5):
+            pool.add_task_for_model({"model": index}, model_last_modified=old_time)
+        self.assertEqual(100, client.active)
+        with self.assertRaises(ModelHubAPIError):
+            pool.add_task_for_model({"model": "test-101"}, model_last_modified=old_time)
+        self.assertEqual(1, pool.weekly_calibration_status()["completedAccounts"])
+        self.assertEqual([95], pool.old_model_queue_thresholds())
+        pool.configure_weekly_calibration(self.monday_calibration_time(11))
+        with self.assertRaises(OldModelQueuePolicyError):
+            pool.add_task_for_model({"model": "after"}, model_last_modified=old_time)
+        self.assertEqual(100, client.active)
+
+    def test_weekly_calibration_stops_when_window_deadline_passes(self) -> None:
+        client = DynamicCapacityClient(active=95, limit=200)
+        pool = ModelHubClientPool([client], active_task_cap=100)  # type: ignore[list-item]
+        pool.active_task_counts()
+        pool.configure_weekly_calibration(self.monday_calibration_time())
+        pool._weekly_calibration_deadline = 0.0
+        with self.assertRaises(OldModelQueuePolicyError):
+            pool.add_task_for_model(
+                {"model": "too-late"},
+                model_last_modified=datetime(2026, 9, 1, tzinfo=timezone.utc),
+            )
+        self.assertEqual(95, client.active)
+
+    def test_weekly_calibration_discovers_200_and_persists_per_account_completion(self) -> None:
+        client = DynamicCapacityClient(active=95, limit=200)
+        old_time = datetime(2026, 9, 1, tzinfo=timezone.utc)
+        with tempfile.TemporaryDirectory() as temporary_dir:
+            state_path = Path(temporary_dir) / "account_capacity.json"
+            pool = ModelHubClientPool(
+                [client],  # type: ignore[list-item]
+                active_task_cap=100,
+                capacity_state_path=state_path,
+            )
+            pool.active_task_counts()
+            pool.configure_weekly_calibration(self.monday_calibration_time())
+            for index in range(105):
+                pool.add_task_for_model({"model": index}, model_last_modified=old_time)
+            with self.assertRaises(ModelHubAPIError):
+                pool.add_task_for_model({"model": "test-201"}, model_last_modified=old_time)
+            self.assertEqual([200], pool.account_capacity_limits())
+            self.assertEqual([195], pool.old_model_queue_thresholds())
+            stored = read_json(state_path)
+            account = next(iter(stored["accounts"].values()))
+            self.assertEqual("saturated", account["calibrationResult"])
+
+            restarted = ModelHubClientPool(
+                [client],  # type: ignore[list-item]
+                active_task_cap=100,
+                capacity_state_path=state_path,
+            )
+            restarted.active_task_counts()
+            restarted.configure_weekly_calibration(self.monday_calibration_time())
+            self.assertEqual([200], restarted.account_capacity_limits())
+            self.assertIsNone(restarted.weekly_calibration_status()["targetAccount"])
+
+    def test_weekly_calibration_reads_full_active_count_not_cap_plus_one(self) -> None:
+        client = DynamicCapacityClient(active=200, limit=200)
+        pool = ModelHubClientPool([client], active_task_cap=100)  # type: ignore[list-item]
+        pool.configure_weekly_calibration(self.monday_calibration_time())
+        self.assertEqual([200], pool.active_task_counts())
+        self.assertEqual([200], pool.account_capacity_limits())
+        with self.assertRaises(ModelHubAPIError):
+            pool.add_task_for_model(
+                {"model": "full-capacity"},
+                model_last_modified=datetime(2026, 9, 1, tzinfo=timezone.utc),
+            )
+        self.assertEqual(1, pool.weekly_calibration_status()["completedAccounts"])
+
+    def test_weekly_calibration_targets_accounts_separately(self) -> None:
+        first = DynamicCapacityClient(active=2, limit=2, token="first")
+        second = DynamicCapacityClient(active=1, limit=3, token="second")
+        pool = ModelHubClientPool(
+            [first, second],  # type: ignore[list-item]
+            active_task_cap=2,
+            recent_model_reserve_slots=1,
+        )
+        pool.active_task_counts()
+        pool.configure_weekly_calibration(self.monday_calibration_time())
+        self.assertEqual(1, pool.weekly_calibration_status()["targetAccount"])
+        pool.add_task_for_model(
+            {"model": "try-first-then-second"},
+            model_last_modified=datetime(2026, 9, 1, tzinfo=timezone.utc),
+        )
+        self.assertEqual(2, first.active)
+        self.assertEqual(2, second.active)
+        self.assertEqual(2, pool.weekly_calibration_status()["targetAccount"])
+        pool.add_task_for_model(
+            {"model": "second-capacity"},
+            model_last_modified=datetime(2026, 9, 1, tzinfo=timezone.utc),
+        )
+        with self.assertRaises(ModelHubAPIError):
+            pool.add_task_for_model(
+                {"model": "test-second-full"},
+                model_last_modified=datetime(2026, 9, 1, tzinfo=timezone.utc),
+            )
+        self.assertEqual([2, 3], pool.account_capacity_limits())
+        self.assertEqual([1, 2], pool.old_model_queue_thresholds())
+        self.assertEqual(2, pool.weekly_calibration_status()["completedAccounts"])
+
+    def test_weekly_calibration_skips_unknown_counts_and_retries_no_candidates(self) -> None:
+        class UnreadableCountClient(DynamicCapacityClient):
+            def count_active_tasks(self, **_kwargs) -> int:
+                raise TimeoutError("cannot read active count")
+
+        unreadable = UnreadableCountClient(active=95, limit=100, token="unreadable")
+        readable = DynamicCapacityClient(active=95, limit=100, token="readable")
+        pool = ModelHubClientPool(
+            [unreadable, readable],  # type: ignore[list-item]
+            active_task_cap=100,
+        )
+        pool.active_task_counts()
+        pool.configure_weekly_calibration(self.monday_calibration_time())
+        self.assertEqual(2, pool.weekly_calibration_status()["targetAccount"])
+        pool.pause_weekly_calibration_for_no_candidates()
+        self.assertEqual(0, pool.weekly_calibration_status()["completedAccounts"])
+        self.assertIsNone(pool.weekly_calibration_status()["targetAccount"])
+        pool._weekly_no_candidate_retry_at = 0.0
+        self.assertEqual(2, pool.weekly_calibration_status()["targetAccount"])
+
+    def test_rate_limit_does_not_finish_weekly_capacity_calibration(self) -> None:
+        class RateLimitedClient(DynamicCapacityClient):
+            def add_task(self, _payload: dict) -> dict:
+                raise ModelHubAPIError("rate limited", code=429)
+
+        client = RateLimitedClient(active=100, limit=200)
+        pool = ModelHubClientPool([client], active_task_cap=100)  # type: ignore[list-item]
+        pool.active_task_counts()
+        pool.configure_weekly_calibration(self.monday_calibration_time())
+        with self.assertRaises(ModelHubAPIError):
+            pool.add_task_for_model(
+                {"model": "rate-limited"},
+                model_last_modified=datetime(2026, 9, 1, tzinfo=timezone.utc),
+            )
+        self.assertEqual([100], pool.account_capacity_limits())
+        self.assertEqual(1, pool.weekly_calibration_status()["targetAccount"])
+
+    def test_rate_limited_submission_is_deferred_not_failed_or_capacity(self) -> None:
+        class RateLimitedClient:
+            @staticmethod
+            def add_task(_payload: dict) -> dict:
+                raise ModelHubAPIError("rate limited", code=429)
+
+        result = submit_candidate(
+            {
+                "repoId": "owner/model",
+                "modelAddress": "https://modelscope.cn/models/owner/model",
+                "targetGpu": "gpu-a",
+                "framework": "vllm",
+                "taskType": "text-generation",
+                "configParams": "safe",
+            },
+            RateLimitedClient(),  # type: ignore[arg-type]
+        )
+        self.assertEqual("rate_limit_deferred", result["outcome"])
+
     def test_old_models_use_another_account_and_reserve_last_five_positions(self) -> None:
         below_threshold = FakeClient(active_count=94)
         at_threshold = FakeClient(active_count=95)