diff --git a/.changeset/mean-tips-dig.md b/.changeset/mean-tips-dig.md new file mode 100644 index 0000000000..429d1fcb04 --- /dev/null +++ b/.changeset/mean-tips-dig.md @@ -0,0 +1,5 @@ +--- +"gradio": patch +--- + +fix:Drain run-history writes on shutdown diff --git a/gradio/route_utils.py b/gradio/route_utils.py index 87b364ede1..d5fffadc58 100644 --- a/gradio/route_utils.py +++ b/gradio/route_utils.py @@ -1157,6 +1157,34 @@ async def _delete_state_handler(app: App): await _cancel_background_task(task) +HISTORY_DRAIN_TIMEOUT = 5.0 + + +@asynccontextmanager +async def _drain_history_tasks(app: App, timeout: float = HISTORY_DRAIN_TIMEOUT): + """On shutdown, give in-flight run-history writes a bounded chance to finish. + + A record is filed as a detached task so a prediction never waits on the Hub, + and nothing else ever awaits those tasks. Without this, a shutdown cancels + whatever is still in `record_run`, and a record that has not yet reached the + Hub call is lost outright. + """ + try: + yield + finally: + # `state` itself is absent on the stand-in apps the lifespan tests use. + state = getattr(app, "state", None) + tasks = { + task + for task in getattr(state, "history_tasks", None) or () + if not task.done() + } + if tasks: + _, pending = await asyncio.wait(tasks, timeout=timeout) + for task in pending: + task.cancel() + + def create_lifespan_handler( user_lifespan: Callable[[App], AbstractAsyncContextManager] | None, frequency: int | None = 1, @@ -1173,6 +1201,9 @@ async def _handler(app: App): await stack.enter_async_context(_lifespan_handler(app, frequency, age)) if user_lifespan is not None: state = await stack.enter_async_context(user_lifespan(app)) + # Entered last so it unwinds first: the records are drained while + # the rest of the app is still up. + await stack.enter_async_context(_drain_history_tasks(app)) yield state return _handler diff --git a/test/test_history.py b/test/test_history.py index 3fbe4fcc15..fbf83feaf4 100644 --- a/test/test_history.py +++ b/test/test_history.py @@ -458,6 +458,43 @@ def test_a_hub_failure_does_not_fail_the_prediction(self, recording_app): assert r.json()["data"] == ["hello world"] +def test_a_record_still_in_flight_survives_shutdown(monkeypatch): + """Records are filed as detached tasks, so a shutdown that does not wait for + them loses whatever has not reached the Hub yet. The delay goes in the + prelude rather than in the Hub call on purpose: `anyio.to_thread.run_sync` + cannot interrupt its worker thread, so a write that has already started + finishes either way and would not tell the two behaviours apart.""" + monkeypatch.delenv("GRADIO_HISTORY_BUCKET", raising=False) + hub = FakeHub() + original = history_mod.externalize_assets + + async def slow_externalize(*args, **kwargs): + await asyncio.sleep(0.3) + return await original(*args, **kwargs) + + def greet(name): + return f"hello {name}" + + io = gr.Interface(greet, "text", "text", api_name="greet") + app, _, _ = io.launch(prevent_thread_lock=True) + with ( + patch("gradio.history.resolve_token", return_value="tok"), + patch("gradio.history.HfApi", return_value=hub), + patch("gradio.history.externalize_assets", slow_externalize), + TestClient(app) as client, + ): + r = client.post( + "/gradio_api/run/greet", + json={"data": ["world"]}, + headers={"X-Gradio-History-Bucket": "alice/hist"}, + ) + assert r.status_code == 200, r.text + # Leaving the block runs the lifespan shutdown, which is where the drain is. + assert len(hub.files) == 1, "the record was dropped by the shutdown" + io.close() + close_all() + + # ------------------------------------------------- end-to-end: workflow canvas