ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Close video single-file, training reservation, and image mount-resume gaps

Browse source 
Route on-device single-checkpoint video folders through the single_file loader:
a bare local .safetensors directory (no model_index.json) is advertised as a
pipeline with no filename, so validation rejected it before it could load.
Reinterpret the pick as a single_file load of the sole checkpoint, mirroring the
image load route.

Treat a reserved-but-not-yet-spawned LLM training start as active in
is_training_active() so /images/load, /video/load, and /diffusion/start cannot
race the reserved run for VRAM during the pre-spawn free window. Mirrors the
diffusion training service reservation.

Resume an in-flight image generation on the Images page mount: probe
generate-progress, re-enter the poll loop, and refresh the gallery on completion
so a run started elsewhere is reflected and its saved image appears without a
manual refresh. Seed resident image defaults from the resolved base_repo rather
than a possibly path-shaped repo_id so the first resident generation uses the
right recipe.
This commit is contained in:
Daniel Han 2026-07-12 12:41:16 +00:00
commit 23a71b1c2c
5 changed files with 179 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

9
studio/backend/core/training/training.py
View file

@ -1141,6 +1141,15 @@ class TrainingBackend:
# training invisibly behind a frozen UI. Cheap enough for per-second polls.
self._ensure_pump_alive()
with self._lock:
# A start reserved via the compare-and-set guard in start_training but not yet
# spawned (before_spawn frees residents, then GPU auto-selection, then proc.start())
# is already "active": the load/start guards read this to refuse a concurrent
# /images/load, /video/load, or /diffusion/start, so treating the pre-spawn window
# as idle would let another pipeline race the reserved run for VRAM. Mirrors the
# diffusion training service's reserve()/is_active().
if self._start_in_progress:
return True
if self._proc is not None and self._proc.is_alive():
return True

21
studio/backend/routes/video.py
View file

