tests/test_estimator.py
Rama main · réplica de lectura del repositorio original.
"""Unit tests for the v1 wall-time estimator (jobd.estimator)."""
from __future__ import annotations
import json
from datetime import UTC, datetime, timedelta
import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from jobd.db import Job, init_db, migrate
from jobd.estimator import (
CLIP_THRESHOLD,
HISTORY_WINDOW,
MIN_SAMPLES,
ColdStart,
WallPrediction,
_quantile,
cmd_head,
make_predict_cache,
predict_wall,
queue_start_eta,
remaining_for_running,
)
from jobd.models import JobState
# ---------------- cmd_head ----------------
@pytest.mark.parametrize(
"cmd,expected",
[
([], ""),
(["python", "train.py"], "python train.py"),
(["python", "train.py", "--seed", "0"], "python train.py"),
(["python3", "/abs/path/run.py"], "python3 run.py"),
(["/usr/bin/python3", "/abs/path/run.py"], "python3 run.py"),
(["python", "-m", "torch.distributed.run"], "python -m torch.distributed.run"),
(["python", "-c", "print(1)"], "python -c"),
(["Rscript", "pipeline.R"], "Rscript pipeline.R"),
(["./train.sh"], "train.sh"),
(["bash", "-c", "echo hi"], "bash -c"),
(["bash", "build.sh"], "bash build.sh"),
(["uv", "run", "pytest"], "uv run"),
(["foo"], "foo"),
(["/usr/local/bin/foo"], "foo"),
# Skip --flag --flag value before settling on the script.
(["python", "--unbuffered", "train.py"], "python train.py"),
],
)
def test_cmd_head_cases(cmd, expected):
assert cmd_head(cmd) == expected
def test_cmd_head_caps_at_80_chars():
long = ["python", "x" * 200]
out = cmd_head(long)
assert len(out) <= 80
# ---------------- quantile math ----------------
def test_quantile_basic():
assert _quantile([10.0], 0.5) == 10.0
assert _quantile([10.0, 20.0], 0.5) == 15.0
# 100 evenly spaced 1..100; p50 is interpolated → 50.5
vals = [float(i) for i in range(1, 101)]
assert _quantile(vals, 0.5) == pytest.approx(50.5)
assert _quantile(vals, 0.9) == pytest.approx(90.1)
def test_quantile_empty():
assert _quantile([], 0.5) == 0.0
# ---------------- predict_wall ----------------
@pytest.fixture
def session(tmp_path):
engine = create_engine(f"sqlite:///{tmp_path}/est.db", future=True)
init_db(engine)
migrate(engine)
Session = sessionmaker(engine, future=True)
with Session() as s:
yield s
def _add_completed(
session,
project: str,
cmd: list[str],
wall_s: float,
*,
max_wall_s: int | None = None,
base_id: int | None = None,
):
"""Insert a completed Job row with deterministic timing."""
started = datetime(2026, 4, 1, tzinfo=UTC) + timedelta(seconds=base_id or 0)
finished = started + timedelta(seconds=wall_s)
job = Job(
id=base_id,
project=project,
host_pin="any",
priority=50,
state=JobState.COMPLETED.value,
cmd_json=json.dumps(cmd),
cwd="/tmp",
env_json="{}",
preemptible=False,
vram_gb=0,
ram_gb=0,
cpus=1,
submitted_at=started,
started_at=started,
finished_at=finished,
max_wall_s=max_wall_s,
)
session.add(job)
session.commit()
return job
def test_predict_wall_cold_start(session):
pred = predict_wall(session, "p1", "python train.py")
assert isinstance(pred, ColdStart)
assert pred.n_samples == 0
assert pred.basis == "insufficient-history-N=0"
def test_predict_wall_below_min_samples(session):
for i in range(MIN_SAMPLES - 1):
_add_completed(session, "p1", ["python", "t.py"], 100.0, base_id=i + 1)
pred = predict_wall(session, "p1", "python t.py")
assert isinstance(pred, ColdStart)
assert pred.n_samples == MIN_SAMPLES - 1
def test_predict_wall_quantiles_and_basis(session):
# 10 walls: 60..600s, p50 ≈ 330, p90 ≈ 546
for i, w in enumerate(range(60, 601, 60)):
_add_completed(session, "p1", ["python", "t.py"], float(w), base_id=i + 1)
pred = predict_wall(session, "p1", "python t.py")
assert isinstance(pred, WallPrediction)
assert pred.n_samples == 10
assert pred.p50_s == pytest.approx(330.0)
assert pred.p90_s == pytest.approx(546.0)
assert pred.basis == "history-N=10"
assert pred.clipped is False
def test_predict_wall_clip_detection(session):
# 5 walls: 4 short, 1 within 95% of max_wall_s=100 (i.e. wall=98)
for i, w in enumerate([10.0, 12.0, 15.0, 20.0]):
_add_completed(session, "p1", ["python", "t.py"], w, base_id=i + 1)
_add_completed(session, "p1", ["python", "t.py"], 98.0, max_wall_s=100, base_id=99)
pred = predict_wall(session, "p1", "python t.py")
assert isinstance(pred, WallPrediction)
assert pred.clipped is True
def test_predict_wall_does_not_mix_projects(session):
for i in range(5):
_add_completed(session, "p1", ["python", "t.py"], 100.0, base_id=i + 1)
for i in range(5):
_add_completed(session, "p2", ["python", "t.py"], 999.0, base_id=i + 100)
pred = predict_wall(session, "p1", "python t.py")
assert isinstance(pred, WallPrediction)
assert pred.p50_s == pytest.approx(100.0)
def test_predict_wall_does_not_mix_cmd_heads(session):
for i in range(5):
_add_completed(session, "p1", ["python", "train.py"], 100.0, base_id=i + 1)
for i in range(5):
_add_completed(session, "p1", ["python", "infer.py"], 999.0, base_id=i + 100)
pred = predict_wall(session, "p1", "python train.py")
assert isinstance(pred, WallPrediction)
assert pred.p50_s == pytest.approx(100.0)
def test_predict_wall_history_window_caps_samples(session):
# Insert 2x HISTORY_WINDOW completions; only the most recent N are sampled.
n = HISTORY_WINDOW * 2
for i in range(n):
wall = 10.0 if i < HISTORY_WINDOW else 1000.0
_add_completed(session, "p1", ["python", "t.py"], wall, base_id=i + 1)
pred = predict_wall(session, "p1", "python t.py")
assert isinstance(pred, WallPrediction)
# Most recent (highest IDs) had wall=1000s
assert pred.n_samples == HISTORY_WINDOW
assert pred.p50_s == pytest.approx(1000.0)
# ---------------- remaining_for_running ----------------
def test_remaining_for_running_subtracts_elapsed():
pred = WallPrediction(p50_s=600.0, p90_s=900.0, n_samples=10, clipped=False)
started = datetime(2026, 4, 1, tzinfo=UTC)
now = started + timedelta(seconds=300)
job = Job(
project="p1",
host_pin="any",
priority=50,
state=JobState.RUNNING.value,
cmd_json="[]",
cwd="/tmp",
env_json="{}",
submitted_at=started,
started_at=started,
)
rem_p50, rem_p90 = remaining_for_running(job, pred, now)
assert rem_p50 == pytest.approx(300.0)
assert rem_p90 == pytest.approx(600.0)
def test_remaining_for_running_floors_at_zero():
pred = WallPrediction(p50_s=60.0, p90_s=120.0, n_samples=10, clipped=False)
started = datetime(2026, 4, 1, tzinfo=UTC)
now = started + timedelta(seconds=99999)
job = Job(
project="p1",
host_pin="any",
priority=50,
state=JobState.RUNNING.value,
cmd_json="[]",
cwd="/tmp",
env_json="{}",
submitted_at=started,
started_at=started,
)
rem_p50, rem_p90 = remaining_for_running(job, pred, now)
assert rem_p50 == 0.0
assert rem_p90 == 0.0
# ---------------- queue_start_eta ----------------
def test_queue_start_eta_no_competition_returns_zero(session):
cache = make_predict_cache(session)
target = Job(
id=1,
project="p1",
host_pin="any",
priority=50,
state=JobState.QUEUED.value,
cmd_json=json.dumps(["python", "t.py"]),
cwd="/tmp",
env_json="{}",
submitted_at=datetime(2026, 4, 1, tzinfo=UTC),
)
now = datetime(2026, 4, 1, tzinfo=UTC)
eta = queue_start_eta(target, [], [], cache, now)
assert eta == 0.0
def test_queue_start_eta_sums_running_remaining(session):
# Build history so prediction exists for the running job's (project, head).
for i in range(5):
_add_completed(session, "p1", ["python", "t.py"], 600.0, base_id=i + 1)
cache = make_predict_cache(session)
started = datetime(2026, 4, 1, tzinfo=UTC)
running = Job(
id=10,
project="p1",
host_pin="any",
priority=50,
state=JobState.RUNNING.value,
cmd_json=json.dumps(["python", "t.py"]),
cwd="/tmp",
env_json="{}",
submitted_at=started,
started_at=started,
)
target = Job(
id=11,
project="p1",
host_pin="any",
priority=50,
state=JobState.QUEUED.value,
cmd_json=json.dumps(["python", "t.py"]),
cwd="/tmp",
env_json="{}",
submitted_at=started,
)
now = started + timedelta(seconds=200)
eta = queue_start_eta(target, [target], [running], cache, now)
# 600 - 200 elapsed = 400s remaining on the running job.
assert eta == pytest.approx(400.0)
def test_queue_start_eta_skips_lower_priority_queued(session):
for i in range(5):
_add_completed(session, "p1", ["python", "t.py"], 100.0, base_id=i + 1)
cache = make_predict_cache(session)
target = Job(
id=20,
project="p1",
host_pin="any",
priority=50,
state=JobState.QUEUED.value,
cmd_json=json.dumps(["python", "t.py"]),
cwd="/tmp",
env_json="{}",
submitted_at=datetime(2026, 4, 1, tzinfo=UTC),
)
lower = Job(
id=21,
project="p1",
host_pin="any",
priority=10, # lower than target
state=JobState.QUEUED.value,
cmd_json=json.dumps(["python", "t.py"]),
cwd="/tmp",
env_json="{}",
submitted_at=datetime(2026, 4, 1, tzinfo=UTC),
)
now = datetime(2026, 4, 1, tzinfo=UTC)
eta = queue_start_eta(target, [target, lower], [], cache, now)
# Lower-priority queued job is behind target → not counted.
assert eta == 0.0
def test_queue_start_eta_respects_host_pin(session):
for i in range(5):
_add_completed(session, "p1", ["python", "t.py"], 600.0, base_id=i + 1)
cache = make_predict_cache(session)
started = datetime(2026, 4, 1, tzinfo=UTC)
running_other = Job(
id=30,
project="p1",
host_pin="laptop", # different specific host
priority=50,
state=JobState.RUNNING.value,
cmd_json=json.dumps(["python", "t.py"]),
cwd="/tmp",
env_json="{}",
submitted_at=started,
started_at=started,
)
target = Job(
id=31,
project="p1",
host_pin="desktop",
priority=50,
state=JobState.QUEUED.value,
cmd_json=json.dumps(["python", "t.py"]),
cwd="/tmp",
env_json="{}",
submitted_at=started,
)
eta = queue_start_eta(target, [target], [running_other], cache, started)
# Different non-any host pins → not competing → no wait.
assert eta == 0.0
def test_constants_sane():
assert MIN_SAMPLES >= 3
assert 0 < CLIP_THRESHOLD <= 1.0
assert HISTORY_WINDOW >= MIN_SAMPLES