Plan the pre-cast text encoder, and make the cross-trainer GPU admission atomic
An fp8 text-encoder request loads a hosted PRE-CAST checkpoint, but the image download plan never received text_encoder_quant, so the manager staged the base repo's dense encoder (FLUX.2-dev's Mistral-24B is ~48 GB, Qwen-Image's Qwen2.5-VL ~16.6 GB) and the load then pulled the pre-cast file inline, outside the manager's progress and disk preflight. The plan now takes the field, resolves the hosted artifact with the same resolver the injection uses, stages that file, and drops only those components' dense weight shards. The load's own prefetch takes the same treatment, since it paid the same cost. Only a checkpoint that really resolves on the Hub earns the drop, so a gated or renamed artifact still stages the dense encoder the load will fall back to. The two trainers admitted each other with independent check-then-act guards: the diffusion route checks the LLM backend several network-bound preflights before it reserves, and the LLM route checks the diffusion service well before it spawns, so two near-simultaneous starts could both pass and train on one GPU. reserve() now re-tests the LLM backend under its own lock, and the LLM route holds the diffusion service's gpu_load_admission across its spawn, so exactly one of the two wins. Both halves fail open, so a chat-only install still trains.
This commit is contained in:
parent
a2342f80df
commit
794e43ffa1
7 changed files with 381 additions and 10 deletions
|
|
@ -891,6 +891,12 @@ class DiffusionBackend:
|
|||
kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token")
|
||||
)
|
||||
kwargs["base_repo"] = base
|
||||
# The pre-cast encoder injection replaces these components' weights, so prefetching the
|
||||
# dense shards for them downloads tens of GB nothing opens. Same resolver the injection
|
||||
# uses, so the prefetch and the load can never disagree about which encoders are dense.
|
||||
te_prequant_files = self._te_prequant_plan_files(
|
||||
fam, kwargs.get("text_encoder_quant"), kwargs.get("hf_token")
|
||||
)
|
||||
expected, base_files = self._estimate_download_bytes(
|
||||
kwargs["repo_id"],
|
||||
kwargs.get("gguf_filename"),
|
||||
|
|
@ -901,6 +907,7 @@ class DiffusionBackend:
|
|||
# Pull the base shards here rather than inside the locked, unpreemptable finalize.
|
||||
include_transformer = kind == "gguf"
|
||||
and self._dense_quant_prefetch_needed(fam, kwargs),
|
||||
skip_te_components = tuple(te_prequant_files),
|
||||
)
|
||||
with self._lock:
|
||||
# Stamp progress only if this load is still current (a superseder has its own token).
|
||||
|
|
@ -974,6 +981,33 @@ class DiffusionBackend:
|
|||
return ()
|
||||
return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
|
||||
|
||||
@staticmethod
|
||||
def _te_prequant_plan_files(
|
||||
fam: Any, text_encoder_quant: Optional[str], hf_token: Optional[str]
|
||||
) -> dict[str, tuple[str, list[tuple[str, int]]]]:
|
||||
"""``{component: (repo_id, [(rfilename, size)])}`` for the text encoders this pick will
|
||||
take PRE-CAST from a hosted checkpoint instead of the base repo's dense weights.
|
||||
|
||||
Empty unless the request asked for a scheme with a hosted artifact AND that artifact
|
||||
really resolves, so a plan can never drop a dense encoder the load still wants."""
|
||||
try:
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
from .diffusion_te_prequant import te_prequant_hub_files, te_prequant_sources
|
||||
|
||||
sources = te_prequant_sources(
|
||||
fam,
|
||||
te_quant_mode = text_encoder_quant,
|
||||
target = resolve_diffusion_device_target(),
|
||||
)
|
||||
if not sources:
|
||||
return {}
|
||||
files = te_prequant_hub_files(sources, HfApi(token = hf_token or None), logger)
|
||||
return {c: (sources[c].location, f) for c, f in files.items()}
|
||||
except Exception as exc: # noqa: BLE001 -- an unresolvable pre-cast keeps the dense encoder
|
||||
logger.warning("diffusion.te_prequant_plan_failed: %s", exc)
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def _estimate_download_bytes(
|
||||
repo_id: str,
|
||||
|
|
@ -985,6 +1019,7 @@ class DiffusionBackend:
|
|||
single_file_is_pipeline: bool = False,
|
||||
include_transformer: bool = False,
|
||||
sizes_out: Optional[dict[str, int]] = None,
|
||||
skip_te_components: tuple[str, ...] = (),
|
||||
) -> tuple[int, list[str]]:
|
||||
"""Total download size for the progress bar, plus the base-repo files to
|
||||
fetch (the prefetch reuses this list, so the base is listed only once).
|
||||
|
|
@ -997,16 +1032,35 @@ class DiffusionBackend:
|
|||
single-file paths, where the transformer is the single file and the base repo
|
||||
supplies only the companions. For a ``single_file_is_pipeline`` family (SDXL) the
|
||||
single file is the WHOLE pipeline, so the base repo supplies only config/tokenizer
|
||||
(no weights) and its weight files are skipped."""
|
||||
(no weights) and its weight files are skipped.
|
||||
|
||||
``skip_te_components`` names the text encoders this pick loads PRE-CAST from a hosted
|
||||
checkpoint, so their dense weight shards are not counted or fetched: staging the dense
|
||||
encoder for a pre-cast load wastes tens of GB (FLUX.2-dev's Mistral-24B is ~48 GB,
|
||||
Qwen-Image's Qwen2.5-VL ~16.6 GB) and nothing ever opens them. Everything else in the
|
||||
component folder (config, shard index, tokenizer) is kept -- the pre-cast loader
|
||||
meta-inits the encoder from the base repo's config."""
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
from .diffusion_te_prequant import is_prequant_covered_weight
|
||||
|
||||
api = HfApi()
|
||||
total = 0
|
||||
base_files: list[str] = []
|
||||
|
||||
def _dense_te_shard(rfilename: str) -> bool:
|
||||
return bool(skip_te_components) and is_prequant_covered_weight(
|
||||
rfilename, skip_te_components
|
||||
)
|
||||
|
||||
try:
|
||||
if kind == "pipeline":
|
||||
info = api.model_info(repo_id, files_metadata = True, token = hf_token)
|
||||
picked = [s for s in info.siblings if _pipeline_file_downloaded(s.rfilename)]
|
||||
picked = [
|
||||
s
|
||||
for s in info.siblings
|
||||
if _pipeline_file_downloaded(s.rfilename) and not _dense_te_shard(s.rfilename)
|
||||
]
|
||||
# diffusers prefers safetensors: drop a .bin whose dir also has a picked .safetensors.
|
||||
st_dirs = {
|
||||
s.rfilename.rsplit("/", 1)[0]
|
||||
|
|
@ -1040,7 +1094,7 @@ class DiffusionBackend:
|
|||
base_info = api.model_info(base_repo, files_metadata = True, token = hf_token)
|
||||
base_bytes = 0
|
||||
for s in base_info.siblings:
|
||||
if base_filter(s.rfilename):
|
||||
if base_filter(s.rfilename) and not _dense_te_shard(s.rfilename):
|
||||
base_files.append(s.rfilename)
|
||||
base_bytes += s.size or 0
|
||||
total += base_bytes
|
||||
|
|
@ -1059,6 +1113,7 @@ class DiffusionBackend:
|
|||
family_override: Optional[str] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
hf_token: Optional[str] = None,
|
||||
text_encoder_quant: Optional[str] = None,
|
||||
**load_kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""The repos + exact files this pick needs, so the Hub download manager can fetch
|
||||
|
|
@ -1067,13 +1122,21 @@ class DiffusionBackend:
|
|||
A plain snapshot_download would also pull what the loader deliberately skips (the
|
||||
packaged root single, transformer/ shards, fp16 twins) -- tens of GB per FLUX repo.
|
||||
Resolves family/kind/base exactly as ``_run_load`` does, so the plan and the load
|
||||
agree. Local paths are already on disk and yield no entries."""
|
||||
agree. Local paths are already on disk and yield no entries.
|
||||
|
||||
``text_encoder_quant`` is read for the same reason as the DiT quant: an fp8 request
|
||||
loads a hosted PRE-CAST encoder, so the base repo's dense encoder shards must not be
|
||||
staged and the pre-cast checkpoint must be. Without it the manager stages the dense
|
||||
encoder (tens of GB the load never opens) and the load then pulls the pre-cast file
|
||||
inline, outside the manager's progress and disk preflight."""
|
||||
fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
|
||||
kind = resolve_model_kind(gguf_filename, model_kind)
|
||||
if kind == "pipeline":
|
||||
base = repo_id # the full pipeline IS the repo
|
||||
else:
|
||||
base = _resolve_base_repo(repo_id, base_repo, fam, hf_token)
|
||||
# Only a checkpoint that really resolves on the Hub earns the right to drop dense shards.
|
||||
te_files = self._te_prequant_plan_files(fam, text_encoder_quant, hf_token)
|
||||
sizes: dict[str, int] = {}
|
||||
total, base_files = self._estimate_download_bytes(
|
||||
repo_id,
|
||||
|
|
@ -1085,8 +1148,19 @@ class DiffusionBackend:
|
|||
include_transformer = kind == "gguf"
|
||||
and self._dense_quant_prefetch_needed(fam, load_kwargs),
|
||||
sizes_out = sizes,
|
||||
skip_te_components = tuple(te_files),
|
||||
)
|
||||
entries: list[dict[str, Any]] = []
|
||||
for repo, files in te_files.values():
|
||||
entries.append(
|
||||
{
|
||||
"repo_id": repo,
|
||||
"files": [name for name, _size in files],
|
||||
"bytes": int(sum(size for _name, size in files)),
|
||||
"gguf_filename": None,
|
||||
}
|
||||
)
|
||||
total += int(sum(size for _name, size in files))
|
||||
if gguf_filename and not Path(repo_id).expanduser().exists():
|
||||
entries.append(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -418,6 +418,35 @@ def _validate_checkpoint(ckpt: Any, scheme: str, component: str, base: str, logg
|
|||
return True
|
||||
|
||||
|
||||
def te_prequant_hub_files(
|
||||
sources: dict[str, "TePrequantSource"], api: Any, logger: Any = None
|
||||
) -> dict[str, list[tuple[str, int]]]:
|
||||
"""``{component: [(rfilename, size)]}`` for every hosted pre-cast checkpoint that really
|
||||
resolves on the Hub.
|
||||
|
||||
Only a component listed here may have its dense weights dropped from a plan or a prefetch:
|
||||
an unpublished / gated / renamed artifact keeps its dense encoder, exactly as the load's own
|
||||
fallback does. Checked per source so one missing repo cannot sink the whole plan. A local
|
||||
path override is already on disk and is never staged."""
|
||||
found: dict[str, list[tuple[str, int]]] = {}
|
||||
for component, source in sources.items():
|
||||
if getattr(source, "kind", None) != "repo" or not getattr(source, "filename", None):
|
||||
continue
|
||||
try:
|
||||
info = api.model_info(source.location, files_metadata = True)
|
||||
except Exception as exc: # noqa: BLE001 -- unavailable pre-cast means the dense encoder
|
||||
_warn(logger, f"hub_files:{source.location}", exc)
|
||||
continue
|
||||
files = [
|
||||
(s.rfilename, int(getattr(s, "size", 0) or 0))
|
||||
for s in (info.siblings or [])
|
||||
if s.rfilename == source.filename
|
||||
]
|
||||
if files:
|
||||
found[component] = files
|
||||
return found
|
||||
|
||||
|
||||
def _has_meta_tensors(module: Any) -> bool:
|
||||
"""True if any parameter or buffer is still on the meta device after loading."""
|
||||
from itertools import chain
|
||||
|
|
|
|||
|
|
@ -78,6 +78,16 @@ class DatasetMutationInFlight(RuntimeError):
|
|||
"""A start was refused because a dataset mutation is open (route: 409)."""
|
||||
|
||||
|
||||
def _llm_training_active() -> bool:
|
||||
"""Whether the LLM trainer holds the GPU. Best-effort: an import/health failure must not
|
||||
block a diffusion start, exactly as the route's own reciprocal check fails open."""
|
||||
try:
|
||||
from core.training import get_training_backend
|
||||
return bool(get_training_backend().is_training_active())
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
|
||||
|
||||
# ── persisted run history ──────────────────────────────────────────────────────
|
||||
# Every terminal run is recorded as one JSON file (summary + scrubbed config + metric logs) so
|
||||
# the Train tab can show history. JSON, not the LLM sqlite, so diffusion runs stay off the LLM
|
||||
|
|
@ -275,6 +285,16 @@ class DiffusionTrainingService:
|
|||
"A model is being loaded onto the GPU right now. Wait for that to finish, "
|
||||
"then start the run."
|
||||
)
|
||||
# The LLM trainer under the SAME lock, not just at the route's earlier check: that check
|
||||
# and this reservation are separated by several network-bound preflights, so an LLM start
|
||||
# could spawn in between and both trainers would allocate on one GPU. The LLM route holds
|
||||
# gpu_load_admission() across its own spawn, so between the two either this raises or
|
||||
# that one does -- never neither.
|
||||
if _llm_training_active():
|
||||
raise RuntimeError(
|
||||
"An LLM training job is already running. "
|
||||
"Stop it before starting diffusion (Images) training."
|
||||
)
|
||||
self._reserved = True
|
||||
|
||||
def unreserve(self) -> None:
|
||||
|
|
|
|||
|
|
@ -16095,6 +16095,11 @@ async def diffusion_download_plan(
|
|||
model_kind = kind,
|
||||
hf_token = request.hf_token,
|
||||
transformer_quant = request.transformer_quant,
|
||||
# An fp8 encoder request loads a hosted pre-cast checkpoint, so the plan has to stage
|
||||
# that file instead of the base repo's dense encoder shards -- otherwise the manager
|
||||
# downloads tens of GB the load never opens and then pulls the pre-cast file inline,
|
||||
# outside its progress and disk preflight.
|
||||
text_encoder_quant = request.text_encoder_quant,
|
||||
speed_mode = request.speed_mode,
|
||||
# The dense-quant prefetch decision reads the memory policy, the prequant path and the adapter
|
||||
# selection too (an offload policy never runs the dense build; a baked LoRA always does), so the
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
Training API routes
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
|
@ -541,12 +542,28 @@ async def start_training(
|
|||
# generation locks until an in-flight denoise step hits its cancel callback (and the export
|
||||
# subprocess teardown can take seconds), which would otherwise freeze every concurrent
|
||||
# status/cancel/UI request. Overlapping starts are serialized by the backend's own guard.
|
||||
success = await asyncio.to_thread(
|
||||
backend.start_training,
|
||||
job_id = job_id,
|
||||
before_spawn = _free_vram_for_training,
|
||||
resume_source_run_id = resume_run["id"] if resume_run else None,
|
||||
**training_kwargs,
|
||||
#
|
||||
# The diffusion admission is held ACROSS the spawn so the cross-trainer decision is
|
||||
# atomic. The _diffusion_training_active() check above is separated from this point by
|
||||
# dataset validation and memory coordination, and the diffusion route likewise checks
|
||||
# this backend well before it reserves, so two near-simultaneous starts of different
|
||||
# types could both pass their checks and train on the same GPU. Entering this context
|
||||
# re-tests the diffusion state under the service's own lock, and while it is held
|
||||
# reserve() refuses -- so exactly one of the two wins.
|
||||
with _diffusion_gpu_admission():
|
||||
success = await asyncio.to_thread(
|
||||
backend.start_training,
|
||||
job_id = job_id,
|
||||
before_spawn = _free_vram_for_training,
|
||||
resume_source_run_id = resume_run["id"] if resume_run else None,
|
||||
**training_kwargs,
|
||||
)
|
||||
except _DiffusionStartInFlight as exc:
|
||||
return TrainingJobResponse(
|
||||
job_id = "",
|
||||
status = "error",
|
||||
message = str(exc),
|
||||
error = "Diffusion training already active",
|
||||
)
|
||||
except SidecarSwapInProgress as exc:
|
||||
# Expected loss of the race against a sidecar install: a retryable
|
||||
|
|
@ -1160,6 +1177,46 @@ def _diffusion_training_active() -> bool:
|
|||
return False
|
||||
|
||||
|
||||
class _DiffusionStartInFlight(RuntimeError):
|
||||
"""An LLM start lost the race to a diffusion start (route: refuse, don't spawn)."""
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _diffusion_gpu_admission():
|
||||
"""Hold the diffusion service's GPU admission across the LLM spawn.
|
||||
|
||||
Makes the cross-trainer admission atomic: entering re-tests the diffusion state under the
|
||||
service's own lock and raises if a diffusion run is reserved or active, and while it is held
|
||||
the diffusion ``reserve()`` refuses. So of two near-simultaneous starts of different types,
|
||||
exactly one proceeds. Fails OPEN on an import/health failure, like every other guard here: a
|
||||
chat-only install has no diffusion service and must still be able to train."""
|
||||
try:
|
||||
from core.training.diffusion_training_service import (
|
||||
TrainingActiveError,
|
||||
get_diffusion_training_service,
|
||||
)
|
||||
service = get_diffusion_training_service()
|
||||
except Exception: # noqa: BLE001 -- no diffusion stack: nothing to coordinate with
|
||||
yield
|
||||
return
|
||||
try:
|
||||
cm = service.gpu_load_admission()
|
||||
except Exception: # noqa: BLE001
|
||||
yield
|
||||
return
|
||||
try:
|
||||
cm.__enter__()
|
||||
except TrainingActiveError as exc:
|
||||
raise _DiffusionStartInFlight(
|
||||
"A diffusion (Images) LoRA training job is already running. "
|
||||
"Stop it before starting an LLM training run."
|
||||
) from exc
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
cm.__exit__(None, None, None)
|
||||
|
||||
|
||||
def _require_diffusion_dataset_mutable() -> None:
|
||||
"""Reject a dataset mutation while a diffusion run is active.
|
||||
|
||||
|
|
|
|||
|
|
@ -3868,3 +3868,118 @@ def test_download_plan_is_empty_for_a_local_path(tmp_path, monkeypatch):
|
|||
|
||||
plan = DiffusionBackend().download_plan(str(local), gguf_filename = "weights.gguf")
|
||||
assert plan["entries"] == []
|
||||
|
||||
|
||||
def test_download_plan_stages_the_precast_encoder_instead_of_the_dense_one(monkeypatch):
|
||||
# An fp8 text-encoder request loads a hosted PRE-CAST checkpoint, so the plan must stage that
|
||||
# file and NOT the base repo's dense encoder shards. Without this the manager downloaded tens of
|
||||
# GB the load never opens, and the load then pulled the pre-cast file inline, outside the
|
||||
# manager's progress and disk preflight.
|
||||
_fake_hf_api(
|
||||
monkeypatch,
|
||||
{
|
||||
"unsloth/FLUX.1-dev-GGUF": [_FakeSibling("flux1-dev-Q4_K_M.gguf", 7 * GB)],
|
||||
"black-forest-labs/FLUX.1-dev": _FLUX_BASE_SIBLINGS,
|
||||
"unsloth/FLUX.1-schnell-FP8": [
|
||||
_FakeSibling("text_encoder_2-fp8.pt", 1 * GB),
|
||||
_FakeSibling("README.md", 100),
|
||||
],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion._resolve_base_repo",
|
||||
lambda *a, **k: "black-forest-labs/FLUX.1-dev",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
DiffusionBackend, "_dense_quant_prefetch_needed", lambda self, fam, kwargs: False
|
||||
)
|
||||
# The pick resolves one hosted pre-cast encoder for text_encoder_2 (flux.1 hosts its T5-XXL).
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion_te_prequant.te_prequant_sources",
|
||||
lambda fam, *, te_quant_mode, target: (
|
||||
{
|
||||
"text_encoder_2": types.SimpleNamespace(
|
||||
kind = "repo",
|
||||
location = "unsloth/FLUX.1-schnell-FP8",
|
||||
filename = "text_encoder_2-fp8.pt",
|
||||
)
|
||||
}
|
||||
if te_quant_mode == "fp8"
|
||||
else {}
|
||||
),
|
||||
)
|
||||
|
||||
plan = DiffusionBackend().download_plan(
|
||||
"unsloth/FLUX.1-dev-GGUF",
|
||||
gguf_filename = "flux1-dev-Q4_K_M.gguf",
|
||||
text_encoder_quant = "fp8",
|
||||
)
|
||||
by_repo = {e["repo_id"]: e for e in plan["entries"]}
|
||||
assert "unsloth/FLUX.1-schnell-FP8" in by_repo
|
||||
assert by_repo["unsloth/FLUX.1-schnell-FP8"]["files"] == ["text_encoder_2-fp8.pt"]
|
||||
base = by_repo["black-forest-labs/FLUX.1-dev"]
|
||||
# text_encoder_2's dense weights are gone; text_encoder (no hosted artifact here) stays, and so
|
||||
# do the non-weight files the pre-cast loader still meta-inits from.
|
||||
assert not any(f.startswith("text_encoder_2/") and f.endswith(".safetensors") for f in base["files"])
|
||||
assert "text_encoder/model.safetensors" in base["files"]
|
||||
assert "model_index.json" in base["files"]
|
||||
assert plan["total_bytes"] == sum(e["bytes"] for e in plan["entries"])
|
||||
|
||||
|
||||
def test_download_plan_keeps_the_dense_encoder_without_an_fp8_request(monkeypatch):
|
||||
# No fp8 request -> no hosted checkpoint -> the dense encoder is exactly as before.
|
||||
_fake_hf_api(
|
||||
monkeypatch,
|
||||
{
|
||||
"unsloth/FLUX.1-dev-GGUF": [_FakeSibling("flux1-dev-Q4_K_M.gguf", 7 * GB)],
|
||||
"black-forest-labs/FLUX.1-dev": _FLUX_BASE_SIBLINGS,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion._resolve_base_repo",
|
||||
lambda *a, **k: "black-forest-labs/FLUX.1-dev",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
DiffusionBackend, "_dense_quant_prefetch_needed", lambda self, fam, kwargs: False
|
||||
)
|
||||
plan = DiffusionBackend().download_plan(
|
||||
"unsloth/FLUX.1-dev-GGUF", gguf_filename = "flux1-dev-Q4_K_M.gguf"
|
||||
)
|
||||
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 len(plan["entries"]) == 2
|
||||
|
||||
|
||||
def test_download_plan_keeps_the_dense_encoder_when_the_precast_repo_is_unavailable(monkeypatch):
|
||||
# A gated / renamed / unpublished artifact must NOT cost the dense encoder: the load falls back
|
||||
# to the dense weights, so the plan has to stage them.
|
||||
_fake_hf_api(
|
||||
monkeypatch,
|
||||
{
|
||||
"unsloth/FLUX.1-dev-GGUF": [_FakeSibling("flux1-dev-Q4_K_M.gguf", 7 * GB)],
|
||||
"black-forest-labs/FLUX.1-dev": _FLUX_BASE_SIBLINGS,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion._resolve_base_repo",
|
||||
lambda *a, **k: "black-forest-labs/FLUX.1-dev",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
DiffusionBackend, "_dense_quant_prefetch_needed", lambda self, fam, kwargs: False
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion_te_prequant.te_prequant_sources",
|
||||
lambda fam, *, te_quant_mode, target: {
|
||||
"text_encoder": types.SimpleNamespace(
|
||||
kind = "repo", location = "unsloth/does-not-exist", filename = "te-fp8.pt"
|
||||
)
|
||||
},
|
||||
)
|
||||
plan = DiffusionBackend().download_plan(
|
||||
"unsloth/FLUX.1-dev-GGUF",
|
||||
gguf_filename = "flux1-dev-Q4_K_M.gguf",
|
||||
text_encoder_quant = "fp8",
|
||||
)
|
||||
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"])
|
||||
|
|
|
|||
|
|
@ -1976,3 +1976,74 @@ def test_route_rejects_cond_cache_dir_for_sdxl(client):
|
|||
r = client.post("/api/train/diffusion/start", json = {**_BODY, "model_family": "sdxl"})
|
||||
assert r.status_code == 200, r.text
|
||||
assert client._fake.started_with["cond_cache_dir"] is None
|
||||
|
||||
|
||||
def test_service_reserve_refuses_while_the_llm_trainer_holds_the_gpu(monkeypatch):
|
||||
# The route's reciprocal check runs several network-bound preflights before reserve(), so an LLM
|
||||
# start could spawn inside that window and both trainers would allocate on one GPU. reserve()
|
||||
# re-tests the LLM backend under its own lock, which is the half of the interlock that closes.
|
||||
import types
|
||||
|
||||
import core.training.diffusion_training_service as dts
|
||||
from core.training.diffusion_training_service import DiffusionTrainingService
|
||||
|
||||
svc = DiffusionTrainingService()
|
||||
monkeypatch.setattr(dts, "_llm_training_active", lambda: True)
|
||||
with pytest.raises(RuntimeError, match = "LLM training job is already running"):
|
||||
svc.reserve()
|
||||
assert svc.is_active() is False # the refused start left no claim behind
|
||||
|
||||
monkeypatch.setattr(dts, "_llm_training_active", lambda: False)
|
||||
svc.reserve()
|
||||
assert svc.is_active() is True
|
||||
|
||||
|
||||
def test_llm_active_probe_fails_open(monkeypatch):
|
||||
# A chat-only install (or a wedged backend) must not block a diffusion start.
|
||||
import sys
|
||||
|
||||
import core.training.diffusion_training_service as dts
|
||||
|
||||
monkeypatch.setitem(sys.modules, "core.training", None) # import raises
|
||||
assert dts._llm_training_active() is False
|
||||
|
||||
|
||||
def test_llm_start_holds_the_diffusion_admission_across_its_spawn(monkeypatch):
|
||||
# The other half: while the LLM route is spawning, a diffusion reserve() must lose. The route
|
||||
# holds gpu_load_admission() across start_training, and reserve() already refuses while an
|
||||
# admission is open.
|
||||
from core.training.diffusion_training_service import DiffusionTrainingService
|
||||
import routes.training as tr
|
||||
|
||||
svc = DiffusionTrainingService()
|
||||
monkeypatch.setattr(
|
||||
"core.training.diffusion_training_service.get_diffusion_training_service", lambda: svc
|
||||
)
|
||||
with tr._diffusion_gpu_admission():
|
||||
with pytest.raises(RuntimeError, match = "being loaded onto the GPU"):
|
||||
svc.reserve()
|
||||
svc.reserve() # released once the spawn is done
|
||||
|
||||
|
||||
def test_llm_start_admission_refuses_when_diffusion_is_already_reserved(monkeypatch):
|
||||
from core.training.diffusion_training_service import DiffusionTrainingService
|
||||
import routes.training as tr
|
||||
|
||||
svc = DiffusionTrainingService()
|
||||
monkeypatch.setattr(
|
||||
"core.training.diffusion_training_service.get_diffusion_training_service", lambda: svc
|
||||
)
|
||||
svc.reserve()
|
||||
with pytest.raises(tr._DiffusionStartInFlight):
|
||||
with tr._diffusion_gpu_admission():
|
||||
pass
|
||||
|
||||
|
||||
def test_llm_start_admission_fails_open_without_a_diffusion_stack(monkeypatch):
|
||||
import sys
|
||||
|
||||
import routes.training as tr
|
||||
|
||||
monkeypatch.setitem(sys.modules, "core.training.diffusion_training_service", None)
|
||||
with tr._diffusion_gpu_admission():
|
||||
pass # a chat-only install still trains
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue