Do not let a queued generation outlive the model, and three scan fixes
Five items from the latest review; four were real. An unload or arbiter eviction only cancels the generation holding _generate_lock. A second request queued behind it holds no cancel event yet, and Python locks are not FIFO, so it could take the lock the instant the active denoise released it, still see a loaded pipeline, and run a whole new denoise after the model was told to go away: the eviction then waits minutes for it and an image lands after the eject. Unload and a superseding load now raise a fence under _lock before they queue, and a generation that wins the lock while one is pending refuses instead. The cached-model scan judged pipeline completeness across every revision, so a repo holding an older complete snapshot plus a newer companion-only one read as complete while the snapshot from_pretrained actually opens has no transformer. Both scans now look at the revision the loader will open. Deleting a dataset image deleted its caption sidecar unconditionally, which for cat.jpg alongside cat.png removed the caption the survivor still resolves to. The sidecar now goes only with the last image of that stem, matching what the thumbnail cleanup beside it already did. Importing an example into a folder that holds no images but does hold files fell back to promoting the staging dir one file at a time, so an interruption left a partial dataset that the image_count check accepts as complete on retry. Those files are folded into the staging dir instead and the promotion stays a single atomic rename. The MPS generator report does not apply: torch.Generator(device="mps") has worked since PyTorch 2.0 (pytorch/pytorch#91348) and the studio installer pins torch>=2.4.
This commit is contained in:
Daniel Han
parent
5434520e67
commit
181adc6703
7 changed files with 361 additions and 26 deletions
|
|
@ -554,6 +554,13 @@ class DiffusionBackend:
|
|||
self._cancel_event = threading.Event()
|
||||
# Cancel Event of the in-flight generation; per-generation so a cancel can't be lost or leak.
|
||||
self._active_generate_cancel: Optional[threading.Event] = None
|
||||
# How many unloads / superseding loads are waiting on _generate_lock to free this pipeline.
|
||||
# A generation queued behind the active one holds no cancel event yet, so the cancel they
|
||||
# signal cannot reach it: without this fence it could win the lock as the active denoise
|
||||
# released it, see a still-loaded _state, and run a whole new denoise after the model was
|
||||
# told to go away -- stalling a chat/video GPU handoff for minutes and painting an image
|
||||
# after an eject. A count, not a flag, so concurrent teardowns each own their own release.
|
||||
self._teardown_waiters = 0
|
||||
# Written by the callback, read lock-free by generate_progress().
|
||||
self._gen: Optional[_GenState] = None
|
||||
# img2img/inpaint pipes built via from_pipe (shared modules, no extra VRAM); cleared on unload.
|
||||
|
|
@ -1368,14 +1375,22 @@ class DiffusionBackend:
|
|||
raise RuntimeError("Diffusion load was cancelled.")
|
||||
if self._active_generate_cancel is not None:
|
||||
self._active_generate_cancel.set()
|
||||
# Same fence unload() takes: a generation queued behind the active denoise would
|
||||
# otherwise slip in here and run on the pipeline this load is about to free.
|
||||
self._teardown_waiters += 1
|
||||
with self._generate_lock:
|
||||
with self._lock:
|
||||
# Re-check: a newer load/unload may have superseded this one while we waited.
|
||||
if _load_token is not None and _load_token != self._load_token:
|
||||
raise RuntimeError("Diffusion load was cancelled.")
|
||||
try:
|
||||
# Re-check: a newer load/unload may have superseded this one while we waited.
|
||||
if _load_token is not None and _load_token != self._load_token:
|
||||
raise RuntimeError("Diffusion load was cancelled.")
|
||||
|
||||
# Free the old pipeline before allocating the new one (never two in VRAM).
|
||||
self._unload_locked()
|
||||
# Free the old pipeline before allocating the new one (never two in VRAM).
|
||||
self._unload_locked()
|
||||
finally:
|
||||
# Released here, not at the end of the load: the old pipe is gone (or this load
|
||||
# bailed), and the rest of the load holds _generate_lock anyway.
|
||||
self._teardown_waiters -= 1
|
||||
|
||||
# Single-file kinds resolve a checkpoint path; the pipeline kind has none.
|
||||
single_file_path = (
|
||||
|
|
@ -2816,6 +2831,11 @@ class DiffusionBackend:
|
|||
cancel = threading.Event()
|
||||
with self._generate_lock:
|
||||
with self._lock:
|
||||
# An unload / superseding load signalled the active denoise and is waiting for this
|
||||
# lock. Python locks are not FIFO, so this request can get in first; refuse instead
|
||||
# of starting a denoise on a pipeline that is already being torn down.
|
||||
if self._teardown_waiters:
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
state = self._state
|
||||
if state is None:
|
||||
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
|
||||
|
|
@ -3280,6 +3300,9 @@ class DiffusionBackend:
|
|||
# Abort an in-flight denoise via ITS cancel event.
|
||||
if self._active_generate_cancel is not None:
|
||||
self._active_generate_cancel.set()
|
||||
# Fence queued generations too: they hold no cancel event yet, so the signal above
|
||||
# cannot reach them.
|
||||
self._teardown_waiters += 1
|
||||
# Cancel any in-flight load (its worker checks this token) and drop the marker.
|
||||
self._load_token += 1
|
||||
self._loading = None
|
||||
|
|
@ -3289,6 +3312,9 @@ class DiffusionBackend:
|
|||
with self._generate_lock:
|
||||
with self._lock:
|
||||
self._unload_locked()
|
||||
# Teardown is done: _state is None, so the fence has nothing left to protect and
|
||||
# the next generation gets the plain not-loaded message.
|
||||
self._teardown_waiters -= 1
|
||||
return self.status()
|
||||
|
||||
def _unload_locked(self) -> None:
|
||||
|
|
|
|||
|
|
@ -457,6 +457,38 @@ def is_snapshot_partial(
|
|||
)
|
||||
|
||||
|
||||
def _current_revisions(repo_info):
|
||||
"""The revisions to judge a cached repo by: just the one the loader will actually open.
|
||||
|
||||
``from_pretrained`` resolves the newest snapshot by mtime (:func:`latest_snapshot_dir`), so a
|
||||
repo cached twice -- an older complete snapshot plus a newer companion-only scoped one -- must
|
||||
be judged on the newer one. Scanning every revision let the old snapshot's denoiser satisfy the
|
||||
completeness check, so the row read as complete while the snapshot the loader picks has no
|
||||
transformer/unet: offline loads then fail and online loads silently pull the multi-GB weight.
|
||||
|
||||
Falls back to the newest revision by ``last_modified``, then to every revision, so a cache
|
||||
layout this cannot resolve behaves as before rather than reporting nothing.
|
||||
"""
|
||||
revisions = list(getattr(repo_info, "revisions", ()) or ())
|
||||
if len(revisions) <= 1:
|
||||
return revisions
|
||||
repo_path = getattr(repo_info, "repo_path", None)
|
||||
if repo_path is not None:
|
||||
latest = latest_snapshot_dir(Path(repo_path))
|
||||
if latest is not None:
|
||||
scoped = [
|
||||
rev for rev in revisions
|
||||
if getattr(rev, "snapshot_path", None) is not None
|
||||
and Path(rev.snapshot_path) == latest
|
||||
]
|
||||
if scoped:
|
||||
return scoped
|
||||
dated = [rev for rev in revisions if getattr(rev, "last_modified", None) is not None]
|
||||
if dated:
|
||||
return [max(dated, key = lambda rev: rev.last_modified)]
|
||||
return revisions
|
||||
|
||||
|
||||
def repo_has_pipeline_index(repo_info) -> bool:
|
||||
"""Whether the cached snapshot carries a ROOT model_index.json, i.e. is loadable
|
||||
as a full diffusers pipeline (from_pretrained reads only the repo root). A nested
|
||||
|
|
@ -465,7 +497,7 @@ def repo_has_pipeline_index(repo_info) -> bool:
|
|||
a name match alone would also claim nested copies -- scope by file_path when the
|
||||
scan provides it."""
|
||||
try:
|
||||
for rev in repo_info.revisions:
|
||||
for rev in _current_revisions(repo_info):
|
||||
snapshot = getattr(rev, "snapshot_path", None)
|
||||
for f in rev.files:
|
||||
name = str(getattr(f, "file_name", "") or "")
|
||||
|
|
@ -494,7 +526,7 @@ def repo_pipeline_missing_denoiser(repo_info) -> bool:
|
|||
_DENOISER_DIRS = ("transformer", "unet")
|
||||
_WEIGHT_SUFFIXES = (".safetensors", ".bin")
|
||||
try:
|
||||
for rev in repo_info.revisions:
|
||||
for rev in _current_revisions(repo_info):
|
||||
snapshot = getattr(rev, "snapshot_path", None)
|
||||
for f in rev.files:
|
||||
name = str(getattr(f, "file_name", "") or "")
|
||||
|
|
|
|||
|
|
@ -2342,8 +2342,20 @@ async def delete_diffusion_dataset_image(
|
|||
import glob as _glob
|
||||
|
||||
image_path.unlink(missing_ok = True)
|
||||
for ext in (".txt", ".caption"):
|
||||
image_path.with_suffix(ext).unlink(missing_ok = True)
|
||||
# Sidecars are keyed on the STEM, so cat.jpg and cat.png share cat.txt: both the trainer's
|
||||
# pair discovery and the labeling grid resolve either image to it. Deleting it with one of
|
||||
# them would silently strip the survivor's caption and change what the next run trains on.
|
||||
# New collisions are refused at upload, but hand-made and legacy folders still have them.
|
||||
stem_still_used = any(
|
||||
p.is_file()
|
||||
and p != image_path
|
||||
and p.stem == image_path.stem
|
||||
and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS
|
||||
for p in folder.iterdir()
|
||||
)
|
||||
if not stem_still_used:
|
||||
for ext in (".txt", ".caption"):
|
||||
image_path.with_suffix(ext).unlink(missing_ok = True)
|
||||
thumbs_dir = folder / _THUMBS_DIRNAME
|
||||
if thumbs_dir.is_dir():
|
||||
# Thumbs are keyed on the full filename (stem + extension), so match that here too; a stem-only
|
||||
|
|
@ -2681,18 +2693,39 @@ async def import_diffusion_dataset_example(
|
|||
status_code = 502,
|
||||
detail = f"No images found in '{entry['repo']}'.",
|
||||
)
|
||||
# Promote the fully-materialized staging dir as a UNIT. A per-file move loop is not atomic: a hard
|
||||
# process death mid-loop would leave SOME images, which the image_count>0 idempotency check would
|
||||
# accept as complete on retry. The folder was created empty here, so a single same-filesystem
|
||||
# rename is atomic. If it holds unrelated non-image files (rmdir refuses), fall back to a per-file
|
||||
# move rather than abort.
|
||||
# Promote the fully-materialized staging dir as a UNIT: a same-filesystem rename is
|
||||
# atomic, so a hard process death leaves either the old folder or the finished
|
||||
# import, never a half-filled one that the image_count>0 check above would accept as
|
||||
# complete on retry. rmdir needs an empty target, and an image-empty folder can still
|
||||
# hold files (a .thumbs cache, or a metadata.jsonl / captions from an earlier
|
||||
# upload), so fold those INTO the staging dir first and keep one atomic promotion.
|
||||
# Moving them one by one into a live folder instead -- the old fallback -- gave up
|
||||
# exactly the atomicity this whole staging dance exists for.
|
||||
for p in sorted(folder.iterdir()):
|
||||
dest = staging / p.name
|
||||
if dest.exists():
|
||||
# Same name in both: the import's own file wins, exactly as the previous
|
||||
# per-file move did by overwriting it. Drop the old one so the folder can
|
||||
# still be emptied for the rename.
|
||||
if p.is_dir():
|
||||
shutil.rmtree(p, ignore_errors = True)
|
||||
else:
|
||||
p.unlink(missing_ok = True)
|
||||
continue
|
||||
shutil.move(str(p), str(dest))
|
||||
try:
|
||||
os.rmdir(folder)
|
||||
except OSError:
|
||||
for p in staging.iterdir():
|
||||
shutil.move(str(p), str(folder / p.name))
|
||||
else:
|
||||
os.replace(str(staging), str(folder))
|
||||
except OSError as e:
|
||||
# Something landed in the folder in the meantime. Fail with the dataset
|
||||
# untouched rather than promoting it piecemeal.
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
f"'{folder.name}' changed while the example was being imported "
|
||||
f"({e.strerror or e}). Nothing was written; try again."
|
||||
),
|
||||
)
|
||||
os.replace(str(staging), str(folder))
|
||||
finally:
|
||||
shutil.rmtree(staging, ignore_errors = True)
|
||||
return _import_response(entry, folder, imported = imported)
|
||||
|
|
|
|||
|
|
@ -1654,6 +1654,56 @@ def test_repo_has_pipeline_index_requires_root_model_index(tmp_path):
|
|||
assert models_route._repo_has_pipeline_index(repo_root) is True
|
||||
|
||||
|
||||
def test_pipeline_scans_read_the_snapshot_the_loader_will_open(tmp_path):
|
||||
# A repo cached twice -- an older complete snapshot plus a newer companion-only one, the shape a
|
||||
# GGUF load leaves when it prefetches the base repo's VAE / text encoder and skips the
|
||||
# transformer -- must be judged on the snapshot from_pretrained resolves, i.e. the newest by
|
||||
# mtime. Scanning every revision let the OLD snapshot's transformer satisfy completeness, so the
|
||||
# row read as on-device while the load would fail offline or silently pull multi-GB weights.
|
||||
import os
|
||||
|
||||
import hub.utils.inventory_scan as scan
|
||||
|
||||
repo_dir = tmp_path / "models--Org--Repo"
|
||||
old_snap = repo_dir / "snapshots" / "old"
|
||||
new_snap = repo_dir / "snapshots" / "new"
|
||||
for d in (old_snap / "transformer", new_snap / "vae"):
|
||||
d.mkdir(parents = True)
|
||||
(old_snap / "model_index.json").write_text("{}", encoding = "utf-8")
|
||||
(new_snap / "model_index.json").write_text("{}", encoding = "utf-8")
|
||||
# Make "new" unambiguously newer than "old" for the mtime rule both this and the loader use.
|
||||
os.utime(old_snap, (1_000_000, 1_000_000))
|
||||
os.utime(new_snap, (2_000_000, 2_000_000))
|
||||
|
||||
def _rev(snap, files):
|
||||
return SimpleNamespace(
|
||||
snapshot_path = snap,
|
||||
last_modified = float(snap.stat().st_mtime),
|
||||
files = [
|
||||
SimpleNamespace(file_name = Path(f).name, file_path = snap / f) for f in files
|
||||
],
|
||||
)
|
||||
|
||||
info = SimpleNamespace(
|
||||
repo_id = "Org/Repo",
|
||||
repo_path = repo_dir,
|
||||
revisions = [
|
||||
_rev(old_snap, ["model_index.json", "transformer/diffusion_pytorch_model.safetensors"]),
|
||||
_rev(new_snap, ["model_index.json", "vae/diffusion_pytorch_model.safetensors"]),
|
||||
],
|
||||
)
|
||||
assert scan.repo_has_pipeline_index(info) is True
|
||||
assert scan.repo_pipeline_missing_denoiser(info) is True
|
||||
|
||||
# The reverse cache (the complete snapshot is the newer one) still reports complete.
|
||||
os.utime(old_snap, (3_000_000, 3_000_000))
|
||||
info.revisions = [
|
||||
_rev(old_snap, ["model_index.json", "transformer/diffusion_pytorch_model.safetensors"]),
|
||||
_rev(new_snap, ["model_index.json", "vae/diffusion_pytorch_model.safetensors"]),
|
||||
]
|
||||
assert scan.repo_pipeline_missing_denoiser(info) is False
|
||||
|
||||
|
||||
def test_list_cached_models_flags_single_file_diffusion_repos(monkeypatch, tmp_path):
|
||||
# A diffusion-tagged repo with NO top-level model_index.json is a single-file checkpoint, so it
|
||||
# carries single_file=True; a full pipeline repo and a chat repo carry no flag.
|
||||
|
|
|
|||
|
|
@ -3985,3 +3985,82 @@ def test_download_plan_keeps_the_dense_encoder_when_the_precast_repo_is_unavaila
|
|||
base = next(e for e in plan["entries"] if e["repo_id"] == "black-forest-labs/FLUX.1-dev")
|
||||
assert "text_encoder/model.safetensors" in base["files"]
|
||||
assert not any(e["repo_id"] == "unsloth/does-not-exist" for e in plan["entries"])
|
||||
|
||||
|
||||
# ── teardown fence ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_unload_fences_queued_generations_while_it_waits(fake_runtime, tmp_path):
|
||||
# A generation queued behind the active one holds no cancel event, so unload's signal cannot
|
||||
# reach it. Python locks are not FIFO, so when the active denoise released _generate_lock the
|
||||
# queued request could get in ahead of the unload, see the still-loaded pipeline, and run a
|
||||
# whole new denoise after the model was told to go away.
|
||||
(tmp_path / "model.gguf").write_bytes(b"weights")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo",
|
||||
family_override = "z-image",
|
||||
)
|
||||
|
||||
seen: list[int] = []
|
||||
real_unload_locked = backend._unload_locked
|
||||
|
||||
def _record_then_unload():
|
||||
# Sampled at the moment unload holds both locks, i.e. exactly the window a queued
|
||||
# generation could have slipped through.
|
||||
seen.append(backend._teardown_waiters)
|
||||
real_unload_locked()
|
||||
|
||||
backend._unload_locked = _record_then_unload
|
||||
backend.unload()
|
||||
|
||||
assert seen == [1] # the fence was up for the whole wait
|
||||
assert backend._teardown_waiters == 0 # and released once the pipeline was gone
|
||||
|
||||
|
||||
def test_generation_refuses_while_a_teardown_is_waiting(fake_runtime, tmp_path):
|
||||
# The fence's effect: with a teardown waiting on _generate_lock, a generation that wins the
|
||||
# lock refuses instead of denoising on a pipeline that is being freed.
|
||||
(tmp_path / "model.gguf").write_bytes(b"weights")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo",
|
||||
family_override = "z-image",
|
||||
)
|
||||
assert backend.generate(prompt = "before", steps = 2)["images"]
|
||||
|
||||
backend._teardown_waiters = 1
|
||||
with pytest.raises(RuntimeError, match = "cancelled"):
|
||||
backend.generate(prompt = "during", steps = 2)
|
||||
# Still loaded: the refusal is about the pending teardown, not a missing model.
|
||||
assert backend._state is not None
|
||||
|
||||
backend._teardown_waiters = 0
|
||||
assert backend.generate(prompt = "after", steps = 2)["images"]
|
||||
|
||||
|
||||
def test_a_superseding_load_fences_queued_generations_too(fake_runtime, tmp_path):
|
||||
# begin_load frees the old pipeline behind the same barrier, so it needs the same fence: a
|
||||
# queued generation would otherwise run on the pipe the new load is about to drop.
|
||||
(tmp_path / "model.gguf").write_bytes(b"weights")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo",
|
||||
family_override = "z-image",
|
||||
)
|
||||
|
||||
seen: list[int] = []
|
||||
real_unload_locked = backend._unload_locked
|
||||
|
||||
def _record_then_unload():
|
||||
seen.append(backend._teardown_waiters)
|
||||
real_unload_locked()
|
||||
|
||||
backend._unload_locked = _record_then_unload
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo",
|
||||
family_override = "z-image",
|
||||
)
|
||||
|
||||
assert seen == [1]
|
||||
assert backend._teardown_waiters == 0
|
||||
|
|
|
|||
|
|
@ -215,6 +215,47 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
|
|||
assert not list((folder / ".thumbs").glob("x.png_*.jpg"))
|
||||
|
||||
|
||||
def test_delete_keeps_a_caption_a_same_stem_sibling_still_uses(client, ds_root):
|
||||
# cat.jpg and cat.png share cat.txt: the trainer's pair discovery and the labeling grid both
|
||||
# resolve either image to it. Deleting one image must not strip the survivor's caption and
|
||||
# silently change what the next run trains on.
|
||||
folder = ds_root / "d"
|
||||
folder.mkdir()
|
||||
_write_png(folder / "cat.png")
|
||||
Image.new("RGB", (8, 8), (9, 9, 9)).save(folder / "cat.jpg", format = "JPEG")
|
||||
(folder / "cat.txt").write_text("a cat", encoding = "utf-8")
|
||||
|
||||
r = client.delete("/api/train/diffusion/dataset/d/image/cat.jpg")
|
||||
assert r.status_code == 200, r.text
|
||||
assert not (folder / "cat.jpg").exists()
|
||||
assert (folder / "cat.txt").read_text(encoding = "utf-8") == "a cat"
|
||||
# The survivor is still reported as captioned.
|
||||
info = client.get("/api/train/diffusion/info").json()
|
||||
row = [d for d in info["datasets"] if d["name"] == "d"][0]
|
||||
assert (row["image_count"], row["caption_count"]) == (1, 1)
|
||||
|
||||
# Deleting the last image with that stem does take the sidecar.
|
||||
r = client.delete("/api/train/diffusion/dataset/d/image/cat.png")
|
||||
assert r.status_code == 200, r.text
|
||||
assert not (folder / "cat.txt").exists()
|
||||
|
||||
|
||||
def test_delete_removes_a_caption_no_other_image_shares(client, ds_root):
|
||||
# A same-stem NON-image file must not keep the sidecar alive.
|
||||
folder = ds_root / "d"
|
||||
folder.mkdir()
|
||||
_write_png(folder / "cat.png")
|
||||
(folder / "cat.txt").write_text("a cat", encoding = "utf-8")
|
||||
(folder / "cat.caption").write_text("also a cat", encoding = "utf-8")
|
||||
(folder / "cat.json").write_text("{}", encoding = "utf-8")
|
||||
|
||||
r = client.delete("/api/train/diffusion/dataset/d/image/cat.png")
|
||||
assert r.status_code == 200, r.text
|
||||
assert not (folder / "cat.txt").exists()
|
||||
assert not (folder / "cat.caption").exists()
|
||||
assert (folder / "cat.json").exists()
|
||||
|
||||
|
||||
def test_thumb_cache_key_distinguishes_same_stem_extensions(client, ds_root):
|
||||
# sample.png and sample.jpg share a stem; each must get its OWN thumbnail cache file, so the grid
|
||||
# never serves one image's thumbnail for the other.
|
||||
|
|
@ -714,13 +755,18 @@ def test_delete_image_with_glob_chars_only_removes_own_thumbs(client, ds_root):
|
|||
|
||||
|
||||
def test_import_preserves_unrelated_files_when_folder_not_empty(client, ds_root, monkeypatch):
|
||||
# If the target folder already holds unrelated NON-image files (so image_count is still 0 and the
|
||||
# import runs), the atomic rmdir refuses and the code falls back to a per-file move: the images
|
||||
# are imported AND the pre-existing file is preserved.
|
||||
# A target folder holding unrelated NON-image files still has image_count 0, so the import runs.
|
||||
# Those files are folded into the staging dir and promoted with it: the images are imported, the
|
||||
# pre-existing file survives, and the promotion stays a single atomic rename (it used to fall
|
||||
# back to a per-file move into the live folder, which is what makes a partial import possible).
|
||||
_install_fake_load_dataset(monkeypatch, n_rows = 3)
|
||||
folder = ds_root / "my-tux"
|
||||
folder.mkdir(parents = True)
|
||||
(folder / "notes.md").write_text("keep me", encoding = "utf-8")
|
||||
# The promoted folder is the staging dir renamed into place, so its inode changes. A per-file
|
||||
# move into the live folder would keep the original directory, which is how this pins that the
|
||||
# promotion really was one atomic rename rather than a loop.
|
||||
inode_before = folder.stat().st_ino
|
||||
r = client.post(
|
||||
"/api/train/diffusion/dataset/import-example",
|
||||
json = {"id": "tuxemon", "name": "my-tux"},
|
||||
|
|
@ -729,6 +775,45 @@ def test_import_preserves_unrelated_files_when_folder_not_empty(client, ds_root,
|
|||
assert r.json()["imported"] == 3
|
||||
assert sorted(p.name for p in folder.glob("*.png")) == [f"img_{i:04d}.png" for i in range(3)]
|
||||
assert (folder / "notes.md").read_text(encoding = "utf-8") == "keep me"
|
||||
assert folder.stat().st_ino != inode_before
|
||||
# No staging dir left behind, so the folder that is there is the promoted one.
|
||||
assert [d.name for d in ds_root.glob(".my-tux.import-*")] == []
|
||||
|
||||
|
||||
def test_import_promotes_atomically_over_a_thumbs_cache(client, ds_root, monkeypatch):
|
||||
# .thumbs is the case that actually shows up: a folder whose images were deleted keeps the
|
||||
# thumbnail cache, which used to force the non-atomic per-file promote.
|
||||
_install_fake_load_dataset(monkeypatch, n_rows = 2)
|
||||
folder = ds_root / "my-tux"
|
||||
(folder / ".thumbs").mkdir(parents = True)
|
||||
(folder / ".thumbs" / "old.png_32.jpg").write_bytes(b"stale")
|
||||
r = client.post(
|
||||
"/api/train/diffusion/dataset/import-example",
|
||||
json = {"id": "tuxemon", "name": "my-tux"},
|
||||
)
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["imported"] == 2
|
||||
assert sorted(p.name for p in folder.glob("*.png")) == ["img_0000.png", "img_0001.png"]
|
||||
assert (folder / ".thumbs" / "old.png_32.jpg").read_bytes() == b"stale"
|
||||
assert [d.name for d in ds_root.glob(".my-tux.import-*")] == []
|
||||
|
||||
|
||||
def test_import_replaces_a_pre_existing_file_the_import_also_writes(client, ds_root, monkeypatch):
|
||||
# Same name on both sides -- here a stray caption sidecar with no image, so image_count is still
|
||||
# 0 and the import runs. The imported file wins, which is what the old per-file move did by
|
||||
# overwriting; the point is that the outcome did not change with the atomic promote.
|
||||
_install_fake_load_dataset(monkeypatch, n_rows = 2)
|
||||
folder = ds_root / "my-tux"
|
||||
folder.mkdir(parents = True)
|
||||
(folder / "img_0000.txt").write_text("stale caption", encoding = "utf-8")
|
||||
r = client.post(
|
||||
"/api/train/diffusion/dataset/import-example",
|
||||
json = {"id": "tuxemon", "name": "my-tux"},
|
||||
)
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["imported"] == 2
|
||||
assert (folder / "img_0000.txt").read_text(encoding = "utf-8") == "caption 0"
|
||||
assert [d.name for d in ds_root.glob(".my-tux.import-*")] == []
|
||||
|
||||
|
||||
def test_a_second_concurrent_import_of_the_same_name_is_refused(client, ds_root, monkeypatch):
|
||||
|
|
|
|||
|
|
@ -134,7 +134,14 @@ const galleryCache: {
|
|||
hasMore: boolean;
|
||||
selectedId: string | null;
|
||||
quant: string | null;
|
||||
srcById: Map<string, string>;
|
||||
// id -> the signed link and when it was minted. The link is short-lived (the backend expires it,
|
||||
// and its signing secret is per-process, so a server restart invalidates every outstanding one),
|
||||
// while this cache deliberately survives navigation -- so an entry has to be re-mintable rather
|
||||
// than final, or playback, seeking and Save would 401 until a full page reload.
|
||||
srcById: Map<string, { url: string; mintedAt: number }>;
|
||||
// Ids re-minted once after a media error already, so a clip that is broken for any other reason
|
||||
// cannot spin in a mint/error loop.
|
||||
refreshed: Set<string>;
|
||||
// Ids with a mint in flight, so concurrent ensureSrc calls don't double-request.
|
||||
inflight: Set<string>;
|
||||
// Ids deleted while their link was still being minted, so a reply that lands after the delete
|
||||
|
|
@ -148,11 +155,16 @@ const galleryCache: {
|
|||
selectedId: null,
|
||||
quant: null,
|
||||
srcById: new Map(),
|
||||
refreshed: new Set(),
|
||||
inflight: new Set(),
|
||||
deleted: new Set(),
|
||||
epoch: 0,
|
||||
};
|
||||
|
||||
// Re-mint a cached link once it is this old. Comfortably inside the backend's own expiry, so a
|
||||
// long-lived tab keeps working without waiting for a 401 to tell it the link died.
|
||||
const VIDEO_LINK_REFRESH_MS = 6 * 60 * 60 * 1000;
|
||||
|
||||
// Videos loaded per infinite-scroll page.
|
||||
const PAGE_SIZE = 50;
|
||||
|
||||
|
|
@ -574,7 +586,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
if (!active) previewRef.current?.pause();
|
||||
}, [active]);
|
||||
const [srcById, setSrcById] = useState<Record<string, string>>(() =>
|
||||
Object.fromEntries(galleryCache.srcById),
|
||||
Object.fromEntries([...galleryCache.srcById].map(([id, e]) => [id, e.url])),
|
||||
);
|
||||
// Guards a "load more" so a fast scroll can't fire several at once.
|
||||
const loadingMore = useRef(false);
|
||||
|
|
@ -696,7 +708,9 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
// which streams ranges as it plays, so a clip starts on its first seconds instead of after a
|
||||
// full download and seeking works.
|
||||
const ensureSrc = useCallback(async (video: GalleryVideo) => {
|
||||
if (galleryCache.srcById.has(video.id) || galleryCache.inflight.has(video.id)) return;
|
||||
const cached = galleryCache.srcById.get(video.id);
|
||||
if (cached && Date.now() - cached.mintedAt < VIDEO_LINK_REFRESH_MS) return;
|
||||
if (galleryCache.inflight.has(video.id)) return;
|
||||
galleryCache.inflight.add(video.id);
|
||||
const epochAtStart = galleryCache.epoch;
|
||||
try {
|
||||
|
|
@ -704,7 +718,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
// The record can be deleted (or the gallery cleared) while the link is being minted;
|
||||
// caching it then would leave an entry for a card that no longer exists.
|
||||
if (galleryCache.deleted.has(video.id) || galleryCache.epoch !== epochAtStart) return;
|
||||
galleryCache.srcById.set(video.id, url);
|
||||
galleryCache.srcById.set(video.id, { url, mintedAt: Date.now() });
|
||||
// The URL is cached above either way; skip the state update after unmount
|
||||
// (matches the other async callbacks in this file).
|
||||
if (isMounted.current) setSrcById((prev) => ({ ...prev, [video.id]: url }));
|
||||
|
|
@ -715,6 +729,18 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
}
|
||||
}, []);
|
||||
|
||||
// A media error on a clip that was playing means its link died early -- the server restarted, so
|
||||
// its signing secret changed. Drop the entry and mint a fresh one, once per clip per session.
|
||||
const remintSrc = useCallback(
|
||||
(video: GalleryVideo) => {
|
||||
if (galleryCache.refreshed.has(video.id)) return;
|
||||
galleryCache.refreshed.add(video.id);
|
||||
galleryCache.srcById.delete(video.id);
|
||||
void ensureSrc(video);
|
||||
},
|
||||
[ensureSrc],
|
||||
);
|
||||
|
||||
// A card's poster frame only appears once its src lands, and each src costs a request, so a
|
||||
// full gallery page (PAGE_SIZE records) minted up front would queue PAGE_SIZE requests ahead
|
||||
// of the one clip the user is actually waiting on. Mint a card's link as it nears the viewport
|
||||
|
|
@ -832,6 +858,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
return;
|
||||
}
|
||||
galleryCache.srcById.delete(id);
|
||||
galleryCache.refreshed.delete(id);
|
||||
// A mint still in flight for this id must throw its link away rather than cache it.
|
||||
galleryCache.deleted.add(id);
|
||||
setSrcById((prev) => {
|
||||
|
|
@ -851,6 +878,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
return;
|
||||
}
|
||||
galleryCache.srcById.clear();
|
||||
galleryCache.refreshed.clear();
|
||||
// Every mint in flight now belongs to a cleared gallery, so their links are discarded on
|
||||
// arrival. The epoch covers ids this page never listed too.
|
||||
galleryCache.epoch += 1;
|
||||
|
|
@ -1788,6 +1816,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
void e.currentTarget.play();
|
||||
}
|
||||
}}
|
||||
onError={() => remintSrc(selected)}
|
||||
className="max-h-full max-w-full rounded-xl object-contain shadow-sm"
|
||||
/>
|
||||
{selected.has_audio && (
|
||||
|
|
@ -1916,6 +1945,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
muted
|
||||
playsInline
|
||||
preload="metadata"
|
||||
onError={() => remintSrc(video)}
|
||||
className="absolute inset-0 size-full object-cover"
|
||||
/>
|
||||
) : (
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue