diff --git a/.changeset/quiet-sessions-return.md b/.changeset/quiet-sessions-return.md new file mode 100644 index 00000000000..34cb2d2b7fb --- /dev/null +++ b/.changeset/quiet-sessions-return.md @@ -0,0 +1,7 @@ +--- +"@gradio/client": minor +"gradio": minor +"gradio_client": minor +--- + +feat: Resume queued jobs and restore undelivered output after temporary disconnects or page refreshes diff --git a/client/js/src/client.ts b/client/js/src/client.ts index 75ea0f46eac..1fe5bd02da5 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, @@ -37,6 +42,7 @@ import { open_stream, readable_stream, close_stream } from "./utils/stream"; import { API_INFO_ERROR_MSG, APP_ID_URL, + CLOSE_URL, CONFIG_ERROR_MSG, HEARTBEAT_URL, COMPONENT_SERVER_URL @@ -57,14 +63,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; @@ -181,6 +191,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, @@ -215,6 +226,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); @@ -232,10 +258,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(); @@ -293,13 +328,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}` @@ -321,7 +367,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..a4720d3b959 100644 --- a/client/js/src/constants.ts +++ b/client/js/src/constants.ts @@ -15,6 +15,8 @@ export const HEARTBEAT_URL = `heartbeat`; export const COMPONENT_SERVER_URL = `component_server`; export const RESET_URL = `reset`; export const CANCEL_URL = `cancel`; +export const ACK_URL = `queue/ack`; +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 6ea916ad113..9c99f4bdd26 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 type { SpaceStatus, diff --git a/client/js/src/test/init.test.ts b/client/js/src/test/init.test.ts index 17f8ec50916..9ca870e9220 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 { 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,59 @@ 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 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..8400bb190c6 --- /dev/null +++ b/client/js/src/test/session.test.ts @@ -0,0 +1,87 @@ +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("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 36240d81c78..d1c063867cf 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; } @@ -338,6 +339,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..983843c7d68 --- /dev/null +++ b/client/js/src/utils/session.ts @@ -0,0 +1,149 @@ +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[]; + expires_at: number; +} + +const STORAGE_KEY = "gradio_active_session"; +const COOKIE_NAME = "gradio_active_session"; +const RESUME_TTL_SECONDS = 600; + +function read_session(): ResumableSession | null { + if (typeof sessionStorage === "undefined") return null; + + try { + const value = sessionStorage.getItem(STORAGE_KEY); + if (!value) return null; + const session = JSON.parse(value) as ResumableSession; + if (session.expires_at <= Date.now()) { + sessionStorage.removeItem(STORAGE_KEY); + return null; + } + return session; + } 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)}; Max-Age=${RESUME_TTL_SECONDS}; 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], + expires_at: Date.now() + RESUME_TTL_SECONDS * 1000 + }); +} + +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 db642d823a3..440d78d85f2 100644 --- a/client/js/src/utils/submit.ts +++ b/client/js/src/utils/submit.ts @@ -25,9 +25,11 @@ import { SSE_DATA_URL, RESET_URL, CANCEL_URL, + ACK_URL, 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"; export function submit( @@ -37,7 +39,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; @@ -72,7 +75,9 @@ 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 stream: EventSource | null; let protocol = config.protocol ?? "ws"; @@ -137,6 +142,23 @@ 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}/${ACK_URL}`, { + headers: { "Content-Type": "application/json" }, + method: "POST", + body: JSON.stringify({ + event_id: event_id_final, + session_hash + }) + }) + .catch(() => null); + if (response?.ok) clear_resumable_event(event_id_final); } const resolve_heartbeat = async (config: Config): Promise => { @@ -168,11 +190,179 @@ 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); + } + 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, @@ -490,146 +680,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); } }); } @@ -725,7 +782,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..199a9c8f31f 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 = 600 -@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, session_hash) 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,25 @@ 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, session_hash: str | None = None + ) -> None: + self.pending_event_ids.discard(event_id) + try: + response = httpx.post( + urllib.parse.urljoin(self.src_prefixed, utils.ACK_URL), + json={ + "event_id": event_id, + "session_hash": session_hash or self.session_hash, + }, + 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 +547,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 +590,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 +628,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 +649,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 +1345,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 +1566,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 +1585,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 +1593,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 +1601,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 +1655,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/gradio_client/utils.py b/client/python/gradio_client/utils.py index ecc531c4c10..71a4449b347 100644 --- a/client/python/gradio_client/utils.py +++ b/client/python/gradio_client/utils.py @@ -58,6 +58,7 @@ SPACE_URL = "https://hf.space/{}" HEARTBEAT_URL = "heartbeat/{session_hash}" CANCEL_URL = "cancel" +ACK_URL = "queue/ack" STATE_COMPONENT = "state" INVALID_RUNTIME = [ diff --git a/client/python/test/test_client.py b/client/python/test/test_client.py index 78c6c6062de..33ff073b15a 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", "session-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 9dd7912ff84..78cc233d8f4 100644 --- a/gradio/data_classes.py +++ b/gradio/data_classes.py @@ -53,6 +53,15 @@ class CancelBody(BaseModel): event_id: str +class QueueAckBody(BaseModel): + session_hash: str + 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 3d67f79ecad..b74a9e66fd6 100644 --- a/gradio/queueing.py +++ b/gradio/queueing.py @@ -16,7 +16,7 @@ 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, @@ -33,6 +33,7 @@ ProgressMessage, ProgressUnit, ServerMessage, + UnexpectedErrorMessage, ) from gradio.utils import ( LRUCache, @@ -126,6 +127,7 @@ def __init__( LRUCache(2000) ) self.pending_event_ids_session: dict[str, set[str]] = {} + self.message_history_per_session: dict[str, list[EventMessage]] = {} self.event_ids_to_events: dict[str, Event] = {} self.pending_message_lock = safe_get_lock() self.event_queue_per_concurrency_id: dict[str, EventQueue] = {} @@ -152,6 +154,10 @@ def __init__( self.event_analytics: dict[str, dict[str, float | str | None]] = {} self.cached_event_analytics_summary = {"functions": {}} self.event_count_at_last_cache = 0 + self.detached_session_expirations: dict[str, float] = {} + self.active_session_streams: defaultdict[str, int] = defaultdict(int) + self.resume_ttl = float(os.getenv("GRADIO_QUEUE_SESSION_RESUME_TTL", "600")) + 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") ) @@ -244,9 +250,167 @@ def send_message( if not event.alive: return event_message.event_id = event._id + self.message_history_per_session.setdefault(event.session_hash, []).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() + history = self.message_history_per_session.get(session_hash, []) + 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, session_hash: str, event_id: str) -> None: + pending_ids = self.pending_event_ids_session.get(session_hash) + if pending_ids is not None: + pending_ids.discard(event_id) + self.event_ids_to_events.pop(event_id, None) + self.message_history_per_session[session_hash] = [ + message + for message in self.message_history_per_session.get(session_hash, []) + if message.event_id != event_id + ] + if not pending_ids: + self.pending_event_ids_session.pop(session_hash, None) + await self.delete_session(session_hash) + + def mark_session_attached(self, session_hash: str) -> None: + self.active_session_streams[session_hash] += 1 + self.detached_session_expirations.pop(session_hash, None) + + def mark_session_closing(self, session_hash: str) -> None: + if not self.pending_event_ids_session.get(session_hash): + return + expiration = time.monotonic() + self.close_grace_period + current_expiration = self.detached_session_expirations.get(session_hash) + self.detached_session_expirations[session_hash] = ( + min(expiration, current_expiration) + if current_expiration is not None + else expiration + ) + + def close_session_stream(self, session_hash: str) -> None: + if self.active_session_streams.get(session_hash, 0) > 0: + self.active_session_streams[session_hash] -= 1 + if self.active_session_streams[session_hash] == 0: + self.active_session_streams.pop(session_hash, None) + + async def mark_session_detached( + self, + session_hash: str, + *, + stream_closed: bool = True, + ) -> None: + if stream_closed: + self.close_session_stream(session_hash) + + if self.active_session_streams.get(session_hash, 0) > 0: + return + + if self.pending_event_ids_session.get(session_hash): + expiration = time.monotonic() + self.resume_ttl + current_expiration = self.detached_session_expirations.get(session_hash) + self.detached_session_expirations[session_hash] = ( + min(expiration, current_expiration) + if current_expiration is not None + else expiration + ) + 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.detached_session_expirations.pop(session_hash, None) + self.active_session_streams.pop(session_hash, None) + self.pending_messages_per_session.pop(session_hash, None) + self.message_history_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, expires_at in self.detached_session_expirations.items() + if 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: @@ -520,6 +684,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 @@ -631,11 +796,19 @@ def log_message( async def clean_events( self, *, session_hash: str | None = None, event_id: str | None = None ) -> None: + removed_ids = { + current_event_id + for current_event_id, event in self.event_ids_to_events.items() + if event.session_hash == session_hash or current_event_id == event_id + } + if session_hash: + removed_ids.update(self.pending_event_ids_session.get(session_hash, set())) 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 + removed_ids.add(job._id) async with self.delete_lock: events_to_remove: list[Event] = [] @@ -643,18 +816,25 @@ async def clean_events( for event in event_queue.queue: if event.session_hash == session_hash or event._id == event_id: events_to_remove.append(event) + removed_ids.add(event._id) for event in events_to_remove: self.event_queue_per_concurrency_id[event.concurrency_id].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) + for removed_id in removed_ids: + self.event_ids_to_events.pop(removed_id, 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 764326eb524..77cd8e7e264 100644 --- a/gradio/routes.py +++ b/gradio/routes.py @@ -81,6 +81,8 @@ JsonData, PredictBody, PredictBodyInternal, + QueueAckBody, + QueueCloseBody, ResetBody, SimplePredictBody, UserProvidedPath, @@ -1228,6 +1230,14 @@ 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 + ): + await app.get_blocks()._queue.mark_session_detached( + session_hash, stream_closed=False + ) + return + req = Request(request, username, session_hash=session_hash) root_path = route_utils.get_root_url( request=request, @@ -1456,6 +1466,18 @@ async def cancel_event(body: CancelBody): app.iterators_to_reset.add(body.event_id) return {"success": True} + @router.post("/queue/ack", dependencies=[Depends(login_check)]) + async def acknowledge_queue_event(body: QueueAckBody): + await app.get_blocks()._queue.acknowledge_event( + body.session_hash, 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)] ) @@ -1494,20 +1516,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) @@ -1518,11 +1553,14 @@ async def heartbeat(): await queue.put(HeartbeatMessage()) async def sse_stream(request: fastapi.Request): + blocks._queue.mark_session_attached(session_hash) heartbeat_task = asyncio.create_task(heartbeat()) + delivered_completed_event_ids: set[str] = set() + remaining_resume_event_ids = set(resume_event_ids or []) try: while True: if await request.is_disconnected(): - await blocks._queue.clean_events(session_hash=session_hash) + await blocks._queue.mark_session_detached(session_hash) heartbeat_task.cancel() return @@ -1552,14 +1590,22 @@ 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 ): + delivered_completed_event_ids.add(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 ( + if not acknowledgements and ( message.event_id in ( blocks._queue.pending_event_ids_session[ @@ -1570,32 +1616,55 @@ async def sse_stream(request: fastapi.Request): 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 ( + blocks._queue.pending_event_ids_session.get( + session_hash, set() + ).issubset(delivered_completed_event_ids) + if acknowledgements + 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.close_session_stream(session_hash) + blocks._queue.message_history_per_session.pop( + session_hash, None + ) + heartbeat_task.cancel() + return + except asyncio.CancelledError: + await blocks._queue.mark_session_detached(session_hash) + 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 blocks._queue.mark_session_detached(session_hash) 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..2ef051c7cee 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**: `600` +- **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..dc8be110734 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 (10 minutes 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..5f21d581574 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 server keeps disconnected jobs for up to `GRADIO_QUEUE_SESSION_RESUME_TTL` seconds (10 minutes 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 45b3b977868..b8755c33c0e 100644 --- a/js/app/src/routes/[...catchall]/+page.svelte +++ b/js/app/src/routes/[...catchall]/+page.svelte @@ -267,8 +267,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 0cdd18225de..c07dbef3727 100644 --- a/js/core/src/Blocks.svelte +++ b/js/core/src/Blocks.svelte @@ -370,9 +370,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 = @@ -384,10 +381,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 @@ -450,11 +444,6 @@ } 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); } @@ -471,7 +460,13 @@ app_tree.ready.then(() => { ready = true; - 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 c1c31398ffe..be7962b7c94 100644 --- a/js/spa/src/Index.svelte +++ b/js/spa/src/Index.svelte @@ -346,11 +346,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 b47e327fad9..557001dbdfb 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 @@ -307,6 +308,160 @@ def slow(): demo.close() +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(0.05) + finished.set() + return "done" + + start.click(slow, None, output) + + 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={ + "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 in demo._queue.pending_event_ids_session["resume_session"] + assert demo._queue.message_history_per_session["resume_session"] + + response = test_client.post( + f"{API_PREFIX}/queue/ack", + json={"session_hash": "resume_session", "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.message_history_per_session + assert "resume_session" not in demo._queue.detached_session_expirations + 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: + 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={ + "data": [], + "fn_index": 0, + "event_data": None, + "session_hash": "expired_session", + "trigger_id": None, + }, + ) + event_id = response.json()["event_id"] + assert started.wait(timeout=1) + + response = test_client.post( + f"{API_PREFIX}/queue/close", + json={"session_hash": "expired_session"}, + ) + assert response.status_code == 200 + assert ( + demo._queue.detached_session_expirations["expired_session"] + <= time.monotonic() + 5 + ) + + demo._queue.detached_session_expirations["expired_session"] = 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() + + def test_analytics_summary(monkeypatch): """Test that the analytics summary endpoint is correctly being computed every N requests, where N is set by the GRADIO_ANALYTICS_CACHE_FREQUENCY environment variable.""" diff --git a/test/test_routes.py b/test/test_routes.py index 3a499325e9a..c81b2257ca3 100644 --- a/test/test_routes.py +++ b/test/test_routes.py @@ -989,6 +989,28 @@ def test_monitoring_route(self): response = client.get("/monitoring/summary") assert response.status_code == 401 + def test_queue_ack_route(self): + io = Interface(lambda x: x, "text", "text") + app, _, _ = io.launch( + auth=("test", "correct_password"), + prevent_thread_lock=True, + ) + client = TestClient(app) + body = {"session_hash": "session", "event_id": "event"} + + try: + response = client.post(f"{API_PREFIX}/queue/ack", json=body) + assert response.status_code == 401 + + client.post( + "/login", + data={"username": "test", "password": "correct_password"}, + ) + response = client.post(f"{API_PREFIX}/queue/ack", json=body) + assert response.status_code == 200 + finally: + io.close() + class TestQueueRoutes: @pytest.mark.asyncio 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()