| """Tests for src/queue_eta — slot-simulation ETA calculation.""" |
|
|
| import json |
| import math |
| import os |
| import tempfile |
|
|
| import pytest |
|
|
| from src.queue_eta import ( |
| _LARGE_MODEL_HOURS, |
| _SIZE_THRESHOLD_B, |
| _SMALL_MODEL_HOURS, |
| compute_single_eta, |
| estimate_task_hours, |
| format_eta, |
| _get_params, |
| ) |
|
|
|
|
| |
| |
| |
|
|
| def _write(directory, owner, filename, data): |
| d = os.path.join(directory, owner) |
| os.makedirs(d, exist_ok=True) |
| fpath = os.path.join(d, filename) |
| with open(fpath, "w") as f: |
| json.dump(data, f) |
| return fpath |
|
|
|
|
| |
| |
| |
|
|
| class TestEstimateTaskHours: |
| def test_small_model(self): |
| assert estimate_task_hours(7.0) == _SMALL_MODEL_HOURS |
|
|
| def test_at_threshold(self): |
| assert estimate_task_hours(_SIZE_THRESHOLD_B) == _SMALL_MODEL_HOURS |
|
|
| def test_large_model(self): |
| assert estimate_task_hours(70.0) == _LARGE_MODEL_HOURS |
|
|
| def test_none_defaults_small(self): |
| assert estimate_task_hours(None) == _SMALL_MODEL_HOURS |
|
|
| def test_zero_defaults_small(self): |
| assert estimate_task_hours(0) == _SMALL_MODEL_HOURS |
|
|
| def test_negative_defaults_small(self): |
| assert estimate_task_hours(-5) == _SMALL_MODEL_HOURS |
|
|
|
|
| |
| |
| |
|
|
| class TestFormatEta: |
| def test_zero(self): |
| assert format_eta(0) == "< 1h" |
|
|
| def test_negative(self): |
| assert format_eta(-1) == "< 1h" |
|
|
| def test_fraction(self): |
| assert format_eta(0.5) == "~1h" |
|
|
| def test_exact_hours(self): |
| assert format_eta(6.0) == "~6h" |
|
|
| def test_fractional_hours(self): |
| assert format_eta(6.3) == "~7h" |
|
|
| def test_exact_day(self): |
| assert format_eta(24.0) == "~1d" |
|
|
| def test_day_and_hours(self): |
| assert format_eta(27.0) == "~1d 3h" |
|
|
| def test_multi_days(self): |
| assert format_eta(50.0) == "~2d 2h" |
|
|
|
|
| |
| |
| |
|
|
| class TestGetParams: |
| def test_model_params(self): |
| assert _get_params({"model_params": 7.0}) == 7.0 |
|
|
| def test_params_fallback(self): |
| assert _get_params({"params": 14}) == 14.0 |
|
|
| def test_model_params_preferred(self): |
| assert _get_params({"model_params": 7, "params": 14}) == 7.0 |
|
|
| def test_none_when_missing(self): |
| assert _get_params({}) is None |
|
|
| def test_string_value(self): |
| assert _get_params({"model_params": "7.5"}) == 7.5 |
|
|
| def test_bad_string(self): |
| assert _get_params({"model_params": "unknown"}) is None |
|
|
|
|
| |
| |
| |
|
|
| class TestComputeSingleEta: |
|
|
| def setup_method(self): |
| self._tmpdir = tempfile.mkdtemp(prefix="test_single_eta_") |
|
|
| def test_empty_queue(self): |
| """No running, no pending — new model is position 1.""" |
| eta = compute_single_eta(self._tmpdir, model_params=7.0, concurrency=2) |
| |
| |
| assert eta == 3.0 |
|
|
| def test_with_2r_1p(self): |
| """2 running + 1 pending → new model at pos 2.""" |
| _write(self._tmpdir, "a", "r1.json", { |
| "model": "a/r1", "status": "Running", "script": "auto_quant", |
| "model_params": 7.0, |
| }) |
| _write(self._tmpdir, "b", "r2.json", { |
| "model": "b/r2", "status": "Running", "script": "auto_eval", |
| "model_params": 7.0, |
| }) |
| _write(self._tmpdir, "c", "p1.json", { |
| "model": "c/p1", "status": "Pending", "script": "auto_quant", |
| "model_params": 7.0, "submitted_time": "2026-04-01T00:00:00Z", |
| }) |
|
|
| eta = compute_single_eta(self._tmpdir, model_params=7.0, concurrency=2) |
| |
| |
| assert eta == 6.0 |
|
|
| def test_large_model(self): |
| eta = compute_single_eta(self._tmpdir, model_params=70.0, concurrency=2) |
| |
| |
| assert eta == 5.0 |
|
|
| def test_nonexistent_dir(self): |
| eta = compute_single_eta("/tmp/nonexistent_eta_dir", model_params=7.0, concurrency=2) |
| assert eta == 3.0 |
|
|
