diff --git a/.changeset/whole-radios-lick.md b/.changeset/whole-radios-lick.md new file mode 100644 index 00000000000..887b660fad8 --- /dev/null +++ b/.changeset/whole-radios-lick.md @@ -0,0 +1,10 @@ +--- +"@gradio/client": minor +"@gradio/core": minor +"@self/app": minor +"@self/spa": minor +"gradio": minor +"gradio_client": minor +--- + +feat:Resume sessions diff --git a/client/js/src/client.ts b/client/js/src/client.ts index a43e04bf856..f6c4dab5eb7 100644 --- a/client/js/src/client.ts +++ b/client/js/src/client.ts @@ -22,6 +22,11 @@ import { post_data } from "./utils/post_data"; import { predict } from "./utils/predict"; import { duplicate } from "./utils/duplicate"; import { submit } from "./utils/submit"; +import { + get_resumable_events, + get_resumable_session_hash, + type ResumableJob +} from "./utils/session"; import { RE_SPACE_NAME, process_endpoint } from "./helpers/api_info"; import { map_names_to_ids, @@ -38,6 +43,7 @@ import { clear_run_history } from "./utils/run_history"; import { API_INFO_ERROR_MSG, APP_ID_URL, + CLOSE_URL, CONFIG_ERROR_MSG, HEARTBEAT_URL, COMPONENT_SERVER_URL @@ -58,14 +64,18 @@ export class Client { last_status: Record = {}; private cookies: string | null = null; + private restored_session_hash = false; // streaming stream_status = { open: false }; closed = false; - pending_stream_messages: Record = {}; + pending_stream_messages: Record = {}; pending_diff_streams: Record = {}; event_callbacks: Record Promise> = {}; unclosed_events: Set = new Set(); + events_to_resume: Set = new Set(); + stream_reconnect_attempts = 0; + stream_reconnect_timer: ReturnType | null = null; heartbeat_event: EventSource | null = null; abort_controller: AbortController | null = null; stream_instance: EventSource | null = null; @@ -182,6 +192,7 @@ export class Client { trigger_id?: number | null, all_events?: boolean ) => SubmitIterable; + resume_jobs: (jobs?: ResumableJob[]) => SubmitIterable[]; predict: ( endpoint: string | number, data: unknown[] | Record | undefined, @@ -216,6 +227,21 @@ export class Client { this.handle_blob = handle_blob.bind(this); this.post_data = post_data.bind(this); this.submit = submit.bind(this); + this.resume_jobs = (jobs = this.get_resumable_events()) => { + this.options.resume_sessions = true; + return jobs.map(({ fn_index, event_id }) => + submit.call( + this, + fn_index, + {}, + undefined, + undefined, + undefined, + undefined, + event_id + ) + ); + }; this.predict = predict.bind(this) as typeof this.predict; this.open_stream = open_stream.bind(this); this.resolve_config = resolve_config.bind(this); @@ -233,10 +259,19 @@ export class Client { await this.resolve_cookies(); } - await this._resolve_config().then( - (res: { config: Config } | undefined) => - res?.config && this._resolve_heartbeat(res.config) - ); + const resolved = (await this._resolve_config()) as + | { config: Config } + | undefined; + if (resolved?.config) { + if ( + this.restored_session_hash && + typeof sessionStorage !== "undefined" && + get_resumable_events(resolved.config, this.session_hash).length === 0 + ) { + this.session_hash = Math.random().toString(36).substring(2); + } + await this._resolve_heartbeat(resolved.config); + } try { this.api_info = await this.view_api(); @@ -294,13 +329,24 @@ export class Client { } ): Promise { const client = new this(app_reference, options); // this refers to the class itself, not the instance - if (options.session_hash) { - client.session_hash = options.session_hash; + const session_hash = + options.session_hash || + (options.resume_sessions + ? get_resumable_session_hash(options.cookies) + : null); + if (session_hash) { + client.session_hash = session_hash; + client.restored_session_hash = !options.session_hash; } await client.init(); return client; } + get_resumable_events(): ResumableJob[] { + if (!this.options.resume_sessions || !this.config) return []; + return get_resumable_events(this.config, this.session_hash); + } + async reconnect(): Promise<"connected" | "broken" | "changed"> { const app_id_url = new URL( `${this.config!.root}${this.api_prefix}/${APP_ID_URL}` @@ -322,7 +368,31 @@ export class Client { } close(): void { + if ( + !this.closed && + this.options.resume_sessions && + this.config && + typeof window !== "undefined" + ) { + const headers: Record = { + "Content-Type": "application/json" + }; + if (this.options.token) { + headers.Authorization = `Bearer ${this.options.token}`; + } + void this.fetch(`${this.config.root}${this.api_prefix}/${CLOSE_URL}`, { + method: "POST", + headers, + body: JSON.stringify({ session_hash: this.session_hash }), + credentials: this.options.credentials ?? "same-origin", + keepalive: true + }).catch(() => {}); + } this.closed = true; + if (this.stream_reconnect_timer) { + clearTimeout(this.stream_reconnect_timer); + this.stream_reconnect_timer = null; + } close_stream(this.stream_status, this.abort_controller); } diff --git a/client/js/src/constants.ts b/client/js/src/constants.ts index 4da0d13ed43..a6f6762cdfc 100644 --- a/client/js/src/constants.ts +++ b/client/js/src/constants.ts @@ -15,6 +15,7 @@ export const HEARTBEAT_URL = `heartbeat`; export const COMPONENT_SERVER_URL = `component_server`; export const RESET_URL = `reset`; export const CANCEL_URL = `cancel`; +export const CLOSE_URL = `queue/close`; export const APP_ID_URL = `app_id`; export const RAW_API_INFO_URL = `info?serialize=False`; diff --git a/client/js/src/index.ts b/client/js/src/index.ts index 0c3185edf6e..d5a73940972 100644 --- a/client/js/src/index.ts +++ b/client/js/src/index.ts @@ -5,6 +5,7 @@ export { submit } from "./utils/submit"; export { upload_files } from "./utils/upload_files"; export { FileData, upload, prepare_files } from "./upload"; export { handle_file } from "./helpers/data"; +export type { ResumableJob } from "./utils/session"; export { clear_run_history, consume_run_history_replay, diff --git a/client/js/src/test/init.test.ts b/client/js/src/test/init.test.ts index 17f8ec50916..211a500c4a2 100644 --- a/client/js/src/test/init.test.ts +++ b/client/js/src/test/init.test.ts @@ -16,6 +16,8 @@ import { } from "./test_data"; import { initialise_server } from "./server"; import { SPACE_NOT_FOUND_MSG } from "../constants"; +import { get_resumable_events, track_resumable_event } from "../utils/session"; +import { http, HttpResponse } from "msw"; const app_reference = "hmb/hello_world"; const broken_app_reference = "hmb/bye_world"; @@ -65,6 +67,23 @@ describe("Client class", () => { expect(app.config).toEqual(config_response); }); + test.skipIf(typeof sessionStorage === "undefined")( + "does not restore a session from a different app", + async () => { + track_resumable_event( + { ...config_response, app_id: "another-app" }, + "restored-session", + { event_id: "event-id", fn_index: 0 } + ); + + const app = await Client.connect(direct_app_reference, { + resume_sessions: true + }); + + expect(app.session_hash).not.toBe("restored-session"); + } + ); + test("connecting successfully to a private running app with a space reference", async () => { const app = await Client.connect("hmb/secret_world", { token: "hf_123" @@ -122,6 +141,85 @@ describe("Client class", () => { }); }); + describe("resume_jobs", () => { + test("reattaches each job to its existing queue event", async () => { + const app = await Client.connect(direct_app_reference); + app.stream_status.open = true; + + const submissions = app.resume_jobs([ + { event_id: "event-1", fn_index: 0 }, + { event_id: "event-2", fn_index: 0 } + ]); + + await expect(submissions[0].wait_for_id()).resolves.toBe("event-1"); + await expect(submissions[1].wait_for_id()).resolves.toBe("event-2"); + expect(app.event_callbacks["event-1"]).toBeDefined(); + expect(app.event_callbacks["event-2"]).toBeDefined(); + expect(app.options.resume_sessions).toBe(true); + + await Promise.all(submissions.map((submission) => submission.return())); + }); + + test.skipIf(typeof sessionStorage === "undefined")( + "clears a resumable event rejected by the server", + async () => { + const app = await Client.connect(direct_app_reference); + app.stream_status.open = true; + track_resumable_event(app.config!, app.session_hash, { + event_id: "expired-event", + fn_index: 0 + }); + + const submission = app.resume_jobs([ + { event_id: "expired-event", fn_index: 0 } + ])[0]; + await submission.wait_for_id(); + await app.event_callbacks["expired-event"]({ + msg: "unexpected_error", + event_id: "expired-event", + message: "Session event not found.", + session_not_found: true + }); + + expect(get_resumable_events(app.config!, app.session_hash)).toEqual([]); + await submission.return(); + } + ); + + test.skipIf(typeof window === "undefined")( + "notifies the server when a resumable client closes", + async () => { + const app = await Client.connect(secret_direct_app_reference, { + token: "hf_123", + resume_sessions: true + }); + let received_session_hash: string | undefined; + let received_authorization: string | null = null; + let resolve_request: () => void = () => {}; + const request_received = new Promise((resolve) => { + resolve_request = resolve; + }); + server.resetHandlers( + http.post( + `${secret_direct_app_reference}/queue/close`, + async ({ request }) => { + const body = (await request.json()) as { session_hash: string }; + received_session_hash = body.session_hash; + received_authorization = request.headers.get("Authorization"); + resolve_request(); + return HttpResponse.json({ success: true }); + } + ) + ); + app.close(); + await request_received; + + expect(received_session_hash).toBe(app.session_hash); + expect(received_authorization).toBe("Bearer hf_123"); + } + ); + }); + describe("duplicate", () => { test("backwards compatibility of duplicate using deprecated syntax", async () => { const app = await duplicate("gradio/hello_world", { diff --git a/client/js/src/test/session.test.ts b/client/js/src/test/session.test.ts new file mode 100644 index 00000000000..8cc0bc7e765 --- /dev/null +++ b/client/js/src/test/session.test.ts @@ -0,0 +1,105 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { Config } from "../types"; +import { + clear_resumable_event, + get_resumable_events, + get_resumable_session_hash, + track_resumable_event +} from "../utils/session"; + +class MemoryStorage implements Storage { + values = new Map(); + + get length(): number { + return this.values.size; + } + + clear(): void { + this.values.clear(); + } + + getItem(key: string): string | null { + return this.values.get(key) ?? null; + } + + key(index: number): string | null { + return [...this.values.keys()][index] ?? null; + } + + removeItem(key: string): void { + this.values.delete(key); + } + + setItem(key: string, value: string): void { + this.values.set(key, value); + } +} + +const config = { + app_id: "app-1", + root: + typeof location === "undefined" ? "https://example.com/" : location.origin +} as Config; + +describe("resumable sessions", () => { + beforeEach(() => { + if (typeof sessionStorage === "undefined") { + vi.stubGlobal("sessionStorage", new MemoryStorage()); + } + sessionStorage.clear(); + if (typeof document !== "undefined") { + document.cookie = "gradio_active_session=; Path=/; Max-Age=0"; + } + }); + + it("stores only active events for the current app session", () => { + track_resumable_event(config, "session-1", { + event_id: "event-1", + fn_index: 3 + }); + + expect(get_resumable_session_hash()).toBe("session-1"); + expect(get_resumable_events(config, "session-1")).toEqual([ + { event_id: "event-1", fn_index: 3 } + ]); + + clear_resumable_event("event-1"); + expect(get_resumable_events(config, "session-1")).toEqual([]); + }); + + it("drops events when the app has changed", () => { + track_resumable_event(config, "session-1", { + event_id: "event-1", + fn_index: 3 + }); + + expect( + get_resumable_events({ ...config, app_id: "app-2" }, "session-1") + ).toEqual([]); + expect(get_resumable_session_hash()).toBeNull(); + }); + + it("keeps active events until the server resolves them", () => { + vi.useFakeTimers(); + try { + track_resumable_event(config, "session-1", { + event_id: "event-1", + fn_index: 3 + }); + + vi.advanceTimersByTime(2 * 60 * 60 * 1000); + + expect(get_resumable_events(config, "session-1")).toEqual([ + { event_id: "event-1", fn_index: 3 } + ]); + } finally { + vi.useRealTimers(); + } + }); + + it("reads the active session cookie during server rendering", () => { + expect( + get_resumable_session_hash("other=value; gradio_active_session=session-2") + ).toBe("session-2"); + }); +}); diff --git a/client/js/src/test/stream.test.ts b/client/js/src/test/stream.test.ts index 649052c6702..ffb5c287e08 100644 --- a/client/js/src/test/stream.test.ts +++ b/client/js/src/test/stream.test.ts @@ -37,7 +37,8 @@ describe("open_stream", () => { }); afterEach(() => { - vi.clearAllMocks(); + vi.restoreAllMocks(); + vi.useRealTimers(); }); it("should throw an error if config is not defined", async () => { @@ -62,10 +63,10 @@ describe("open_stream", () => { if (!app.stream_instance?.onmessage || !app.stream_instance?.onerror) { throw new Error("stream instance is not defined"); } - + const stream = app.stream_instance; const message = { msg: "hello jerry" }; - app.stream_instance.onmessage({ + stream.onmessage({ data: JSON.stringify(message) } as MessageEvent); expect(app.stream_status.open).toBe(true); @@ -74,14 +75,73 @@ describe("open_stream", () => { expect(app.pending_stream_messages).toEqual({}); const close_stream_message = { msg: "close_stream" }; - app.stream_instance.onmessage({ + await stream.onmessage({ data: JSON.stringify(close_stream_message) } as MessageEvent); expect(app.stream_status.open).toBe(false); + expect(app.stream_instance).toBeNull(); + stream.onerror?.({ + data: JSON.stringify("closed stream") + } as MessageEvent); + expect(app.stream_status.open).toBe(false); + expect(app.stream).toHaveBeenCalledTimes(1); + }); + + it("should open another stream when jobs remain after a normal close", async () => { + const callback = vi.fn().mockResolvedValue(undefined); + app.event_callbacks["event-1"] = callback; + app.unclosed_events.add("event-1"); + + await app.open_stream(); + + const stream = app.stream_instance; + if (!stream?.onmessage) { + throw new Error("stream instance is not defined"); + } + + await stream.onmessage({ + data: JSON.stringify({ msg: "close_stream" }) + } as MessageEvent); + + expect(app.stream).toHaveBeenCalledTimes(2); + expect(app.stream_status.open).toBe(true); + + stream.onerror?.({ + data: JSON.stringify("closed stream") + } as MessageEvent); + expect(app.stream).toHaveBeenCalledTimes(2); + }); + + it("should reconnect the SSE stream after an error when jobs are active", async () => { + vi.useFakeTimers(); + vi.spyOn(console, "error").mockImplementation(() => {}); + const callback = vi.fn().mockResolvedValue(undefined); + app.event_callbacks["event-1"] = callback; + app.unclosed_events.add("event-1"); + + await app.open_stream(); + + if (!app.stream_instance?.onerror) { + throw new Error("stream instance is not defined"); + } app.stream_instance.onerror({ - data: JSON.stringify("404") + data: JSON.stringify("network error") } as MessageEvent); + expect(app.stream_status.open).toBe(false); + expect(callback).not.toHaveBeenCalled(); + expect(app.stream).toHaveBeenCalledTimes(1); + expect(app.events_to_resume).toContain("event-1"); + + await vi.advanceTimersByTimeAsync(500); + + expect(app.stream).toHaveBeenCalledTimes(2); + expect(app.stream_status.open).toBe(true); + expect( + (app.stream as Mock).mock.calls[1][0].searchParams.getAll( + "resume_event_id" + ) + ).toEqual(["event-1"]); }); }); diff --git a/client/js/src/types.ts b/client/js/src/types.ts index 1836dcc1e47..a031f2f5030 100644 --- a/client/js/src/types.ts +++ b/client/js/src/types.ts @@ -112,6 +112,7 @@ export interface SubmitIterable extends AsyncIterable { event_id: () => string; send_chunk: (payload: Record) => void; wait_for_id: () => Promise; + acknowledge: () => Promise; close_stream: () => void; } @@ -344,6 +345,7 @@ export interface ClientOptions { headers?: Record | Headers; query_params?: Record; session_hash?: string; + resume_sessions?: boolean; cookies?: string; credentials?: RequestCredentials; /** diff --git a/client/js/src/utils/session.ts b/client/js/src/utils/session.ts new file mode 100644 index 00000000000..a6166e7d560 --- /dev/null +++ b/client/js/src/utils/session.ts @@ -0,0 +1,141 @@ +import type { Config } from "../types"; + +export interface ResumableJob { + event_id: string; + fn_index: number; +} + +interface ResumableSession { + app_id?: string; + root: string; + session_hash: string; + events: ResumableJob[]; +} + +const STORAGE_KEY = "gradio_active_session"; +const COOKIE_NAME = "gradio_active_session"; + +function read_session(): ResumableSession | null { + if (typeof sessionStorage === "undefined") return null; + + try { + const value = sessionStorage.getItem(STORAGE_KEY); + if (!value) return null; + return JSON.parse(value) as ResumableSession; + } catch { + return null; + } +} + +function cookie_value(cookies: string): string | null { + const cookie = cookies + .split(";") + .map((value) => value.trim()) + .find((value) => value.startsWith(`${COOKIE_NAME}=`)); + const value = cookie?.slice(COOKIE_NAME.length + 1).split(";")[0]; + return value ? decodeURIComponent(value) : null; +} + +function cookie_path(root: string): string { + try { + return new URL(root).pathname || "/"; + } catch { + return "/"; + } +} + +function write_session(session: ResumableSession): void { + if (typeof sessionStorage === "undefined") return; + + try { + sessionStorage.setItem(STORAGE_KEY, JSON.stringify(session)); + if (typeof document !== "undefined") { + document.cookie = `${COOKIE_NAME}=${encodeURIComponent(session.session_hash)}; Path=${cookie_path(session.root)}; SameSite=Lax`; + } + } catch { + return; + } +} + +function remove_session(session: ResumableSession): void { + try { + if (typeof sessionStorage !== "undefined") { + sessionStorage.removeItem(STORAGE_KEY); + } + if (typeof document !== "undefined") { + document.cookie = `${COOKIE_NAME}=; Path=${cookie_path(session.root)}; Max-Age=0; SameSite=Lax`; + } + } catch { + return; + } +} + +export function get_resumable_session_hash(cookies?: string): string | null { + const session = read_session(); + if (session) { + if (typeof location === "undefined") return session.session_hash; + try { + const root = new URL(session.root); + if ( + root.origin === location.origin && + location.pathname.startsWith(root.pathname) + ) { + return session.session_hash; + } + } catch { + return null; + } + } + + const cookie_string = + cookies ?? (typeof document !== "undefined" ? document.cookie : ""); + return cookie_value(cookie_string); +} + +export function get_resumable_events( + config: Config, + session_hash: string +): ResumableJob[] { + const session = read_session(); + if ( + !session || + session.session_hash !== session_hash || + session.root !== config.root || + (session.app_id && session.app_id !== config.app_id) + ) { + if (session) remove_session(session); + return []; + } + return session.events; +} + +export function track_resumable_event( + config: Config, + session_hash: string, + event: ResumableJob +): void { + const current = read_session(); + const events = + current?.session_hash === session_hash && current.root === config.root + ? current.events.filter(({ event_id }) => event_id !== event.event_id) + : []; + + write_session({ + app_id: config.app_id, + root: config.root, + session_hash, + events: [...events, event] + }); +} + +export function clear_resumable_event(event_id: string): void { + const session = read_session(); + if (!session) return; + + const events = session.events.filter((event) => event.event_id !== event_id); + if (events.length === 0) { + remove_session(session); + } else { + write_session({ ...session, events }); + } +} diff --git a/client/js/src/utils/stream.ts b/client/js/src/utils/stream.ts index 534f36d51aa..c94544d474b 100644 --- a/client/js/src/utils/stream.ts +++ b/client/js/src/utils/stream.ts @@ -1,4 +1,4 @@ -import { BROKEN_CONNECTION_MSG, SSE_URL } from "../constants"; +import { SSE_URL } from "../constants"; import type { Client } from "../client"; import { stream } from "fetch-event-stream"; @@ -13,25 +13,41 @@ export async function open_stream(this: Client): Promise { } = this; const that = this; - if (!config) { throw new Error("Could not resolve app config"); } + if (stream_status.open) { + return; + } + + if (this.stream_reconnect_timer) { + clearTimeout(this.stream_reconnect_timer); + this.stream_reconnect_timer = null; + } + stream_status.open = true; let stream: EventSource | null = null; let params = new URLSearchParams({ - session_hash: this.session_hash - }).toString(); + session_hash: this.session_hash, + acknowledgements: this.options.resume_sessions ? "true" : "false" + }); + this.events_to_resume.forEach((event_id) => { + params.append("resume_event_id", event_id); + }); + this.events_to_resume.clear(); - let url = new URL(`${config.root}${this.api_prefix}/${SSE_URL}?${params}`); + let url = new URL( + `${config.root}${this.api_prefix}/${SSE_URL}?${params.toString()}` + ); if (jwt) { url.searchParams.set("__sign", jwt); } stream = this.stream(url); + const abort_controller = this.abort_controller; if (!stream) { console.warn("Cannot connect to SSE endpoint: " + url.toString()); @@ -39,9 +55,17 @@ export async function open_stream(this: Client): Promise { } stream.onmessage = async function (event: MessageEvent) { + if (that.stream_instance !== stream) { + return; + } + that.stream_reconnect_attempts = 0; let _data = JSON.parse(event.data); if (_data.msg === "close_stream") { - close_stream(stream_status, that.abort_controller); + that.stream_instance = null; + close_stream(stream_status, abort_controller); + if (!that.closed && unclosed_events.size > 0) { + await that.open_stream(); + } return; } const event_id = _data.event_id; @@ -67,10 +91,7 @@ export async function open_stream(this: Client): Promise { typeof document !== "undefined" && document.visibilityState !== "hidden" ) { - // Put the event at the end of the event loop so the browser can refresh - // between callbacks and not freeze during quick generations. Hidden tabs - // throttle timers, so process those messages immediately instead. - setTimeout(fn, 0, _data); // See https://github.com/gradio-app/gradio/pull/7055 + setTimeout(fn, 0, _data); // need to do this to put the event on the end of the event loop, so the browser can refresh between callbacks and not freeze in case of quick generations. See https://github.com/gradio-app/gradio/pull/7055 } else { fn(_data); } @@ -82,15 +103,42 @@ export async function open_stream(this: Client): Promise { } }; stream.onerror = async function (e) { + if (that.stream_instance !== stream) { + return; + } console.error(e); - await Promise.all( - Object.keys(event_callbacks).map((event_id) => - event_callbacks[event_id]({ - msg: "broken_connection", - message: BROKEN_CONNECTION_MSG - }) - ) - ); + close_stream(stream_status, abort_controller); + unclosed_events.forEach((event_id) => { + that.events_to_resume.add(event_id); + delete that.pending_diff_streams[event_id]; + }); + + if ( + that.closed || + that.stream_reconnect_timer || + (unclosed_events.size === 0 && Object.keys(event_callbacks).length === 0) + ) { + return; + } + + const delay = Math.min(500 * 2 ** that.stream_reconnect_attempts, 5000); + that.stream_reconnect_attempts += 1; + + that.stream_reconnect_timer = setTimeout(async () => { + that.stream_reconnect_timer = null; + if ( + that.closed || + (unclosed_events.size === 0 && + Object.keys(event_callbacks).length === 0) + ) { + return; + } + try { + await that.open_stream(); + } catch (error) { + stream?.onerror?.(error as Event); + } + }, delay); }; } diff --git a/client/js/src/utils/submit.ts b/client/js/src/utils/submit.ts index 4876b3b073f..fb4272a0ef4 100644 --- a/client/js/src/utils/submit.ts +++ b/client/js/src/utils/submit.ts @@ -28,6 +28,7 @@ import { WS_PROTOCOL_MSG } from "../constants"; import { apply_diff_stream, close_stream } from "./stream"; +import { clear_resumable_event, track_resumable_event } from "./session"; import { Client } from "../client"; import { start_run_history, @@ -42,7 +43,8 @@ export function submit( event_data?: unknown, trigger_id?: number | null, all_events?: boolean, - additional_headers?: Record + additional_headers?: Record, + resume_event_id?: string ): SubmitIterable { try { const { token } = this.options; @@ -77,7 +79,10 @@ export function submit( config ); - let resolved_data = map_data_to_params(data, endpoint_info); + let resolved_data = resume_event_id + ? [] + : map_data_to_params(data, endpoint_info); + let protocol = config.protocol ?? "ws"; if (protocol === "ws") { throw new Error(WS_PROTOCOL_MSG); @@ -115,7 +120,7 @@ export function submit( config.run_history !== false && this.options.record_history !== false; const history_scope = { app_id: config.app_id, username: config.username }; const history_run_id = - !history_enabled || !is_documented_endpoint + resume_event_id || !history_enabled || !is_documented_endpoint ? null : start_run_history({ ...history_scope, @@ -196,6 +201,20 @@ export function submit( "The `/reset` endpoint could not be called. Subsequent endpoint results may be unreliable." ); } + await acknowledge(); + } + + async function acknowledge(): Promise { + if (!event_id_final) return; + if (!options.resume_sessions) return; + const response = await that + .fetch(`${config!.root}${api_prefix}/${RESET_URL}`, { + headers: { "Content-Type": "application/json" }, + method: "POST", + body: JSON.stringify({ event_id: event_id_final }) + }) + .catch(() => null); + if (response?.ok) clear_resumable_event(event_id_final); } const resolve_heartbeat = async (config: Config): Promise => { @@ -227,11 +246,182 @@ export function submit( }); } + async function register_queue_event(queue_event_id: string): Promise { + event_id = queue_event_id; + event_id_final = queue_event_id; + let callback = async function (_data: unknown): Promise { + try { + const { type, status, data, original_msg } = handle_message( + _data, + last_status[fn_index] + ); + + if (type == "heartbeat") { + return; + } + + if (type === "update" && status && !complete) { + fire_event({ + type: "status", + endpoint: _endpoint, + fn_index, + time: new Date(), + original_msg, + ...status + }); + } else if (type === "complete") { + complete = status; + } else if ( + type == "unexpected_error" || + type == "broken_connection" + ) { + if (type === "unexpected_error") { + unclosed_events.delete(queue_event_id); + if (status?.session_not_found) { + clear_resumable_event(queue_event_id); + } + } + console.error("Unexpected error", status?.message); + const broken = type === "broken_connection"; + fire_event({ + type: "status", + stage: "error", + message: status?.message || "An Unexpected Error Occurred!", + queue: true, + endpoint: _endpoint, + broken, + session_not_found: status?.session_not_found, + fn_index, + time: new Date() + }); + } else if (type === "log") { + fire_event({ + type: "log", + title: data.title, + log: data.log, + level: data.level, + endpoint: _endpoint, + duration: data.duration, + visible: data.visible, + fn_index + }); + return; + } else if (type === "generating" || type === "streaming") { + fire_event({ + type: "status", + time: new Date(), + ...status, + stage: status?.stage!, + queue: true, + endpoint: _endpoint, + fn_index + }); + if ( + data && + dependency.connection !== "stream" && + ["sse_v2", "sse_v2.1", "sse_v3"].includes(protocol) + ) { + apply_diff_stream(pending_diff_streams, event_id!, data); + } + } + if (data) { + fire_event({ + type: "data", + time: new Date(), + data: handle_payload( + data.data, + dependency, + config!.components, + "output", + options.with_null_state + ), + endpoint: _endpoint, + fn_index + }); + if (data.render_config) { + await handle_render_config(data.render_config); + } + + if (complete) { + fire_event({ + type: "status", + time: new Date(), + ...complete, + stage: status?.stage!, + queue: true, + endpoint: _endpoint, + fn_index + }); + close(); + } + } + + if (status?.stage === "complete" || status?.stage === "error") { + delete event_callbacks[queue_event_id]; + delete pending_diff_streams[queue_event_id]; + close(); + } + } catch (e) { + console.error("Unexpected client exception", e); + fire_event({ + type: "status", + stage: "error", + message: "An Unexpected Error Occurred!", + queue: true, + endpoint: _endpoint, + fn_index, + time: new Date() + }); + if (["sse_v2", "sse_v2.1", "sse_v3"].includes(protocol)) { + close_stream(stream_status, that.abort_controller); + stream_status.open = false; + close(); + } + } + }; + + event_callbacks[queue_event_id] = callback; + unclosed_events.add(queue_event_id); + if (queue_event_id in pending_stream_messages) { + for (const msg of pending_stream_messages[queue_event_id]) { + if ( + msg.msg === "process_completed" && + ["sse", "sse_v1", "sse_v2", "sse_v2.1", "sse_v3"].includes(protocol) + ) { + unclosed_events.delete(queue_event_id); + } + await callback(msg); + } + delete pending_stream_messages[queue_event_id]; + } + if (!stream_status.open && unclosed_events.has(queue_event_id)) { + await that.open_stream(); + } + } + + if (resume_event_id) { + event_id = resume_event_id; + event_id_final = resume_event_id; + this.events_to_resume.add(resume_event_id); + } + const job = this.handle_blob( config.root, resolved_data, endpoint_info ).then(async (_payload) => { + if (resume_event_id) { + fire_event({ + type: "status", + stage: "pending", + queue: true, + endpoint: _endpoint, + fn_index, + time: new Date() + }); + await register_queue_event(resume_event_id); + return; + } let input_data = handle_payload( _payload, dependency, @@ -550,146 +740,13 @@ export function submit( } else { event_id = response.event_id as string; event_id_final = event_id; - let callback = async function (_data: object): Promise { - try { - const { type, status, data, original_msg } = handle_message( - _data, - last_status[fn_index] - ); - - if (type == "heartbeat") { - return; - } - - if (type === "update" && status && !complete) { - // call 'status' listeners - fire_event({ - type: "status", - endpoint: _endpoint, - fn_index, - time: new Date(), - original_msg: original_msg, - ...status - }); - } else if (type === "complete") { - complete = status; - } else if ( - type == "unexpected_error" || - type == "broken_connection" - ) { - console.error("Unexpected error", status?.message); - const broken = type === "broken_connection"; - fire_event({ - type: "status", - stage: "error", - message: status?.message || "An Unexpected Error Occurred!", - queue: true, - endpoint: _endpoint, - broken, - session_not_found: status?.session_not_found, - fn_index, - time: new Date() - }); - } else if (type === "log") { - fire_event({ - type: "log", - title: data.title, - log: data.log, - level: data.level, - endpoint: _endpoint, - duration: data.duration, - visible: data.visible, - fn_index - }); - return; - } else if (type === "generating" || type === "streaming") { - fire_event({ - type: "status", - time: new Date(), - ...status, - stage: status?.stage!, - queue: true, - endpoint: _endpoint, - fn_index - }); - if ( - data && - dependency.connection !== "stream" && - ["sse_v2", "sse_v2.1", "sse_v3"].includes(protocol) - ) { - apply_diff_stream(pending_diff_streams, event_id!, data); - } - } - if (data) { - fire_event({ - type: "data", - time: new Date(), - data: handle_payload( - data.data, - dependency, - config.components, - "output", - options.with_null_state - ), - endpoint: _endpoint, - fn_index - }); - if (data.render_config) { - await handle_render_config(data.render_config); - } - - if (complete) { - fire_event({ - type: "status", - time: new Date(), - ...complete, - stage: status?.stage!, - queue: true, - endpoint: _endpoint, - fn_index - }); - close(); - } - } - - if (status?.stage === "complete" || status?.stage === "error") { - if (event_callbacks[event_id!]) { - delete event_callbacks[event_id!]; - } - if (event_id! in pending_diff_streams) { - delete pending_diff_streams[event_id!]; - } - close(); - } - } catch (e) { - console.error("Unexpected client exception", e); - fire_event({ - type: "status", - stage: "error", - message: "An Unexpected Error Occurred!", - queue: true, - endpoint: _endpoint, - fn_index, - time: new Date() - }); - if (["sse_v2", "sse_v2.1", "sse_v3"].includes(protocol)) { - close_stream(stream_status, that.abort_controller); - stream_status.open = false; - close(); - } - } - }; - - if (event_id in pending_stream_messages) { - pending_stream_messages[event_id].forEach((msg) => callback(msg)); - delete pending_stream_messages[event_id]; - } - // @ts-ignore - event_callbacks[event_id] = callback; - unclosed_events.add(event_id); - if (!stream_status.open) { - await this.open_stream(); + if (options.resume_sessions) { + track_resumable_event(config, session_hash, { + event_id, + fn_index + }); } + await register_queue_event(event_id); } }); } @@ -785,7 +842,8 @@ export function submit( wait_for_id: async () => { await job; return event_id; - } + }, + acknowledge }; return iterator; diff --git a/client/python/gradio_client/__init__.py b/client/python/gradio_client/__init__.py index 037264eeba8..40656945674 100644 --- a/client/python/gradio_client/__init__.py +++ b/client/python/gradio_client/__init__.py @@ -1,4 +1,4 @@ -from gradio_client.client import Client +from gradio_client.client import Client, ResumableJob from gradio_client.data_classes import FileData from gradio_client.utils import __version__, file, handle_file @@ -7,5 +7,6 @@ "file", "handle_file", "FileData", + "ResumableJob", "__version__", ] diff --git a/client/python/gradio_client/client.py b/client/python/gradio_client/client.py index fab17c0570d..ff64efcd575 100644 --- a/client/python/gradio_client/client.py +++ b/client/python/gradio_client/client.py @@ -26,7 +26,7 @@ from functools import partial from pathlib import Path from threading import Lock -from typing import Any, Literal +from typing import Any, Literal, TypedDict import httpx import huggingface_hub @@ -56,9 +56,15 @@ DEFAULT_TEMP_DIR = os.environ.get("GRADIO_TEMP_DIR") or str( Path(tempfile.gettempdir()) / "gradio" ) +SESSION_RESUME_TTL_SECONDS = 3600 -@document("predict", "submit", "view_api", "duplicate") +class ResumableJob(TypedDict): + event_id: str + fn_index: int + + +@document("predict", "submit", "resume_jobs", "view_api", "duplicate") class Client: """ The main Client class for the Python client. This class is used to connect to a remote Gradio app and call its API endpoints. @@ -88,6 +94,8 @@ def __init__( headers: dict[str, str] | None = None, download_files: str | Path | Literal[False] = DEFAULT_TEMP_DIR, ssl_verify: bool = True, + session_hash: str | None = None, + resume_sessions: bool = False, _skip_components: bool = True, # internal parameter to skip values certain components (e.g. State) that do not need to be displayed to users. analytics_enabled: bool = True, oauth_token: str | None = None, @@ -101,6 +109,8 @@ def __init__( headers: additional headers to send to the remote Gradio app on every request. By default only the HF authorization and user-agent headers are sent. This parameter will override the default headers if they have the same keys. download_files: directory where the client should download output files on the local machine from the remote API. By default, uses the value of the GRADIO_TEMP_DIR environment variable which, if not set by the user, is a temporary directory on your machine. If False, the client does not download files and returns a FileData dataclass object with the filepath on the remote machine instead. ssl_verify: if False, skips certificate validation which allows the client to connect to Gradio apps that are using self-signed certificates. + session_hash: Session hash from a previous Client whose queued jobs should be resumed. + resume_sessions: If True, active queued jobs are retained and automatically reattached after temporary network interruptions. httpx_kwargs: additional keyword arguments to pass to `httpx.Client`, `httpx.stream`, `httpx.get` and `httpx.post`. This can be used to set timeouts, proxies, http auth, etc. analytics_enabled: Whether to allow basic telemetry. If None, will use GRADIO_ANALYTICS_ENABLED environment variable or default to True. oauth_token: optional Hugging Face token for the app to act on your behalf, for endpoints whose function takes a `gr.OAuthToken`. Unlike `token`, which only authenticates you to the app, this is passed to the app's code, so it is sent only to endpoints that declare they need it — `view_api()` marks those. It is never sent anywhere else, and is not inferred from your locally saved token. @@ -191,7 +201,8 @@ def __init__( self.reset_url = urllib.parse.urljoin(self.src_prefixed, utils.RESET_URL) self.app_version = version.parse(self.config.get("version", "2.0")) self._info = self._get_api_info() - self.session_hash = str(uuid.uuid4()) + self.session_hash = session_hash or str(uuid.uuid4()) + self.resume_sessions = resume_sessions self.endpoints = { dependency.get("id", fn_index): Endpoint( @@ -218,8 +229,10 @@ def __init__( self.streaming_future: Future | None = None self.pending_messages_per_event: dict[str, deque[Message | None]] = {} self.pending_event_ids: set[str] = set() + self._closed = False def close(self): + self._closed = True self._kill_heartbeat.set() self.heartbeat.join(timeout=1) @@ -250,34 +263,46 @@ def stream_messages( self, protocol: Literal["sse_v1", "sse_v2", "sse_v2.1", "sse_v3"], session_hash: str, + resume_event_ids: list[str] | None = None, ) -> None: - try: - httpx_kwargs = self.httpx_kwargs.copy() - httpx_kwargs.setdefault("timeout", httpx.Timeout(timeout=None)) - with httpx.Client( - verify=self.ssl_verify, - **httpx_kwargs, - ) as client: - with client.stream( - "GET", - self.sse_url, - params={"session_hash": session_hash}, - headers=self.headers, - cookies=self.cookies, - ) as response: - buffer = b"" - for chunk in response.iter_bytes(): - buffer += chunk - while b"\n\n" in buffer: - line, buffer = buffer.split(b"\n\n", 1) - line = line.decode("utf-8").rstrip("\n") - if not len(line): - continue - if line.startswith("data:"): + reconnect_deadline = None + while not self._closed: + connected = False + try: + httpx_kwargs = self.httpx_kwargs.copy() + httpx_kwargs.setdefault("timeout", httpx.Timeout(timeout=None)) + with httpx.Client( + verify=self.ssl_verify, + **httpx_kwargs, + ) as client: + with client.stream( + "GET", + self.sse_url, + params={ + "session_hash": session_hash, + "resume_event_id": resume_event_ids or [], + "acknowledgements": self.resume_sessions, + }, + headers=self.headers, + cookies=self.cookies, + ) as response: + response.raise_for_status() + connected = True + resume_event_ids = None + buffer = b"" + for chunk in response.iter_bytes(): + buffer += chunk + while b"\n\n" in buffer: + line, buffer = buffer.split(b"\n\n", 1) + line = line.decode("utf-8").rstrip("\n") + if not len(line): + continue + if not line.startswith("data:"): + raise ValueError(f"Unexpected SSE line: '{line}'") resp = json.loads(line[5:]) if resp["msg"] == ServerMessage.heartbeat: continue - elif ( + if ( resp.get("message", "") == ServerMessage.server_stopped ): @@ -286,7 +311,7 @@ def stream_messages( ) in self.pending_messages_per_event.values(): pending_messages.append(resp) return - elif resp["msg"] == ServerMessage.close_stream: + if resp["msg"] == ServerMessage.close_stream: self.stream_open = False return event_id = resp["event_id"] @@ -294,27 +319,30 @@ def stream_messages( self.pending_messages_per_event[event_id] = deque() self.pending_messages_per_event[event_id].append(resp) if resp["msg"] == ServerMessage.process_completed: - self.pending_event_ids.remove(event_id) + self.pending_event_ids.discard(event_id) + if self.resume_sessions: + self._acknowledge_event(event_id) if ( len(self.pending_event_ids) == 0 and protocol != "sse_v3" ): self.stream_open = False return - else: - raise ValueError(f"Unexpected SSE line: '{line}'") - except BaseException as e: - # If the job is cancelled the stream will close so we - # should not raise this httpx exception that comes from the - # stream abruply closing - if isinstance(e, httpx.RemoteProtocolError): - return - import traceback + except httpx.TransportError: + if not self.resume_sessions: + return - traceback.print_exc() - raise e + if not self.pending_event_ids: + return + if not self.resume_sessions: + return + if connected or reconnect_deadline is None: + reconnect_deadline = time.monotonic() + SESSION_RESUME_TTL_SECONDS + if time.monotonic() >= reconnect_deadline: + return + time.sleep(1) - def send_data(self, data, hash_data, protocol, request_headers): + def send_data(self, data, hash_data, request_headers): headers = self.add_zero_gpu_headers(self.headers) if request_headers is not None: headers = {**request_headers, **headers} @@ -332,14 +360,22 @@ def send_data(self, data, hash_data, protocol, request_headers): raise ValidationError(validation_message) req.raise_for_status() resp = req.json() - event_id = resp["event_id"] + return resp["event_id"] + def _start_stream( + self, + protocol: Literal["sse_v1", "sse_v2", "sse_v2.1", "sse_v3"], + session_hash: str, + resume_event_ids: list[str] | None = None, + ) -> None: if not self.stream_open: self.stream_open = True def open_stream(): return self.stream_messages( - protocol, session_hash=hash_data["session_hash"] + protocol, + session_hash=session_hash, + resume_event_ids=resume_event_ids, ) def close_stream(_): @@ -351,7 +387,20 @@ def close_stream(_): self.streaming_future = self.executor.submit(open_stream) self.streaming_future.add_done_callback(close_stream) - return event_id + def _acknowledge_event(self, event_id: str) -> None: + self.pending_event_ids.discard(event_id) + try: + response = httpx.post( + self.reset_url, + json={"event_id": event_id}, + headers=self.headers, + cookies=self.cookies, + verify=self.ssl_verify, + **self.httpx_kwargs, + ) + response.raise_for_status() + except httpx.HTTPError: + pass @classmethod def duplicate( @@ -493,7 +542,7 @@ def predict( The result of the API call. Will be a Tuple if the API has multiple outputs. Example: from gradio_client import Client - client = Client(src="gradio/calculator") + client = Client(src="gradio/calculator", resume_sessions=True) client.predict(5, "add", 4, api_name="/predict") >> 9.0 """ @@ -536,7 +585,7 @@ def submit( A Job object that can be used to retrieve the status and result of the remote API call. Example: from gradio_client import Client - client = Client(src="gradio/calculator") + client = Client(src="gradio/calculator", resume_sessions=True) job = client.submit(5, "add", 4, api_name="/predict") job.status() >> @@ -574,6 +623,7 @@ def submit( verbose=self.verbose, space_id=self.space_id, _cancel_fn=cancel_fn, + fn_index=inferred_fn_index, ) if result_callbacks: @@ -594,6 +644,91 @@ def fn(future): return job + def resume_jobs( + self, + jobs: list[ResumableJob], + *, + session_hash: str | None = None, + ) -> list[Job]: + """ + Resumes queued jobs created by the same Gradio app and session. + + Parameters: + jobs: The event ID and function index for each queued job to resume. + session_hash: The session hash that originally submitted the jobs. Uses this client's current session hash when omitted. + Returns: + A list of Job objects in the same order as the provided jobs. + Example: + from gradio_client import Client + client = Client(src="gradio/calculator") + job = client.submit(5, "add", 4, api_name="/predict") + resumable_job = {"event_id": job.wait_for_id(), "fn_index": job.fn_index} + session_hash = client.session_hash + resumed_job = Client(src="gradio/calculator").resume_jobs( + [resumable_job], session_hash=session_hash + )[0] + resumed_job.result() + >> 9.0 + """ + if self.protocol not in ("sse_v1", "sse_v2", "sse_v2.1", "sse_v3"): + raise ValueError(f"Cannot resume jobs using protocol: {self.protocol}") + if self.stream_open: + raise ValueError("Cannot resume jobs while the queue stream is active.") + if session_hash is not None and session_hash != self.session_hash: + self.session_hash = session_hash + self._refresh_heartbeat.set() + self.resume_sessions = True + resumed_session_hash = self.session_hash + + resumable_jobs: list[tuple[Endpoint, Communicator, str, int]] = [] + for job_info in jobs: + event_id = job_info["event_id"] + fn_index = job_info["fn_index"] + endpoint = self.endpoints.get(fn_index) + if not isinstance(endpoint, Endpoint): + raise ValueError(f"No endpoint found for fn_index {fn_index}.") + helper = self.new_helper(fn_index, headers={"x-gradio-user": "api"}) + helper.event_id = event_id + self.pending_event_ids.add(event_id) + self.pending_messages_per_event[event_id] = deque() + resumable_jobs.append((endpoint, helper, event_id, fn_index)) + + result_jobs = [] + for endpoint, helper, event_id, fn_index in resumable_jobs: + + def resume_job( + endpoint: Endpoint = endpoint, + helper: Communicator = helper, + event_id: str = event_id, + ): + result = endpoint._sse_fn_v1plus(helper, event_id, self.protocol) + output = endpoint.process_result(result) + predictions = endpoint.process_predictions(*output) + with helper.lock: + if not helper.job.outputs: + helper.job.outputs.append(predictions) + return predictions + + future = self.executor.submit(copy_context().run, resume_job) + result_jobs.append( + Job( + future, + communicator=helper, + verbose=self.verbose, + space_id=self.space_id, + _cancel_fn=endpoint.make_cancel(helper), + fn_index=fn_index, + ) + ) + + if resumable_jobs: + self._start_stream( + self.protocol, + resumed_session_hash, + [event_id for _, _, event_id, _ in resumable_jobs], + ) + return result_jobs + def _get_api_info(self): api_info_url = urllib.parse.urljoin(self.src_prefixed, utils.RAW_API_INFO_URL) if self.app_version > version.Version("3.36.1"): @@ -1205,46 +1340,50 @@ def _predict(*data, **kwargs) -> tuple: result = self._sse_fn_v0(data, hash_data, helper) # type: ignore elif self.protocol in ("sse_v1", "sse_v2", "sse_v2.1", "sse_v3"): event_id = self.client.send_data( - data, hash_data, self.protocol, helper.request_headers + data, hash_data, helper.request_headers ) self.client.pending_event_ids.add(event_id) self.client.pending_messages_per_event[event_id] = deque() helper.event_id = event_id + self.client._start_stream(self.protocol, hash_data["session_hash"]) result = self._sse_fn_v1plus(helper, event_id, self.protocol) else: raise ValueError(f"Unsupported protocol: {self.protocol}") - if "error" in result: - if result["error"] is None: - raise AppError( - "The upstream Gradio app has raised an exception but has not enabled " - "verbose error reporting. To enable, set show_error=True in launch()." - ) - else: - message = result.pop("error") - raise AppError(message=message, **result) + return self.process_result(result) - try: - output = result["data"] - except KeyError as ke: - is_public_space = ( - self.client.space_id - and not huggingface_hub.space_info(self.client.space_id).private + return _predict + + def process_result(self, result: dict[str, Any]) -> tuple: + if "error" in result: + if result["error"] is None: + raise AppError( + "The upstream Gradio app has raised an exception but has not enabled " + "verbose error reporting. To enable, set show_error=True in launch()." ) - if "error" in result and "429" in result["error"] and is_public_space: - raise utils.TooManyRequestsError( - f"Too many requests to the API, please try again later. To avoid being rate-limited, " - f"please duplicate the Space using Client.duplicate({self.client.space_id}) " - f"and pass in your Hugging Face token." - ) from None - elif "error" in result: - raise ValueError(result["error"]) from None - raise KeyError( - f"Could not find 'data' key in response. Response received: {result}" - ) from ke - return tuple(output) + else: + message = result.pop("error") + raise AppError(message=message, **result) - return _predict + try: + output = result["data"] + except KeyError as ke: + is_public_space = ( + self.client.space_id + and not huggingface_hub.space_info(self.client.space_id).private + ) + if "error" in result and "429" in result["error"] and is_public_space: + raise utils.TooManyRequestsError( + f"Too many requests to the API, please try again later. To avoid being rate-limited, " + f"please duplicate the Space using Client.duplicate({self.client.space_id}) " + f"and pass in your Hugging Face token." + ) from None + elif "error" in result: + raise ValueError(result["error"]) from None + raise KeyError( + f"Could not find 'data' key in response. Response received: {result}" + ) from ke + return tuple(output) def insert_empty_state(self, *data) -> tuple: data = list(data) @@ -1422,7 +1561,7 @@ def _sse_fn_v1plus( ) -@document("result", "outputs", "status") +@document("result", "outputs", "status", "wait_for_id") class Job(Future): """ A Job is a wrapper over the Future class that represents a prediction call that has been @@ -1441,6 +1580,7 @@ def __init__( verbose: bool = True, space_id: str | None = None, _cancel_fn: Callable[[], None] | None = None, + fn_index: int | None = None, ): """ Parameters: @@ -1448,6 +1588,7 @@ def __init__( communicator: The communicator object that is used to communicate between the client and the background thread running the job verbose: Whether to print any status-related messages to the console space_id: The space ID corresponding to the Client object that created this Job object + fn_index: The function index of the API endpoint for this job """ self.future = future self.communicator = communicator @@ -1455,6 +1596,7 @@ def __init__( self.verbose = verbose self.space_id = space_id self.cancel_fn = _cancel_fn + self.fn_index = fn_index def __iter__(self) -> Job: return self @@ -1508,6 +1650,27 @@ def result(self, timeout: float | None = None) -> Any: """ return super().result(timeout=timeout) + def wait_for_id(self, timeout: float | None = None) -> str: + """ + Wait until the server assigns an event ID to this job. + + Parameters: + timeout: The number of seconds to wait. If None, there is no limit. + Returns: + The server-assigned event ID. + """ + if not self.communicator: + raise ValueError("This job does not have an event ID.") + deadline = None if timeout is None else time.monotonic() + timeout + while self.communicator.event_id is None: + if self.done(): + self.result() + raise ValueError("This job did not receive an event ID.") + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError() + time.sleep(0.01) + return self.communicator.event_id + def outputs(self) -> list[tuple | Any]: """ Returns a list containing the latest outputs from the Job. diff --git a/client/python/test/test_client.py b/client/python/test/test_client.py index 78c6c6062de..a34e74d868c 100644 --- a/client/python/test/test_client.py +++ b/client/python/test/test_client.py @@ -5,6 +5,7 @@ import threading import time import uuid +from collections import deque from concurrent.futures import CancelledError, TimeoutError, wait from contextlib import contextmanager from datetime import datetime, timedelta @@ -237,6 +238,76 @@ def test_job_status(self, calculator_demo): s.code for s in statuses if s ] + def test_job_wait_for_id(self, calculator_demo): + with connect(calculator_demo) as client: + job = client.submit(5, "add", 4, api_name="/predict") + + assert job.wait_for_id(timeout=5) + assert job.fn_index == 0 + assert job.result() == 9 + + def test_resume_jobs(self, calculator_demo): + with connect(calculator_demo) as client: + with patch.object(client, "_start_stream") as start_stream: + job = client.resume_jobs( + [{"event_id": "event-1", "fn_index": 0}], + session_hash="existing-session", + )[0] + client.pending_messages_per_event["event-1"].append( + { + "msg": "process_completed", + "event_id": "event-1", + "output": {"data": [9]}, + "success": True, + } + ) + + assert job.result(timeout=5) == 9 + assert job.wait_for_id() == "event-1" + assert job.fn_index == 0 + assert client.session_hash == "existing-session" + start_stream.assert_called_once_with( + client.protocol, "existing-session", ["event-1"] + ) + + def test_queue_stream_reconnects_for_active_jobs(self): + requests = [] + + def handle_request(request): + requests.append(request) + if len(requests) == 1: + raise httpx.ConnectError("offline", request=request) + return httpx.Response( + 200, + content=( + b'data: {"msg":"process_completed","event_id":"event-1",' + b'"output":{"data":[9]},"success":true}\n\n' + ), + ) + + client = Client.__new__(Client) + client._closed = False + client.resume_sessions = True + client.httpx_kwargs = {"transport": httpx.MockTransport(handle_request)} + client.ssl_verify = True + client.sse_url = "https://example.test/queue/data" + client.headers = {} + client.cookies = {} + client.stream_open = True + client.pending_event_ids = {"event-1"} + client.pending_messages_per_event = {"event-1": deque()} + client._acknowledge_event = MagicMock() + + with patch("gradio_client.client.time.sleep"): + client.stream_messages("sse_v2", "session-1") + + assert len(requests) == 2 + assert requests[1].url.params["session_hash"] == "session-1" + assert requests[1].url.params["acknowledgements"] == "true" + assert not client.pending_event_ids + assert client.pending_messages_per_event["event-1"][0]["success"] is True + client._acknowledge_event.assert_called_once_with("event-1") + @pytest.mark.flaky def test_intermediate_outputs(self, count_generator_demo): with connect(count_generator_demo) as client: diff --git a/gradio/data_classes.py b/gradio/data_classes.py index db26578ec10..7b071b01cd4 100644 --- a/gradio/data_classes.py +++ b/gradio/data_classes.py @@ -53,6 +53,10 @@ class CancelBody(BaseModel): event_id: str +class QueueCloseBody(BaseModel): + session_hash: str + + class SimplePredictBody(BaseModel): data: list[Any] session_hash: str | None = None diff --git a/gradio/queueing.py b/gradio/queueing.py index 53b1961740b..b789a9a58f2 100644 --- a/gradio/queueing.py +++ b/gradio/queueing.py @@ -10,13 +10,14 @@ import uuid from asyncio import Queue as AsyncQueue from collections import defaultdict +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast import fastapi import numpy as np from anyio.to_thread import run_sync -from gradio import route_utils, routes +from gradio import caching, route_utils, routes from gradio.caching import CacheMissError, ProbeCache from gradio.data_classes import ( PredictBodyInternal, @@ -40,6 +41,7 @@ ProgressMessage, ProgressUnit, ServerMessage, + UnexpectedErrorMessage, ) from gradio.utils import ( LRUCache, @@ -107,6 +109,18 @@ def __init__(self, concurrency_id: str, concurrency_limit: int | None): ) +@dataclass +class ResumableSession: + history: list[EventMessage] = field(default_factory=list) + expires_at: float | None = None + active_streams: int = 0 + closing: bool = False + + def expire_after(self, seconds: float) -> None: + deadline = time.monotonic() + seconds + self.expires_at = min(deadline, self.expires_at or deadline) + + class ProcessTime: def __init__(self): self.process_time = 0 @@ -133,6 +147,7 @@ def __init__( LRUCache(2000) ) self.pending_event_ids_session: dict[str, set[str]] = {} + self.resumable_sessions: dict[str, ResumableSession] = {} self.event_ids_to_events: dict[str, Event] = {} self.pending_message_lock = safe_get_lock() self.event_queue_per_concurrency_id: dict[str, EventQueue] = {} @@ -164,6 +179,8 @@ def __init__( self.cached_event_analytics_summary = {"functions": {}} self.events_recorded = 0 self.event_count_at_last_cache = 0 + self.resume_ttl = float(os.getenv("GRADIO_QUEUE_SESSION_RESUME_TTL", "3600")) + self.close_grace_period = max(0.0, min(5.0, self.resume_ttl)) self.ANAYLTICS_CACHE_FREQUENCY = int( os.getenv("GRADIO_ANALYTICS_CACHE_FREQUENCY", "1") ) @@ -256,9 +273,170 @@ def send_message( if not event.alive: return event_message.event_id = event._id + self.resumable_sessions.setdefault( + event.session_hash, ResumableSession() + ).history.append(event_message) messages = self.pending_messages_per_session[event.session_hash] messages.put_nowait(event_message) + def resume_session(self, session_hash: str, event_ids: list[str]) -> None: + requested_ids = set(event_ids) + messages: AsyncQueue[EventMessage] = AsyncQueue() + resumable = self.resumable_sessions.setdefault(session_hash, ResumableSession()) + history = resumable.history + known_ids = { + message.event_id for message in history if message.event_id is not None + } + known_ids.update( + event_id + for event_id, event in self.event_ids_to_events.items() + if event.session_hash == session_hash + ) + for message in history: + if message.event_id in requested_ids: + messages.put_nowait(message) + for event_id in requested_ids - known_ids: + messages.put_nowait( + UnexpectedErrorMessage( + event_id=event_id, + message="Session event not found.", + session_not_found=True, + ) + ) + self.pending_messages_per_session[session_hash] = messages + + async def acknowledge_event(self, event_id: str) -> None: + event = self.event_ids_to_events.get(event_id) + session_hash = ( + event.session_hash + if event is not None + else next( + ( + current_session_hash + for current_session_hash, resumable in self.resumable_sessions.items() + if event_id + in self.pending_event_ids_session.get(current_session_hash, set()) + or any( + message.event_id == event_id for message in resumable.history + ) + ), + None, + ) + ) + if session_hash is None: + return + pending_ids = self.pending_event_ids_session.get(session_hash) + if pending_ids is not None: + pending_ids.discard(event_id) + if resumable := self.resumable_sessions.get(session_hash): + resumable.history = [ + message for message in resumable.history if message.event_id != event_id + ] + has_unacknowledged_history = resumable and any( + message.event_id for message in resumable.history + ) + if not pending_ids and not has_unacknowledged_history: + self.pending_event_ids_session.pop(session_hash, None) + await self.delete_session(session_hash) + else: + await self.clean_events(event_id=event_id) + + def mark_session_attached(self, session_hash: str) -> None: + resumable = self.resumable_sessions.setdefault(session_hash, ResumableSession()) + resumable.active_streams += 1 + resumable.expires_at = None + resumable.closing = False + + def mark_session_closing(self, session_hash: str) -> None: + if not self.pending_event_ids_session.get(session_hash): + return + resumable = self.resumable_sessions.setdefault(session_hash, ResumableSession()) + resumable.closing = True + if not resumable.active_streams: + resumable.expire_after(self.close_grace_period) + + async def mark_session_detached(self, session_hash: str) -> None: + resumable = self.resumable_sessions.setdefault(session_hash, ResumableSession()) + resumable.active_streams = max(0, resumable.active_streams - 1) + if resumable.active_streams: + return + if self.pending_event_ids_session.get(session_hash) or resumable.history: + resumable.expire_after( + self.close_grace_period if resumable.closing else self.resume_ttl + ) + else: + await self.delete_session(session_hash) + + async def delete_session( + self, session_hash: str, *, run_unload: bool = False + ) -> None: + if run_unload: + await self._run_unload(session_hash) + self.resumable_sessions.pop(session_hash, None) + self.pending_messages_per_session.pop(session_hash, None) + await self.clean_events(session_hash=session_hash) + + async def _run_unload(self, session_hash: str) -> None: + app = self.server_app + if app is None: + return + event = next( + ( + event + for event in self.event_ids_to_events.values() + if event.session_hash == session_hash + ), + None, + ) + if event is not None: + body = PredictBodyInternal( + session_hash=session_hash, data=[], request=event.request + ) + gr_request = route_utils.Request( + event.request, event.username, session_hash=session_hash + ) + root_path = route_utils.get_root_url( + request=event.request, + route_path=route_utils.get_api_call_path(request=event.request), + root_path=app.root_path, + ) + unload_fns = [ + fn + for fn in app.get_blocks().fns.values() + if any(target[1] == "unload" for target in fn.targets) + ] + await asyncio.gather( + *( + route_utils.call_process_api( + app=app, + body=body, + gr_request=gr_request, + fn=fn, + root_path=root_path, + ) + for fn in unload_fns + ), + return_exceptions=True, + ) + if session_hash in app.state_holder.session_data: + app.state_holder.session_data[session_hash].is_closed = True + caching.clear_session_caches(session_hash) + for event_id in self.pending_event_ids_session.get(session_hash, set()): + pending_event = self.event_ids_to_events.get(event_id) + if pending_event is not None: + pending_event.run_time = float("inf") + pending_event.signal.set() + + async def clean_expired_detached_sessions(self) -> None: + now = time.monotonic() + expired_sessions = [ + session_hash + for session_hash, resumable in self.resumable_sessions.items() + if resumable.expires_at is not None and resumable.expires_at <= now + ] + for session_hash in expired_sessions: + await self.delete_session(session_hash, run_unload=True) + def _resolve_concurrency_limit( self, default_concurrency_limit: int | None | Literal["not_set"] ) -> int | None: @@ -536,6 +714,7 @@ def get_events(self) -> tuple[list[Event], bool, str] | None: async def start_processing(self) -> None: try: while not self.stopped: + await self.clean_expired_detached_sessions() if len(self) == 0: await asyncio.sleep(self.sleep_when_free) continue @@ -649,30 +828,35 @@ def log_message( async def clean_events( self, *, session_hash: str | None = None, event_id: str | None = None ) -> None: - for job_set in self.active_jobs: - if job_set: - for job in job_set: - if job.session_hash == session_hash or job._id == event_id: - job.alive = False + events = { + event._id: event + for current_event_id, event in self.event_ids_to_events.items() + if event.session_hash == session_hash or current_event_id == event_id + } + for job in self.active_jobs: + for event in job or []: + if event.session_hash == session_hash or event._id == event_id: + events[event._id] = event + + for event in events.values(): + event.alive = False async with self.delete_lock: - events_to_remove: list[Event] = [] - for event_queue in self.event_queue_per_concurrency_id.values(): - for event in event_queue.queue: - if event.session_hash == session_hash or event._id == event_id: - events_to_remove.append(event) - - for event in events_to_remove: - self.event_queue_per_concurrency_id[event.concurrency_id].queue.remove( - event - ) + for event in events.values(): + queue = self.event_queue_per_concurrency_id[event.concurrency_id].queue + if event in queue: + queue.remove(event) self.event_ids_to_events.pop(event._id, None) - if session_hash and session_hash in self.pending_event_ids_session: - removed_ids = {e._id for e in events_to_remove} - self.pending_event_ids_session[session_hash] -= removed_ids - if not self.pending_event_ids_session[session_hash]: - self.pending_event_ids_session.pop(session_hash, None) + if session_hash: + self.pending_event_ids_session.pop(session_hash, None) + elif event_id: + for current_session_hash, pending_ids in list( + self.pending_event_ids_session.items() + ): + pending_ids.discard(event_id) + if not pending_ids: + self.pending_event_ids_session.pop(current_session_hash) async def notify_clients(self) -> None: """ diff --git a/gradio/routes.py b/gradio/routes.py index 688bf237df7..2a15b94f512 100644 --- a/gradio/routes.py +++ b/gradio/routes.py @@ -81,6 +81,7 @@ JsonData, PredictBody, PredictBodyInternal, + QueueCloseBody, ResetBody, SimplePredictBody, UserProvidedPath, @@ -1230,10 +1231,10 @@ async def _(session_hash: str, run: int, component_id: int): async def file_deprecated(path: str, request: fastapi.Request): return await file(path, request) - @router.post("/reset/") - @router.post("/reset") - async def reset_iterator(body: ResetBody): # noqa: ARG001 - # No-op, all the cancelling/reset logic handled by /cancel + @router.post("/reset/", dependencies=[Depends(login_check)]) + @router.post("/reset", dependencies=[Depends(login_check)]) + async def reset_iterator(body: ResetBody): + await app.get_blocks()._queue.acknowledge_event(body.event_id) return {"success": True} @router.get("/heartbeat/{session_hash}") @@ -1266,6 +1267,11 @@ async def iterator(): if not stop_stream_task.done(): stop_stream_task.cancel() + if app.get_blocks()._queue.pending_event_ids_session.get( + session_hash + ): + return + req = Request(request, username, session_hash=session_hash) root_path = route_utils.get_root_url( request=request, @@ -1503,6 +1509,11 @@ async def cancel_event(body: CancelBody): app.iterators_to_reset.add(body.event_id) return {"success": True} + @router.post("/queue/close", dependencies=[Depends(login_check)]) + async def close_queue_session(body: QueueCloseBody): + app.get_blocks()._queue.mark_session_closing(body.session_hash) + return {"success": True} + @router.get( "/call/v2/{api_name}/{event_id}", dependencies=[Depends(login_check)] ) @@ -1558,20 +1569,33 @@ def process_msg(message: EventMessage) -> str | None: async def queue_data( request: fastapi.Request, session_hash: str, + resume_event_id: list[str] = fastapi.Query(default=[]), + acknowledgements: bool = False, ): def process_msg(message: EventMessage) -> str: return f"data: {orjson.dumps(message.model_dump(), default=str).decode('utf-8')}\n\n" - return await queue_data_helper(request, session_hash, process_msg) + return await queue_data_helper( + request, + session_hash, + process_msg, + resume_event_id, + acknowledgements, + ) async def queue_data_helper( request: fastapi.Request, session_hash: str, process_msg: Callable[[EventMessage], str | None], + resume_event_ids: list[str] | None = None, + acknowledgements: bool = False, ): blocks = app.get_blocks() heartbeat_rate = utils.get_heartbeat_rate() + if resume_event_ids: + blocks._queue.resume_session(session_hash, resume_event_ids) + async def heartbeat(): while blocks.is_running: await asyncio.sleep(heartbeat_rate) @@ -1582,11 +1606,21 @@ async def heartbeat(): await queue.put(HeartbeatMessage()) async def sse_stream(request: fastapi.Request): + if acknowledgements: + blocks._queue.mark_session_attached(session_hash) heartbeat_task = asyncio.create_task(heartbeat()) + remaining_resume_event_ids = set(resume_event_ids or []) + + async def close_session() -> None: + if acknowledgements: + await blocks._queue.mark_session_detached(session_hash) + else: + await blocks._queue.delete_session(session_hash) + try: while True: if await request.is_disconnected(): - await blocks._queue.clean_events(session_hash=session_hash) + await close_session() heartbeat_task.cancel() return @@ -1616,50 +1650,71 @@ async def sse_stream(request: fastapi.Request): response = process_msg(message) if response is not None: yield response + if ( + isinstance(message, UnexpectedErrorMessage) + and message.session_not_found + and message.event_id + ): + remaining_resume_event_ids.discard(message.event_id) if ( isinstance(message, ProcessCompletedMessage) and message.event_id ): + remaining_resume_event_ids.discard(message.event_id) # It's possible that the event_id has already been removed # for example, the user sent two duplicate `/cancel` requests. # The first one would have removed the event_id from pending_event_ids_session if ( message.event_id - in ( - blocks._queue.pending_event_ids_session[ - session_hash - ] + in blocks._queue.pending_event_ids_session.get( + session_hash, set() ) ): blocks._queue.pending_event_ids_session[ session_hash ].remove(message.event_id) - if message.msg == ServerMessage.server_stopped or ( + all_events_delivered = ( + not remaining_resume_event_ids + if resume_event_ids + else not blocks._queue.pending_event_ids_session.get( + session_hash, set() + ) + ) + if message.msg == ServerMessage.server_stopped or ( + all_events_delivered + and ( message.msg == ServerMessage.process_completed - and ( - len( - blocks._queue.pending_event_ids_session[ - session_hash - ] - ) - == 0 + or ( + isinstance(message, UnexpectedErrorMessage) + and message.session_not_found ) - ): - message = CloseStreamMessage() - response = process_msg(message) - if response is not None: - yield response - heartbeat_task.cancel() - return + ) + ): + message = CloseStreamMessage() + response = process_msg(message) + if response is not None: + yield response + if acknowledgements: + await blocks._queue.mark_session_detached( + session_hash + ) + else: + blocks._queue.resumable_sessions.pop( + session_hash, None + ) + heartbeat_task.cancel() + return + except asyncio.CancelledError: + await close_session() + heartbeat_task.cancel() + raise except BaseException as e: message = UnexpectedErrorMessage( message=str(e), session_not_found=isinstance(e, HTTPException), ) response = process_msg(message) - if isinstance(e, asyncio.CancelledError): - del blocks._queue.pending_messages_per_session[session_hash] - await blocks._queue.clean_events(session_hash=session_hash) + await close_session() if response is not None: yield response heartbeat_task.cancel() diff --git a/guides/04_additional-features/10_environment-variables.md b/guides/04_additional-features/10_environment-variables.md index 3e9a7808046..ad2a3184c92 100644 --- a/guides/04_additional-features/10_environment-variables.md +++ b/guides/04_additional-features/10_environment-variables.md @@ -16,7 +16,7 @@ Environment variables in Gradio provide a way to customize your applications and ### 2. `GRADIO_SERVER_NAME` - **Description**: Defines the host name for the Gradio server. To make Gradio accessible from any IP address, set this to `"0.0.0.0"` -- **Default**: `"127.0.0.1"` +- **Default**: `"127.0.0.1"` - **Example**: ```bash export GRADIO_SERVER_NAME="0.0.0.0" @@ -33,7 +33,7 @@ Environment variables in Gradio provide a way to customize your applications and ### 4. `GRADIO_ANALYTICS_ENABLED` -- **Description**: Whether Gradio should provide +- **Description**: Whether Gradio should provide - **Default**: `"True"` - **Options**: `"True"`, `"False"` - **Example**: @@ -124,7 +124,6 @@ Environment variables in Gradio provide a way to customize your applications and export GRADIO_CACHE_EXAMPLES="true" ``` - ### 14. `GRADIO_CACHE_MODE` - **Description**: How to cache examples. Only applies if `cache_examples` is set to `True` either via enviornment variable or by an explicit parameter, AND no no explicit argument is passed for the `cache_mode` parameter in `gr.Interface()`, `gr.ChatInterface()` or in `gr.Examples()`. Can be set to either the strings "lazy" or "eager." If "lazy", examples are cached after their first use for all users of the app. If "eager", all examples are cached at app launch. @@ -135,17 +134,15 @@ Environment variables in Gradio provide a way to customize your applications and export GRADIO_CACHE_MODE="lazy" ``` - ### 15. `GRADIO_EXAMPLES_CACHE` -- **Description**: If you set `cache_examples=True` in `gr.Interface()`, `gr.ChatInterface()` or in `gr.Examples()`, Gradio will run your prediction function and save the results to disk. By default, this is in the `.gradio/cached_examples//` subdirectory within your app's working directory. You can customize the location of cached example files created by Gradio by setting the environment variable `GRADIO_EXAMPLES_CACHE` to an absolute path or a path relative to your working directory. +- **Description**: If you set `cache_examples=True` in `gr.Interface()`, `gr.ChatInterface()` or in `gr.Examples()`, Gradio will run your prediction function and save the results to disk. By default, this is in the `.gradio/cached_examples//` subdirectory within your app's working directory. You can customize the location of cached example files created by Gradio by setting the environment variable `GRADIO_EXAMPLES_CACHE` to an absolute path or a path relative to your working directory. - **Default**: `".gradio/cached_examples/"` - **Example**: ```sh export GRADIO_EXAMPLES_CACHE="custom_cached_examples/" ``` - ### 16. `GRADIO_SSR_MODE` - **Description**: Controls whether server-side rendering (SSR) is enabled. When enabled, the initial HTML is rendered on the server rather than the client, which can improve initial page load performance and SEO. @@ -177,7 +174,7 @@ Environment variables in Gradio provide a way to customize your applications and ### 19. `GRADIO_RESET_EXAMPLES_CACHE` -- **Description**: If set to "True", Gradio will delete and recreate the examples cache directory when the app starts instead of reusing the cached example if they already exist. +- **Description**: If set to "True", Gradio will delete and recreate the examples cache directory when the app starts instead of reusing the cached example if they already exist. - **Default**: `"False"` - **Options**: `"True"`, `"False"` - **Example**: @@ -224,7 +221,6 @@ Environment variables in Gradio provide a way to customize your applications and export GRADIO_MCP_SERVER="True" ``` - ### 24. `GRADIO_NUM_WORKERS` - **Description**: Number of multiple workers to launch in the background to offload traffic for file I/O and static assets from the main Gradio server. Only works when SSR mode is set. @@ -244,6 +240,15 @@ Environment variables in Gradio provide a way to customize your applications and export GRADIO_HEARTBEAT_INTERVAL=5 ``` +### 26. `GRADIO_QUEUE_SESSION_RESUME_TTL` + +- **Description**: Sets how many seconds an active queued job remains available for reconnection after an unexpected connection loss. Set this to `0` to cancel disconnected jobs instead of allowing them to resume. Deliberately closing or navigating away from a page uses a shorter grace period so a refresh can reconnect without leaving abandoned jobs running. +- **Default**: `3600` +- **Example**: + ```sh + export GRADIO_QUEUE_SESSION_RESUME_TTL=60 + ``` + ## How to Set Environment Variables To set environment variables in your terminal, use the `export` command followed by the variable name and its value. For example: @@ -260,6 +265,3 @@ GRADIO_SERVER_NAME="localhost" ``` Then, use a tool like `dotenv` to load these variables when running your application. - - - diff --git a/guides/09_gradio-clients-and-lite/01_getting-started-with-the-python-client.md b/guides/09_gradio-clients-and-lite/01_getting-started-with-the-python-client.md index 70b8ce4949b..54b99363703 100644 --- a/guides/09_gradio-clients-and-lite/01_getting-started-with-the-python-client.md +++ b/guides/09_gradio-clients-and-lite/01_getting-started-with-the-python-client.md @@ -252,6 +252,27 @@ job.status() _Note_: The `Job` class also has a `.done()` instance method which returns a boolean indicating whether the job has completed. +## Resuming Jobs After a Disconnect + +Pass `resume_sessions=True` to keep active queued jobs available when the client temporarily loses its connection. Save the job's event ID, function index, and the client's session hash if you need to recreate the client process and retrieve the same job later: + +```py +from gradio_client import Client + +client = Client("gradio/calculator", resume_sessions=True) +job = client.submit(5, "add", 4, api_name="/predict") + +saved_job = {"event_id": job.wait_for_id(), "fn_index": job.fn_index} +saved_session = client.session_hash + +# The same values can be loaded from your own persistent storage after restart. +client = Client("gradio/calculator") +job = client.resume_jobs([saved_job], session_hash=saved_session)[0] +print(job.result()) # 9.0 +``` + +The server retains an interrupted job for up to `GRADIO_QUEUE_SESSION_RESUME_TTL` seconds (1 hour by default). The result is acknowledged and removed from the resume buffer after the resumed client receives it. + ## Cancelling Jobs The `Job` class also has a `.cancel()` instance method that cancels jobs that have been queued but not started. For example, if you run: diff --git a/guides/09_gradio-clients-and-lite/02_getting-started-with-the-js-client.md b/guides/09_gradio-clients-and-lite/02_getting-started-with-the-js-client.md index df6fca2c999..8eed3082df4 100644 --- a/guides/09_gradio-clients-and-lite/02_getting-started-with-the-js-client.md +++ b/guides/09_gradio-clients-and-lite/02_getting-started-with-the-js-client.md @@ -294,6 +294,28 @@ for await (const message of job) { } ``` +## Resuming Jobs After a Disconnect + +Set `resume_sessions: true` when connecting from a browser. The client stores active queued event IDs in browser session storage, reconnects the stream after temporary network failures, and exposes unfinished jobs after a page reload: + +```js +import { Client } from "@gradio/client"; + +const app = await Client.connect(location.origin, { + resume_sessions: true, + events: ["status", "data"] +}); + +const resumed_jobs = app.resume_jobs(); +for (const job of resumed_jobs) { + for await (const message of job) { + console.log(message); + } +} +``` + +New calls made with `app.submit()` are tracked automatically. Call `app.resume_jobs()` after reconnecting the page to reattach to those same event IDs rather than submitting duplicate work. The browser keeps active event IDs for the current tab until the server acknowledges or rejects them, while the server keeps disconnected jobs for up to `GRADIO_QUEUE_SESSION_RESUME_TTL` seconds (1 hour by default). + ## Cancelling Jobs The job instance also has a `.cancel()` method that cancels jobs that have been queued but not started. For example, if you run: diff --git a/guides/09_gradio-clients-and-lite/08_server-mode.md b/guides/09_gradio-clients-and-lite/08_server-mode.md index 047696664d6..08ac7d9c7ca 100644 --- a/guides/09_gradio-clients-and-lite/08_server-mode.md +++ b/guides/09_gradio-clients-and-lite/08_server-mode.md @@ -121,6 +121,37 @@ Then open `http://localhost:7860` in your browser. The custom HTML page uses the Note: if your `Server` app uses ZeroGPU, you _must_ call Gradio API endpoints through `@gradio/client` from the browser. The JavaScript client forwards the Hugging Face iframe auth headers needed for ZeroGPU quota handling. +## Resuming Queued Work + +Server-mode endpoints use the same resumable queue as `Blocks` apps. For a custom browser frontend, connect with `resume_sessions: true` and call `resume_jobs()` when the page loads. This restores any active submissions saved in the browser session instead of starting them again: + +```js +const app = await Client.connect(location.origin, { + resume_sessions: true, + events: ["status", "data"] +}); + +for (const job of app.resume_jobs()) { + for await (const message of job) { + console.log(message); + } +} +``` + +Python services can persist the event ID, function index, and session hash, then recreate the job with `Client.resume_jobs()`: + +```python +client = Client("http://localhost:7860", resume_sessions=True) +job = client.submit("hello", api_name="/generate") +saved_job = {"event_id": job.wait_for_id(), "fn_index": job.fn_index} +saved_session = client.session_hash + +client = Client("http://localhost:7860") +job = client.resume_jobs([saved_job], session_hash=saved_session)[0] +``` + +When storing the values separately, preserve the original client's `session_hash` and pass that value to the new client. Resumed generator endpoints replay their buffered chunks in order before continuing with live output. + ## Concurrency and Streaming `app.api()` supports all of the same concurrency and streaming options as `gr.api()`: diff --git a/js/app/src/routes/[...catchall]/+page.svelte b/js/app/src/routes/[...catchall]/+page.svelte index 8359a8ea3be..93eab85be4c 100644 --- a/js/app/src/routes/[...catchall]/+page.svelte +++ b/js/app/src/routes/[...catchall]/+page.svelte @@ -265,8 +265,8 @@ } onMount(async () => { - window.addEventListener("beforeunload", () => { - app?.close(); + window.addEventListener("pagehide", (event) => { + if (!event.persisted) app?.close(); }); //@ts-ignore config = data.config; diff --git a/js/app/src/routes/[...catchall]/+page.ts b/js/app/src/routes/[...catchall]/+page.ts index d4c467bf966..09875d2c3bf 100644 --- a/js/app/src/routes/[...catchall]/+page.ts +++ b/js/app/src/routes/[...catchall]/+page.ts @@ -95,6 +95,7 @@ export async function load({ try { app = await Client.connect(api_url, { with_null_state: true, + resume_sessions: true, events: ["data", "log", "status", "render"], query_params: deepLink ? { deep_link: deepLink } : undefined, headers, diff --git a/js/core/src/Blocks.svelte b/js/core/src/Blocks.svelte index a7718bcd1a0..7bf1b0a60fa 100644 --- a/js/core/src/Blocks.svelte +++ b/js/core/src/Blocks.svelte @@ -391,9 +391,6 @@ const MESSAGE_QUOTE_RE = /^'([^]+)'$/; const DUPLICATE_MESSAGE = $reactive_formatter("blocks.long_requests_queue"); - const MOBILE_QUEUE_WARNING = $reactive_formatter( - "blocks.connection_can_break" - ); const LOST_CONNECTION_MESSAGE = "Connection to the server was lost. Attempting reconnection..."; const CHANGED_CONNECTION_MESSAGE = @@ -405,10 +402,7 @@ "blocks.waiting_for_inputs" ); const SHOW_DUPLICATE_MESSAGE_ON_ETA = 15; - const SHOW_MOBILE_QUEUE_WARNING_ON_ETA = 10; - let is_mobile_device = false; let showed_duplicate_message = false; - let showed_mobile_warning = false; let inputs_waiting: number[] = []; // as state updates are not synchronous, we need to ensure updates are flushed before triggering any requests @@ -464,10 +458,9 @@ } onMount(() => { - is_mobile_device = - /Android|webOS|iPhone|iPad|iPod|BlackBerry|IEMobile|Opera Mini/i.test( - navigator.userAgent - ); + if ("parentIFrame" in window) { + window.parentIFrame?.autoResize(false); + } refresh_run_count(); const unsubscribe_run_history = on_run_history_change(refresh_run_count); @@ -491,7 +484,13 @@ ready = true; reset_resize_growth(resize_state); void settled().then(handle_resize); - dep_manager.dispatch_load_events(); + const resumable_events = dep_manager.get_resumable_events(); + dep_manager.dispatch_load_events( + new Set(resumable_events.map(({ fn_index }) => fn_index)) + ); + if (resumable_events.length > 0) { + void dep_manager.resume(resumable_events); + } }); if (vibe_mode) { diff --git a/js/core/src/dependency.test.ts b/js/core/src/dependency.test.ts index 4ed36e66516..80fe0a72a68 100644 --- a/js/core/src/dependency.test.ts +++ b/js/core/src/dependency.test.ts @@ -83,3 +83,26 @@ describe("DependencyManager.reload", () => { expect(dependency_manager.loading_stati.fn_outputs[0]).toEqual([32]); }); }); + +describe("DependencyManager.dispatch_load_events", () => { + test("dispatches load handlers that are not already being resumed", () => { + const resumed_load = dependency(0, "resumed_load", []); + resumed_load.targets = [[1, "load"]]; + const normal_load = dependency(1, "normal_load", []); + normal_load.targets = [[2, "load"]]; + const dependency_manager = manager([resumed_load, normal_load]); + const dispatch = vi + .spyOn(dependency_manager, "dispatch") + .mockResolvedValue(undefined); + + dependency_manager.dispatch_load_events(new Set([0])); + + expect(dispatch).toHaveBeenCalledTimes(1); + expect(dispatch).toHaveBeenCalledWith({ + type: "fn", + fn_index: 1, + event_data: null, + target_id: 2 + }); + }); +}); diff --git a/js/core/src/dependency.ts b/js/core/src/dependency.ts index bac38d18fd7..92bf7451363 100644 --- a/js/core/src/dependency.ts +++ b/js/core/src/dependency.ts @@ -6,7 +6,7 @@ import type { Payload } from "./types.js"; import { AsyncFunction } from "./init_utils"; -import { Client, type client_return } from "@gradio/client"; +import { Client, type client_return, type ResumableJob } from "@gradio/client"; import { LoadingStatusState, type LoadingStatusArgs @@ -600,6 +600,7 @@ export class DependencyManager { ); }); unset_args.forEach((fn) => fn()); + await dep_submission.data.acknowledge(); this.clear_submission(dep.id, dep_submission.data); if (this.queue.has(dep.id)) { this.queue.delete(dep.id); @@ -677,6 +678,7 @@ export class DependencyManager { this.handle_log(result); } } + await dep_submission.data.acknowledge(); all.forEach((dep_id) => { this.dispatch({ type: "fn", @@ -694,6 +696,7 @@ export class DependencyManager { } } } catch (error) { + await this.submissions.get(dep.id)?.acknowledge(); this.loading_stati.update({ status: "error", fn_index: dep.id, @@ -718,6 +721,66 @@ export class DependencyManager { return; } + get_resumable_events(): ResumableJob[] { + return this.client + .get_resumable_events() + .filter(({ fn_index }) => this.dependencies_by_fn.has(fn_index)); + } + + async resume( + events: ResumableJob[] = this.get_resumable_events() + ): Promise { + const submissions = this.client.resume_jobs(events); + await Promise.all( + events.map(async ({ fn_index }, index) => { + const dep = this.dependencies_by_fn.get(fn_index); + if (!dep) return; + + const submission = submissions[index]; + this.submissions.set(fn_index, submission); + this.loading_stati.update({ + status: "pending", + fn_index, + stream_state: null + }); + await this.update_loading_stati_state(); + + try { + for await (const result of submission) { + if (result.type === "data") { + await this.handle_data(dep.outputs, result.data); + } else if (result.type === "log") { + this.handle_log(result); + } else if (result.type === "status") { + this.loading_stati.update({ + status: result.stage, + fn_index, + stream_state: null, + queue: result.queue, + size: result.size, + position: result.position, + eta: result.eta, + progress_data: result.progress_data, + time_limit: result.time_limit, + used_cache: result.used_cache, + cache_duration: result.cache_duration, + avg_time: result.avg_time + }); + await this.update_loading_stati_state(); + if (result.stage === "complete" || result.stage === "error") { + this.dispatch_state_change_events(result); + break; + } + } + } + } finally { + await submission.acknowledge(); + this.clear_submission(fn_index, submission); + } + }) + ); + } + /** * Creates a map of dependencies for easy lookup * @@ -917,8 +980,9 @@ export class DependencyManager { } } - dispatch_load_events() { + dispatch_load_events(excluded_fn_indices: Set = new Set()) { this.dependencies_by_fn.forEach((dep) => { + if (excluded_fn_indices.has(dep.id)) return; dep.targets.forEach(([target_id, event_name]) => { if (event_name === "load") { this.dispatch({ diff --git a/js/spa/src/Index.svelte b/js/spa/src/Index.svelte index f938afee233..73cbe9e3725 100644 --- a/js/spa/src/Index.svelte +++ b/js/spa/src/Index.svelte @@ -409,11 +409,12 @@ app = await Client.connect(api_url, { status_callback: handle_status, with_null_state: true, + resume_sessions: true, events: ["data", "log", "status", "render"], query_params }); - window.addEventListener("beforeunload", () => { - app.close(); + window.addEventListener("pagehide", (event) => { + if (!event.persisted) app.close(); }); if (!app.config && !config?.auth_required) { diff --git a/js/spa/test/cancel_events.spec.ts b/js/spa/test/cancel_events.spec.ts index 65e44ce2984..7d5e83cf367 100644 --- a/js/spa/test/cancel_events.spec.ts +++ b/js/spa/test/cancel_events.spec.ts @@ -1,5 +1,4 @@ import { test, expect } from "@self/tootils"; -import { readFileSync } from "fs"; test("when using an iterative function the UI should update over time as iteration results are received", async ({ page @@ -45,27 +44,3 @@ test("when using an iterative function it should be possible to cancel the funct "Cancel follow-up ran" ); }); - -test("when using an iterative function and the user closes the page, the python function should stop running", async ({ - page -}) => { - const start_button = await page.locator("button", { - hasText: /Start Iterating/ - }); - - await start_button.click(); - await page.waitForTimeout(300); - await page.close(); - - // wait for the duration of the entire iteration - // check that the final value did not get written - // to the log file. That's our proof python stopped - // running - await new Promise((resolve) => setTimeout(resolve, 2000)); - const data = readFileSync( - "../../demo/cancel_events/cancel_events_output_log.txt", - "utf-8" - ); - expect(data).toContain("Current step: 0"); - expect(data).not.toContain("Current step: 8"); -}); diff --git a/test/test_queueing.py b/test/test_queueing.py index c7bd934022a..08961189a33 100644 --- a/test/test_queueing.py +++ b/test/test_queueing.py @@ -1,6 +1,7 @@ import asyncio import json import sys +import threading import time from unittest.mock import patch @@ -231,14 +232,16 @@ def tracking_create_task(coro, **kwargs): demo.close() -def test_cancel_removes_pending_event_from_queue(): - """Cancelling a queued (not yet running) event should remove it from the queue.""" +def test_detached_queue_session_can_resume(): + finished = threading.Event() + with gr.Blocks() as demo: start = gr.Button() output = gr.Textbox() def slow(): - time.sleep(2) + time.sleep(0.05) + finished.set() return "done" start.click(slow, None, output) @@ -247,62 +250,137 @@ def slow(): app, _, _ = demo.launch(prevent_thread_lock=True) test_client = TestClient(app) - join_payload = { - "data": [], - "fn_index": 0, - "event_data": None, - "session_hash": "sess1", - "trigger_id": None, - } + try: + response = test_client.post( + f"{API_PREFIX}/queue/join", + json={ + "data": [], + "fn_index": 0, + "event_data": None, + "session_hash": "resume_session", + "trigger_id": None, + }, + ) + assert response.status_code == 200 + event_id = response.json()["event_id"] + + asyncio.run(demo._queue.mark_session_detached("resume_session")) + + assert event_id in demo._queue.event_ids_to_events + assert event_id in demo._queue.pending_event_ids_session["resume_session"] + assert "resume_session" in demo._queue.pending_messages_per_session + assert finished.wait(timeout=1) + + deadline = time.monotonic() + 1 + while any(demo._queue.active_jobs) and time.monotonic() < deadline: + time.sleep(0.01) + assert not any(demo._queue.active_jobs) + + response = test_client.get( + f"{API_PREFIX}/queue/data?session_hash=resume_session" + f"&resume_event_id={event_id}&acknowledgements=true" + ) + completed_data = None + for line in response.iter_lines(): + if "data" not in line: + continue + data = json.loads(line[5:]) + if data["msg"] == "process_completed": + completed_data = data["output"]["data"] + + assert completed_data == ["done"] + assert event_id not in demo._queue.pending_event_ids_session["resume_session"] + assert demo._queue.resumable_sessions["resume_session"].history + + response = test_client.post( + f"{API_PREFIX}/reset", + json={"event_id": event_id}, + ) + assert response.status_code == 200 + assert "resume_session" not in demo._queue.pending_event_ids_session + assert "resume_session" not in demo._queue.resumable_sessions + finally: + demo.close() + + +def test_missing_resumed_event_closes_stream(): + with gr.Blocks() as demo: + gr.Button().click(lambda: None) + + demo.queue() + app, _, _ = demo.launch(prevent_thread_lock=True) + test_client = TestClient(app) try: - first = test_client.post(f"{API_PREFIX}/queue/join", json=join_payload) - second = test_client.post(f"{API_PREFIX}/queue/join", json=join_payload) - third = test_client.post(f"{API_PREFIX}/queue/join", json=join_payload) - assert first.status_code == 200 - assert second.status_code == 200 - assert third.status_code == 200 - - second_event_id = second.json()["event_id"] - third_event_id = third.json()["event_id"] - - # First event gets picked up by the worker; second and third are queued. - # The worker dequeues asynchronously, so wait for it to settle before - # asserting (avoids a race on slower CI runners where all three are - # momentarily still in the queue). - for _ in range(50): - if len(demo._queue) == 2: - break - time.sleep(0.1) - assert len(demo._queue) == 2 - assert second_event_id in demo._queue.event_ids_to_events - assert second_event_id in demo._queue.pending_event_ids_session["sess1"] - - # Cancel the second (pending/queued) event - resp = test_client.post( - f"{API_PREFIX}/cancel", + response = test_client.get( + f"{API_PREFIX}/queue/data?session_hash=missing_session" + "&resume_event_id=missing_event&acknowledgements=true" + ) + messages = [json.loads(line[5:]) for line in response.iter_lines() if line] + + assert messages[0]["session_not_found"] is True + assert messages[-1]["msg"] == "close_stream" + finally: + demo.close() + + +def test_expired_detached_session_is_fully_cleaned_up(): + started = threading.Event() + unloaded = threading.Event() + + with gr.Blocks() as demo: + start = gr.Button() + + def slow(): + started.set() + time.sleep(0.5) + + start.click(slow) + demo.unload(unloaded.set) + + demo.queue(default_concurrency_limit=1) + app, _, _ = demo.launch(prevent_thread_lock=True) + test_client = TestClient(app) + + try: + response = test_client.post( + f"{API_PREFIX}/queue/join", json={ - "session_hash": "sess1", + "data": [], "fn_index": 0, - "event_id": second_event_id, + "event_data": None, + "session_hash": "expired_session", + "trigger_id": None, }, ) - assert resp.status_code == 200 - assert third_event_id in demo._queue.event_ids_to_events + event_id = response.json()["event_id"] + assert started.wait(timeout=1) - assert len(demo._queue) == 1 - r = test_client.get(f"{API_PREFIX}/queue/data?session_hash=sess1") + response = test_client.post( + f"{API_PREFIX}/queue/close", + json={"session_hash": "expired_session"}, + ) + assert response.status_code == 200 + assert ( + demo._queue.resumable_sessions["expired_session"].expires_at + <= time.monotonic() + 5 + ) - # Verify we got a process_completed message - got_completed = False - for line in r.iter_lines(): - if "data" in line: - data = json.loads(line[5:]) - if data["msg"] == "process_completed": - got_completed = True - assert got_completed - assert second_event_id not in demo._queue.pending_event_ids_session["sess1"] - assert second_event_id not in demo._queue.event_ids_to_events + demo._queue.resumable_sessions["expired_session"].expires_at = 0 + asyncio.run(demo._queue.clean_expired_detached_sessions()) + + assert unloaded.is_set() + assert app.state_holder.session_data["expired_session"].is_closed is True + assert event_id not in demo._queue.event_ids_to_events + assert "expired_session" not in demo._queue.pending_event_ids_session + + response = test_client.get( + f"{API_PREFIX}/queue/data?session_hash=expired_session" + f"&resume_event_id={event_id}&acknowledgements=true" + ) + messages = [json.loads(line[5:]) for line in response.iter_lines() if line] + assert messages[0]["session_not_found"] is True + assert messages[-1]["msg"] == "close_stream" finally: demo.close() diff --git a/test/test_routes.py b/test/test_routes.py index 630b256e5fa..8096ea592a2 100644 --- a/test/test_routes.py +++ b/test/test_routes.py @@ -1131,6 +1131,28 @@ def test_monitoring_route(self): response = client.get("/monitoring/summary") assert response.status_code == 401 + def test_queue_reset_route(self): + io = Interface(lambda x: x, "text", "text") + app, _, _ = io.launch( + auth=("test", "correct_password"), + prevent_thread_lock=True, + ) + client = TestClient(app) + body = {"event_id": "event"} + + try: + response = client.post(f"{API_PREFIX}/reset", json=body) + assert response.status_code == 401 + + client.post( + "/login", + data={"username": "test", "password": "correct_password"}, + ) + response = client.post(f"{API_PREFIX}/reset", json=body) + assert response.status_code == 200 + finally: + io.close() + class TestConfigUsername: """`/config` reports who is logged in. Resolving the user is async, so a diff --git a/test/test_server.py b/test/test_server.py index f57acb2c26d..b4634fe0cb8 100644 --- a/test/test_server.py +++ b/test/test_server.py @@ -1,4 +1,6 @@ import inspect +import time +from collections.abc import Iterator from contextlib import asynccontextmanager import pytest @@ -58,6 +60,47 @@ def echo(x: str) -> str: finally: gr.close_all() + def test_server_stream_can_resume_with_python_client(self): + server = gr.Server() + + @server.api(name="count") + def count(limit: int) -> Iterator[str]: + for value in range(limit): + time.sleep(0.02) + yield str(value) + + app, local_url, _ = server.launch(prevent_thread_lock=True) + session_hash = "server-resume-session" + try: + with TestClient(app) as test_client: + response = test_client.post( + f"{API_PREFIX}/queue/join", + json={ + "data": [3], + "fn_index": 0, + "event_data": None, + "session_hash": session_hash, + "trigger_id": None, + }, + ) + event_id = response.json()["event_id"] + + client = Client(local_url) + try: + job = client.resume_jobs( + [{"event_id": event_id, "fn_index": 0}], + session_hash=session_hash, + )[0] + deadline = time.monotonic() + 5 + while not job.done() and time.monotonic() < deadline: + time.sleep(0.01) + assert job.done() + assert job.outputs() == ["0", "1", "2"] + finally: + client.close() + finally: + gr.close_all() + def test_server_launch_enables_mcp(self): server = gr.Server()