25 lines
1.0 KiB
Python
25 lines
1.0 KiB
Python
|
from app.domain.gpu import expand_gpu_aliases, parse_aliases
|
||
|
|
from app.domain.model_type import engine_for_model
|
||
|
|
from app.storage.repositories import Repository
|
||
|
|
from app.scheduler.budget import Budget
|
||
|
|
|
||
|
|
|
||
|
|
class Planner:
|
||
|
|
def __init__(self, repo: Repository, budget: Budget, default_aliases: list[str]):
|
||
|
|
self.repo = repo
|
||
|
|
self.budget = budget
|
||
|
|
self.default_aliases = default_aliases
|
||
|
|
|
||
|
|
def materialize_tasks(self) -> int:
|
||
|
|
created = 0
|
||
|
|
for candidate in self.repo.list_candidates():
|
||
|
|
if engine_for_model(candidate["model_id"]) != "vllm":
|
||
|
|
continue
|
||
|
|
aliases = parse_aliases(candidate["target_gpu_aliases"], self.default_aliases)
|
||
|
|
for gpu_alias, gpu_type in expand_gpu_aliases(aliases):
|
||
|
|
if not self.budget.can_create_task():
|
||
|
|
return created
|
||
|
|
if self.repo.insert_task_if_absent(candidate["model_id"], gpu_alias, gpu_type, candidate["priority"]):
|
||
|
|
created += 1
|
||
|
|
return created
|