Address review findings: dataset preflight, sd.cpp unload barrier, caption tombstone, ControlNet cache race

- Run the trainer's caption discovery in the start route BEFORE freeing GPU
  residents, so a missing or uncaptionable dataset 400s without evicting the
  loaded chat/Images model.
- sd.cpp unload now waits out a cancelled one-shot generation on the generate
  lock before reporting the device free, matching the diffusers backend.
- Clearing a caption that came from metadata.jsonl writes an empty sidecar
  tombstone instead of unlinking (both readers treat an existing sidecar as
  authoritative), so the cleared label cannot resurface.
- The ControlNet wrapper pipe is only cached while its load is still current,
  closing the unload race the model cache already handled.

studio/backend/core/inference/diffusion.py

@@ -1582,7 +1582,14 @@ class DiffusionBackend:
             pipe = getattr(diffusers, pipe_cls_name).from_pipe(
                 state.pipe, controlnet = cn_model, torch_dtype = None
             )
-            self._cn_pipes[key] = pipe
+            with self._lock:
+                # Same race as the model cache above: an unload/superseding load may
+                # have cleared _cn_pipes while from_pipe ran; caching now would pin a
+                # pipeline built around the UNLOADED base and hand it to the next load.
+                if cancel.is_set() or self._state is not state:
+                    del pipe
+                    raise RuntimeError(DIFFUSION_CANCELLED_MSG)
+                self._cn_pipes[key] = pipe
         return pipe
 
     @staticmethod

studio/backend/core/inference/sd_cpp_backend.py

@@ -1096,6 +1096,13 @@ class SdCppDiffusionBackend:
             state.server.stop()
         if pending is not None and pending is not (state.server if state else None):
             pending.stop()
+        # Wait for a signalled one-shot generation to actually exit before reporting
+        # unloaded: callers (the GPU arbiter, training cleanup) treat this return as
+        # "the device is free", but a one-shot sd-cli child killed by the cancel above
+        # unwinds under _generate_lock. A bare acquire is the exit barrier (never taken
+        # while holding _lock; same pattern as DiffusionBackend.unload).
+        with self._generate_lock:
+            pass
         return self.status()
 
     def status(self) -> dict[str, Any]:

studio/backend/routes/training.py

@@ -1234,6 +1234,21 @@ async def start_diffusion_training(
     # user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
     _preflight_gated_base(config.get("base_model", ""), config.get("hf_token"))
 
+    # Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise
+    # fails inside the spawned trainer AFTER the user's chat/Images model was
+    # evicted. Same discovery the trainer runs, so the two cannot disagree.
+    from core.training import diffusion_train_common as _dtc
+
+    try:
+        await asyncio.to_thread(
+            _dtc.discover_image_caption_pairs,
+            config["data_dir"],
+            instance_prompt = config.get("instance_prompt") or None,
+            caption_column = config.get("caption_column") or "text",
+        )
+    except (FileNotFoundError, ValueError) as e:
+        raise HTTPException(status_code = 400, detail = str(e))
+
     # Free resident GPU workloads (export / Images pipeline / chat) before the trainer
     # loads its own pipeline.
     _free_gpu_for_diffusion_training()
@@ -1636,12 +1651,24 @@ async def set_diffusion_dataset_caption(
         sidecar = image_path.with_suffix(".txt")
         if caption:
             sidecar.write_text(caption, encoding = "utf-8")
-        else:
-            # Blank clears the sidecar; also drop a stale .caption so the image reads as
-            # uncaptioned afterwards.
-            sidecar.unlink(missing_ok = True)
             image_path.with_suffix(".caption").unlink(missing_ok = True)
-        return _image_record(folder, image_path, _load_metadata_captions(folder))
+            return _image_record(folder, image_path, _load_metadata_captions(folder))
+        # Blank must actually clear. Unlinking alone would resurface this image's
+        # metadata.jsonl / captions.jsonl caption (the fallback source), so when one
+        # exists write an EMPTY sidecar instead: both the record reader and the
+        # trainer's discovery treat an existing sidecar as authoritative even when
+        # empty, which makes it a tombstone. No metadata caption -> plain cleanup.
+        meta = _load_metadata_captions(folder)
+        try:
+            rel = image_path.relative_to(folder).as_posix()
+        except ValueError:
+            rel = image_path.name
+        if image_path.name in meta or rel in meta:
+            sidecar.write_text("", encoding = "utf-8")
+        else:
+            sidecar.unlink(missing_ok = True)
+        image_path.with_suffix(".caption").unlink(missing_ok = True)
+        return _image_record(folder, image_path, meta)
 
     return await asyncio.to_thread(write)
 

studio/backend/tests/test_diffusion_controlnet.py

@@ -239,6 +239,9 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
     monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
     b = DiffusionBackend()
     st = _state()
+    # The pipe cache only commits while ``st`` is the CURRENT load (an unload racing
+    # from_pipe must not repopulate the cache), so mirror the loaded invariant.
+    b._state = st
     resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
     p1 = b._controlnet_pipe(st, resolved, threading.Event())
     assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel)
@@ -259,3 +262,18 @@ def test_controlnet_pipe_rejects_family_without_classes():
     st.family.controlnet_pipeline_class = None
     with pytest.raises(ValueError, match = "not supported"):
         b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event())
+
+def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
+    # An unload that lands while from_pipe is assembling must not let the wrapper
+    # repopulate the cache around the torn-down base pipe.
+    import threading
+
+    from core.inference.diffusion import DiffusionBackend
+
+    monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
+    b = DiffusionBackend()
+    st = _state()  # never committed to b._state: the load is already gone
+    resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
+    with pytest.raises(RuntimeError, match = "cancelled"):
+        b._controlnet_pipe(st, resolved, threading.Event())
+    assert b._cn_pipes == {}

studio/backend/tests/test_diffusion_training.py

@@ -265,6 +265,13 @@ def client(monkeypatch):
 
     monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False))
     monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
+    # The dataset preflight runs the trainer's discovery against _BODY's fake
+    # data_dir; stub it here so wiring tests pass, and let the dedicated preflight
+    # tests below re-point it at a real tmp dataset.
+    monkeypatch.setattr(
+        "core.training.diffusion_train_common.discover_image_caption_pairs",
+        lambda data_dir, **kw: [("img.png", "caption")],
+    )
     app = FastAPI()
     app.include_router(training_router, prefix = "/api/train")
     app.dependency_overrides[get_current_subject] = lambda: "test-user"