@ -77,13 +77,28 @@ def _guard_video_load_against_training() -> None:
async def load_video_model(
request: VideoLoadRequest, current_subject: str = Depends(get_current_subject)
):
from core.inference.diffusion import resolve_local_single_file
from core.inference.diffusion_device import resolve_diffusion_device_target
from core.inference.gpu_arbiter import VIDEO, acquire_for, release
from core.inference.video import get_video_backend
from core.inference.video import get_video_backend, resolve_video_model_kind
from utils.native_path_leases import redact_native_paths
backend = get_video_backend()
try:
# Resolve the load kind once (gguf / single_file / pipeline) so validation and the
# load agree; a bad explicit kind raises here -> 400.
kind = resolve_video_model_kind(request.gguf_filename, request.model_kind)
# A local On-Device pick can be a bare single-file .safetensors directory (no
# model_index.json): the scanner advertises it as a text-to-video model, but the
# local picker starts it as a pipeline with no filename, so a pipeline load would
# 400 on the missing model_index.json and the advertised model is unusable. If the
# directory holds exactly one checkpoint, reinterpret the pick as a single_file load
# of it, so validation and the load agree. Mirrors the image load route.
if kind == "pipeline" and not request.gguf_filename:
sole = await asyncio.to_thread(resolve_local_single_file, request.model_path)
if sole is not None:
request.gguf_filename = sole
kind = resolve_video_model_kind(sole, None)
# Validate cheaply BEFORE touching the GPU so an unloadable pick can't evict chat then 400.
await asyncio.to_thread(
backend.validate_load_request,
@ -91,7 +106,7 @@ async def load_video_model(
gguf_filename = request.gguf_filename,
base_repo = request.base_repo,
family_override = request.family_override,
model_kind = request.model_kind,
model_kind = kind,
transformer_quant = request.transformer_quant,
text_encoder_quant = request.text_encoder_quant,
)
@ -118,7 +133,7 @@ async def load_video_model(
transformer_cache_threshold = request.transformer_cache_threshold,
transformer_quant = request.transformer_quant,
text_encoder_quant = request.text_encoder_quant,
model_kind = request.model_kind,
model_kind = kind,
)
return VideoStatusResponse(**status_dict)
except (ValueError, FileNotFoundError) as exc:

36
studio/backend/tests/test_training_start_offload.py
View file

@ -113,3 +113,39 @@ def test_backend_start_guard_blocks_overlapping_starts():
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
def test_is_training_active_true_during_start_reservation():
# While start_training holds the compare-and-set reservation but has not yet spawned
# (before_spawn frees residents, then GPU auto-selection, then proc.start()), the LLM
# training run must already read as active: /images/load, /video/load, and
# /diffusion/start all gate on is_training_active(), so an idle reading in this window
# would let another pipeline race the reserved run for the just-freed VRAM.
from core.training.training import TrainingBackend
backend = TrainingBackend()
# Not reserved yet: idle.
assert backend.is_training_active() is False
entered = threading.Event()
release = threading.Event()
captured = {}
def _slow_impl(job_id, *, before_spawn = None, **kwargs):
entered.set()
release.wait(timeout = 5.0)
return True
backend._start_training_impl = _slow_impl
t = threading.Thread(target = lambda: backend.start_training("job-a"), daemon = True)
t.start()
assert entered.wait(timeout = 5.0)
# Inside the pre-spawn window: reserved, so active even though no proc/progress is set.
captured["in_window"] = backend.is_training_active()
release.set()
t.join(timeout = 5.0)
assert captured["in_window"] is True
# Reservation cleared once the start returns; with no live proc it reads idle again.
assert backend.is_training_active() is False

44
studio/backend/tests/test_video_routes.py
View file

@ -106,10 +106,17 @@ class _FakeBackend(video_module.VideoBackend):
):
# Mirror the real backend's cheap validation so the route's
# validate-before-evict ordering is exercised.
from pathlib import Path
kind = (model_kind or ("gguf" if gguf_filename else "pipeline")).lower()
if kind in ("gguf", "single_file") and not gguf_filename:
raise ValueError("A gguf/single_file load needs the checkpoint filename.")
if kind != "gguf" and not model_path.lower().startswith(("unsloth/", "lightricks/")):
# Non-GGUF loads are gated to unsloth/* repos, the official bases, and local paths
# (the real backend trusts an existing local path via _is_trusted_video_repo).
trusted = model_path.lower().startswith(("unsloth/", "lightricks/")) or (
Path(model_path).expanduser().exists()
)
if kind != "gguf" and not trusted:
raise ValueError(
f"Non-GGUF video loads are limited to unsloth/* repos, the official family "
f"base repos, and local paths; '{model_path}' is neither."
@ -378,6 +385,41 @@ def test_load_progress_route(client):
assert ready.json()["phase"] == "ready"
def test_load_local_single_file_dir_routes_through_single_file(client, tmp_path):
# An On-Device pick of a local directory named for a video family that holds exactly one
# .safetensors and no model_index.json arrives as a pipeline with no filename. The route
# reinterprets it as a single_file load of the sole checkpoint (mirrors the image route),
# so it is loadable instead of 400ing on the missing model_index.json.
d = tmp_path / "ltx-2.3-local"
d.mkdir()
(d / "ltx-dit.safetensors").write_bytes(b"0")
resp = client.post(
"/api/inference/video/load",
json = {"model_path": str(d), "model_kind": "pipeline"},
)
assert resp.status_code == 200
kwargs = video_module.get_video_backend().last_load_kwargs
assert kwargs["model_kind"] == "single_file"
assert kwargs["gguf_filename"] == "ltx-dit.safetensors"
def test_load_local_pipeline_dir_stays_pipeline(client, tmp_path):
# A real diffusers directory (has model_index.json) is left as a pipeline load:
# resolve_local_single_file returns None for it, so the pick is not rewritten.
d = tmp_path / "ltx-2.3-pipeline"
d.mkdir()
(d / "model_index.json").write_text("{}")
(d / "diffusion_pytorch_model.safetensors").write_bytes(b"0")
resp = client.post(
"/api/inference/video/load",
json = {"model_path": str(d), "model_kind": "pipeline"},
)
assert resp.status_code == 200
kwargs = video_module.get_video_backend().last_load_kwargs
assert kwargs["model_kind"] == "pipeline"
assert not kwargs.get("gguf_filename")
def test_generate_happy_path_persists_and_reports_record(client):
client.post(
"/api/inference/video/load",

76
studio/frontend/src/features/images/images-page.tsx
View file

@ -1450,6 +1450,56 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
pollTimer.current = setTimeout(() => void pollLoadProgress(), 1000);
}, [dismissLoadToast, refreshStatus]);
// Re-enter the per-step generation poll for a run already in flight on the backend --
// one this page did not start (another client called /images/generate, or this browser
// reloaded mid-generate). The backend runs generations under a serialising lock and keeps
// reporting progress, so track it here instead of showing a stale idle view. The images
// generate-progress carries only per-step progress (no terminal record), so on completion
// refresh the gallery to merge any image saved after the mount fetch. Kept separate from
// handleGenerate's own loop, which prepends its synchronous results directly.
const resumeGeneratePoll = useCallback(() => {
if (genPollTimer.current) clearInterval(genPollTimer.current);
if (genVisibilityListener.current)
document.removeEventListener("visibilitychange", genVisibilityListener.current);
let pollInFlight = false;
const pollGenerateOnce = async () => {
if (pollInFlight) return;
pollInFlight = true;
try {
const p = await getGenerateProgress();
if (!p.active) {
if (genPollTimer.current) clearInterval(genPollTimer.current);
genPollTimer.current = null;
if (genVisibilityListener.current) {
document.removeEventListener("visibilitychange", genVisibilityListener.current);
genVisibilityListener.current = null;
}
if (!isMounted.current) return;
setBusy(null);
setGenStep(null);
// The finished run saved its image(s) to the backend gallery; re-fetch the first
// page to merge them (a full replace, so deduped) and resync status.
void loadGallery();
void refreshStatus();
return;
}
setGenStep((prev) => {
if (prev && prev.step === p.step && prev.eta_seconds === p.eta_seconds) return prev;
return p;
});
} catch {
// transient; keep polling
} finally {
pollInFlight = false;
}
};
genVisibilityListener.current = () => {
if (document.visibilityState === "visible") void pollGenerateOnce();
};
document.addEventListener("visibilitychange", genVisibilityListener.current);
genPollTimer.current = setInterval(() => void pollGenerateOnce(), 300);
}, [loadGallery, refreshStatus]);
useEffect(() => {
void (async () => {
await refreshStatus();
@ -1468,6 +1518,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
} catch {
// Resume is best-effort; a failed probe just leaves the idle view.
}
// A generation started elsewhere -- another client, or this browser before a reload --
// keeps running on the backend under a serialising lock. Resume tracking it so the page
// reflects the in-flight run (progress bar + disabled Generate) instead of a stale idle
// view, then refreshes the gallery when it finishes to pick up an image saved after the
// mount fetch. Mirrors the video page's mount resume.
try {
const g = await getGenerateProgress();
if (g.active) {
setBusy("generating");
setGenStep(g);
resumeGeneratePoll();
}
} catch {
// Resume is best-effort; a failed probe just leaves the idle view.
}
})();
// Stop polling if the page unmounts mid-load / mid-generate, and dismiss the
// load toast — its poll loop is gone, so it would otherwise hang forever
@ -1481,7 +1546,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
}
dismissLoadToast();
};
}, [refreshStatus, dismissLoadToast, pollLoadProgress]);
}, [refreshStatus, dismissLoadToast, pollLoadProgress, resumeGeneratePoll]);
// Seed the generation sliders to a resident model's recipe when the page discovers one it
// did not load itself (left loaded by a prior session or another route). refreshStatus() only
@ -1496,7 +1561,12 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
if (lastLoad.current) return;
if (seededResident.current === repoId) return;
seededResident.current = repoId;
const d = defaultsFor(repoId);
// Seed from the resolved base_repo, not repo_id: a GGUF/single_file resident (or one
// loaded from a bare local path like /models/checkpoint.safetensors) carries a repo_id
// with no family substring, so defaultsFor(repoId) would fall back to the distilled
// few-step/no-CFG recipe and the first resident generation would run with wrong defaults.
// base_repo is the resolved diffusers base (it holds the family), so prefer it when set.
const d = defaultsFor(status?.base_repo ?? repoId);
setSteps(d.steps);
setGuidance(d.guidance);
// Wire "Reapply" to the resident model too, so an advanced-option reload works without
@ -1508,7 +1578,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
if (kind === "pipeline") {
lastLoad.current = { repoId, kind };
}
}, [status?.loaded, status?.repo_id, status?.model_kind]);
}, [status?.loaded, status?.repo_id, status?.base_repo, status?.model_kind]);
const handleLoad = useCallback(
// Resolves true when the background load STARTED (callers may revert