fix: verify capacity from dual paginated active-task scans

This commit is contained in:
CoolBoy
2026-10-28 00:38:09 +08:00
parent 0ba4977e88
commit 5dead8172c
5 changed files with 148 additions and 53 deletions

View File

@@ -15,10 +15,11 @@ class CapacityRefreshTests(unittest.TestCase):
def client(self, total=97, observed=94):
client = ModelHubClient(token="test")
def page(*, status, page_size, **kwargs):
def page(*, status, page_size, current=1, **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)]}}
{"taskId": str(i), "status": status} for i in range(
(current - 1) * page_size, min(current * page_size, observed if status == "waiting" else 0))]}}
client.list_tasks_page = page
return client
@@ -26,8 +27,8 @@ class CapacityRefreshTests(unittest.TestCase):
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.assertEqual([94], pool.refresh_active_counts())
self.assertEqual(6, pool.available_submit_slots())
self.assertTrue(client.active_count_details[0]["mismatch"])
def test_error_is_unknown_and_recovery_reopens_slots(self):
@@ -40,22 +41,82 @@ class CapacityRefreshTests(unittest.TestCase):
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())
self.assertEqual(6, pool.available_submit_slots())
def test_incomplete_or_repeated_pages_do_not_open_slots(self):
client = self.client(total=200, observed=100)
client.list_tasks_page = lambda **kwargs: {"data": {"records": [
{"taskId": "same", "status": kwargs["status"]}]}}
with self.assertRaises(ModelHubAPIError):
client.count_active_tasks_exact()
client.list_tasks_page = lambda **kwargs: {"data": {"records": []}}
client.list_tasks_page = lambda **kwargs: {"data": {}}
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": [
client.list_tasks_page = lambda **kwargs: {"data": {"total": 1, "records": [] if kwargs['current'] > 1 else [
{"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())
def test_inflated_total_does_not_hide_capacity(self):
self.assertEqual(100, self.client(total=773, observed=100).count_active_tasks_exact())
def test_underreported_pages_do_not_truncate_scan(self):
client = self.client(total=58, observed=175)
original = client.list_tasks_page
def page(**kwargs):
result = original(**kwargs)
result['data']['pages'] = 1
return result
client.list_tasks_page = page
self.assertEqual(175, client.count_active_tasks_exact())
def test_different_confirmation_ids_fail_closed_even_with_equal_counts(self):
client = self.client(total=1, observed=1)
original = client.list_tasks_page
def page(**kwargs):
result = original(**kwargs)
if kwargs['page_size'] == 50:
for row in result['data']['records']:
row['taskId'] = 'replacement'
return result
client.list_tasks_page = page
with self.assertRaisesRegex(ModelHubAPIError, 'active_snapshot_changed'):
client.count_active_tasks_exact()
def test_confirmation_error_does_not_open_slots(self):
client = self.client(total=773, observed=94)
original = client.list_tasks_page
def page(**kwargs):
if kwargs['page_size'] == 50:
raise ModelHubAPIError('temporary failure')
return original(**kwargs)
client.list_tasks_page = page
pool = ModelHubClientPool([client], active_task_cap=100)
pool.refresh_active_counts()
self.assertEqual(0, pool.available_submit_slots())
self.assertFalse(pool.active_count_status()[0]['known'])
def test_empty_queue_confirmed_despite_stale_total(self):
self.assertEqual(0, self.client(total=773, observed=0).count_active_tasks_exact())
def test_short_page_is_not_end_of_scan(self):
client = ModelHubClient(token='test')
def page(**kwargs):
rows = [{'taskId': str(kwargs['current']), 'status': 'waiting'}] if (
kwargs['status'] == 'waiting' and kwargs['current'] <= 3) else []
return {'data': {'records': rows}}
client.list_tasks_page = page
self.assertEqual(3, client.count_active_tasks_exact())
def test_scan_budget_is_bounded(self):
client = ModelHubClient(token='test')
client.list_tasks_page = lambda **kw: {'data': {'records': [
{'taskId': str(kw['current']), 'status': kw['status']}]}}
with self.assertRaisesRegex(ModelHubAPIError, 'active_pagination_budget_exhausted'):
client.count_active_tasks_exact()

View File

@@ -290,16 +290,17 @@ class ClientPoolConcurrencyTests(unittest.TestCase):
client = ModelHubClient(token="test-token")
calls: list[str] = []
def page(*, status: str, current: int = 1, **_kwargs) -> dict: # noqa: ANN003
def page(*, status: str, current: int = 1, page_size: int = 100, **_kwargs) -> dict: # noqa: ANN003
calls.append(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))]}}
for i in range((current - 1) * page_size, min(current * page_size, total))]}}
client.list_tasks_page = page # type: ignore[method-assign]
self.assertEqual(200, client.count_active_tasks_exact())
self.assertEqual(["waiting", "waiting", "running"], calls)
self.assertEqual(8, calls.count("waiting"))
self.assertEqual(4, calls.count("running"))
@staticmethod
def monday_calibration_time(hour: int = 10) -> datetime: