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

@@ -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)} />