Studio: defer chat GPU handoff, unload video before training, stage example imports, name-aware Wan GGUF video tag, and diffusion upload body passthrough
- Chat load defers the CHAT arbiter handoff until after identifier / gpu_ids / training-memory validation, so a doomed chat load (bad id, unsupported gpu_ids on GGUF, or a training 409) no longer evicts a resident Images/Video pipeline and then errors. The already-loaded fast paths re-assert CHAT ownership themselves. Mirrors the image and video loaders, which validate before acquire_for. - Both training-start GPU cleanups now unload a resident Video pipeline and release the VIDEO arbiter owner, not just Images/DIFFUSION, so starting LLM or diffusion training after a video generation session no longer competes with the still-resident video model and OOMs the run. - Example dataset import materializes into a private staging dir and promotes into the dataset folder only after the whole import succeeds. A materialize that fails partway no longer leaves a partial dataset that a retry would treat as complete (imported=0), stranding the user with a truncated dataset. Hidden dirs are skipped by the dataset scan so the staging dir never surfaces as a dataset. - Cached/local Wan GGUFs are now classified name-aware: arch "wan" alone is ambiguous between the loadable single-DiT TI2V-5B and the dual-expert A14B MoE the loader refuses, so _arch_to_task falls back to the repo/file name (as the loader's own detect_video_family does) and tags only a non-MoE match text-to-video, surfacing loadable Wan GGUFs in the Video picker without surfacing unloadable A14B files. - The diffusion dataset upload route is added to the MaxBodyMiddleware upload passthrough, so its own get_upload_limit_bytes() cap (plus multipart overhead, and a raised max_upload_size_mb) applies instead of the default body limit rejecting near-limit batches with 413 before the handler runs. Adds regression tests for each: the chat handoff not evicting on a doomed load, the video unload on diffusion-training start, the atomic import leaving no partial dataset, the name-aware Wan GGUF classification (TI2V-5B video, A14B unsupported, bare arch unsupported), and the diffusion upload passthrough cap.
This commit is contained in:
Daniel Han
parent
d36aeb54f3
commit
b6ab866c46
8 changed files with 289 additions and 38 deletions
|
|
@ -743,16 +743,26 @@ _DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload"
|
|||
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = (
|
||||
"/api/data-recipe/seed/upload-unstructured-file"
|
||||
)
|
||||
# The diffusion dataset upload route (POST /api/train/diffusion/dataset) is a multipart
|
||||
# image upload under the protected /api/train prefix. Like /api/datasets/upload it enforces
|
||||
# its own get_upload_limit_bytes() cap, so it must bypass the default body cap here or the
|
||||
# middleware would 413 near-limit batches (and ignore a raised max_upload_size_mb) before the
|
||||
# handler runs. Its small-JSON sub-routes (import-example, caption) merely inherit the more
|
||||
# generous cap, which is harmless.
|
||||
_DIFFUSION_DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/train/diffusion/dataset"
|
||||
_BODY_UPLOAD_PASSTHROUGH_PREFIXES = (
|
||||
_DATASET_UPLOAD_PASSTHROUGH_PREFIX,
|
||||
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX,
|
||||
_DIFFUSION_DATASET_UPLOAD_PASSTHROUGH_PREFIX,
|
||||
)
|
||||
|
||||
|
||||
def _get_upload_passthrough_request_max_bytes(path: str) -> int:
|
||||
if path.startswith(_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX):
|
||||
return upload_request_limit_bytes(UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES)
|
||||
if path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX):
|
||||
if path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX) or path.startswith(
|
||||
_DIFFUSION_DATASET_UPLOAD_PASSTHROUGH_PREFIX
|
||||
):
|
||||
return upload_request_limit_bytes()
|
||||
return default_request_body_limit_bytes()
|
||||
|
||||
|
|
|
|||
|
|
@ -3136,14 +3136,14 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
|
|||
user_override = request.chat_template_override,
|
||||
)
|
||||
|
||||
# Reclaim the GPU from the diffusion (Images) backend before any chat
|
||||
# load path — including the already-loaded fast path below, so chat
|
||||
# ownership is asserted without depending on chat/diffusion exclusivity
|
||||
# holding. No-op when diffusion isn't loaded.
|
||||
# Reclaim the GPU for chat (evicting a resident Images/Video pipeline) only once the
|
||||
# load is known viable: the already-loaded fast paths below re-assert CHAT ownership
|
||||
# themselves, and the real handoff is deferred past identifier / gpu_ids / training-memory
|
||||
# validation so a doomed chat load (bad id, unsupported gpu_ids on GGUF, or a training
|
||||
# 409) can't evict a working image/video model and then error. Mirrors the image/video
|
||||
# loaders, which validate before acquire_for.
|
||||
from core.inference.gpu_arbiter import acquire_for, CHAT
|
||||
|
||||
await asyncio.to_thread(acquire_for, CHAT)
|
||||
|
||||
# ── Already-loaded check: skip reload if the exact model is active ──
|
||||
backend = get_inference_backend()
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
|
|
@ -3177,6 +3177,10 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
|
|||
|
||||
_gguf_audio = getattr(llama_backend, "_audio_type", None)
|
||||
_gguf_is_audio = getattr(llama_backend, "_is_audio", False)
|
||||
# The requested GGUF chat model is already resident: assert CHAT ownership (a
|
||||
# no-op when it already holds it) so a drifted arbiter owner is corrected. This
|
||||
# is a guaranteed-success path, not a doomed load, so evicting here is correct.
|
||||
await asyncio.to_thread(acquire_for, CHAT)
|
||||
return LoadResponse(
|
||||
status = "already_loaded",
|
||||
model = model_log_label
|
||||
|
|
@ -3229,6 +3233,10 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
|
|||
_sf_flags = _detect_safetensors_features(backend, _chat_template)
|
||||
_sf_supports_reasoning = _sf_flags["supports_reasoning"]
|
||||
_sf_reasoning_style = _sf_flags["reasoning_style"]
|
||||
# The requested chat model is already resident: assert CHAT ownership (no-op when
|
||||
# it already holds it) to correct a drifted arbiter owner. Guaranteed-success
|
||||
# path, not a doomed load, so evicting here is correct.
|
||||
await asyncio.to_thread(acquire_for, CHAT)
|
||||
return LoadResponse(
|
||||
status = "already_loaded",
|
||||
model = model_log_label if native_grant_backed else backend.active_model_name,
|
||||
|
|
@ -3303,6 +3311,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
|
|||
n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1),
|
||||
)
|
||||
|
||||
# The load is now known viable (valid identifier, gpu_ids ok, fits alongside any active
|
||||
# training): reclaim the GPU for chat, evicting a resident Images/Video pipeline. Doing
|
||||
# this only here -- not before the validation above -- is what keeps a doomed chat load
|
||||
# from evicting a working image/video model and then erroring. No-op when chat already
|
||||
# owns the GPU.
|
||||
await asyncio.to_thread(acquire_for, CHAT)
|
||||
|
||||
# ── GGUF path: load via llama-server ──────────────────────
|
||||
if config.is_gguf:
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
|
|
|
|||
|
|
@ -3206,20 +3206,31 @@ def _gguf_architecture(path: str) -> Optional[str]:
|
|||
return arch.strip() if isinstance(arch, str) and arch.strip() else None
|
||||
|
||||
|
||||
def _arch_to_task(arch: Optional[str]) -> Optional[str]:
|
||||
def _arch_to_task(arch: Optional[str], name_hints: tuple[Optional[str], ...] = ()) -> Optional[str]:
|
||||
if arch is None:
|
||||
return None
|
||||
a = arch.lower()
|
||||
if a in _DIFFUSION_GGUF_ARCHS:
|
||||
return "text-to-image"
|
||||
if a in _VIDEO_GGUF_ARCHS:
|
||||
# Only advertise as loadable video when a VideoFamily is actually registered for this arch
|
||||
# (the video registry differs across the stacked video PRs -- wan is pre-registered here but
|
||||
# its family lands later). An unregistered video arch would 400 on load, so fall it through
|
||||
# to the unsupported bucket (hidden from chat, not surfaced in the Video/Images pickers)
|
||||
# until its family exists, rather than advertising a GGUF that cannot load.
|
||||
# Advertise as loadable video only when a VideoFamily actually resolves. Some archs map
|
||||
# straight from the arch (ltxv); others are ambiguous at the arch level -- bare "wan"
|
||||
# covers both the single-DiT TI2V-5B (GGUF-loadable) and the dual-expert A14B MoE whose
|
||||
# single file the loader refuses -- so when the bare arch does not resolve, fall back to
|
||||
# the repo/file names like the loader's own detect_video_family does (each tried
|
||||
# separately, since it matches on name segments not substrings), and surface only a
|
||||
# non-MoE (loadable) match. Without a name we cannot disambiguate, so a bare-arch Wan
|
||||
# GGUF (which the loader also cannot resolve) stays in the unsupported bucket rather than
|
||||
# advertising a GGUF that would 400 on load.
|
||||
from core.inference.video_families import detect_video_family
|
||||
if detect_video_family("", override = a) is not None:
|
||||
fam = detect_video_family("", override = a)
|
||||
if fam is None:
|
||||
for hint in name_hints:
|
||||
if hint:
|
||||
fam = detect_video_family(hint)
|
||||
if fam is not None:
|
||||
break
|
||||
if fam is not None and not getattr(fam, "is_moe", False):
|
||||
return _VIDEO_GEN_TASK
|
||||
return _UNSUPPORTED_DIFFUSION_TASK
|
||||
# A diffusion arch the backend can't assemble: hide it from chat (it would die
|
||||
|
|
@ -3234,11 +3245,14 @@ def _repo_gguf_task(repo_info) -> Optional[str]:
|
|||
'text-to-image' for a loadable diffusion arch, the non-loadable diffusion tag
|
||||
for a recognized-but-unsupported image arch, else 'text-generation' (None if
|
||||
unreadable)."""
|
||||
repo_id = getattr(repo_info, "repo_id", None)
|
||||
try:
|
||||
for path in _iter_gguf_paths(Path(repo_info.repo_path)):
|
||||
if _is_mmproj_filename(path.name):
|
||||
continue
|
||||
task = _arch_to_task(_gguf_architecture(str(path)))
|
||||
task = _arch_to_task(
|
||||
_gguf_architecture(str(path)), name_hints = (repo_id, path.name)
|
||||
)
|
||||
if task is not None:
|
||||
return task
|
||||
except Exception:
|
||||
|
|
@ -3255,15 +3269,20 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]:
|
|||
models get the 'text-to-image' tag instead of being dropped as task=null; the load
|
||||
path accepts these as a local pipeline."""
|
||||
path = model.path
|
||||
_id_hints = (model.model_id, model.display_name, model.id)
|
||||
if model.model_format == "gguf":
|
||||
try:
|
||||
p = Path(path)
|
||||
if p.suffix.lower() == ".gguf" and p.is_file():
|
||||
return _arch_to_task(_gguf_architecture(str(p)))
|
||||
return _arch_to_task(
|
||||
_gguf_architecture(str(p)), name_hints = _id_hints + (p.name,)
|
||||
)
|
||||
for f in _iter_gguf_paths(p):
|
||||
if _is_mmproj_filename(f.name):
|
||||
continue
|
||||
task = _arch_to_task(_gguf_architecture(str(f)))
|
||||
task = _arch_to_task(
|
||||
_gguf_architecture(str(f)), name_hints = _id_hints + (f.name,)
|
||||
)
|
||||
if task is not None:
|
||||
return task
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -443,6 +443,24 @@ async def start_training(
|
|||
except Exception as e:
|
||||
logger.warning("Could not unload diffusion model for training: %s", e)
|
||||
|
||||
try:
|
||||
# A resident or in-flight Video pipeline holds GPU memory the training run
|
||||
# needs too, and it loads under the VIDEO arbiter owner the diffusion teardown
|
||||
# above never touches. Tear it down the same way (unload is a no-op when nothing
|
||||
# is loaded and preempts an in-flight load) and release VIDEO, so starting
|
||||
# training while a generated-video session is resident can't OOM the run. Must
|
||||
# precede the chat block, which early-returns.
|
||||
from core.inference import gpu_arbiter
|
||||
from core.inference.video import get_video_backend
|
||||
|
||||
video = get_video_backend()
|
||||
if video.status().get("loaded"):
|
||||
logger.info("Unloading Video model to free GPU memory for training")
|
||||
video.unload()
|
||||
gpu_arbiter.release(gpu_arbiter.VIDEO)
|
||||
except Exception as e:
|
||||
logger.warning("Could not unload video model for training: %s", e)
|
||||
|
||||
try:
|
||||
from routes.training_vram import (
|
||||
can_keep_chat_during_training,
|
||||
|
|
@ -1130,6 +1148,21 @@ def _free_gpu_for_diffusion_training() -> None:
|
|||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("Could not unload Images pipeline for diffusion training: %s", e)
|
||||
|
||||
try:
|
||||
# A resident Video pipeline loads under the VIDEO arbiter owner, which the Images
|
||||
# teardown above does not free; unload it too (no-op when nothing is loaded) and release
|
||||
# VIDEO so a generated-video session left resident can't OOM the diffusion trainer.
|
||||
from core.inference import gpu_arbiter
|
||||
from core.inference.video import get_video_backend
|
||||
|
||||
video = get_video_backend()
|
||||
if video.status().get("loaded"):
|
||||
logger.info("Unloading resident Video pipeline to free GPU memory for training")
|
||||
video.unload() # no-op when nothing is loaded; also preempts an in-flight load
|
||||
gpu_arbiter.release(gpu_arbiter.VIDEO)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("Could not unload Video pipeline for diffusion training: %s", e)
|
||||
|
||||
try:
|
||||
# The SDXL trainer's footprint can't be cheaply sized against a resident chat
|
||||
# model, so free chat unconditionally (same conservative choice the LLM path
|
||||
|
|
@ -1445,7 +1478,11 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub
|
|||
root = datasets_root()
|
||||
found: list[DiffusionDatasetSummary] = []
|
||||
try:
|
||||
children = sorted(p for p in root.iterdir() if p.is_dir())
|
||||
# Skip hidden dirs: they are never user datasets, and an in-progress example
|
||||
# import stages into a dot-prefixed sibling that must not surface as a dataset.
|
||||
children = sorted(
|
||||
p for p in root.iterdir() if p.is_dir() and not p.name.startswith(".")
|
||||
)
|
||||
except OSError:
|
||||
children = []
|
||||
for child in children:
|
||||
|
|
@ -2068,28 +2105,46 @@ async def import_diffusion_dataset_example(
|
|||
folder = _resolve_dataset_folder(body.name or entry["id"], must_exist = False)
|
||||
|
||||
def do_import() -> DiffusionDatasetImportResponse:
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
folder.mkdir(parents = True, exist_ok = True)
|
||||
existing = _diffusion_dataset_summary(folder)
|
||||
imported = 0
|
||||
if existing.image_count == 0:
|
||||
cap = int(entry["image_cap"])
|
||||
# Materialize into a private staging dir and promote into the dataset folder only
|
||||
# after the whole import succeeds. A materialize that fails partway (a transient
|
||||
# fetch/copy error after writing some images) then leaves only the staging dir,
|
||||
# never a half-filled dataset -- otherwise the image_count>0 idempotency check
|
||||
# above would treat that partial result as complete on the next retry (imported=0)
|
||||
# and strand the user with a truncated dataset (there is no dataset-delete flow).
|
||||
# Staged as a hidden sibling on the same filesystem so promotion is an atomic rename.
|
||||
staging = Path(
|
||||
tempfile.mkdtemp(dir = folder.parent, prefix = f".{folder.name}.import-")
|
||||
)
|
||||
try:
|
||||
if entry["loader"] == "imagefolder_jsonl":
|
||||
imported = _materialize_imagefolder_jsonl(entry, folder, cap)
|
||||
else:
|
||||
imported = _materialize_hf_dataset(entry, folder, cap)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 -- surface a readable fetch/parse failure
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = f"Could not import '{entry['repo']}': {e}",
|
||||
)
|
||||
if imported == 0:
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = f"No images found in '{entry['repo']}'.",
|
||||
)
|
||||
try:
|
||||
if entry["loader"] == "imagefolder_jsonl":
|
||||
imported = _materialize_imagefolder_jsonl(entry, staging, cap)
|
||||
else:
|
||||
imported = _materialize_hf_dataset(entry, staging, cap)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 -- surface a readable fetch/parse failure
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = f"Could not import '{entry['repo']}': {e}",
|
||||
)
|
||||
if imported == 0:
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = f"No images found in '{entry['repo']}'.",
|
||||
)
|
||||
for p in staging.iterdir():
|
||||
shutil.move(str(p), str(folder / p.name))
|
||||
finally:
|
||||
shutil.rmtree(staging, ignore_errors = True)
|
||||
summary = _diffusion_dataset_summary(folder)
|
||||
return DiffusionDatasetImportResponse(
|
||||
name = folder.name,
|
||||
|
|
|
|||
|
|
@ -796,12 +796,27 @@ def test_arch_to_task_hides_unsupported_diffusion_from_chat():
|
|||
# LTX-2.x GGUFs ship general.architecture "ltxv" and ltx-2 is registered).
|
||||
assert models_route._arch_to_task("ltxv") == models_route._VIDEO_GEN_TASK
|
||||
assert models_route._arch_to_task("ltxv") not in ("text-generation", "text-to-image")
|
||||
# A video arch that is pre-registered in _VIDEO_GGUF_ARCHS but has NO VideoFamily yet ("wan"
|
||||
# in this build) must NOT be advertised as loadable video -- it would 400 on load. It falls to
|
||||
# the unsupported bucket (hidden from chat and from the Video/Images pickers) until its family
|
||||
# lands, rather than surfacing a GGUF that cannot load.
|
||||
# A video arch that does not resolve from the bare arch alone ("wan" is ambiguous -- it
|
||||
# covers both the loadable single-DiT TI2V-5B and the A14B MoE whose single file the loader
|
||||
# refuses) stays unsupported when no repo/file name is available to disambiguate, rather than
|
||||
# surfacing a GGUF that might 400 on load.
|
||||
assert models_route._arch_to_task("wan") == models_route._UNSUPPORTED_DIFFUSION_TASK
|
||||
assert models_route._arch_to_task("wan") not in ("text-generation", "text-to-image")
|
||||
# With a repo/file name hint, the loadable TI2V-5B Wan GGUF resolves to the Video task (so it
|
||||
# surfaces in the Video On-Device picker), while the A14B MoE (single file refused by the
|
||||
# loader) stays in the unsupported bucket -- matching the loader's own name-aware detection.
|
||||
assert (
|
||||
models_route._arch_to_task("wan", ("QuantStack/Wan2.2-TI2V-5B-GGUF",))
|
||||
== models_route._VIDEO_GEN_TASK
|
||||
)
|
||||
assert (
|
||||
models_route._arch_to_task("wan", (None, "Wan2.2-TI2V-5B-Q4_K_M.gguf"))
|
||||
== models_route._VIDEO_GEN_TASK
|
||||
)
|
||||
assert (
|
||||
models_route._arch_to_task("wan", ("QuantStack/Wan2.2-T2V-A14B-GGUF",))
|
||||
== models_route._UNSUPPORTED_DIFFUSION_TASK
|
||||
)
|
||||
# Drift guard: every diffusion arch llama.cpp rejects as a chat model must be
|
||||
# classified here as some non-chat task (image, video, or unsupported).
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
|
|
|||
|
|
@ -774,6 +774,71 @@ def test_diffusion_dataset_upload_rejects_unsupported_files(client, dataset_root
|
|||
assert "Unsupported file" in r.json()["detail"]
|
||||
|
||||
|
||||
def test_free_gpu_for_diffusion_training_unloads_video(monkeypatch):
|
||||
# A resident Video pipeline loads under the VIDEO arbiter owner, which the Images teardown
|
||||
# does not free; starting diffusion training must unload it too or the trainer OOMs against
|
||||
# the still-resident video model.
|
||||
import routes.training as tr
|
||||
from core.inference import gpu_arbiter
|
||||
|
||||
class _Exp:
|
||||
current_checkpoint = None
|
||||
|
||||
def is_export_active(self):
|
||||
return False
|
||||
|
||||
class _Diff:
|
||||
is_loaded = False
|
||||
|
||||
def unload(self):
|
||||
pass
|
||||
|
||||
unloaded = {"video": False}
|
||||
|
||||
class _Vid:
|
||||
def status(self):
|
||||
return {"loaded": True}
|
||||
|
||||
def unload(self):
|
||||
unloaded["video"] = True
|
||||
|
||||
released = []
|
||||
monkeypatch.setattr("core.export.get_export_backend", lambda: _Exp())
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion_engine_router.get_active_diffusion_engine", lambda: _Diff()
|
||||
)
|
||||
monkeypatch.setattr("core.inference.video.get_video_backend", lambda: _Vid())
|
||||
monkeypatch.setattr(gpu_arbiter, "release", lambda owner: released.append(owner))
|
||||
|
||||
tr._free_gpu_for_diffusion_training()
|
||||
|
||||
assert unloaded["video"] is True
|
||||
assert gpu_arbiter.VIDEO in released
|
||||
|
||||
|
||||
def test_import_example_partial_failure_leaves_no_partial_dataset(
|
||||
client, dataset_roots, monkeypatch
|
||||
):
|
||||
# A materialize that writes some images then fails must not leave a partial dataset: it stages
|
||||
# into a discarded temp dir, so the target folder stays empty and a retry re-materializes
|
||||
# instead of the image_count>0 idempotency check treating a truncated result as complete.
|
||||
import routes.training as tr
|
||||
|
||||
ds_root, _ = dataset_roots
|
||||
|
||||
def _boom(entry, dest, cap):
|
||||
(dest / "img_0000.png").write_bytes(b"x") # partial write into staging
|
||||
raise RuntimeError("transient copy error")
|
||||
|
||||
monkeypatch.setattr(tr, "_materialize_hf_dataset", _boom)
|
||||
r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "dreambooth-dog"})
|
||||
assert r.status_code == 502
|
||||
folder = ds_root / "dreambooth-dog"
|
||||
assert not folder.exists() or not any(folder.iterdir())
|
||||
# And no leftover staging dir surfaces as a dataset.
|
||||
assert not any(p.name.startswith(".dreambooth-dog.import-") for p in ds_root.iterdir())
|
||||
|
||||
|
||||
def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypatch):
|
||||
# A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed,
|
||||
# so a bad pick never unloads the user's working chat/Images model.
|
||||
|
|
|
|||
|
|
@ -904,6 +904,60 @@ class TestRouteErrors(unittest.TestCase):
|
|||
self.assertEqual(exc_info.exception.status_code, 400)
|
||||
self.assertIn("GGUF", exc_info.exception.detail)
|
||||
|
||||
def test_inference_route_defers_gpu_handoff_until_after_validation(self):
|
||||
# A doomed chat load (GGUF + gpu_ids -> 400) must NOT reclaim the CHAT arbiter owner
|
||||
# first: the handoff is deferred past validation, so a resident Images/Video pipeline is
|
||||
# never evicted for a load that then errors.
|
||||
import core.inference.gpu_arbiter as arb
|
||||
|
||||
inference_route = _load_route_module(
|
||||
"inference_route_module_for_handoff_test",
|
||||
"routes/inference.py",
|
||||
)
|
||||
request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1])
|
||||
model_config = SimpleNamespace(
|
||||
is_gguf = True,
|
||||
is_lora = False,
|
||||
gguf_hf_repo = None,
|
||||
gguf_file = "/tmp/test.gguf",
|
||||
gguf_mmproj_file = None,
|
||||
gguf_variant = None,
|
||||
identifier = "unsloth/test.gguf",
|
||||
display_name = "unsloth/test.gguf",
|
||||
is_vision = False,
|
||||
is_audio = False,
|
||||
audio_type = None,
|
||||
has_audio_input = False,
|
||||
)
|
||||
acquired = []
|
||||
with (
|
||||
patch.object(
|
||||
inference_route,
|
||||
"ModelConfig",
|
||||
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
||||
),
|
||||
patch.object(
|
||||
inference_route, "_guard_chat_load_against_training", return_value = None
|
||||
),
|
||||
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
||||
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
||||
patch.object(arb, "acquire_for", lambda owner: acquired.append(owner)),
|
||||
):
|
||||
with self.assertRaises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
inference_route._load_model_impl(
|
||||
request,
|
||||
SimpleNamespace(
|
||||
app = SimpleNamespace(
|
||||
state = SimpleNamespace(llama_parallel_slots = 1),
|
||||
),
|
||||
),
|
||||
current_subject = "test-user",
|
||||
)
|
||||
)
|
||||
self.assertEqual(exc_info.exception.status_code, 400)
|
||||
self.assertEqual(acquired, []) # no CHAT handoff before the doomed load errored
|
||||
|
||||
def test_training_route_returns_400_for_invalid_gpu_ids(self):
|
||||
training_route = _load_route_module(
|
||||
"training_route_module_for_test",
|
||||
|
|
|
|||
|
|
@ -165,6 +165,24 @@ class TestMaxBodyMiddleware:
|
|||
assert r.status_code == 200
|
||||
assert r.json()["total"] == 512
|
||||
|
||||
def test_diffusion_dataset_upload_in_body_passthrough(self, main_module):
|
||||
# The diffusion dataset upload route lives under the protected /api/train prefix, so it
|
||||
# must be in the REAL passthrough allowlist with the DB-aware + multipart-overhead cap;
|
||||
# otherwise MaxBodyMiddleware would 413 near-limit batches (and ignore a raised
|
||||
# max_upload_size_mb) before the handler's own get_upload_limit_bytes() check runs.
|
||||
from utils.upload_limits import (
|
||||
default_request_body_limit_bytes,
|
||||
upload_request_limit_bytes,
|
||||
)
|
||||
|
||||
path = "/api/train/diffusion/dataset"
|
||||
assert any(
|
||||
path.startswith(p) for p in main_module._BODY_UPLOAD_PASSTHROUGH_PREFIXES
|
||||
)
|
||||
cap = main_module._get_upload_passthrough_request_max_bytes(path)
|
||||
assert cap == upload_request_limit_bytes() # DB-aware cap + multipart overhead
|
||||
assert cap > default_request_body_limit_bytes() # not the plain default body cap
|
||||
|
||||
def test_upload_passthrough_rejects_declared_body_over_dedicated_cap(self, main_module):
|
||||
app = _make_protected_app(
|
||||
128,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue