Spaces:
Running on Zero
Running on Zero
| """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:<id>` 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, | |
| ) | |
| ) | |
| def provider(tmp_path): | |
| return LocalDiffusionImageProvider(str(tmp_path)) | |
| 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" | |
| 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" | |
| 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] | |