Fix training start blocking, dataset resolution shadowing, duplicate uploads, partial-caption gate

Run backend.start_training off the event loop with asyncio.to_thread so the
synchronous diffusion/video unload calls (which wait on engine generation
locks) cannot freeze concurrent requests; guard against overlapping starts
with a _start_in_progress compare-and-set under the service lock.

Resolve bare diffusion dataset names directly under datasets_root() before
falling back to the generic resolver, so an unrelated LLM upload file or
recipe folder sharing the name cannot shadow the image dataset.

Reject exact duplicate filenames within one multipart upload batch: two
parts staged to the same destination would let the later tmp.replace
silently discard the earlier file. Case variants stay exempt per the
existing stem-guard contract.

Require an instance prompt in the train panel when only some images have
captions, since backend discovery silently skips uncaptioned images.

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)

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 <stem>.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

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.

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
 

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

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;
     }