"""The automatic model-fallback loop for one image-edit request. What is under test is the ORDER, not the policy: `model_fallback` owns what counts as recoverable, and these tests pin that the provider honours it - one model per attempt, never the same model twice, a release between attempts, a hard stop at three, and a final report that names the whole chain. The selector (`_plan_attempt`) and the GPU boundaries are scripted, so nothing here downloads a weight or needs a card. The loop itself is real: it is the same object the service runs, deciding what to try next and when to stop. """ from __future__ import annotations import asyncio import pytest from PIL import Image from app.services import model_fallback from app.services.providers import LocalDiffusionImageProvider from app.services.providers import image_provider as ip KLEIN = "black-forest-labs/FLUX.2-klein-4b-fp8" SDXL = "stabilityai/stable-diffusion-xl-inpainting-1.0" KONTEXT = "black-forest-labs/FLUX.1-Kontext-dev" QWEN = "Qwen/Qwen-Image-Edit-2511" OOM = "CUDA out of memory. Tried to allocate 20.00 GiB" BAD_INPUT = "reference image avatar.png is not a readable file" def _script(prov, order, failures, ran, monkeypatch): """Make `prov` plan `order` and fail `failures`, recording every attempt. `order` is what the selector will name, in sequence; `failures` maps a model id to the error it raises inside its boundary. `ran` collects, in order, each model that entered a boundary and each `released:` cleanup. """ pending = list(order) def plan(self, torch, *, editing): # The real selector gates `SelectionRequest.exclude_model_ids`, which # `_select_edit_model` fills from `self._failed_model_ids`. Honouring it # here is the point: it proves the loop actually hands the selector the # models this request has already failed on. eligible = [m for m in pending if m not in self._failed_model_ids] if not eligible: raise ip.ModelSelectionError( "no registered image-editing model can serve this request" ) return eligible[0], 4 def boundary_for(model_id): def boundary(job): ran.append(job["selected_model"]) error = failures.get(model_id) if error is not None: raise error return "rendered" return boundary prov._plan_attempt = plan.__get__(prov, LocalDiffusionImageProvider) prov._release_attempt_resources = lambda failed: ran.append("released:%s" % failed) monkeypatch.setattr(ip, "_boundary_for", boundary_for) return prov def _edit(prov, reference): return asyncio.run( prov.generate_image( prompt="a red coat", character_name="Aria", character_details="dark hair", aspect_ratio="1:1", output_path="out.png", reference_image_path=reference, ) ) @pytest.fixture def provider(tmp_path): return LocalDiffusionImageProvider(str(tmp_path)) @pytest.fixture def reference(tmp_path): """A REAL reference file: the provider checks readability before planning.""" path = tmp_path / "avatar.png" Image.new("RGB", (32, 32), (200, 160, 140)).save(path) return str(path) def test_a_recoverable_failure_falls_back_to_the_next_compatible_model( provider, reference, monkeypatch ): """The whole point: one model dies on OOM, the next one serves the request.""" ran = [] _script(provider, [KLEIN, SDXL], {KLEIN: RuntimeError(OOM)}, ran, monkeypatch) assert _edit(provider, reference) == "rendered" assert ran == [KLEIN, "released:%s" % KLEIN, SDXL] def test_a_model_that_already_failed_is_never_tried_again( provider, reference, monkeypatch ): """A repeat run must not re-enter a model this request already failed on.""" ran = [] _script( provider, [KLEIN, SDXL, SDXL], {KLEIN: RuntimeError(OOM), SDXL: RuntimeError(OOM)}, ran, monkeypatch, ) with pytest.raises(RuntimeError): _edit(provider, reference) model_runs = [item for item in ran if not item.startswith("released:")] assert model_runs == [KLEIN, SDXL], ran def test_three_attempts_is_the_hard_ceiling(provider, reference, monkeypatch): """One initial model plus two fallbacks, never a fourth.""" ran = [] # Four DISTINCT models are on offer, and the loop must still stop after three. _script( provider, [KLEIN, SDXL, KONTEXT, QWEN], {KLEIN: RuntimeError(OOM), SDXL: RuntimeError(OOM), KONTEXT: RuntimeError(OOM)}, ran, monkeypatch, ) with pytest.raises(RuntimeError) as failure: _edit(provider, reference) model_runs = [item for item in ran if not item.startswith("released:")] assert model_runs == [KLEIN, SDXL, KONTEXT], ran # The fourth was available and was never entered. assert QWEN not in ran, ran assert str(failure.value).count("attempt ") == model_fallback.MAX_ATTEMPTS def test_a_permanent_failure_is_never_retried(provider, reference, monkeypatch): """A missing reference image fails identically for every model.""" ran = [] _script( provider, [KLEIN, SDXL, SDXL], {KLEIN: RuntimeError(BAD_INPUT)}, ran, monkeypatch ) with pytest.raises(RuntimeError) as failure: _edit(provider, reference) assert ran == [KLEIN], "no second model was tried: %r" % (ran,) assert BAD_INPUT in str(failure.value) def test_the_final_error_names_the_original_failure_and_every_model( provider, reference, monkeypatch ): """Requirement: the report leads with the cause and lists every attempt.""" ran = [] _script( provider, [KLEIN, SDXL], { KLEIN: RuntimeError(OOM), SDXL: RuntimeError("SDXL pipeline initialization failed"), }, ran, monkeypatch, ) with pytest.raises(RuntimeError) as failure: _edit(provider, reference) message = str(failure.value) assert OOM in message, "the original failure must come first" assert "attempt 1" in message and "attempt 2" in message assert KLEIN in message and SDXL in message assert "pipeline initialization failed" in message def test_a_reference_request_is_never_turned_into_text_to_image( provider, reference, monkeypatch ): """A failed edit must not silently become a brand new picture.""" entered = [] def plan(self, torch, *, editing): entered.append("editing" if editing else "text-to-image") return (SDXL if editing else ip.GENERATION_MODEL_ID), 4 def boundary_for(model_id): def boundary(job): entered.append(job["selected_model"]) raise RuntimeError(OOM) return boundary provider._plan_attempt = plan.__get__(provider, LocalDiffusionImageProvider) provider._release_attempt_resources = lambda failed: None monkeypatch.setattr(ip, "_boundary_for", boundary_for) with pytest.raises(RuntimeError): _edit(provider, reference) assert entered, "nothing ran" # Every attempt planned the EDIT path; none reached the text-to-image boundary. assert "text-to-image" not in entered, entered assert ip.GENERATION_MODEL_ID not in entered, entered def test_a_plain_success_never_falls_back(provider, reference, monkeypatch): """No failure means no fallback: the first compatible model is the only one.""" ran = [] _script(provider, [KLEIN, SDXL], {}, ran, monkeypatch) assert _edit(provider, reference) == "rendered" assert ran == [KLEIN] def test_the_budget_is_three_and_no_more(): """The policy constants themselves, so a later edit cannot quietly widen them.""" assert model_fallback.MAX_FALLBACKS == 2 assert model_fallback.MAX_ATTEMPTS == 3 # --------------------------------------------------------------------------- # Provider-wide outages # # Every model on this service is served by the SAME GPU provider, so when the # provider is out, swapping checkpoints cannot help - it only repeats the same # wait once per model. These pin that the loop stops instead. # --------------------------------------------------------------------------- def test_a_zero_gpu_allocation_failure_is_a_provider_outage(): """Recoverable, yes - but NOT fixable by trying another model.""" error = _gradio_error(NO_GPU_60S) # Still recoverable: a later request may well get a card. assert model_fallback.classify(error) is True # ...and that is precisely why the two must be asked separately. assert model_fallback.is_provider_outage(error) is True def test_a_model_level_failure_is_not_a_provider_outage(): """OOM on a working card is about the MODEL, so the fallback should proceed.""" oom = RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB") assert model_fallback.classify(oom) is True assert model_fallback.is_provider_outage(oom) is False def test_a_missing_provider_is_a_provider_outage(): """The fast-fail raised before any boundary is entered.""" error = RuntimeError( "%s: no local CUDA device was found" % model_fallback.PROVIDER_UNAVAILABLE_MESSAGE ) assert model_fallback.is_provider_outage(error) is True def test_a_provider_outage_does_not_walk_through_every_model( provider, reference, monkeypatch ): """The headline behaviour: one attempt, not 60s x 3. A ZeroGPU scheduler that is out of capacity will refuse the next model exactly as it refused this one, so the loop stops after the first. """ ran = [] _script( provider, [KLEIN, SDXL, KONTEXT], {KLEIN: _gradio_error(NO_GPU_60S), SDXL: _gradio_error(NO_GPU_60S)}, ran, monkeypatch, ) with pytest.raises(Exception) as failure: _edit(provider, reference) model_runs = [item for item in ran if not item.startswith("released:")] assert model_runs == [KLEIN], "only one model may be tried: %r" % (ran,) # The real reason survives to the caller, not a "all models failed" summary. assert NO_GPU_60S in str(failure.value) def test_a_model_level_failure_still_falls_back(provider, reference, monkeypatch): """The existing model fallback is untouched by the outage short-circuit.""" ran = [] _script( provider, [KLEIN, SDXL], {KLEIN: RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB")}, ran, monkeypatch, ) assert _edit(provider, reference) == "rendered" assert ran == [KLEIN, "released:%s" % KLEIN, SDXL] # --------------------------------------------------------------------------- # ZeroGPU temporary GPU-allocation failures # # `spaces` raises a generic `gradio.exceptions.Error` for these, so the exception # class says nothing at all: a busy scheduler and a corrupt input are the SAME # class. These pin that the MESSAGE decides, in both directions. # --------------------------------------------------------------------------- NO_GPU_60S = ( "No GPU was available after 60s. Try re-running outside of examples if it " "happened after clicking one" ) def _gradio_error(message): """The exception class ZeroGPU actually raises. Used so the tests exercise the real class rather than a stand-in: that shared class is exactly why message matching is required here. """ from gradio.exceptions import Error return Error(message) def test_a_busy_zerogpu_scheduler_is_recoverable_and_named(): """The reported bug: a full queue looked like a permanent failure.""" error = _gradio_error(NO_GPU_60S) assert model_fallback.is_gpu_allocation_failure(error) is True assert model_fallback.classify(error) is True assert model_fallback.error_type(error) == "gpu_allocation" @pytest.mark.parametrize( "message", [ NO_GPU_60S, "ZeroGPU quota exceeded", "Waiting for a GPU to become available", ], ) def test_every_zerogpu_scheduling_message_is_recoverable(message): assert model_fallback.classify(_gradio_error(message)) is True assert model_fallback.error_type(_gradio_error(message)) == "gpu_allocation" @pytest.mark.parametrize( "message", [ "Invalid prompt: the reference image is corrupt", "401 Unauthorized: invalid token", "the input image is not a readable file", "TypeError: unexpected keyword argument 'foo'", ], ) def test_other_gradio_errors_are_not_mistaken_for_a_busy_gpu(message): """The exception class alone must NEVER make a failure recoverable. Every message here is raised as the same `gradio.exceptions.Error` as the scheduler refusal above. If the class were the test, all of these would be retried against other models - which cannot fix any of them. """ error = _gradio_error(message) assert model_fallback.is_gpu_allocation_failure(error) is False assert model_fallback.error_type(error) == "Error" assert model_fallback.classify(error) is False def test_the_gpu_allocation_classification_survives_any_exception_class(): """It is the message, not the class: a plain RuntimeError classifies the same.""" assert model_fallback.classify(RuntimeError(NO_GPU_60S)) is True assert model_fallback.error_type(RuntimeError(NO_GPU_60S)) == "gpu_allocation" def test_a_permanent_phrase_inside_a_scheduler_message_still_wins(): """Non-recoverable is checked first, as everywhere else.""" error = _gradio_error("ZeroGPU quota exceeded: the reference image is corrupt") assert model_fallback.classify(error) is False assert model_fallback.error_type(error) == "gpu_allocation" def test_a_busy_scheduler_continues_to_the_next_model(provider, reference, monkeypatch): """End to end: Klein is refused a GPU, SDXL serves the request.""" ran = [] _script( provider, [KLEIN, SDXL], {KLEIN: _gradio_error(NO_GPU_60S)}, ran, monkeypatch ) assert _edit(provider, reference) == "rendered" assert ran == [KLEIN, "released:%s" % KLEIN, SDXL]