Add GPU selection to job launcher

This commit is contained in:
2026-06-30 15:57:16 +08:00
parent 4b3d750df9
commit 73d15e9dce
5 changed files with 155 additions and 14 deletions

View File

@@ -1,7 +1,10 @@
from fastapi.testclient import TestClient
from app import jobs
from app.commands import CommandSpec
from app.jobs import default_conda_env_for_job
from app.main import _job_with_progress, app
from app.schemas import JobCreate
def test_mmseg_jobs_use_mmseg_conda_env_by_default():
@@ -9,6 +12,16 @@ def test_mmseg_jobs_use_mmseg_conda_env_by_default():
assert default_conda_env_for_job("segmodel.train") == "seg_smp"
def test_build_task_sets_cuda_visible_devices(tmp_path, monkeypatch):
def fake_build_module_task(job_type, params, conda_env):
return CommandSpec(["python", "-c", "print('ok')"], tmp_path, "fake task")
monkeypatch.setattr(jobs, "build_module_task", fake_build_module_task)
spec = jobs._build_task(JobCreate(type="mock.echo", params={}, gpus=[2, 0]))
assert spec.env["CUDA_VISIBLE_DEVICES"] == "2,0"
def test_job_progress_reports_log_size(tmp_path):
log_path = tmp_path / "job.log"
log_path.write_text("line one\nline two\n", encoding="utf-8")