From fb94a793375c59e7ec96c39faadf9c0108985db6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 02:16:56 +0000 Subject: [PATCH] Studio: fix dataset upload data loss, caption over-count, and custom-root sd.cpp uninstall - Diffusion dataset upload now streams each file into a sibling temp file and atomically os.replace()s it into place only after the whole file is written and within the size cap. A mid-batch 413 (or any abort) removes the temp, never an example already stored under the same name, so re-uploading a too-large batch can no longer truncate or delete a previously uploaded image. - _diffusion_dataset_summary counts an image as captioned only when it resolves to a non-empty caption via the same sidecar-over-metadata precedence the trainer uses. An empty (tombstone) sidecar shadows a metadata row and makes the trainer skip the image, so counting it over-reported caption_count and mislabeled an effectively uncaptioned dataset as captioned. - uninstall.sh/.ps1 now remove a custom/env-mode Studio's native diffusion build that installs beside the root as a stable-diffusion.cpp sibling (find_sd_cpp_binary resolves it from the Studio home's parent), guarded by the same unsafe-path check, and stop processes locking the default-mode stable-diffusion.cpp before removing it. Adds regression tests for the upload data-loss and caption-count paths and a hermetic shell test for the custom-root stable-diffusion.cpp removal. --- scripts/uninstall.ps1 | 12 ++- scripts/uninstall.sh | 10 +++ studio/backend/routes/training.py | 53 ++++++++++-- .../backend/tests/test_diffusion_training.py | 56 ++++++++++++ tests/sh/test_uninstall_sd_cpp_custom_root.sh | 85 +++++++++++++++++++ 5 files changed, 206 insertions(+), 10 deletions(-) create mode 100755 tests/sh/test_uninstall_sd_cpp_custom_root.sh diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 5d60bacf90..10c71ecb21 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -366,7 +366,7 @@ function Uninstall-UnslothStudio { _StopStudioProcesses -KnownRoots $knownRoots # Also stop anything holding a handle on the exact paths we delete (llama-server, # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused. - _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode)) + _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultSdCpp, $defaultCache, $defaultNode)) # ── Remove custom-root install trees ── _Step "Removing data and install directories..." @@ -380,6 +380,16 @@ function Uninstall-UnslothStudio { continue } _RemovePath $r + # Native diffusion (stable-diffusion.cpp) for a custom/env-mode Studio installs beside + # the root at \stable-diffusion.cpp -- find_sd_cpp_binary resolves it from + # UNSLOTH_STUDIO_HOME.parent (sd_cpp_engine.py) -- so removing only the root leaves the + # build behind. Derive and remove the sibling, guarding the parent path the same way. + $customSdCpp = Join-Path (Split-Path -LiteralPath $r -Parent) "stable-diffusion.cpp" + if (_IsUnsafeRoot $customSdCpp) { + _Substep "refusing to remove unsafe path: $customSdCpp" "Yellow" + } else { + _RemovePath $customSdCpp + } } # Default install dir (always at %USERPROFILE%\.unsloth\studio when present). if ($defaultStudioHome) { _RemovePath $defaultStudioHome } diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index f68fd37131..4e3e34136b 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -210,6 +210,16 @@ _custom_studio_roots | while IFS= read -r _custom_root; do continue fi _remove_path "$_custom_root" + # Native diffusion (stable-diffusion.cpp) for a custom/env-mode Studio installs beside + # the root at /stable-diffusion.cpp -- find_sd_cpp_binary resolves it from + # UNSLOTH_STUDIO_HOME.parent (sd_cpp_engine.py) -- so removing only the root leaves the + # build behind. Derive and remove the sibling, guarding the parent path the same way. + _custom_sd_cpp="$(dirname "$_custom_root")/stable-diffusion.cpp" + if _is_unsafe_root "$_custom_sd_cpp"; then + echo " refusing to remove unsafe path: $_custom_sd_cpp" >&2 + else + _remove_path "$_custom_sd_cpp" + fi done _remove_path "$HOME/.unsloth/studio" # Default-mode shared llama.cpp build + cache are siblings of studio (not removed diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c9197559c9..318de3d85d 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1388,20 +1388,45 @@ _DIFFUSION_DATASET_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} _DIFFUSION_DATASET_TEXT_EXTS = {".txt", ".caption", ".jsonl"} +def _resolve_dataset_caption( + folder: Path, image_path: Path, meta_captions: dict[str, str] +) -> Optional[str]: + """Resolve an image's caption using the same sidecar > metadata precedence the trainer + applies in ``discover_image_caption_pairs``. A per-image .txt/.caption sidecar wins and + is stripped, so an empty (tombstone) sidecar shadows metadata and yields "" -- the + trainer then skips that image (``if caption:``), so it must not count as captioned.""" + caption: Optional[str] = None + for ext in (".txt", ".caption"): + sidecar = image_path.with_suffix(ext) + if sidecar.is_file(): + try: + caption = sidecar.read_text(encoding = "utf-8").strip() + except OSError: + caption = None + break + if caption is None: + try: + rel = image_path.relative_to(folder).as_posix() + except ValueError: + rel = None + caption = meta_captions.get(image_path.name) or ( + meta_captions.get(rel) if rel is not None else None + ) + return caption + + def _diffusion_dataset_summary(folder: Path) -> DiffusionDatasetSummary: - # Count an image as captioned when a metadata/captions.jsonl row or a per-image - # sidecar (.txt / .caption) resolves a caption for it -- the same sources the - # trainer reads. Counting metadata-only captions here keeps a metadata-captioned - # dataset from reporting caption_count=0 and being treated as uncaptioned. + # Count an image as captioned only when it resolves to a NON-EMPTY caption via the same + # sidecar > metadata precedence the trainer uses -- an empty tombstone sidecar shadows a + # metadata row and makes the trainer skip the image, so counting it here would over-report + # caption_count and mislabel an effectively-uncaptioned dataset as captioned. meta_captions = _load_metadata_captions(folder) images = captions = 0 for f in folder.iterdir(): if not f.is_file() or f.suffix.lower() not in _DIFFUSION_DATASET_IMAGE_EXTS: continue images += 1 - if f.name in meta_captions or any( - f.with_suffix(ext).is_file() for ext in (".txt", ".caption") - ): + if _resolve_dataset_caption(folder, f, meta_captions): captions += 1 return DiffusionDatasetSummary( name = folder.name, path = str(folder), image_count = images, caption_count = captions @@ -1475,6 +1500,9 @@ async def upload_diffusion_dataset( named folder under the Studio datasets root, creating it if needed. Repeat uploads into the same name accumulate, so large datasets can arrive in batches. The returned name can be passed directly as ``data_dir`` to /diffusion/start.""" + import os + import tempfile + from utils.paths import datasets_root from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label @@ -1496,9 +1524,15 @@ async def upload_diffusion_dataset( detail = f"Unsupported file '{f.filename}'. Allowed: {exts}", ) dest = folder / filename + # Stream into a sibling temp file and only atomically promote it once the whole file + # is written and within the limit. A mid-stream 413 (or any abort) then removes the + # TEMP file, never dest, so re-uploading a batch that trips the limit can no longer + # truncate/delete an example that was already stored under the same name. + fd, tmp_name = tempfile.mkstemp(dir = folder, prefix = ".upload-", suffix = ext) + tmp = Path(tmp_name) complete = False try: - with open(dest, "wb") as out: + with os.fdopen(fd, "wb") as out: while chunk := await f.read(1024 * 1024): total_bytes += len(chunk) if total_bytes > limit_bytes: @@ -1511,11 +1545,12 @@ async def upload_diffusion_dataset( ), ) out.write(chunk) + os.replace(tmp, dest) complete = True finally: if not complete: try: - dest.unlink(missing_ok = True) + tmp.unlink(missing_ok = True) except OSError: pass uploaded += 1 diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 729c975d82..9f477a3545 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -698,6 +698,62 @@ def test_diffusion_dataset_upload_accumulates(client, dataset_roots): assert r.json()["image_count"] == 3 +def test_diffusion_dataset_upload_over_cap_keeps_existing_example( + client, dataset_roots, monkeypatch +): + # A re-upload that trips the size cap mid-write must not destroy an example already + # stored under the same name: the write goes to a sibling temp file and only atomically + # replaces the original on success, so the 413 leaves the prior good bytes intact. + import utils.upload_limits as ul + + ds_root, _ = dataset_roots + folder = ds_root / "my style" + folder.mkdir() + (folder / "cat.png").write_bytes(b"ORIGINAL-CAT-BYTES") + + monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 8) + monkeypatch.setattr(ul, "get_upload_limit_label", lambda: "8B") + + r = client.post( + "/api/train/diffusion/dataset", + data = {"name": "my style"}, + files = [("files", ("cat.png", b"x" * 64, "image/png"))], + ) + assert r.status_code == 413, r.text + # The pre-existing example survives untouched, and no temp file is left behind. + assert (folder / "cat.png").read_bytes() == b"ORIGINAL-CAT-BYTES" + assert sorted(p.name for p in folder.iterdir()) == ["cat.png"] + + +def test_diffusion_info_empty_sidecar_shadows_metadata_caption(client, dataset_roots): + # An empty (tombstone) .txt sidecar shadows a metadata row -- the trainer strips it and + # skips the image -- so the summary must not count it as captioned (which would report a + # dataset as captioned that the trainer would reject as having no captioned images). + ds_root, _ = dataset_roots + folder = ds_root / "tombstoned" + folder.mkdir() + (folder / "a.png").write_bytes(b"x") + (folder / "b.png").write_bytes(b"x") + (folder / "c.png").write_bytes(b"x") + (folder / "metadata.jsonl").write_text( + json.dumps({"file_name": "a.png", "text": "cap a"}) + + "\n" + + json.dumps({"file_name": "c.png", "text": "cap c"}) + + "\n", + encoding = "utf-8", + ) + # a.png: metadata caption but an empty sidecar tombstone -> uncaptioned. + (folder / "a.txt").write_text(" ", encoding = "utf-8") + # b.png: real sidecar caption. c.png: metadata only. Both captioned. + (folder / "b.txt").write_text("cap b", encoding = "utf-8") + + r = client.get("/api/train/diffusion/info") + assert r.status_code == 200, r.text + summary = next(d for d in r.json()["datasets"] if d["name"] == "tombstoned") + assert summary["image_count"] == 3 + assert summary["caption_count"] == 2 + + def test_diffusion_dataset_upload_rejects_traversal_names(client, dataset_roots): for bad in ("../evil", "a/b", ".hidden", " "): r = client.post( diff --git a/tests/sh/test_uninstall_sd_cpp_custom_root.sh b/tests/sh/test_uninstall_sd_cpp_custom_root.sh new file mode 100755 index 0000000000..202b046ecf --- /dev/null +++ b/tests/sh/test_uninstall_sd_cpp_custom_root.sh @@ -0,0 +1,85 @@ +#!/bin/bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit test for custom/env-mode stable-diffusion.cpp removal in scripts/uninstall.sh. +# +# A custom Studio (UNSLOTH_STUDIO_HOME=) installs its native diffusion build beside +# the root at /stable-diffusion.cpp -- find_sd_cpp_binary resolves it from +# UNSLOTH_STUDIO_HOME.parent (sd_cpp_engine.py). Uninstall must remove that sibling too, or +# a stale build lingers and a fresh install's finder can pick it up. Tested hermetically: +# the real custom-root removal loop + its helpers are extracted from uninstall.sh and run +# against per-test fixtures. Follows the extract-via-sed pattern of test_uninstall_shared_icon.sh. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +UNINSTALL_SH="$SCRIPT_DIR/../../scripts/uninstall.sh" +PASS=0 +FAIL=0 + +_TMP_ROOT=$(mktemp -d) +trap 'rm -rf "$_TMP_ROOT"' EXIT +# Deterministic deny-list checks: keep $HOME clear of the fixture trees. +HOME="$_TMP_ROOT/home" +mkdir -p "$HOME" + +assert_nodir() { _l="$1"; [ -d "$2" ] && { echo " FAIL: $_l (still present: $2)"; FAIL=$((FAIL+1)); } || { echo " PASS: $_l"; PASS=$((PASS+1)); }; } +assert_dir() { _l="$1"; [ -d "$2" ] && { echo " PASS: $_l"; PASS=$((PASS+1)); } || { echo " FAIL: $_l (missing dir $2)"; FAIL=$((FAIL+1)); }; } + +# Extract the helpers the loop depends on, plus the real custom-root removal loop. +HELPERS_FILE=$(mktemp -p "$_TMP_ROOT") +{ + sed -n '/^_remove_path() {/,/^}/p' "$UNINSTALL_SH" + sed -n '/^_is_studio_root() {/,/^}/p' "$UNINSTALL_SH" + sed -n '/^_is_unsafe_root() {/,/^}/p' "$UNINSTALL_SH" +} > "$HELPERS_FILE" +LOOP_FILE=$(mktemp -p "$_TMP_ROOT") +sed -n '/^_custom_studio_roots | while IFS= read -r _custom_root; do/,/^done/p' "$UNINSTALL_SH" > "$LOOP_FILE" + +# shellcheck disable=SC1090 +. "$HELPERS_FILE" + +# make_studio : a valid custom Studio root (share/studio.conf owner marker) plus its +# sibling /stable-diffusion.cpp build, each with a file so removal is observable. +make_studio() { + mkdir -p "$1/share" + : > "$1/share/studio.conf" + _sib="$(dirname "$1")/stable-diffusion.cpp" + mkdir -p "$_sib" + : > "$_sib/sd-cli" +} +run_loop() { + # shellcheck disable=SC1090 + . "$LOOP_FILE" +} + +# 1. Single custom root -> root AND its sibling stable-diffusion.cpp both removed. +p1="$_TMP_ROOT/inst1" +make_studio "$p1/studioA" +: > "$p1/keep.txt" # unrelated sibling file must be untouched +_custom_studio_roots() { printf '%s\n' "$p1/studioA"; } +run_loop +assert_nodir "single custom root removed" "$p1/studioA" +assert_nodir "custom-root sibling stable-diffusion.cpp removed" "$p1/stable-diffusion.cpp" +[ -f "$p1/keep.txt" ] && { echo " PASS: unrelated sibling file kept"; PASS=$((PASS+1)); } || { echo " FAIL: unrelated sibling file removed"; FAIL=$((FAIL+1)); } + +# 2. Two custom roots sharing a parent share one sd.cpp -> all removed, no error on the +# second (already-gone) removal. +p2="$_TMP_ROOT/inst2" +make_studio "$p2/studioB" +make_studio "$p2/studioC" # same parent -> same sibling sd.cpp +_custom_studio_roots() { printf '%s\n%s\n' "$p2/studioB" "$p2/studioC"; } +run_loop +assert_nodir "shared-parent root B removed" "$p2/studioB" +assert_nodir "shared-parent root C removed" "$p2/studioC" +assert_nodir "shared sibling stable-diffusion.cpp removed" "$p2/stable-diffusion.cpp" + +# 3. Default-mode sd.cpp (a bare ~/.unsloth/stable-diffusion.cpp with no custom root) is NOT +# touched by the custom-root loop -- it is removed by the separate default-mode line. +mkdir -p "$HOME/.unsloth/stable-diffusion.cpp" +_custom_studio_roots() { printf '%s\n' "$p1/studioA"; } # a now-removed root -> guard skips +run_loop +assert_dir "default-mode sd.cpp untouched by custom loop" "$HOME/.unsloth/stable-diffusion.cpp" + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" = 0 ]