| """CPU execution of the installed image classes (including 0031c), without SO1. |
| |
| Only dependency imports, GPU construction, KV allocation and delivery are faked. |
| No scheduler statement, Request body or priority queue is replaced. |
| """ |
| from __future__ import annotations |
|
|
| from types import SimpleNamespace |
|
|
| import pytest |
|
|
| from d2.scheduler_cpu import load_engine, make_request, make_scheduler |
|
|
|
|
| def engine() -> SimpleNamespace: |
| result = load_engine() |
| result.scheduler._DSV41_THROTTLE_MIN_PRIO = 5 |
| return result |
|
|
|
|
| @pytest.mark.parametrize('asynchronous', [False, True]) |
| @pytest.mark.parametrize('aux_priority,expected', [(5, ['ndm', 'd2', 'legal']), |
| (20, ['ndm', 'legal', 'd2'])]) |
| def test_actual_waiting_admission_order(asynchronous: bool, aux_priority: int, |
| expected: list[str]) -> None: |
| e = engine() |
| s = make_scheduler(e, asynchronous) |
| for rid, priority in [('d2', aux_priority), ('legal', 10), ('ndm', 0)]: |
| r = make_request(e, rid, 100, priority=priority) |
| s.requests[rid] = r |
| s.waiting.add_request(r) |
| assert list(s.schedule().num_scheduled_tokens) == expected |
| assert [(r.request_id, r.priority) for r in s.running] == [ |
| (rid, {'ndm': 0, 'legal': 10, 'd2': aux_priority}[rid]) for rid in expected] |
|
|
|
|
| @pytest.mark.parametrize('asynchronous', [False, True]) |
| @pytest.mark.parametrize('foreground,throttled', [(0, True), (10, False)]) |
| def test_0031_does_not_protect_priority10_from_auxiliary(asynchronous: bool, |
| foreground: int, throttled: bool) -> None: |
| e = engine() |
| s = make_scheduler(e, asynchronous) |
| for rid, priority in [('d2', 20), ('legal', foreground)]: |
| r = make_request(e, rid, 32000 if rid == 'd2' else 256, priority=priority) |
| r.status = e.request.RequestStatus.RUNNING |
| r.is_prefill_chunk = True |
| s.requests[rid] = r |
| s.running.append(r) |
| before = {r.request_id: r.num_computed_tokens for r in s.running} |
| counts = s.schedule(throttle_prefills=True).num_scheduled_tokens |
| if throttled: |
| assert counts == {'legal': 256} |
| assert s.requests['d2'].num_computed_tokens == before['d2'] |
| else: |
| assert counts == {'d2': 1024, 'legal': 256} |
| |
| assert s.schedule(throttle_prefills=False).num_scheduled_tokens['d2'] == 1024 |
| s.prefill_capacity_bound = True |
| assert s.schedule(throttle_prefills=True).num_scheduled_tokens['d2'] == 1024 |
|
|
|
|
| @pytest.mark.parametrize('asynchronous', [False, True]) |
| @pytest.mark.parametrize('draft_slots', [0, 4]) |
| def test_three_auxiliary_prefills_leave_admission_budget(asynchronous: bool, draft_slots: int) -> None: |
| e = engine() |
| s = make_scheduler(e, asynchronous) |
| s.vllm_config.speculative_config = SimpleNamespace(max_num_new_slots_for_drafting=draft_slots) |
| for j in range(3): |
| r = make_request(e, f'd2-{j}', 32000, priority=20) |
| s.requests[r.request_id] = r |
| s.waiting.add_request(r) |
| s.schedule() |
| legal = make_request(e, 'legal', 32000, priority=10, arrival=1) |
| s.requests['legal'] = legal |
| s.waiting.add_request(legal) |
| counts = s.schedule(throttle_prefills=True).num_scheduled_tokens |
| assert counts == {'d2-0': 1024, 'd2-1': 1024, 'd2-2': 1024, 'legal': 1024 - 4 * draft_slots} |
| assert len(s.running) == 4 and not list(s.waiting) |
|
|
