diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 2ddda19951..0481bb90ba 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -739,6 +739,10 @@ class TrainingBackend: # True while a pump thread should be running; cleared on intended exits. # Left True after an abnormal death so _ensure_pump_alive spots a crash. self._pump_running: bool = False + # True from the start_training() guard passing until its spawn attempt + # finishes; blocks a second concurrent start (routes call start_training + # from a worker thread, so overlapping requests are possible). + self._start_in_progress: bool = False self._lock = threading.Lock() # Progress state (updated by pump thread from subprocess events) @@ -800,11 +804,32 @@ class TrainingBackend: still letting auto-selection place training against the freed memory. Hook failures never block the start. """ + # Compare-and-set start guard: the route runs this whole method on a worker + # thread (asyncio.to_thread), so two overlapping /train/start requests can + # reach it concurrently. Without the flag both would pass the alive-check + # below (the proc is only assigned at the end) and double-spawn. Mirrors the + # diffusion training service's reserve(). with self._lock: + if self._start_in_progress: + logger.warning("Training start already in progress") + return False if self._proc is not None and self._proc.is_alive(): logger.warning("Training subprocess already running") return False + self._start_in_progress = True + try: + return self._start_training_impl(job_id, before_spawn = before_spawn, **kwargs) + finally: + with self._lock: + self._start_in_progress = False + def _start_training_impl( + self, + job_id: str, + *, + before_spawn = None, + **kwargs, + ) -> bool: # Join prior pump thread — refuse to start if it won't die if self._pump_thread is not None and self._pump_thread.is_alive(): self._pump_thread.join(timeout = 5.0) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 0ef854cb70..01718c33ac 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -506,8 +506,18 @@ async def start_training( logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e) # The hook runs only once start guards pass -> VRAM freed iff training starts. - success = backend.start_training( - job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs + # Offloaded to a worker thread: the hook's diffusion/video unload() waits on the + # engines' generation locks until an in-flight denoise step reaches its cancel + # callback (and the export subprocess teardown can take seconds), which would + # otherwise block the event loop and freeze every concurrent status/cancel/UI + # request -- the same reason start_diffusion_training runs + # _free_gpu_for_diffusion_training via asyncio.to_thread. Overlapping starts are + # serialized by the backend's own start-in-progress guard. + success = await asyncio.to_thread( + backend.start_training, + job_id = job_id, + before_spawn = _free_vram_for_training, + **training_kwargs, ) if not success: @@ -1213,6 +1223,29 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None: return +def _resolve_diffusion_data_dir(raw: str) -> Path: + """Resolve a diffusion-training ``data_dir``. The upload/labeling routes create and + manage image datasets directly under ``datasets_root()`` and the UI passes the bare + folder name back as ``data_dir``, but the generic :func:`resolve_dataset_path` + searches the LLM uploads and recipe dataset roots FIRST -- so an unrelated upload + file or recipe folder sharing that name would shadow the just-uploaded image + dataset (preflight 400 "not a directory", or training the wrong data). Prefer the + image dataset root for a bare single-component name that exists there; everything + else (explicit "uploads/..." / "recipes/..." prefixes, absolute paths, missing + names) resolves exactly as before.""" + from utils.paths import datasets_root + + value = str(raw or "").strip() + if value and "\x00" not in value: + p = Path(value) + # Single component and not ".." -> joining under datasets_root() cannot escape it. + if not p.is_absolute() and len(p.parts) == 1 and p.parts[0] != "..": + direct = datasets_root() / value + if direct.is_dir(): + return direct + return resolve_dataset_path(raw) + + @router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse) async def start_diffusion_training( body: DiffusionTrainingStartRequest, @@ -1259,8 +1292,8 @@ async def start_diffusion_training( # trainer subprocess otherwise resolves them relative to its own cwd. config = body.model_dump() try: - from utils.paths import resolve_dataset_path, resolve_output_dir - config["data_dir"] = str(resolve_dataset_path(config["data_dir"])) + from utils.paths import resolve_output_dir + config["data_dir"] = str(_resolve_diffusion_data_dir(config["data_dir"])) config["output_dir"] = str(resolve_output_dir(config["output_dir"])) except ValueError as e: raise HTTPException(status_code = 400, detail = str(e)) @@ -1600,6 +1633,25 @@ async def upload_diffusion_dataset( status_code = 400, detail = f"Unsupported file '{f.filename}'. Allowed: {exts}", ) + # Reject an EXACT duplicate name within THIS batch (two cat.png dragged from + # different folders, or an API client repeating a part). The same-name exemption + # below exists for SEPARATE repeat uploads, where re-sending a name is a + # deliberate overwrite of the file on disk; inside one batch the two parts are + # distinct files staged to the same destination on EVERY filesystem, so the later + # tmp.replace(dest) in the commit loop would silently discard the earlier one + # while `uploaded` still counts both. Exact match only: a case VARIANT pair + # (pic.png vs Pic.png) stays exempt like the stem guard documents -- one file / + # an overwrite on case-insensitive filesystems, two files on Linux. + fname_cf = filename.casefold() + if filename in names: + raise HTTPException( + status_code = 400, + detail = ( + f"Duplicate file '{filename}' appears more than once in this upload. " + "Files sharing a name would overwrite each other; rename one before " + "uploading." + ), + ) # Reject a second IMAGE that shares this one's stem but differs by extension (sample.png # vs sample.jpg): both resolve to the same .txt caption sidecar (the kohya/diffusers # convention the reader, editor, and delete paths all use), so keeping both would silently @@ -1616,7 +1668,6 @@ async def upload_diffusion_dataset( # one caption. Casefolding the name guard too keeps a same-name case variant # (sample.png vs Sample.png, one file / an overwrite on those filesystems) exempt. stem_cf = stem.casefold() - fname_cf = filename.casefold() clash = next( ( p.name diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index 52c525d084..e3d30b93ef 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -392,6 +392,31 @@ def test_upload_same_stem_collision_within_one_batch(client, ds_root): assert "Duplicate image name" in r.json()["detail"] +def test_upload_rejects_exact_duplicate_name_within_one_batch(client, ds_root): + # Two parts with the SAME name in ONE multipart batch are distinct files (dragged from + # different folders, or an API client repeating a part); the staged commit would let the + # later tmp.replace(dest) silently discard the earlier one while `uploaded` still counts + # both. The batch must be rejected whole. Re-sending a name in a SEPARATE upload stays a + # deliberate overwrite (test_upload_allows_exact_name_overwrite_and_caption_sidecar). + r = _upload( + client, + "styleset", + [("sample.png", _png_bytes((10, 20, 30))), ("sample.png", _png_bytes((90, 90, 90)))], + ) + assert r.status_code == 400 + assert "more than once" in r.json()["detail"] + assert not (ds_root / "styleset" / "sample.png").exists() # all-or-nothing + # Caption files collide at one destination the same way. + r = _upload(client, "styleset", [("sample.txt", b"a"), ("sample.txt", b"b")]) + assert r.status_code == 400 + assert "more than once" in r.json()["detail"] + # A case VARIANT pair (Cat.png vs cat.png) stays exempt, matching the stem-guard + # contract: it is one file / an overwrite on case-insensitive filesystems and two + # files on Linux, not silent same-destination data loss. + r = _upload(client, "styleset", [("Cat.png", _png_bytes()), ("cat.png", _png_bytes())]) + assert r.status_code == 200 + + def test_upload_allows_exact_name_overwrite_and_caption_sidecar(client, ds_root): # Re-uploading the EXACT same name (stem AND extension) is an allowed overwrite, and a .txt # caption for the same stem is the intended kohya flow -- neither is a same-stem image collision. diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index fb947e6e14..46a4a16862 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -655,6 +655,30 @@ def test_route_start_rejects_uncontained_paths(client): assert r.status_code == 400 +def test_route_start_resolves_bare_name_under_image_dataset_root(client, monkeypatch, tmp_path): + # The upload/labeling routes manage image datasets directly under datasets_root() and + # the UI passes the bare folder name back as data_dir. The generic resolve_dataset_path + # searches the LLM uploads and recipe roots FIRST, so an unrelated upload file or recipe + # folder sharing the name would shadow the just-uploaded image dataset (preflight 400 + # "not a directory", or training the wrong data). The route must prefer the image + # dataset root for a bare name that exists there. + import utils.paths as up + + ds_root = tmp_path / "assets" / "datasets" + img_ds = ds_root / "my-photos" + img_ds.mkdir(parents = True) + (img_ds / "a.png").write_bytes(b"x") + # Shadowing entries the generic resolver would pick first. + (ds_root / "uploads").mkdir() + (ds_root / "uploads" / "my-photos").write_text("an LLM dataset upload, not a folder") + (ds_root / "recipes" / "my-photos").mkdir(parents = True) + monkeypatch.setattr(up, "datasets_root", lambda: ds_root) + + r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "my-photos"}) + assert r.status_code == 200, r.text + assert client._fake.started_with["data_dir"] == str(img_ds) + + def test_route_start_blocked_by_active_llm_training(client, monkeypatch): import routes.training as tr diff --git a/studio/backend/tests/test_training_start_offload.py b/studio/backend/tests/test_training_start_offload.py new file mode 100644 index 0000000000..d38c7b1aa5 --- /dev/null +++ b/studio/backend/tests/test_training_start_offload.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""/api/train/start must run backend.start_training off the event loop. + +start_training() runs the _free_vram_for_training before_spawn hook inline, and that +hook's diffusion/video unload() blocks on the engines' generation locks until an +in-flight denoise step reaches its cancel callback (seconds to tens of seconds for +video). Executed inline in the async route it would freeze every concurrent +status/cancel/UI request -- the same reason start_diffusion_training offloads +_free_gpu_for_diffusion_training via asyncio.to_thread. The backend guards the +overlapping-starts window this offload opens with a compare-and-set flag. +""" + +import asyncio +import threading + +import routes.training as tr +from models import TrainingStartRequest + + +class _FakeBackend: + def __init__(self, result = True): + self._result = result + self.start_thread = None + self.hook = None + self.current_job_id = None + + def is_training_active(self): + return False + + def start_training(self, job_id, *, before_spawn = None, **kwargs): + # The real backend runs before_spawn synchronously inside this call, so the + # thread this method runs on is the thread the blocking VRAM hook runs on. + self.start_thread = threading.current_thread() + self.hook = before_spawn + self.current_job_id = job_id + return self._result + + +def _request() -> TrainingStartRequest: + return TrainingStartRequest( + model_name = "unsloth/tiny-model", + training_type = "LoRA/QLoRA", + format_type = "alpaca", + hf_dataset = "org/data", + # Skip the YAML trust_remote_code lookup (needs the model catalog on disk). + trust_remote_code = True, + ) + + +def test_start_route_offloads_blocking_start(monkeypatch): + fake = _FakeBackend() + monkeypatch.setattr(tr, "get_training_backend", lambda: fake) + monkeypatch.setattr(tr, "_diffusion_training_active", lambda: False) + + async def _run(): + return threading.current_thread(), await tr.start_training( + request = _request(), current_subject = "test-user", via_api_key = False + ) + + loop_thread, resp = asyncio.run(_run()) + + assert resp.status == "queued", resp + # The VRAM-freeing hook was wired in and the blocking call left the loop thread. + assert fake.hook is not None + assert fake.start_thread is not None + assert fake.start_thread is not loop_thread + + +def test_backend_start_guard_blocks_overlapping_starts(): + # With the route offloaded to worker threads, two overlapping /train/start requests + # can reach TrainingBackend.start_training concurrently; the compare-and-set + # _start_in_progress flag must let exactly one of them spawn. + from core.training.training import TrainingBackend + + backend = TrainingBackend() + first_entered = threading.Event() + release_first = threading.Event() + results = {} + + def _slow_impl(job_id, *, before_spawn = None, **kwargs): + first_entered.set() + release_first.wait(timeout = 5.0) + return True + + backend._start_training_impl = _slow_impl + + def _first(): + results["first"] = backend.start_training("job-a") + + t = threading.Thread(target = _first, daemon = True) + t.start() + assert first_entered.wait(timeout = 5.0) + # Second start while the first is still inside the impl: refused by the guard, + # without ever entering the impl. + results["second"] = backend.start_training("job-b") + release_first.set() + t.join(timeout = 5.0) + + assert results["first"] is True + assert results["second"] is False + # The flag is cleared once the winning start returns, so a later start may proceed. + assert backend._start_in_progress is False diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index ed440f277f..fa11c01590 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -612,10 +612,21 @@ export function DiffusionTrainPanel({ toast.error("Name the adapter (this becomes its folder under Studio outputs)."); return; } - if (selectedDataset && selectedDataset.caption_count === 0 && !instancePrompt.trim()) { + // Require a trigger prompt whenever ANY image lacks a caption, not only when none + // have one: without an instance_prompt the backend discovery silently skips every + // uncaptioned image, so a partially captioned dataset would train on a subset. + if ( + selectedDataset && + selectedDataset.caption_count < selectedDataset.image_count && + !instancePrompt.trim() + ) { toast.error( - "These images have no captions - add a trigger prompt so the trainer knows " + - "what to learn (it becomes the caption for every image).", + selectedDataset.caption_count === 0 + ? "These images have no captions - add a trigger prompt so the trainer knows " + + "what to learn (it becomes the caption for every image)." + : `Only ${selectedDataset.caption_count} of ${selectedDataset.image_count} images ` + + "have captions - the rest would be silently skipped. Add a trigger prompt " + + "(it becomes their caption) or caption every image.", ); return; }