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

@@ -67,6 +67,10 @@ without changing the original training scripts. Starting any web job or
dataset YOLO shortcut automatically opens its live log; the SSE stream resumes
from the current log size after the initial tail so existing lines are not
duplicated in the panel.
The task builder also reads `GET /api/system/gpus` and lets an operator choose
CPU or one or more GPUs before launch. Selected GPUs are passed to the backend
as `gpus`, exported as `CUDA_VISIBLE_DEVICES`, and reflected into YOLO/visual
`device` parameters and MMSeg config-generation `gpu_count/gpu_ids`.
The coverage panel calls `GET /api/coverage` and verifies that the user-facing
scripts from the existing `Seg/` workspace are mapped to web jobs. MMSeg

View File

@@ -28,6 +28,7 @@ def evaluate_project() -> dict:
"""Return product/implementation suggestions for the current web platform."""
frontend = settings.project_root / "frontend" / "src" / "main.tsx"
backend = settings.project_root / "backend" / "app" / "main.py"
jobs_backend = settings.project_root / "backend" / "app" / "jobs.py"
readme = settings.project_root / "README.md"
catalog = get_catalog()
coverage = get_coverage_report()
@@ -36,6 +37,7 @@ def evaluate_project() -> dict:
frontend_text = frontend.read_text(encoding="utf-8") if frontend.exists() else ""
backend_text = backend.read_text(encoding="utf-8") if backend.exists() else ""
jobs_text = jobs_backend.read_text(encoding="utf-8") if jobs_backend.exists() else ""
acceptance_text = (settings.project_root / "backend" / "app" / "acceptance.py").read_text(encoding="utf-8")
readme_text = readme.read_text(encoding="utf-8") if readme.exists() else ""
@@ -59,6 +61,13 @@ def evaluate_project() -> dict:
and "setResults(resultsNext)" in frontend_text
and "slice(0, 240)" not in frontend_text,
"job_progress_ui": "JobProgressBar" in frontend_text and "progressTrack" in frontend_text,
"gpu_task_selection_ui": "selectedGpuIds" in frontend_text
and "gpuSelector" in frontend_text
and "jobPayload" in frontend_text
and "selectedGpuDevice" in frontend_text
and "gpu_count" in frontend_text
and "gpus" in frontend_text
and "CUDA_VISIBLE_DEVICES" in jobs_text,
"live_log_stream_ui": "EventSource" in frontend_text
and "eventSourceRef" in frontend_text
and "log_size" in frontend_text

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")

View File

