diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 09698bf677..41c793d503 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -228,6 +228,10 @@ class DiffusionTrainingService: # A start refuses while any is open and a mutation refuses once a start is reserved, both # decided under _lock, so neither can slip through the other's check-then-act window. self._dataset_mutations = 0 + # GPU load admissions in flight (an image/video/chat load between its training guard and + # the moment it registers with the arbiter). Same two-sided rule as the dataset mutations: + # a start refuses while one is open, and an admission refuses once a start is reserved. + self._gpu_admissions = 0 self._proc: Any = None self._stop_queue: Any = None self._pump: Optional[threading.Thread] = None @@ -262,6 +266,15 @@ class DiffusionTrainingService: "The training images are being changed right now. Wait for that to finish, " "then start the run." ) + if self._gpu_admissions: + # A load already passed its training guard and is about to take the GPU. Reserving + # now would free residents it has not registered yet, so the trainer and a + # brand-new pipeline would allocate together. Refusing is safe: the admission is + # held only across the load's registration, not the load itself. + raise RuntimeError( + "A model is being loaded onto the GPU right now. Wait for that to finish, " + "then start the run." + ) self._reserved = True def unreserve(self) -> None: @@ -295,6 +308,34 @@ class DiffusionTrainingService: with self._lock: self._dataset_mutations = max(0, self._dataset_mutations - 1) + @contextlib.contextmanager + def gpu_load_admission(self): + """Hold the GPU-admission interlock across a load's guard -> arbiter -> registration. + + The load guards read ``is_active()`` and only THEN acquire the arbiter and register the + load, so a start reserving inside that gap freed residents the load had not registered + yet and the trainer came up beside a brand-new pipeline. Registering the admission under + the same lock ``reserve()`` uses closes it from both sides, exactly like + ``dataset_mutation``: this raises once a start is reserved or running, and ``reserve()`` + raises while an admission is open, so neither waits on the other. + + The span is deliberately short. ``begin_load`` returns as soon as the load is registered + (the download and build run on a daemon thread), and from that point + ``_free_gpu_for_diffusion_training`` preempts the in-flight load, so holding this for the + whole load would block starts for minutes to no purpose.""" + with self._lock: + if self._reserved or (self._proc is not None and self._proc.is_alive()): + raise TrainingActiveError( + "Diffusion training is running, so the GPU is in use. Stop the run before " + "loading a model." + ) + self._gpu_admissions += 1 + try: + yield + finally: + with self._lock: + self._gpu_admissions = max(0, self._gpu_admissions - 1) + def start(self, config: dict) -> str: """Validate ``config``, spawn the trainer, and start pumping its events. diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index dcf10573a4..48ee2c4a69 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -23,7 +23,7 @@ from loggers import get_logger import asyncio import threading import weakref -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager import re as _re @@ -16000,6 +16000,30 @@ def _diffusion_training_active() -> bool: return False +@contextmanager +def _diffusion_training_admission(): + """Hold the diffusion trainer's GPU-admission interlock for this load's registration. + + The guards below only cover the instant they run. A load then selects its engine, acquires + the arbiter and registers with the backend, and a ``/train/diffusion/start`` reserving inside + that window frees residents this load has not registered yet, so the trainer comes up beside + a brand-new pipeline. Registering the admission under the same lock ``reserve()`` takes makes + the two mutually exclusive: this raises (409) once a start is reserved, and a start raises + while an admission is open. + + Fails open on an import error, like the guards it complements. Covers the DIFFUSION trainer + only; the LLM trainer admits loads that fit beside it, which is a different contract.""" + try: + from core.training.diffusion_training_service import get_diffusion_training_service + + service = get_diffusion_training_service() + except Exception: # noqa: BLE001 -- unknowable state never blocks a load + yield + return + with service.gpu_load_admission(): + yield + + def _guard_diffusion_load_against_training() -> None: """Refuse loading an image model while a training run is active. Unlike chat, a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the @@ -16178,8 +16202,14 @@ async def load_diffusion_model( if needs_gpu: # Register the in-flight load UNDER the arbiter lock (not after acquire_for returns): otherwise a # competing Video/chat acquire in that gap evicts DIFFUSION before the load is marked in-flight, - # finds nothing to cancel, and both loaders allocate VRAM at once. - status_dict = await asyncio.to_thread(acquire_for, DIFFUSION, _begin_load) + # finds nothing to cancel, and both loaders allocate VRAM at once. The training admission wraps + # the same span for the OTHER competitor: a diffusion-training start reserving here would free + # residents this load has not registered yet (see _diffusion_training_admission). + def _acquire_and_begin(): + with _diffusion_training_admission(): + return acquire_for(DIFFUSION, _begin_load) + + status_dict = await asyncio.to_thread(_acquire_and_begin) else: # A CPU-only native load never touches the GPU, so it neither acquires nor is tracked by the # arbiter. But switching here FROM a previous diffusers/GPU load leaves DIFFUSION still marked as diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 39998e5f94..406dd81e0f 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -183,8 +183,16 @@ async def load_video_model( # Register the in-flight load UNDER the arbiter lock (not after acquire_for returns): otherwise a # competing Images/chat acquire in that gap evicts VIDEO before the load is marked in-flight, # finds nothing to cancel, and both loaders allocate VRAM at once. Mirrors the images/load - # handoff. - status_dict = await asyncio.to_thread(acquire_for, VIDEO, _begin_load) + # handoff. The training admission wraps the same span for the OTHER competitor: a + # diffusion-training start reserving here would free residents this load has not + # registered yet (see _diffusion_training_admission). + from routes.inference import _diffusion_training_admission + + def _acquire_and_begin(): + with _diffusion_training_admission(): + return acquire_for(VIDEO, _begin_load) + + status_dict = await asyncio.to_thread(_acquire_and_begin) else: await asyncio.to_thread(release, VIDEO) status_dict = await asyncio.to_thread(_begin_load) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index ed8aa10156..e2331e3b93 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -1877,3 +1877,37 @@ def test_diffusion_seed_is_bounded_to_torch_range(): # The extremes torch does accept stay valid. for good in (2**64 - 1, -(2**63)): assert DiffusionTrainingStartRequest(**request, seed = good).seed == good + + +def test_gpu_load_admission_and_reserve_exclude_each_other(): + # The load guards read is_active() and only THEN acquire the arbiter and register the load, so a + # start reserving inside that window freed residents the load had not registered yet and the + # trainer came up beside a brand-new pipeline. The admission closes it from both sides, exactly + # like the dataset interlock. + from core.training.diffusion_training_service import TrainingActiveError + + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + + # A start cannot reserve while a load is registering. + with svc.gpu_load_admission(): + with pytest.raises(RuntimeError, match = "loaded onto the GPU"): + svc.reserve() + # ...and the admission is released afterwards, so the start goes through. + svc.reserve() + try: + # A load cannot register while a start is reserved, and it says why. + with pytest.raises(TrainingActiveError, match = "Diffusion training is running"): + with svc.gpu_load_admission(): + pass + finally: + svc.unreserve() + + # Nested/concurrent admissions are counted, not boolean: the first exit must not open the door + # while a second load is still registering. + with svc.gpu_load_admission(): + with svc.gpu_load_admission(): + pass + with pytest.raises(RuntimeError, match = "loaded onto the GPU"): + svc.reserve() + svc.reserve() + svc.unreserve() diff --git a/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx b/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx index 400f0836b8..5cd566885a 100644 --- a/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx +++ b/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx @@ -1208,7 +1208,12 @@ export const IMAGE_GEN_TASKS = [ // Video-generation pipeline tasks: handled by the Video page, never loadable as // chat models. The backend reports "text-to-video" for video-diffusion GGUFs. The // Video page reuses this as its picker's `task` filter, so it lives here. -export const VIDEO_GEN_TASKS = ["text-to-video"] as const; +// image-to-video is here because that is the pipeline_tag Hugging Face gives the LTX-2 family +// (both Lightricks/LTX-2 and unsloth/LTX-2.3-GGUF report it, alongside text-to-video in their +// tag list), so a text-to-video-only filter dropped the flagship audio family out of Video Hub +// search while the rest of the app routed it to Video. Locally cached video GGUFs are unaffected +// either way: the backend tags those text-to-video itself. +export const VIDEO_GEN_TASKS = ["text-to-video", "image-to-video"] as const; // Diffusion GGUF archs the Images backend can't assemble yet (SD/SDXL/PixArt/Wan/...). The // backend tags them with this task so the chat picker hides them (they die with "unknown model @@ -2135,11 +2140,14 @@ export function HubModelPicker({ ), [cachedModels, downloadedSort, loadTimes, task, catalog], ); - // Task-scoped loads put the whole pipeline on one device, so quant fit must use the - // largest device, not the multi-GPU sum. Chat keeps the sum (llama.cpp splits layers). + // Task-scoped loads put the whole pipeline on ONE device, so quant fit must use a single + // device, not the multi-GPU sum. Specifically the device the load will land on (the lowest + // visible ordinal), not the largest one: on a heterogeneous host sizing against the bigger + // card recommends a checkpoint that then OOMs the smaller card it actually loads onto. + // Chat keeps the sum (llama.cpp splits layers). const expanderGpuGb = gpu.available ? task - ? gpu.maxDeviceMemoryGb + ? gpu.loadDeviceMemoryGb || gpu.maxDeviceMemoryGb : gpu.memoryTotalGb : undefined; diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts index 14fe9a9ac5..f3deca42fa 100644 --- a/studio/frontend/src/hooks/use-gpu-info.ts +++ b/studio/frontend/src/hooks/use-gpu-info.ts @@ -10,8 +10,14 @@ export interface GpuInfo { name: string; memoryTotalGb: number; /** Largest single device's VRAM. Image/video loads live on ONE device (no - * tensor split), so their fit math must use this, not the multi-GPU sum. */ + * tensor split), so their fit math must use a single device, not the multi-GPU sum. */ maxDeviceMemoryGb: number; + /** VRAM of the device an image/video load will actually land on: the lowest visible + * ordinal, since resolve_diffusion_device_target() returns a bare "cuda" and torch places + * on the current device. On a heterogeneous host this is NOT maxDeviceMemoryGb, and sizing + * a pick against the larger card would recommend a checkpoint that then OOMs the smaller + * one it loads onto. */ + loadDeviceMemoryGb: number; cpuCore: number; cpuThread: number; systemRamAvailableGb: number; @@ -36,6 +42,7 @@ const DEFAULT_GPU: GpuInfo = { name: "Unknown", memoryTotalGb: 0, maxDeviceMemoryGb: 0, + loadDeviceMemoryGb: 0, cpuCore: 0, cpuThread: 0, systemRamAvailableGb: 0, @@ -83,6 +90,14 @@ function toGpuInfo(data: SystemInfoResponse | null): GpuInfo { name: devices[0]?.name ?? "Unknown", memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), maxDeviceMemoryGb: devices.reduce((max, d) => Math.max(max, d.memory_total_gb ?? 0), 0), + // Lowest visible ordinal = torch's current device = where the pipeline lands. The list is + // already filtered by any CUDA_VISIBLE_DEVICES mask, so the minimum index is right whether + // the reported indices are physical or relative. + loadDeviceMemoryGb: + devices.reduce( + (pick, d) => ((d.index ?? 0) < (pick.index ?? 0) ? d : pick), + devices[0], + )?.memory_total_gb ?? 0, }; }