43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
def test_db_lifecycle_and_basic_queries(tmp_path: Path):
|
|
from argus.service.db import Db
|
|
|
|
db = Db(str(tmp_path / "mvp.sqlite3"))
|
|
db.init()
|
|
|
|
t = db.create_task(
|
|
task_id="t1",
|
|
workload="ppo",
|
|
jobspec_yaml="workload: ppo\ncode_path: /c\nmodel_id: m\ntrain_file: t\n",
|
|
nnodes=2,
|
|
n_gpus_per_node=4,
|
|
)
|
|
assert t["task_id"] == "t1"
|
|
assert db.get_task("t1") is not None
|
|
|
|
q = db.list_queue()
|
|
assert len(q["pending"]) == 1
|
|
|
|
db.set_task_state(task_id="t1", state="PENDING_RESOURCES", next_run_at="2099-01-01T00:00:00Z")
|
|
assert db.pick_next_runnable_task() is None
|
|
|
|
# next_run_at is sticky unless explicitly updated; a future value keeps it non-runnable.
|
|
db.set_task_state(task_id="t1", state="QUEUED", next_run_at=None)
|
|
assert db.pick_next_runnable_task() is None
|
|
|
|
# Allow it by setting next_run_at into the past.
|
|
db.set_task_state(task_id="t1", state="QUEUED", next_run_at="2000-01-01T00:00:00Z")
|
|
assert db.pick_next_runnable_task() is not None
|
|
|
|
db.create_attempt(task_id="t1", attempt_no=1, ray_submission_id="sid1")
|
|
db.update_attempt(task_id="t1", attempt_no=1, ray_status="RUNNING")
|
|
attempts = db.list_attempts("t1")
|
|
assert attempts[-1]["ray_status"] == "RUNNING"
|
|
|
|
# No-op update is allowed
|
|
db.update_attempt(task_id="t1", attempt_no=1)
|