17 lines
742 B
Python
17 lines
742 B
Python
from app.scheduler.budget import Budget
|
|
from app.scheduler.planner import Planner
|
|
from app.storage.db import connect
|
|
from app.storage.repositories import Repository
|
|
|
|
|
|
def test_task_granularity_per_model_gpu(tmp_path):
|
|
conn = connect(str(tmp_path / "state.db"))
|
|
repo = Repository(conn)
|
|
repo.upsert_candidate("owner/model", "HUGGING_FACE", "manual_bounty", 0, True, target_gpu_aliases="k100,p800")
|
|
created = Planner(repo, Budget(repo, 2000), ["k100"]).materialize_tasks()
|
|
assert created == 2
|
|
rows = repo.due_tasks(("pending",), limit=10)
|
|
assert {row["gpu_type"] for row in rows} == {"hygon_k100-ai", "Kunlunxin_p-800"}
|
|
assert Planner(repo, Budget(repo, 2000), ["k100"]).materialize_tasks() == 0
|
|
conn.close()
|