fix: verify capacity from dual paginated active-task scans
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user