From 73d15e9dce0f7828fabab2d88350ed0508fde58e Mon Sep 17 00:00:00 2001 From: admin <572701190@qq.com> Date: Tue, 30 Jun 2026 15:57:16 +0800 Subject: [PATCH] Add GPU selection to job launcher --- README.md | 4 ++ backend/app/agents/evaluation_agent.py | 9 +++ backend/tests/test_jobs.py | 13 +++++ frontend/src/main.tsx | 67 ++++++++++++++++++----- frontend/src/styles.css | 76 ++++++++++++++++++++++++++ 5 files changed, 155 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index f841f96..eac803b 100644 --- a/README.md +++ b/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 diff --git a/backend/app/agents/evaluation_agent.py b/backend/app/agents/evaluation_agent.py index 18833ed..15c12e0 100644 --- a/backend/app/agents/evaluation_agent.py +++ b/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 diff --git a/backend/tests/test_jobs.py b/backend/tests/test_jobs.py index 2b45ff4..55fd921 100644 --- a/backend/tests/test_jobs.py +++ b/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") diff --git a/frontend/src/main.tsx b/frontend/src/main.tsx index 1486aab..d02f56c 100644 --- a/frontend/src/main.tsx +++ b/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([]); const [uploadKind, setUploadKind] = useState<"images" | "labels" | "masks">("images"); const [uploadFiles, setUploadFiles] = useState(null); const [agentValidation, setAgentValidation] = useState(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) { + 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) { + const payload: { type: string; params: Record; gpus?: number[] } = { + type, + params: paramsWithGpuSelection(type, rawParams) + }; + if (selectedGpuIds.length) payload.gpus = selectedGpuIds; + return payload; + } + function datasetParamsForTask(next: string): Record { const base = { ...(catalog?.task_defaults?.[next] ?? defaultParams[next] ?? {}) }; const layout = selectedDataset?.absolute_layout; @@ -523,7 +552,7 @@ function App() { try { const job = await api("/api/jobs", { method: "POST", - body: JSON.stringify({ type: taskType, params: JSON.parse(params) }) + body: JSON.stringify(jobPayload(taskType, JSON.parse(params) as Record)) }); 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("/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("/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("/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() { ))} +
+
+ GPU 设备 + {selectedGpuIds.length ? selectedGpuIds.map((item) => `GPU ${item}`).join(", ") : "CPU"} +
+
+ + {availableGpus.map((gpu) => ( + + ))} +
+