gradio

1
/// <reference lib="webworker" />
2
/* eslint-env worker */
3

4
import type { PyodideInterface } from "pyodide";
5
import type { PyProxy } from "pyodide/ffi";
6
import type {
7
	InMessage,
8
	InMessageInitEnv,
9
	InMessageInitApp,
10
	OutMessage,
11
	ReplyMessageError,
12
	ReplyMessageSuccess
13
} from "../message-types";
14
import {
15
	writeFileWithParents,
16
	renameWithParents,
17
	getAppHomeDir,
18
	resolveAppHomeBasedPath
19
} from "./file";
20
import { verifyRequirements } from "./requirements";
21
import { makeAsgiRequest } from "./asgi";
22
import { generateRandomString } from "./random";
23
import scriptRunnerPySource from "./py/script_runner.py?raw";
24
import unloadModulesPySource from "./py/unload_modules.py?raw";
25

26
importScripts("https://cdn.jsdelivr.net/pyodide/v0.25.0/full/pyodide.js");
27

28
type MessageTransceiver = DedicatedWorkerGlobalScope | MessagePort;
29

30
let pyodide: PyodideInterface;
31
let micropip: PyProxy;
32

33
let call_asgi_app_from_js: (
34
	appId: string,
35
	scope: unknown,
36
	receive: () => Promise<unknown>,
37
	send: (event: any) => Promise<void>
38
) => Promise<void>;
39
let run_code: (
40
	appId: string,
41
	home_dir: string,
42
	code: string,
43
	path?: string
44
) => Promise<void>;
45
let run_script: (
46
	appId: string,
47
	home_dir: string,
48
	path: string
49
) => Promise<void>;
50
let unload_local_modules: (target_dir_path?: string) => void;
51

