fix: respect task page limits and recover stale capacity counts
This commit is contained in:
61
tests/test_capacity_refresh.py
Normal file
61
tests/test_capacity_refresh.py
Normal file
@@ -0,0 +1,61 @@
|
||||
import sys
|
||||
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 modelhub_client import ModelHubClient, ModelHubClientPool, ModelHubAPIError
|
||||
|
||||
|
||||
class CapacityRefreshTests(unittest.TestCase):
|
||||
def test_api_boundary_clamps_all_callers(self):
|
||||
client = ModelHubClient(token="test")
|
||||
with patch.object(client, "_request", return_value={}) as request:
|
||||
client.list_tasks_page(page_size=200)
|
||||
self.assertEqual(100, request.call_args.kwargs["query"]["pageSize"])
|
||||
|
||||
def client(self, total=97, observed=94):
|
||||
client = ModelHubClient(token="test")
|
||||
def page(*, status, page_size, **kwargs):
|
||||
self.assertLessEqual(page_size, 100)
|
||||
return {"data": {"total": total if status == "waiting" else 0, "records": [
|
||||
{"taskId": str(i), "status": status} for i in range(observed if status == "waiting" else 0)]}}
|
||||
client.list_tasks_page = page
|
||||
return client
|
||||
|
||||
def test_full_cache_refreshes_to_free_slots_with_disagreeing_total(self):
|
||||
client = self.client()
|
||||
pool = ModelHubClientPool([client], active_task_cap=100)
|
||||
pool._remote_counts[0] = 100
|
||||
self.assertEqual([97], pool.refresh_active_counts())
|
||||
self.assertEqual(3, pool.available_submit_slots())
|
||||
self.assertTrue(client.active_count_details[0]["mismatch"])
|
||||
|
||||
def test_error_is_unknown_and_recovery_reopens_slots(self):
|
||||
client = self.client()
|
||||
pool = ModelHubClientPool([client], active_task_cap=100)
|
||||
pool.refresh_active_counts()
|
||||
with patch.object(client, "list_tasks_page", side_effect=ModelHubAPIError("分页大小不能超过100")):
|
||||
pool.refresh_active_counts()
|
||||
self.assertEqual(0, pool.available_submit_slots())
|
||||
self.assertIsNone(pool.active_count_status()[0]["active"])
|
||||
self.assertIsNone(pool._reserve_account(set()))
|
||||
pool.refresh_active_counts()
|
||||
self.assertEqual(3, pool.available_submit_slots())
|
||||
|
||||
def test_incomplete_or_repeated_pages_do_not_open_slots(self):
|
||||
client = self.client(total=200, observed=100)
|
||||
with self.assertRaises(ModelHubAPIError):
|
||||
client.count_active_tasks_exact()
|
||||
client.list_tasks_page = lambda **kwargs: {"data": {"records": []}}
|
||||
with self.assertRaises(ModelHubAPIError):
|
||||
client.count_active_tasks_exact()
|
||||
|
||||
def test_duplicate_ids_and_transition_are_counted_once(self):
|
||||
client = ModelHubClient(token="test")
|
||||
client.list_tasks_page = lambda **kwargs: {"data": {"total": 1, "records": [
|
||||
{"taskId": "same", "status": kwargs["status"]},
|
||||
{"taskId": "same", "status": kwargs["status"]}]}}
|
||||
self.assertEqual(1, client.count_active_tasks_exact())
|
||||
|
||||
def test_total_lower_than_observed_uses_larger_count(self):
|
||||
self.assertEqual(99, self.client(total=97, observed=99).count_active_tasks_exact())
|
||||
@@ -290,13 +290,16 @@ class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||
client = ModelHubClient(token="test-token")
|
||||
calls: list[str] = []
|
||||
|
||||
def page(*, status: str, **_kwargs) -> dict: # noqa: ANN003
|
||||
def page(*, status: str, current: int = 1, **_kwargs) -> dict: # noqa: ANN003
|
||||
calls.append(status)
|
||||
return {"data": {"total": {"waiting": 175, "running": 25}[status]}}
|
||||
total = {"waiting": 175, "running": 25}[status]
|
||||
return {"data": {"total": total, "records": [
|
||||
{"taskId": f"{status}-{i}", "status": status}
|
||||
for i in range((current - 1) * 100, min(current * 100, total))]}}
|
||||
|
||||
client.list_tasks_page = page # type: ignore[method-assign]
|
||||
self.assertEqual(200, client.count_active_tasks_exact())
|
||||
self.assertEqual(["waiting", "running"], calls)
|
||||
self.assertEqual(["waiting", "waiting", "running"], calls)
|
||||
|
||||
@staticmethod
|
||||
def monday_calibration_time(hour: int = 10) -> datetime:
|
||||
|
||||
Reference in New Issue
Block a user