Skip to content

Commit 4edbbb9

Browse files
samuelcolvinsamuelcolvin
authored
Reuse pyodide env - much faster repeat execution (#4)
1 parent b50c3f9 commit 4edbbb9

File tree

4 files changed

+150
-89
lines changed

4 files changed

+150
-89
lines changed

‎examples/sandbox.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import asyncio
2+
import time
3+
14
from mcp_run_python import code_sandbox
25

36

@@ -6,9 +9,10 @@ def log_handler(level: str, message: str):
69

710

811
code = """
9-
import numpy
12+
import numpy, asyncio
1013
a = numpy.array([1, 2, 3])
1114
print(a)
15+
await asyncio.sleep(1)
1216
a
1317
"""
1418

@@ -23,8 +27,11 @@ async def main():
2327
else:
2428
print(result['error'])
2529

30+
tic = time.time()
31+
result = await asyncio.gather(*[sandbox.eval(code) for _ in range(10)])
32+
toc = time.time()
33+
print(f'Execution time: {toc - tic:.3f} seconds')
2634

27-
if __name__ == '__main__':
28-
import asyncio
2935

36+
if __name__ == '__main__':
3037
asyncio.run(main())

‎mcp_run_python/deno/src/main.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import { type LoggingLevel, SetLevelRequestSchema } from '@modelcontextprotocol/
1111
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
1212
import { z } from 'zod'
1313

14-
import { asJson, asXml, runCode } from './runCode.ts'
14+
import { asJson, asXml, RunCode } from './runCode.ts'
1515
import { Buffer } from 'node:buffer'
1616

1717
const VERSION = '0.0.13'
@@ -57,6 +57,7 @@ options:
5757
* Create an MCP server with the `run_python_code` tool registered.
5858
*/
5959
function createServer(deps: string[], returnMode: string): McpServer {
60+
const runCode = new RunCode()
6061
const server = new McpServer(
6162
{
6263
name: 'MCP Run Python',
@@ -90,7 +91,7 @@ The code will be executed with Python 3.12.
9091
{ python_code: z.string().describe('Python code to run') },
9192
async ({ python_code }: { python_code: string }) => {
9293
const logPromises: Promise<void>[] = []
93-
const result = await runCode(
94+
const result = await runCode.run(
9495
deps,
9596
{ name: 'main.py', content: python_code },
9697
(level, data) => {
@@ -251,7 +252,8 @@ async function runStdio(deps: string[], returnMode: string) {
251252
* Run pyodide to download and install dependencies.
252253
*/
253254
async function installDeps(deps: string[]) {
254-
const result = await runCode(
255+
const runCode = new RunCode()
256+
const result = await runCode.run(
255257
deps,
256258
undefined,
257259
(level, data) => console.error(`${level}|${data}`),
@@ -275,7 +277,8 @@ a = numpy.array([1, 2, 3])
275277
print('numpy array:', a)
276278
a
277279
`
278-
const result = await runCode(
280+
const runCode = new RunCode()
281+
const result = await runCode.run(
279282
deps,
280283
{ name: 'example.py', content: code },
281284
// use warn to avoid recursion since console.log is patched in runCode

‎mcp_run_python/deno/src/runCode.ts‎

Lines changed: 127 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
// deno-lint-ignore-file no-explicit-any
2-
import { loadPyodide } from 'pyodide'
2+
import { loadPyodide, type PyodideInterface } from 'pyodide'
33
import { preparePythonCode } from './prepareEnvCode.ts'
44
import type { LoggingLevel } from '@modelcontextprotocol/sdk/types.js'
55

@@ -8,90 +8,141 @@ export interface CodeFile {
88
content: string
99
}
1010

11-
export async function runCode(
12-
dependencies: string[],
13-
file: CodeFile | undefined,
14-
log: (level: LoggingLevel, data: string) => void,
15-
): Promise<RunSuccess | RunError> {
16-
// remove once we can upgrade to pyodide 0.27.7 and console.log is no longer used.
17-
const realConsoleLog = console.log
18-
console.log = (...args: any[]) => log('debug', args.join(' '))
19-
20-
const output: string[] = []
21-
const pyodide = await loadPyodide({
22-
stdout: (msg) => {
23-
log('info', msg)
24-
output.push(msg)
25-
},
26-
stderr: (msg) => {
27-
log('warning', msg)
28-
output.push(msg)
29-
},
30-
})
31-
32-
// see https://github.com/pyodide/pyodide/discussions/5512
33-
const origLoadPackage = pyodide.loadPackage
34-
pyodide.loadPackage = (pkgs, options) =>
35-
origLoadPackage(pkgs, {
36-
// stop pyodide printing to stdout/stderr
37-
messageCallback: (msg: string) => log('debug', `loadPackage: ${msg}`),
38-
errorCallback: (msg: string) => {
39-
log('error', `loadPackage: ${msg}`)
40-
output.push(`install error: ${msg}`)
41-
},
42-
...options,
43-
})
44-
45-
await pyodide.loadPackage(['micropip', 'pydantic'])
46-
const sys = pyodide.pyimport('sys')
47-
48-
const dirPath = '/tmp/mcp_run_python'
49-
sys.path.append(dirPath)
50-
const pathlib = pyodide.pyimport('pathlib')
51-
pathlib.Path(dirPath).mkdir()
52-
const moduleName = '_prepare_env'
53-
54-
pathlib.Path(`${dirPath}/${moduleName}.py`).write_text(preparePythonCode)
55-
56-
const preparePyEnv: PreparePyEnv = pyodide.pyimport(moduleName)
11+
interface PrepResult {
12+
pyodide: PyodideInterface
13+
preparePyEnv: PreparePyEnv
14+
sys: any
15+
prepareStatus: PrepareSuccess | PrepareError
16+
}
5717

58-
const prepareStatus = await preparePyEnv.prepare_env(pyodide.toPy(dependencies))
59-
let runResult: RunSuccess | RunError
60-
if (prepareStatus.kind == 'error') {
61-
runResult = {
62-
status: 'install-error',
63-
output,
64-
error: prepareStatus.message,
18+
export class RunCode {
19+
private output: string[] = []
20+
private pyodide?: PyodideInterface
21+
private preparePyEnv?: PreparePyEnv
22+
private prepPromise?: Promise<PrepResult>
23+
24+
async run(
25+
dependencies: string[],
26+
file: CodeFile | undefined,
27+
log: (level: LoggingLevel, data: string) => void,
28+
): Promise<RunSuccess | RunError> {
29+
// remove once we can upgrade to pyodide 0.27.7 and console.log is no longer used.
30+
const realConsoleLog = console.log
31+
console.log = (...args: any[]) => log('debug', args.join(' '))
32+
33+
let pyodide: PyodideInterface
34+
let sys: any
35+
let prepareStatus: PrepareSuccess | PrepareError | undefined
36+
let preparePyEnv: PreparePyEnv
37+
if (this.pyodide && this.preparePyEnv) {
38+
pyodide = this.pyodide
39+
preparePyEnv = this.preparePyEnv
40+
sys = pyodide.pyimport('sys')
41+
} else {
42+
if (!this.prepPromise) {
43+
this.prepPromise = this.prepEnv(dependencies, log)
44+
}
45+
// TODO is this safe if the promise has already been accessed? it seems to work fine
46+
const prep = await this.prepPromise
47+
pyodide = prep.pyodide
48+
preparePyEnv = prep.preparePyEnv
49+
sys = prep.sys
50+
prepareStatus = prep.prepareStatus
6551
}
66-
} else if (file) {
67-
try {
68-
const rawValue = await pyodide.runPythonAsync(file.content, {
69-
globals: pyodide.toPy({ __name__: '__main__' }),
70-
filename: file.name,
71-
})
52+
53+
let runResult: RunSuccess | RunError
54+
if (prepareStatus && prepareStatus.kind == 'error') {
7255
runResult = {
73-
status: 'success',
74-
output,
75-
returnValueJson: preparePyEnv.dump_json(rawValue),
56+
status: 'install-error',
57+
output: this.takeOutput(sys),
58+
error: prepareStatus.message,
7659
}
77-
} catch (err) {
60+
} else if (file) {
61+
try {
62+
const rawValue = await pyodide.runPythonAsync(file.content, {
63+
globals: pyodide.toPy({ __name__: '__main__' }),
64+
filename: file.name,
65+
})
66+
runResult = {
67+
status: 'success',
68+
output: this.takeOutput(sys),
69+
returnValueJson: preparePyEnv.dump_json(rawValue),
70+
}
71+
} catch (err) {
72+
runResult = {
73+
status: 'run-error',
74+
output: this.takeOutput(sys),
75+
error: formatError(err),
76+
}
77+
}
78+
} else {
7879
runResult = {
79-
status: 'run-error',
80-
output,
81-
error: formatError(err),
80+
status: 'success',
81+
output: this.takeOutput(sys),
82+
returnValueJson: null,
8283
}
8384
}
84-
} else {
85-
runResult = {
86-
status: 'success',
87-
output,
88-
returnValueJson: null,
85+
console.log = realConsoleLog
86+
return runResult
87+
}
88+
89+
async prepEnv(
90+
dependencies: string[],
91+
log: (level: LoggingLevel, data: string) => void,
92+
): Promise<PrepResult> {
93+
const pyodide = await loadPyodide({
94+
stdout: (msg) => {
95+
log('info', msg)
96+
this.output.push(msg)
97+
},
98+
stderr: (msg) => {
99+
log('warning', msg)
100+
this.output.push(msg)
101+
},
102+
})
103+
104+
// see https://github.com/pyodide/pyodide/discussions/5512
105+
const origLoadPackage = pyodide.loadPackage
106+
pyodide.loadPackage = (pkgs, options) =>
107+
origLoadPackage(pkgs, {
108+
// stop pyodide printing to stdout/stderr
109+
messageCallback: (msg: string) => log('debug', `loadPackage: ${msg}`),
110+
errorCallback: (msg: string) => {
111+
log('error', `loadPackage: ${msg}`)
112+
this.output.push(`install error: ${msg}`)
113+
},
114+
...options,
115+
})
116+
117+
await pyodide.loadPackage(['micropip', 'pydantic'])
118+
const sys = pyodide.pyimport('sys')
119+
120+
const dirPath = '/tmp/mcp_run_python'
121+
sys.path.append(dirPath)
122+
const pathlib = pyodide.pyimport('pathlib')
123+
pathlib.Path(dirPath).mkdir()
124+
const moduleName = '_prepare_env'
125+
126+
pathlib.Path(`${dirPath}/${moduleName}.py`).write_text(preparePythonCode)
127+
128+
const preparePyEnv: PreparePyEnv = pyodide.pyimport(moduleName)
129+
130+
const prepareStatus = await preparePyEnv.prepare_env(pyodide.toPy(dependencies))
131+
return {
132+
pyodide,
133+
preparePyEnv,
134+
sys,
135+
prepareStatus,
89136
}
90137
}
91-
sys.stdout.flush()
92-
sys.stderr.flush()
93-
console.log = realConsoleLog
94-
return runResult
138+
139+
private takeOutput(sys: any): string[] {
140+
sys.stdout.flush()
141+
sys.stderr.flush()
142+
const output = this.output
143+
this.output = []
144+
return output
145+
}
95146
}
96147

97148
interface RunSuccess {

‎tests/test_sandbox.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -57,12 +57,12 @@ async def test_sandbox(deps: list[str], code: str, expected: Any):
5757

5858
async def test_multiple_commands():
5959
async with code_sandbox() as sandbox:
60-
result = await sandbox.eval('print(1)')
61-
assert result == snapshot({'status': 'success', 'output': ['1'], 'return_value': None})
62-
result = await sandbox.eval('print(2)')
63-
assert result == snapshot({'status': 'success', 'output': ['2'], 'return_value': None})
64-
result = await sandbox.eval('print(3)')
65-
assert result == snapshot({'status': 'success', 'output': ['3'], 'return_value': None})
60+
result = await sandbox.eval('print(1)\n1')
61+
assert result == snapshot({'status': 'success', 'output': ['1'], 'return_value': 1})
62+
result = await sandbox.eval('print(2)\n2')
63+
assert result == snapshot({'status': 'success', 'output': ['2'], 'return_value': 2})
64+
result = await sandbox.eval('print(3)\n3')
65+
assert result == snapshot({'status': 'success', 'output': ['3'], 'return_value': 3})
6666

6767

6868
async def test_multiple_sandboxes():

0 commit comments

Comments
 (0)