52
async function initializeEnvironment(
53
	options: InMessageInitEnv["data"],
54
	updateProgress: (log: string) => void
55
): Promise<void> {
56
	console.debug("Loading Pyodide.");
57
	updateProgress("Loading Pyodide");
58
	pyodide = await loadPyodide({
59
		stdout: console.debug,
60
		stderr: console.error
61
	});
62
	console.debug("Pyodide is loaded.");
63

64
	console.debug("Loading micropip");
65
	updateProgress("Loading micropip");
66
	await pyodide.loadPackage("micropip");
67
	micropip = pyodide.pyimport("micropip");
68
	console.debug("micropip is loaded.");
69

70
	const gradioWheelUrls = [
71
		options.gradioWheelUrl,
72
		options.gradioClientWheelUrl
73
	];
74
	console.debug("Loading Gradio wheels.", gradioWheelUrls);
75
	updateProgress("Loading Gradio wheels");
76
	await micropip.add_mock_package("ffmpy", "0.3.0");
77
	await micropip.add_mock_package("aiohttp", "3.8.4");
78
	await pyodide.loadPackage(["ssl", "setuptools"]);
79
	await micropip.install(["typing-extensions>=4.8.0"]); // Typing extensions needs to be installed first otherwise the versions from the pyodide lockfile is used which is incompatible with the latest fastapi.
80
	await micropip.install(["markdown-it-py[linkify]~=2.2.0"]); // On 3rd June 2023, markdown-it-py 3.0.0 has been released. The `gradio` package depends on its `>=2.0.0` version so its 3.x will be resolved. However, it conflicts with `mdit-py-plugins`'s dependency `markdown-it-py >=1.0.0,<3.0.0` and micropip currently can't resolve it. So we explicitly install the compatible version of the library here.
81
	await micropip.install(["anyio==3.*"]); // `fastapi` depends on `anyio>=3.4.0,<5` so its 4.* can be installed, but it conflicts with the anyio version `httpx` depends on, `==3.*`. Seems like micropip can't resolve it for now, so we explicitly install the compatible version of the library here.
82
	await micropip.add_mock_package("pydantic", "2.4.2"); // PydanticV2 is not supported on Pyodide yet. Mock it here for installing the `gradio` package to pass the version check. Then, install PydanticV1 below.
83
	await micropip.add_mock_package("ruff", "0.2.2"); // `ruff` was added to the requirements of `gradio` for the custom components (https://github.com/gradio-app/gradio/pull/7030), but it's not working on PYodide yet. Also Lite doesn't need it, so mock it here for installing the `gradio` package to pass the version check.
84
	await micropip.install.callKwargs(gradioWheelUrls, {
85
		keep_going: true
86
	});
87
	await micropip.remove_mock_package("pydantic");
88
	await micropip.install(["pydantic==1.*"]); // Pydantic is necessary for `gradio` to run, so install v1 here as a fallback. Some tricks has been introduced in `gradio/data_classes.py` to make it work with v1.
89
	console.debug("Gradio wheels are loaded.");
90

91
	console.debug("Mocking os module methods.");
92
	updateProgress("Mock os module methods");
93
	// `os.link` is used in `aiofiles` (https://github.com/Tinche/aiofiles/blob/v23.1.0/src/aiofiles/os.py#L31),
94
	// which is imported from `gradio.ranged_response` (https://github.com/gradio-app/gradio/blob/v3.32.0/gradio/ranged_response.py#L12).
95
	// However, it's not available on Wasm.
96
	await pyodide.runPythonAsync(`
97
import os
98

99
os.link = lambda src, dst: None
100
`);
101
	console.debug("os module methods are mocked.");
102

103
	console.debug("Importing gradio package.");
104
	updateProgress("Importing gradio package");
105
	// Importing the gradio package takes a long time, so we do it separately.
106
	// This is necessary for accurate performance profiling.
107
	await pyodide.runPythonAsync(`import gradio`);
108
	console.debug("gradio package is imported.");
109

110
	console.debug("Defining a ASGI wrapper function.");
111
	updateProgress("Defining a ASGI wrapper function");
112
	// TODO: Unlike Streamlit, user's code is executed in the global scope,
113
	//       so we should not define this function in the global scope.
114
	await pyodide.runPythonAsync(`
115
# Based on Shiny's App.call_pyodide().
116
# https://github.com/rstudio/py-shiny/blob/v0.3.3/shiny/_app.py#L224-L258
117
async def _call_asgi_app_from_js(app_id, scope, receive, send):
118
	# TODO: Pretty sure there are objects that need to be destroy()'d here?
119
	scope = scope.to_py()
120

121
	# ASGI requires some values to be byte strings, not character strings. Those are
122
	# not that easy to create in JavaScript, so we let the JS side pass us strings
123
	# and we convert them to bytes here.
124
	if "headers" in scope:
125
			# JS doesn't have \`bytes\` so we pass as strings and convert here
126
			scope["headers"] = [
127
					[value.encode("latin-1") for value in header]
128
					for header in scope["headers"]
129
			]
130
	if "query_string" in scope and scope["query_string"]:
131
			scope["query_string"] = scope["query_string"].encode("latin-1")
132
	if "raw_path" in scope and scope["raw_path"]:
133
			scope["raw_path"] = scope["raw_path"].encode("latin-1")
134

135
	async def rcv():
136
			event = await receive()
137
			return event.to_py()
138

139
	async def snd(event):
140
			await send(event)
141

142
	app = gradio.wasm_utils.get_registered_app(app_id)
143
	if app is None:
144
		raise RuntimeError("Gradio app has not been launched.")
145

146
	await app(scope, rcv, snd)
147
`);
148
	call_asgi_app_from_js = pyodide.globals.get("_call_asgi_app_from_js");
149
	console.debug("The ASGI wrapper function is defined.");
150

151
	console.debug("Mocking async libraries.");
152
	updateProgress("Mocking async libraries");
153
	// FastAPI uses `anyio.to_thread.run_sync` internally which, however, doesn't work in Wasm environments where the `threading` module is not supported.
154
	// So we mock `anyio.to_thread.run_sync` here not to use threads.
155
	await pyodide.runPythonAsync(`
156
async def mocked_anyio_to_thread_run_sync(func, *args, cancellable=False, limiter=None):
157
	return func(*args)
158

159
import anyio.to_thread
160
anyio.to_thread.run_sync = mocked_anyio_to_thread_run_sync
161
	`);
162
	console.debug("Async libraries are mocked.");
163

164
	console.debug("Setting matplotlib backend.");
165
	updateProgress("Setting matplotlib backend");
166
	// Ref: https://github.com/streamlit/streamlit/blob/1.22.0/lib/streamlit/web/bootstrap.py#L111
167
	// This backend setting is required to use matplotlib in Wasm environment.
168
	await pyodide.runPythonAsync(`
169
import matplotlib
170
matplotlib.use("agg")
171
`);
172
	console.debug("matplotlib backend is set.");
173

174
	console.debug("Setting up Python utility functions.");
175
	updateProgress("Setting up Python utility functions");
176
	await pyodide.runPythonAsync(scriptRunnerPySource);
177
	run_code = pyodide.globals.get("_run_code");
178
	run_script = pyodide.globals.get("_run_script");
179
	await pyodide.runPythonAsync(unloadModulesPySource);
180
	unload_local_modules = pyodide.globals.get("unload_local_modules");
181
	console.debug("Python utility functions are set up.");
182

183
	updateProgress("Initialization completed");
184
}
185

