wehub-resource-sync
40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
import time
|
|
|
|
from omlx.prefill_progress import PrefillProgressTracker
|
|
|
|
|
|
def test_prefill_progress_resets_elapsed_when_phase_changes(monkeypatch):
|
|
tracker = PrefillProgressTracker()
|
|
times = iter([100.0, 110.0, 112.0])
|
|
monkeypatch.setattr(time, "monotonic", lambda: next(times))
|
|
|
|
tracker.update("req", 10, 100, "model", phase="specprefill_scoring")
|
|
tracker.update("req", 0, 50, "model", phase="specprefill_sparse")
|
|
|
|
progress = tracker.get_model_progress("model")[0]
|
|
assert progress["phase"] == "specprefill_sparse"
|
|
assert progress["elapsed"] == 2.0
|
|
|
|
|
|
def test_prefill_progress_extra_cannot_clobber_base_fields():
|
|
tracker = PrefillProgressTracker()
|
|
|
|
tracker.update(
|
|
"req",
|
|
10,
|
|
100,
|
|
"model",
|
|
phase="prefill",
|
|
extra={
|
|
"processed": 999,
|
|
"phase": "wrong",
|
|
"start_time": -1,
|
|
"selected_tokens": 5,
|
|
},
|
|
)
|
|
|
|
progress = tracker.get_model_progress("model")[0]
|
|
assert progress["processed"] == 10
|
|
assert progress["phase"] == "prefill"
|
|
assert progress["selected_tokens"] == 5
|