diff --git a/src/hooks/usePython.ts b/src/hooks/usePython.ts index f1df799d..6725395f 100644 --- a/src/hooks/usePython.ts +++ b/src/hooks/usePython.ts @@ -7,7 +7,7 @@ import { useState } from 'react' import { PythonContext, suppressedMessages } from '../providers/PythonProvider' -import { proxy, Remote, wrap } from 'comlink' +import { Remote } from 'comlink' import useFilesystem from './useFilesystem' import { Packages } from '../types/Packages' @@ -18,25 +18,25 @@ interface UsePythonProps { } export default function usePython(props?: UsePythonProps) { - const { packages = {} } = props ?? {} + const { packages } = props ?? {} - const [isLoading, setIsLoading] = useState(false) - const [pyodideVersion, setPyodideVersion] = useState() + const [runnerId, setRunnerId] = useState() const [isRunning, setIsRunning] = useState(false) const [output, setOutput] = useState([]) const [stdout, setStdout] = useState('') const [stderr, setStderr] = useState('') - const [pendingCode, setPendingCode] = useState() const [hasRun, setHasRun] = useState(false) const { - packages: globalPackages, timeout, lazy, - terminateOnCompletion + terminateOnCompletion, + loading, + getRunner, + run, + terminate } = useContext(PythonContext) - const workerRef = useRef() const runnerRef = useRef>() const { @@ -49,76 +49,13 @@ export default function usePython(props?: UsePythonProps) { watchedModules } = useFilesystem({ runner: runnerRef?.current }) - const createWorker = () => { - const worker = new Worker( - new URL('../workers/python-worker', import.meta.url) - ) - workerRef.current = worker - } - useEffect(() => { - if (!lazy) { - // Spawn worker on mount - createWorker() - } - // Cleanup worker on unmount return () => { cleanup() } }, []) - const allPackages = useMemo(() => { - const official = [ - ...new Set([ - ...(globalPackages.official ?? []), - ...(packages.official ?? []) - ]) - ] - const micropip = [ - ...new Set([ - ...(globalPackages.micropip ?? []), - ...(packages.micropip ?? []) - ]) - ] - return [official, micropip] - }, [globalPackages, packages]) - - const isReady = !isLoading && !!pyodideVersion - - useEffect(() => { - if (workerRef.current && !isReady) { - const init = async () => { - try { - setIsLoading(true) - const runner: Remote = wrap(workerRef.current as Worker) - runnerRef.current = runner - - await runner.init( - proxy((msg: string) => { - // Suppress messages that are not useful for the user - if (suppressedMessages.includes(msg)) { - return - } - setOutput((prev) => [...prev, msg]) - }), - proxy(({ version }) => { - // The runner is ready once the Pyodide version has been set - setPyodideVersion(version) - console.debug('Loaded pyodide version:', version) - }), - allPackages - ) - } catch (error) { - console.error('Error loading Pyodide:', error) - } finally { - setIsLoading(false) - } - } - init() - } - }, [workerRef.current]) - // Immediately set stdout upon receiving new input useEffect(() => { if (output.length > 0) { @@ -126,26 +63,33 @@ export default function usePython(props?: UsePythonProps) { } }, [output]) - // React to ready state and run delayed code if pending - useEffect(() => { - if (pendingCode && isReady) { - const delayedRun = async () => { - await runPython(pendingCode) - setPendingCode(undefined) - } - delayedRun() + const cleanup = () => { + if (runnerId) { + terminate(runnerId) } - }, [pendingCode, isReady]) + setIsRunning(false) + setRunnerId(undefined) + setHasRun(false) + setOutput([]) + setStdout('') + setStderr('') + } // React to run completion and run cleanup if worker should terminate on completion useEffect(() => { if (terminateOnCompletion && hasRun && !isRunning) { cleanup() - setIsRunning(false) - setPyodideVersion(undefined) } }, [terminateOnCompletion, hasRun, isRunning]) + const interruptExecution = () => { + cleanup() + } + + const isReady = !!runnerId + + const isLoading = loading && !isReady + const pythonRunnerCode = ` import sys @@ -172,6 +116,15 @@ def run(code, preamble=''): print() ` + // send a callback to the worker to handle the output + const handleOutput = (msg: string) => { + // Suppress messages that are not useful for the user + if (suppressedMessages.includes(msg)) { + return + } + setOutput((prev) => [...prev, msg]) + } + // prettier-ignore const moduleReloadCode = (modules: Set) => ` import importlib @@ -190,10 +143,17 @@ del sys setStdout('') setStderr('') - if (lazy && !isReady) { - // Spawn worker and set pending code - createWorker() - setPendingCode(code) + let newRunnerId + if (!runnerId || terminateOnCompletion || packages) { + console.log('no runnerId, getting runner') + newRunnerId = await getRunner(handleOutput, packages) + setRunnerId(newRunnerId) + } + + const r = runnerId || newRunnerId + + if (!r) { + console.log('no runnerId, returning') return } @@ -201,18 +161,18 @@ del sys code )}, ${JSON.stringify(preamble)})` - if (!isReady) { - throw new Error('Pyodide is not loaded yet') - } + // if (!isReady) { + // throw new Error('Pyodide is not loaded yet') + // } let timeoutTimer try { setIsRunning(true) setHasRun(true) // Clear output setOutput([]) - if (!isReady || !runnerRef.current) { - throw new Error('Pyodide is not loaded yet') - } + // if (!isReady || !runnerRef.current) { + // throw new Error('Pyodide is not loaded yet') + // } if (timeout > 0) { timeoutTimer = setTimeout(() => { setStdout('') @@ -221,9 +181,9 @@ del sys }, timeout) } if (watchedModules.size > 0) { - await runnerRef.current.run(moduleReloadCode(watchedModules)) + await run(r, moduleReloadCode(watchedModules)) } - await runnerRef.current.run(code) + await run(r, code) // eslint-disable-next-line } catch (error: any) { setStderr('Traceback (most recent call last):\n' + error.message) @@ -232,27 +192,9 @@ del sys clearTimeout(timeoutTimer) } }, - [lazy, isReady, timeout, watchedModules] + [runnerId, lazy, timeout, watchedModules] ) - const interruptExecution = () => { - cleanup() - setIsRunning(false) - setPyodideVersion(undefined) - setOutput([]) - - // Spawn new worker - createWorker() - } - - const cleanup = () => { - if (!workerRef.current) { - return - } - console.debug('Terminating worker') - workerRef.current.terminate() - } - return { runPython, stdout, @@ -269,3 +211,19 @@ del sys unwatchModules } } + +// const allPackages = useMemo(() => { +// const official = [ +// ...new Set([ +// ...(globalPackages.official ?? []), +// ...(packages.official ?? []) +// ]) +// ] +// const micropip = [ +// ...new Set([ +// ...(globalPackages.micropip ?? []), +// ...(packages.micropip ?? []) +// ]) +// ] +// return [official, micropip] +// }, [globalPackages, packages]) diff --git a/src/providers/PythonProvider.tsx b/src/providers/PythonProvider.tsx index dff27156..ab237940 100644 --- a/src/providers/PythonProvider.tsx +++ b/src/providers/PythonProvider.tsx @@ -1,15 +1,28 @@ -import { createContext } from 'react' +import { proxy, Remote, wrap } from 'comlink' +import { createContext, useEffect, useRef, useState } from 'react' import { Packages } from '../types/Packages' +import { PythonRunner } from '../types/Runner' const PythonContext = createContext({ packages: {} as Packages, timeout: 0, lazy: false, - terminateOnCompletion: false + terminateOnCompletion: false, + loading: false, + // eslint-disable-next-line @typescript-eslint/no-unused-vars + getRunner: async (msgCallback: (msg: string) => void, packages?: Packages) => + '', + // eslint-disable-next-line @typescript-eslint/no-unused-vars, @typescript-eslint/no-empty-function + run: async (id: string, code: string) => {}, + output: new Map(), + // eslint-disable-next-line @typescript-eslint/no-unused-vars, @typescript-eslint/no-empty-function + terminate: (id: string) => {} }) export const suppressedMessages = ['Python initialization complete'] +const BUFFER = 2 + interface PythonProviderProps { packages?: Packages timeout?: number @@ -27,13 +40,146 @@ function PythonProvider(props: PythonProviderProps) { terminateOnCompletion = false } = props + const workerRef = useRef>(new Map()) + const runnerRef = useRef>>(new Map()) + const assignedRunners = useRef>(new Set()) + const output = useRef>(new Map()) + + const msgCallbacks = useRef void>>(new Map()) + + const [loading, setLoading] = useState(true) + + useEffect(() => { + if (!lazy) { + console.log('Buffer size:', BUFFER) + + for (let i = 0; i < BUFFER; i++) { + createInstance() + } + } + }, []) + + const createInstance = async () => { + try { + const availableRunners = Array.from(runnerRef.current.keys()).filter( + (id) => !assignedRunners.current.has(id) + ) + if (availableRunners.length === 0) { + setLoading(true) + } + + const id = crypto.randomUUID() + + const worker = new Worker( + new URL('../workers/python-worker', import.meta.url) + ) + workerRef.current.set(id, worker) + + const runner: Remote = wrap(worker) + runnerRef.current.set(id, runner) + + await runner.init( + proxy((msg: string) => { + // Suppress messages that are not useful for the user + if (suppressedMessages.includes(msg)) { + return + } + + const callback = msgCallbacks.current.get(id) + if (callback) { + callback(msg) + } else { + console.warn('No callback found for runner', id) + } + }), + proxy(({ version }) => { + console.debug('Loaded pyodide version:', version) + setLoading(false) + }), + [packages.official ?? [], packages.micropip ?? []] + ) + } catch (error) { + console.error('Error loading Pyodide:', error) + } + } + + const getRunner = async (msgCallback: (msg: string) => void) => { + if (lazy) { + await createInstance() + } + + console.log(runnerRef.current.keys(), assignedRunners) + + // get the first available runner + const availableRunners = Array.from(runnerRef.current.keys()).filter( + (id) => !assignedRunners.current.has(id) + ) + + const id = availableRunners[0] + if (!id) { + throw new Error('No runner available') + } + + assignedRunners.current.add(id) + + // register callback + msgCallbacks.current.set(id, msgCallback) + + // if available is less than buffer, create a new instance + if (availableRunners.length - 1 < BUFFER) { + createInstance() + } + + return id + } + + const run = async (id: string, code: string) => { + const runner = runnerRef.current.get(id) + if (!runner) { + throw new Error('Runner not found') + } + await runner.run(code) + + if (terminateOnCompletion) { + console.log('Python terminateOnCompletion', id) + + runnerRef.current.delete(id) + assignedRunners.current.delete(id) + + const worker = workerRef.current.get(id) + if (!worker) { + throw new Error('Worker not found') + } + worker.terminate() + } + } + + const terminate = (id: string) => { + console.log('Python terminate') + + runnerRef.current.delete(id) + assignedRunners.current.delete(id) + + const worker = workerRef.current.get(id) + if (!worker) { + console.error('Worker not found') + return + } + worker.terminate() + } + return ( diff --git a/website/docs/examples/playground.md b/website/docs/examples/playground.md index f11f7113..64157673 100644 --- a/website/docs/examples/playground.md +++ b/website/docs/examples/playground.md @@ -8,13 +8,81 @@ draft: true Used for testing bugs etc. ```python -with open("/hello.txt", "w") as fh: - fh.write("hello world!") - print("done") +print(1) ``` ```python -with open("/hello.txt", "r") as fh: - data = fh.read() -print(data) +print(2) +``` + +```python +print(3) +``` + +```python +print(4) +``` + +```python +print(5) +``` + +```python +print(6) +``` + +```python +print(7) +``` + +```python +print(8) +``` + +```python +print(9) +``` + +```python +print(10) +``` + +```python +print(11) +``` + +```python +print(12) +``` + +```python +print(13) +``` + +```python +print(14) +``` + +```python +print(15) +``` + +```python +print(16) +``` + +```python +print(17) +``` + +```python +print(18) +``` + +```python +print(19) +``` + +```python +print(20) ``` diff --git a/website/docs/introduction/performance.md b/website/docs/introduction/performance.md new file mode 100644 index 00000000..b1661862 --- /dev/null +++ b/website/docs/introduction/performance.md @@ -0,0 +1,11 @@ +--- +sidebar_position: 6 +--- + +# Performance + +Spawning multiple workers can be expensive, so it's important to understand how to use `react-py` efficiently. + +## `PythonProvider` Provider + +Buffer diff --git a/website/src/components/CodeEditor.tsx b/website/src/components/CodeEditor.tsx index 8340654c..6086ccbb 100644 --- a/website/src/components/CodeEditor.tsx +++ b/website/src/components/CodeEditor.tsx @@ -96,7 +96,7 @@ export default function CodeEditor(props: CodeEditorProps) { mode="python" name="CodeBlock" fontSize="0.9rem" - className="min-h-[4rem] overflow-clip rounded shadow-md" + className="min-h-[4rem] overflow-clip rounded border border-solid border-gray-300/25 shadow-md" theme={colorMode === 'dark' ? 'idle_fingers' : 'textmate'} onChange={(newValue) => setInput(newValue)} width="100%" diff --git a/website/src/components/Console.tsx b/website/src/components/Console.tsx index 2ee860d4..a08fe269 100644 --- a/website/src/components/Console.tsx +++ b/website/src/components/Console.tsx @@ -5,6 +5,7 @@ import { ConsoleState } from '@site/../dist/types/Console' import Controls from './Controls' import { ArrowPathIcon, Bars3BottomLeftIcon } from '@heroicons/react/24/solid' +import clsx from 'clsx' const ps1 = '>>> ' const ps2 = '... ' @@ -106,7 +107,10 @@ export default function Console() { {getPrompt()}