From b6ab866c4696c9b0ca30e42f96d69162aff43b7d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:08:41 +0000 Subject: [PATCH] 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. --- studio/backend/main.py | 12 ++- studio/backend/routes/inference.py | 27 ++++-- studio/backend/routes/models.py | 39 +++++--- studio/backend/routes/training.py | 89 +++++++++++++++---- .../backend/tests/test_cached_gguf_routes.py | 23 ++++- .../backend/tests/test_diffusion_training.py | 65 ++++++++++++++ studio/backend/tests/test_gpu_selection.py | 54 +++++++++++ studio/backend/tests/test_middleware.py | 18 ++++ 8 files changed, 289 insertions(+), 38 deletions(-) diff --git a/studio/backend/main.py b/studio/backend/main.py index e5dc29a505..6297bf21d7 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -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() diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8e06b66784..ed66b46113 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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() diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 005e3f7120..fd584acfdc 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 318de3d85d..0d909241f0 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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, diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index b051dba840..1dc391cd30 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 9f477a3545..a35f6ee1a3 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -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. diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index 69ad560788..1267842b3e 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -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", diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 11aeee6d77..21e8cf1058 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -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,