@@ -427,6 +427,7 @@ function App() {
const [selectedCurvePath, setSelectedCurvePath] = useState("");
const [resultFamilyFilter, setResultFamilyFilter] = useState("all");
const [resultRoleFilter, setResultRoleFilter] = useState("all");
const [selectedGpuIds, setSelectedGpuIds] = useState<number[]>([]);
const [uploadKind, setUploadKind] = useState<"images" | "labels" | "masks">("images");
const [uploadFiles, setUploadFiles] = useState<FileList | null>(null);
const [agentValidation, setAgentValidation] = useState<ValidationAgentPayload | null>(null);
@@ -480,6 +481,8 @@ function App() {
};
}, [curves, results, selectedDataset]);
const selectedYoloWeightReady = Boolean(selectedYoloOutputs.bestWeight);
const availableGpus = gpus?.gpus ?? [];
const selectedGpuDevice = selectedGpuIds.length ? selectedGpuIds.map((_, index) => index).join(",") : "cpu";
const resultFamilyOptions = useMemo(() => ["all", ...Array.from(new Set(results.map((item) => item.family ?? "artifact"))).sort()], [results]);
const resultRoleOptions = useMemo(() => ["all", ...Array.from(new Set(results.map((item) => item.role ?? "artifact"))).sort()], [results]);
const filteredResults = useMemo(
@@ -501,6 +504,32 @@ function App() {
window.location.hash = "jobs";
}
function toggleGpu(index: number) {
setSelectedGpuIds((current) => current.includes(index) ? current.filter((item) => item !== index) : [...current, index].sort((a, b) => a - b));
}
function paramsWithGpuSelection(type: string, rawParams: Record<string, unknown>) {
const next = { ...rawParams };
if (!selectedGpuIds.length) return next;
if (type.startsWith("yolo.") || type.startsWith("visual.")) {
next.device = selectedGpuDevice;
}
if (type.startsWith("mmseg.generate_alg")) {
next.gpu_count = selectedGpuIds.length;
next.gpu_ids = selectedGpuIds;
}
return next;
}
function jobPayload(type: string, rawParams: Record<string, unknown>) {
const payload: { type: string; params: Record<string, unknown>; gpus?: number[] } = {
type,
params: paramsWithGpuSelection(type, rawParams)
};
if (selectedGpuIds.length) payload.gpus = selectedGpuIds;
return payload;
}
function datasetParamsForTask(next: string): Record<string, unknown> {
const base = { ...(catalog?.task_defaults?.[next] ?? defaultParams[next] ?? {}) };
const layout = selectedDataset?.absolute_layout;
@@ -523,7 +552,7 @@ function App() {
try {
const job = await api<Job>("/api/jobs", {
method: "POST",
body: JSON.stringify({ type: taskType, params: JSON.parse(params) })
body: JSON.stringify(jobPayload(taskType, JSON.parse(params) as Record<string, unknown>))
});
await inspectJob(job);
await refresh();
@@ -627,7 +656,7 @@ function App() {
const generated = await createSelectedYoloYaml();
if (!generated) return;
setTaskType("yolo.train_custom");
setParams(JSON.stringify({ model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", project: "var/custom_yolo_runs", name: selectedDataset.name, exist_ok: true }, null, 2));
setParams(JSON.stringify(paramsWithGpuSelection("yolo.train_custom", { model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: selectedGpuDevice, project: "var/custom_yolo_runs", name: selectedDataset.name, exist_ok: true }), null, 2));
await refresh();
} finally {
setBusy(false);
@@ -642,10 +671,7 @@ function App() {
if (!generated) return;
const job = await api<Job>("/api/jobs", {
method: "POST",
body: JSON.stringify({
type: "yolo.train_custom",
params: { model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", project: "var/custom_yolo_runs", name: selectedDataset.name, exist_ok: true }
})
body: JSON.stringify(jobPayload("yolo.train_custom", { model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: selectedGpuDevice, project: "var/custom_yolo_runs", name: selectedDataset.name, exist_ok: true }))
});
await inspectJob(job);
window.location.hash = "jobs";
@@ -661,10 +687,7 @@ function App() {
try {
const job = await api<Job>("/api/jobs", {
method: "POST",
body: JSON.stringify({
type: "yolo.predict_custom",
params: { weights: customYoloWeightPath(selectedDataset), source: selectedDataset.absolute_layout.images, imgsz: 640, conf: 0.25, device: "cpu", project: "var/custom_yolo_runs", name: `${selectedDataset.name}_predict`, exist_ok: true }
})
body: JSON.stringify(jobPayload("yolo.predict_custom", { weights: customYoloWeightPath(selectedDataset), source: selectedDataset.absolute_layout.images, imgsz: 640, conf: 0.25, device: selectedGpuDevice, project: "var/custom_yolo_runs", name: `${selectedDataset.name}_predict`, exist_ok: true }))
});
await inspectJob(job);
window.location.hash = "jobs";
@@ -680,10 +703,7 @@ function App() {
try {
const job = await api<Job>("/api/jobs", {
method: "POST",
body: JSON.stringify({
type: "yolo.heatmap_custom",
params: { weights: customYoloWeightPath(selectedDataset), source: selectedDataset.absolute_layout.images, model_key: "YOLO11n-seg", cam_method: "GradCAM", target_layers: "model.model.model[9]", limit: 3, project: "var/custom_yolo_runs", name: `${selectedDataset.name}_heatmap` }
})
body: JSON.stringify(jobPayload("yolo.heatmap_custom", { weights: customYoloWeightPath(selectedDataset), source: selectedDataset.absolute_layout.images, model_key: "YOLO11n-seg", cam_method: "GradCAM", target_layers: "model.model.model[9]", limit: 3, device: selectedGpuDevice, project: "var/custom_yolo_runs", name: `${selectedDataset.name}_heatmap` }))
});
await inspectJob(job);
window.location.hash = "jobs";
@@ -850,6 +870,25 @@ function App() {
</div>
))}
</div>
<div className="gpuSelector">
<div>
<span>GPU </span>
<strong>{selectedGpuIds.length ? selectedGpuIds.map((item) => `GPU ${item}`).join(", ") : "CPU"}</strong>
</div>
<div className="gpuPicker">
<button type="button" className={!selectedGpuIds.length ? "selected" : ""} onClick={() => setSelectedGpuIds([])}>
<Cpu size={14} />
<span>CPU</span>
</button>
{availableGpus.map((gpu) => (
<button key={gpu.index} type="button" className={selectedGpuIds.includes(gpu.index) ? "selected" : ""} onClick={() => toggleGpu(gpu.index)}>
<Cpu size={14} />
<span>GPU {gpu.index}</span>
<small>{gpu.memory_free_mb} MB</small>
</button>
))}
</div>
</div>
<label className="field">
<span> JSON</span>
<textarea value={params} onChange={(event) => setParams(event.target.value)} />

View File

@@ -398,6 +398,81 @@ h2 {
background: rgba(157, 226, 111, 0.08);
}
.gpuSelector {
display: grid;
gap: 10px;
margin-top: 14px;
padding: 10px;
border: 1px solid var(--line);
border-radius: 7px;
background: #101310;
}
.gpuSelector > div:first-child {
min-width: 0;
display: grid;
grid-template-columns: minmax(0, 1fr) auto;
align-items: center;
gap: 8px;
}
.gpuSelector span {
color: var(--muted);
font-size: 12px;
}
.gpuSelector strong {
min-width: 0;
color: var(--green);
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
.gpuPicker {
display: grid;
grid-template-columns: repeat(4, minmax(0, 1fr));
gap: 8px;
}
.gpuPicker button {
min-width: 0;
min-height: 38px;
display: grid;
grid-template-columns: auto minmax(0, 1fr);
align-items: center;
gap: 6px;
padding: 7px 8px;
border-radius: 6px;
border: 1px solid var(--line);
background: #080a08;
color: var(--muted);
}
.gpuPicker button.selected {
border-color: var(--green);
color: var(--ink);
background: rgba(157, 226, 111, 0.08);
}
.gpuPicker button svg {
color: var(--green);
}
.gpuPicker button span,
.gpuPicker button small {
min-width: 0;
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
.gpuPicker button small {
grid-column: 2;
color: var(--muted);
font-size: 10px;
}
.field {
display: grid;
gap: 8px;
@@ -1373,6 +1448,7 @@ meter {
.sampleStrip,
.datasetOutputLinks,
.datasetOutputPreview,
.gpuPicker,
.taskCheckList,
.agentCheckList,
.qualityStats,