186
async function initializeApp(
187
	appId: string,
188
	options: InMessageInitApp["data"],
189
	updateProgress: (log: string) => void
190
): Promise<void> {
191
	const appHomeDir = getAppHomeDir(appId);
192
	console.debug("Creating a home directory for the app.", {
193
		appId,
194
		appHomeDir
195
	});
196
	pyodide.FS.mkdir(appHomeDir);
197

198
	console.debug("Mounting files.", options.files);
199
	updateProgress("Mounting files");
200
	await Promise.all(
201
		Object.keys(options.files).map(async (path) => {
202
			const file = options.files[path];
203

204
			let data: string | ArrayBufferView;
205
			if ("url" in file) {
206
				console.debug(`Fetch a file from ${file.url}`);
207
				data = await fetch(file.url)
208
					.then((res) => res.arrayBuffer())
209
					.then((buffer) => new Uint8Array(buffer));
210
			} else {
211
				data = file.data;
212
			}
213
			const { opts } = options.files[path];
214

215
			const appifiedPath = resolveAppHomeBasedPath(appId, path);
216
			console.debug(`Write a file "${appifiedPath}"`);
217
			writeFileWithParents(pyodide, appifiedPath, data, opts);
218
		})
219
	);
220
	console.debug("Files are mounted.");
221

222
	console.debug("Installing packages.", options.requirements);
223
	updateProgress("Installing packages");
224
	await micropip.install.callKwargs(options.requirements, { keep_going: true });
225
	console.debug("Packages are installed.");
226
}
227

228
const ctx = self as DedicatedWorkerGlobalScope | SharedWorkerGlobalScope;
229

230
/**
231
 * Set up the onmessage event listener.
232
 */
233
if ("postMessage" in ctx) {
234
	// Dedicated worker
235
	setupMessageHandler(ctx);
236
} else {
237
	// Shared worker
238
	ctx.onconnect = (event: MessageEvent): void => {
239
		const port = event.ports[0];
240

241
		setupMessageHandler(port);
242

243
		port.start();
244
	};
245
}
246

247
// Environment initialization is global and should be done only once, so its promise is managed in a global scope.
248
let envReadyPromise: Promise<void> | undefined = undefined;
249

