Add GPU selection to job launcher

README.md

@@ -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

backend/app/agents/evaluation_agent.py

@@ -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

backend/tests/test_jobs.py

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

frontend/src/main.tsx

@@ -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({
+        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 }))
-          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 }
-        })
       });
       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({
+        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 }))
-          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 }
-        })
       });
       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({
+        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` }))
-          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` }
-        })
       });
       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)} />

frontend/src/styles.css

@@ -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,