250
function setupMessageHandler(receiver: MessageTransceiver): void {
251
	// A concept of "app" is introduced to support multiple apps in a single worker.
252
	// Each app has its own home directory (`getAppHomeDir(appId)`) in a shared single Pyodide filesystem.
253
	// The home directory is used as the current working directory for the app.
254
	// Each frontend app has a connection to the worker which is the `receiver` object passed above
255
	// and it is associated with one app.
256
	// One app also has one Gradio server app which is managed by the `gradio.wasm_utils` module.`
257
	// This multi-app mechanism was introduced for a SharedWorker, but the same mechanism is used for a DedicatedWorker as well.
258
	const appId = generateRandomString(8);
259

260
	console.debug("Set up a new app.", { appId });
261

262
	const updateProgress = (log: string): void => {
263
		const message: OutMessage = {
264
			type: "progress-update",
265
			data: {
266
				log
267
			}
268
		};
269
		receiver.postMessage(message);
270
	};
271

272
	// App initialization is per app or receiver, so its promise is managed in this scope.
273
	let appReadyPromise: Promise<void> | undefined = undefined;
274

275
	receiver.onmessage = async function (
276
		event: MessageEvent<InMessage>
277
	): Promise<void> {
278
		const msg = event.data;
279
		console.debug("worker.onmessage", msg);
280

281
		const messagePort = event.ports[0];
282

283
		try {
284
			if (msg.type === "init-env") {
285
				if (envReadyPromise == null) {
286
					envReadyPromise = initializeEnvironment(msg.data, updateProgress);
287
				} else {
288
					updateProgress(
289
						"Pyodide environment initialization is ongoing in another session"
290
					);
291
				}
292

293
				envReadyPromise
294
					.then(() => {
295
						const replyMessage: ReplyMessageSuccess = {
296
							type: "reply:success",
297
							data: null
298
						};
299
						messagePort.postMessage(replyMessage);
300
					})
301
					.catch((error) => {
302
						const replyMessage: ReplyMessageError = {
303
							type: "reply:error",
304
							error
305
						};
306
						messagePort.postMessage(replyMessage);
307
					});
308
				return;
309
			}
310

311
			if (envReadyPromise == null) {
312
				throw new Error("Pyodide Initialization is not started.");
313
			}
314
			await envReadyPromise;
315

316
			if (msg.type === "init-app") {
317
				appReadyPromise = initializeApp(appId, msg.data, updateProgress);
318

319
				const replyMessage: ReplyMessageSuccess = {
320
					type: "reply:success",
321
					data: null
322
				};
323
				messagePort.postMessage(replyMessage);
324
				return;
325
			}
326

327
			if (appReadyPromise == null) {
328
				throw new Error("App initialization is not started.");
329
			}
330
			await appReadyPromise;
331

332
			switch (msg.type) {
333
				case "echo": {
334
					const replyMessage: ReplyMessageSuccess = {
335
						type: "reply:success",
336
						data: msg.data
337
					};
338
					messagePort.postMessage(replyMessage);
339
					break;
340
				}
341
				case "run-python-code": {
342
					unload_local_modules();
343

344
					await run_code(appId, getAppHomeDir(appId), msg.data.code);
345

346
					const replyMessage: ReplyMessageSuccess = {
347
						type: "reply:success",
348
						data: null // We don't send back the execution result because it's not needed for our purpose, and sometimes the result is of type `pyodide.ffi.PyProxy` which cannot be cloned across threads and causes an error.
349
					};
350
					messagePort.postMessage(replyMessage);
351
					break;
352
				}
353
				case "run-python-file": {
354
					unload_local_modules();
355

356
					await run_script(appId, getAppHomeDir(appId), msg.data.path);
357

358
					const replyMessage: ReplyMessageSuccess = {
359
						type: "reply:success",
360
						data: null
361
					};
362
					messagePort.postMessage(replyMessage);
363
					break;
364
				}
365
				case "asgi-request": {
366
					console.debug("ASGI request", msg.data);
367
					makeAsgiRequest(
368
						call_asgi_app_from_js.bind(null, appId),
369
						msg.data.scope,
370
						messagePort
371
					); // This promise is not awaited because it won't resolves until the HTTP connection is closed.
372
					break;
373
				}
374
				case "file:write": {
375
					const { path, data: fileData, opts } = msg.data;
376

377
					const appifiedPath = resolveAppHomeBasedPath(appId, path);
378

379
					console.debug(`Write a file "${appifiedPath}"`);
380
					writeFileWithParents(pyodide, appifiedPath, fileData, opts);
381

382
					const replyMessage: ReplyMessageSuccess = {
383
						type: "reply:success",
384
						data: null
385
					};
386
					messagePort.postMessage(replyMessage);
387
					break;
388
				}
389
				case "file:rename": {
390
					const { oldPath, newPath } = msg.data;
391

392
					const appifiedOldPath = resolveAppHomeBasedPath(appId, oldPath);
393
					const appifiedNewPath = resolveAppHomeBasedPath(appId, newPath);
394
					console.debug(`Rename "${appifiedOldPath}" to ${appifiedNewPath}`);
395
					renameWithParents(pyodide, appifiedOldPath, appifiedNewPath);
396

397
					const replyMessage: ReplyMessageSuccess = {
398
						type: "reply:success",
399
						data: null
400
					};
401
					messagePort.postMessage(replyMessage);
402
					break;
403
				}
404
				case "file:unlink": {
405
					const { path } = msg.data;
406

407
					const appifiedPath = resolveAppHomeBasedPath(appId, path);
408

409
					console.debug(`Remove "${appifiedPath}`);
410
					pyodide.FS.unlink(appifiedPath);
411

412
					const replyMessage: ReplyMessageSuccess = {
413
						type: "reply:success",
414
						data: null
415
					};
416
					messagePort.postMessage(replyMessage);
417
					break;
418
				}
419
				case "install": {
420
					const { requirements } = msg.data;
421

422
					const micropip = pyodide.pyimport("micropip");
423

424
					console.debug("Install the requirements:", requirements);
425
					verifyRequirements(requirements); // Blocks the not allowed wheel URL schemes.
426
					await micropip.install
427
						.callKwargs(requirements, { keep_going: true })
428
						.then(() => {
429
							if (requirements.includes("matplotlib")) {
430
								// Ref: https://github.com/streamlit/streamlit/blob/1.22.0/lib/streamlit/web/bootstrap.py#L111
431
								// This backend setting is required to use matplotlib in Wasm environment.
432
								return pyodide.runPythonAsync(`
433
									import matplotlib
434
									matplotlib.use("agg")
435
								`);
436
							}
437
						})
438
						.then(() => {
439
							console.debug("Successfully installed");
440

441
							const replyMessage: ReplyMessageSuccess = {
442
								type: "reply:success",
443
								data: null
444
							};
445
							messagePort.postMessage(replyMessage);
446
						});
447
					break;
448
				}
449
			}
450
		} catch (error) {
451
			console.error(error);
452

453
			if (!(error instanceof Error)) {
454
				throw error;
455
			}
456

457
			// The `error` object may contain non-serializable properties such as function (for example Pyodide.FS.ErrnoError which has a `.setErrno` function),
458
			// so it must be converted to a plain object before sending it to the main thread.
459
			// Otherwise, the following error will be thrown:
460
			// `Uncaught (in promise) DOMException: Failed to execute 'postMessage' on 'MessagePort': #<Object> could not be cloned.`
461
			// Also, the JSON.stringify() and JSON.parse() approach like https://stackoverflow.com/a/42376465/13103190
462
			// does not work for Error objects because the Error object is not enumerable.
463
			// So we use the following approach to clone the Error object.
464
			const cloneableError = new Error(error.message);
465
			cloneableError.name = error.name;
466
			cloneableError.stack = error.stack;
467

468
			const replyMessage: ReplyMessageError = {
469
				type: "reply:error",
470
				error: cloneableError
471
			};
472
			messagePort.postMessage(replyMessage);
473
		}
474
	};
475
}
476