diff --git a/Cargo.lock b/Cargo.lock index a2f52d4b..b1c9a276 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1234,6 +1234,8 @@ dependencies = [ "eryx", "reqwest", "rustls", + "serde", + "serde_json", "tokio", "tracing", "tracing-subscriber", diff --git a/book/src/guide/packages.md b/book/src/guide/packages.md index 82844f38..e13e04ca 100644 --- a/book/src/guide/packages.md +++ b/book/src/guide/packages.md @@ -114,6 +114,29 @@ Some packages with native extensions have WASI-compiled versions available: Packages with native C extensions (like `numpy`, `pandas`) require WASI-compiled wheels. Check if WASI builds are available for your specific package. +## Baking Callbacks into the Factory + +Registering callbacks installs a Python wrapper for each of them inside the guest. A fresh sandbox starts from the factory's snapshot, and the snapshot only knows the callbacks it was created with — none, by default — so every `create_sandbox()` re-runs that installation on its first execution (a few milliseconds for a handful of callbacks, more than creating the sandbox itself). + +If every sandbox from a factory registers the same callbacks, pass them to the factory instead. Their declarations (name, description, parameter schema) are baked into the snapshot, and `create_sandbox()` / `create_session()` register them unless given callbacks of their own: + +```python +import eryx + +def get_time(): + import time + return {"timestamp": time.time()} + +factory = eryx.SandboxFactory( + callbacks=[{"name": "get_time", "fn": get_time, "description": "Returns current time"}], +) + +sandbox = factory.create_sandbox() # get_time() is ready; no per-sandbox setup +result = sandbox.execute("print((await get_time())['timestamp'] > 0)") +``` + +A sandbox that is given a different set of callbacks still works; it installs them itself as before. Setup code cannot invoke the callbacks, since nothing answers them while the snapshot is being built. A saved factory keeps the baked declarations but not the Python functions, so pass the same `callbacks=` to `SandboxFactory.load()`. + ## Package Dependencies When installing packages, you must include all dependencies: diff --git a/book/src/guide/precompile.md b/book/src/guide/precompile.md index f4a8a8e2..b502f49f 100644 --- a/book/src/guide/precompile.md +++ b/book/src/guide/precompile.md @@ -148,6 +148,25 @@ eryx-precompile compile runtime.wasm -o jinja2.cwasm \ Each sandbox created from the resulting artifact starts with `env` already defined. Per-request code only needs to handle the request-specific work (deserializing data, compiling the template, rendering). +### Callbacks + +Registering callbacks on a sandbox installs a Python wrapper for each of them inside the guest. A persistent session pays for that once, but a fresh sandbox per request starts from the snapshot, and the snapshot only knows the callbacks it was taken with — none, by default — so every request re-runs the installation (a few milliseconds of Python for a handful of callbacks, more than the instantiation itself). + +If the set of callbacks is fixed per deployment, declare it at pre-initialization and the wrappers are part of the snapshot. A sandbox that registers callbacks with the same names, descriptions and parameter schemas (in any order) then skips the installation; one that registers a different set still works, it just installs as before. + +```bash +# callbacks.json: what the sandboxes will register at runtime +# [ +# {"name": "get_time", "description": "Current time", "parameters": {"type": "object", "properties": {}}}, +# {"name": "fetch", "description": "HTTP GET", "parameters": {"type": "object", "properties": {"url": {"type": "string"}}}} +# ] +eryx-precompile compile runtime.wasm -o runtime.cwasm \ + --preinit --stdlib ./python-stdlib \ + --callbacks callbacks.json +``` + +`parameters` defaults to `{}` when omitted. Setup code cannot invoke the callbacks: there is no host to answer them during pre-initialization. In Rust, pass `PreInitOptions::callbacks` (build the declarations with `eryx::preinit::callback_declaration`); in Python, pass `callbacks=` to `SandboxFactory`. + ### Verification By default, `compile` verifies the output by creating a test sandbox. You can add custom verification: diff --git a/crates/eryx-precompile/Cargo.toml b/crates/eryx-precompile/Cargo.toml index 869dc093..f811ca04 100644 --- a/crates/eryx-precompile/Cargo.toml +++ b/crates/eryx-precompile/Cargo.toml @@ -17,6 +17,8 @@ clap = { workspace = true } #unified eryx = { workspace = true, features = ["preinit"] } reqwest.workspace = true rustls = { workspace = true, features = ["aws-lc-rs"] } +serde.workspace = true # --callbacks file +serde_json.workspace = true tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } tracing.workspace = true tracing-subscriber = { version = "0.3", features = ["env-filter", "registry"] } diff --git a/crates/eryx-precompile/src/main.rs b/crates/eryx-precompile/src/main.rs index fe61ae8c..e6300d34 100644 --- a/crates/eryx-precompile/src/main.rs +++ b/crates/eryx-precompile/src/main.rs @@ -36,7 +36,7 @@ //! --verify-code "import numpy; print(numpy.array(\[1,2,3\]).sum())" //! ``` -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::time::Instant; use anyhow::{Context, Result}; @@ -193,6 +193,19 @@ struct CompileArgs { )] setup_file: Option, + /// JSON file declaring the callbacks sandboxes will register (baked into snapshot) + /// + /// A list of objects with "name", "description" and optional "parameters" + /// (a JSON Schema object, `{}` if omitted). Their Python wrappers are + /// installed during pre-initialization, so a fresh instance whose host + /// registers callbacks with the same names, descriptions and schemas skips + /// the per-execution callback setup. Requires --preinit. + /// + /// Example: `--callbacks callbacks.json` with + /// `[{"name": "get_time", "description": "Current time", "parameters": {"type": "object"}}]` + #[arg(long, value_name = "PATH", requires = "preinit")] + callbacks: Option, + /// Python code to execute during verification /// /// Runs after the standard import verification. Useful for testing that @@ -378,6 +391,40 @@ async fn run_setup(args: SetupArgs) -> Result<()> { Ok(()) } +/// One entry of the `--callbacks` file. +#[derive(serde::Deserialize)] +struct CallbackDeclarationFile { + name: String, + #[serde(default)] + description: String, + /// JSON Schema for the arguments; an empty object when omitted. + #[serde(default = "empty_object")] + parameters: serde_json::Value, +} + +fn empty_object() -> serde_json::Value { + serde_json::Value::Object(serde_json::Map::new()) +} + +/// Read the callback declarations to bake into the snapshot from `path`. +fn read_callback_declarations(path: &Path) -> Result> { + let text = std::fs::read_to_string(path) + .with_context(|| format!("Failed to read callbacks file: {}", path.display()))?; + let entries: Vec = serde_json::from_str(&text) + .with_context(|| format!("Invalid callbacks file: {}", path.display()))?; + entries + .into_iter() + .map(|entry| { + Ok(eryx::preinit::CallbackDeclaration { + name: entry.name, + description: entry.description, + parameters_schema_json: serde_json::to_string(&entry.parameters) + .context("Failed to serialize callback parameters schema")?, + }) + }) + .collect() +} + async fn run_compile(args: CompileArgs) -> Result<()> { // Set up tracing let filter = if args.verbose { @@ -460,6 +507,21 @@ async fn run_compile(args: CompileArgs) -> Result<()> { .map_or("inline code".to_string(), |p| p.display().to_string()) ); } + + let callbacks = match &args.callbacks { + Some(path) => read_callback_declarations(path)?, + None => Vec::new(), + }; + if !callbacks.is_empty() { + println!( + "Callbacks: {} (baked into snapshot)", + callbacks + .iter() + .map(|cb| cb.name.as_str()) + .collect::>() + .join(", ") + ); + } println!(); // Read input WASM @@ -497,22 +559,23 @@ async fn run_compile(args: CompileArgs) -> Result<()> { println!("Native extensions: {}", extensions.len()); } - // Convert imports to &str references - let import_refs: Vec<&str> = args.imports.iter().map(|s| s.as_str()).collect(); - println!(); println!("Step 1: Pre-initializing Python..."); let start = Instant::now(); - let preinit_bytes = eryx::preinit::pre_initialize( - stdlib, - final_site_packages.as_deref(), - &import_refs, - &extensions, - setup_code.as_deref(), - ) - .await - .context("Failed to pre-initialize Python")?; + let mut options = eryx::preinit::PreInitOptions::new(stdlib) + .imports(args.imports.iter().cloned()) + .extensions(extensions) + .callbacks(callbacks); + if let Some(path) = &final_site_packages { + options = options.site_packages(path); + } + if let Some(code) = &setup_code { + options = options.setup_code(code.clone()); + } + let preinit_bytes = eryx::preinit::pre_initialize_with_options(options) + .await + .context("Failed to pre-initialize Python")?; let elapsed = start.elapsed(); println!( diff --git a/crates/eryx-python/python/eryx/_eryx.pyi b/crates/eryx-python/python/eryx/_eryx.pyi index a1cbed07..d65032f2 100644 --- a/crates/eryx-python/python/eryx/_eryx.pyi +++ b/crates/eryx-python/python/eryx/_eryx.pyi @@ -646,6 +646,7 @@ class SandboxFactory: packages: Optional[Sequence[PathLike]] = None, imports: Optional[Sequence[str]] = None, setup_code: Optional[str] = None, + callbacks: Optional[Union[CallbackRegistry, Sequence[CallbackDict]]] = None, cache: bool = True, ) -> None: """Create a new sandbox factory with custom packages. @@ -664,6 +665,13 @@ class SandboxFactory: ``SandboxedEnvironment``) so every sandbox starts with them already in memory. Each sandbox gets its own copy-on-write clone, preserving full isolation. + callbacks: Optional callbacks (a ``CallbackRegistry`` or a list of callback + dicts) whose declarations are baked into the snapshot, so sandboxes + created from this factory skip the per-sandbox callback setup (a few + milliseconds per sandbox). ``create_sandbox()`` and ``create_session()`` + register these callbacks unless given their own; a sandbox given a + different set still works, it just installs them itself. Setup code + cannot invoke the callbacks. cache: Whether to cache the pre-compiled component in the process-global cache. When enabled, a BLAKE3 content hash is computed once during factory construction and subsequent ``create_sandbox()`` calls skip @@ -700,6 +708,7 @@ class SandboxFactory: path: PathLike, *, site_packages: Optional[PathLike] = None, + callbacks: Optional[Union[CallbackRegistry, Sequence[CallbackDict]]] = None, cache: bool = True, ) -> SandboxFactory: """Load a sandbox factory from a file. @@ -711,6 +720,9 @@ class SandboxFactory: path: Path to the saved factory file. site_packages: Optional path to site-packages directory. Required if the factory was saved without embedded packages. + callbacks: The callbacks the factory was created with, if any. The file + holds their baked declarations but not the Python callables, so pass + the same callbacks here to keep the setup-free fast path. cache: Whether to cache the pre-compiled component in the process-global cache. When enabled, a BLAKE3 content hash is computed once during loading and subsequent ``create_sandbox()`` calls skip component diff --git a/crates/eryx-python/src/preinit.rs b/crates/eryx-python/src/preinit.rs index 6e736363..d597441d 100644 --- a/crates/eryx-python/src/preinit.rs +++ b/crates/eryx-python/src/preinit.rs @@ -52,6 +52,10 @@ pub struct SandboxFactory { /// Extracted packages (kept alive to prevent temp dir cleanup). #[allow(dead_code)] extracted_packages: Arc>, + /// Callbacks whose declarations were baked into the snapshot (or, for a + /// loaded factory, the ones the caller says were). `create_sandbox()` and + /// `create_session()` register them when given no callbacks of their own. + callbacks: Option>, } /// Construct a pre-compiled artifact with optional content-safe caching. @@ -76,6 +80,11 @@ impl SandboxFactory { /// These are extracted and their native extensions are linked. /// imports: Optional list of module names to pre-import during initialization. /// Pre-imported modules are immediately available without import overhead. + /// callbacks: Optional callbacks (a CallbackRegistry or a list of callback + /// dicts) whose declarations are baked into the snapshot, so sandboxes + /// created from this factory skip the per-sandbox callback setup. + /// `create_sandbox()` and `create_session()` register these callbacks + /// unless given their own. Setup code cannot invoke them. /// cache: Whether to cache the pre-compiled component in the process-global /// cache. When enabled, a BLAKE3 content hash is computed once during /// factory construction and subsequent ``create_sandbox()`` calls skip @@ -97,13 +106,21 @@ impl SandboxFactory { /// ], /// imports=["jinja2"], /// ) + /// + /// # With the callbacks every sandbox will register + /// factory = SandboxFactory(callbacks=[ + /// {"name": "get_time", "fn": get_time, "description": "Returns current time"} + /// ]) + /// sandbox = factory.create_sandbox() # get_time() available, no setup cost #[new] - #[pyo3(signature = (*, site_packages=None, packages=None, imports=None, setup_code=None, cache=true))] + #[pyo3(signature = (*, site_packages=None, packages=None, imports=None, setup_code=None, callbacks=None, cache=true))] fn new( + py: Python<'_>, site_packages: Option, packages: Option>, imports: Option>, setup_code: Option, + callbacks: Option>, cache: bool, ) -> PyResult { // Create tokio runtime for async pre-initialization @@ -120,23 +137,34 @@ impl SandboxFactory { let (final_site_packages, extensions, extracted_packages) = process_packages(site_packages.as_ref(), packages.as_ref())?; - // Convert imports to the format pre_initialize expects - let import_refs: Vec<&str> = imports - .as_ref() - .map(|v| v.iter().map(|s| s.as_str()).collect()) - .unwrap_or_default(); + // Bake the callbacks' declarations so sandboxes registering the same + // set skip the per-sandbox wrapper installation. + let declarations = match &callbacks { + Some(cbs) => extract_callbacks(py, cbs)? + .iter() + .map(|cb| eryx::preinit::callback_declaration(cb)) + .collect(), + None => Vec::new(), + }; + + let mut options = eryx::preinit::PreInitOptions::new(&stdlib_path) + .imports(imports.unwrap_or_default()) + .extensions(extensions) + .callbacks(declarations); + if let Some(path) = &final_site_packages { + options = options.site_packages(path); + } + if let Some(code) = setup_code { + options = options.setup_code(code); + } // Run pre-initialization let preinit_bytes = runtime.block_on(async { - eryx::preinit::pre_initialize( - &stdlib_path, - final_site_packages.as_deref(), - &import_refs, - &extensions, - setup_code.as_deref(), - ) - .await - .map_err(|e| InitializationError::new_err(format!("pre-initialization failed: {e}"))) + eryx::preinit::pre_initialize_with_options(options) + .await + .map_err(|e| { + InitializationError::new_err(format!("pre-initialization failed: {e}")) + }) })?; // Pre-compile to native code for faster instantiation @@ -149,6 +177,7 @@ impl SandboxFactory { stdlib_path, site_packages_path: final_site_packages, extracted_packages: Arc::new(extracted_packages), + callbacks: callbacks.map(Bound::unbind), }) } @@ -159,6 +188,10 @@ impl SandboxFactory { /// /// Args: /// path: Path to the saved factory file. + /// callbacks: The callbacks the factory was created with, if any. The + /// file holds their baked declarations but not the Python callables, + /// so pass the same callbacks here to get the setup-free fast path; + /// a different set still works, it just installs per sandbox. /// cache: Whether to cache the pre-compiled component in the process-global /// cache. When enabled, a BLAKE3 content hash is computed once during /// loading and subsequent ``create_sandbox()`` calls skip component @@ -174,8 +207,13 @@ impl SandboxFactory { /// factory = SandboxFactory.load("/path/to/jinja2-factory.bin") /// sandbox = factory.create_sandbox() #[staticmethod] - #[pyo3(signature = (path, *, site_packages=None, cache=true))] - fn load(path: PathBuf, site_packages: Option, cache: bool) -> PyResult { + #[pyo3(signature = (path, *, site_packages=None, callbacks=None, cache=true))] + fn load( + path: PathBuf, + site_packages: Option, + callbacks: Option>, + cache: bool, + ) -> PyResult { // Get embedded resources for stdlib path let embedded = eryx::embedded::EmbeddedResources::get().map_err(eryx_error_to_py)?; let stdlib_path = embedded.stdlib().to_path_buf(); @@ -194,6 +232,7 @@ impl SandboxFactory { stdlib_path, site_packages_path: site_packages, extracted_packages: Arc::new(Vec::new()), + callbacks: callbacks.map(Bound::unbind), }) } @@ -292,6 +331,7 @@ impl SandboxFactory { max_vfs_bytes: None, }); let limits: eryx::ResourceLimits = (&limits).into(); + let callbacks = self.default_callbacks(py, callbacks); Session::from_executor( py, Arc::new(executor), @@ -393,8 +433,9 @@ impl SandboxFactory { builder = builder.with_network(net.into()); } - // Apply callbacks if provided - if let Some(ref cbs) = callbacks { + // Apply the caller's callbacks, or the factory's own (whose + // declarations are baked into the snapshot) when none are given. + if let Some(ref cbs) = self.default_callbacks(py, callbacks) { let python_callbacks = extract_callbacks(py, cbs)?; for callback in python_callbacks { builder = builder.with_callback(callback); @@ -465,6 +506,18 @@ impl SandboxFactory { } } +impl SandboxFactory { + /// The callbacks a sandbox or session gets when the caller passes none: + /// the factory's own. + fn default_callbacks<'py>( + &self, + py: Python<'py>, + callbacks: Option>, + ) -> Option> { + callbacks.or_else(|| self.callbacks.as_ref().map(|cbs| cbs.bind(py).clone())) + } +} + impl std::fmt::Debug for SandboxFactory { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("SandboxFactory") diff --git a/crates/eryx-python/tests/test_factory_callbacks.py b/crates/eryx-python/tests/test_factory_callbacks.py new file mode 100644 index 00000000..12b55c71 --- /dev/null +++ b/crates/eryx-python/tests/test_factory_callbacks.py @@ -0,0 +1,87 @@ +"""Tests for SandboxFactory(callbacks=...): declarations baked into the snapshot.""" + +import eryx +import pytest + + +def get_time(): + return {"timestamp": 1234} + + +def add(a, b): + return {"sum": a + b} + + +def ping(): + return "pong" + + +CALLBACKS = [ + {"name": "get_time", "fn": get_time, "description": "Returns a fixed time"}, + {"name": "add", "fn": add, "description": "Adds two numbers"}, +] + + +@pytest.fixture(scope="module") +def callback_factory(): + """Factory whose snapshot has get_time and add baked in.""" + return eryx.SandboxFactory(callbacks=CALLBACKS) + + +class TestFactoryCallbacks: + def test_sandbox_registers_factory_callbacks_by_default(self, callback_factory): + sandbox = callback_factory.create_sandbox() + result = sandbox.execute( + "t = await get_time(); r = await add(a=3, b=4); print(t['timestamp'], r['sum'])" + ) + assert result.stdout.strip() == "1234 7" + + def test_introspection_lists_baked_callbacks(self, callback_factory): + sandbox = callback_factory.create_sandbox() + result = sandbox.execute("print(sorted(c['name'] for c in list_callbacks()))") + assert result.stdout.strip() == "['add', 'get_time']" + + def test_repeated_sandboxes_keep_working(self, callback_factory): + for i in range(3): + sandbox = callback_factory.create_sandbox() + result = sandbox.execute(f"print((await add(a={i}, b=1))['sum'])") + assert result.stdout.strip() == str(i + 1) + + def test_explicit_callbacks_override_the_baked_set(self, callback_factory): + sandbox = callback_factory.create_sandbox( + callbacks=[{"name": "ping", "fn": ping, "description": ""}] + ) + result = sandbox.execute( + "print(sorted(c['name'] for c in list_callbacks()), await ping())" + ) + assert result.stdout.strip() == "['ping'] pong" + + def test_session_registers_factory_callbacks_by_default(self, callback_factory): + session = callback_factory.create_session() + result = session.execute("print((await get_time())['timestamp'])") + assert result.stdout.strip() == "1234" + + def test_save_and_load_with_callbacks(self, callback_factory, tmp_path): + path = tmp_path / "factory.bin" + callback_factory.save(path) + + loaded = eryx.SandboxFactory.load(path, callbacks=CALLBACKS) + sandbox = loaded.create_sandbox() + result = sandbox.execute("print((await add(a=20, b=22))['sum'])") + assert result.stdout.strip() == "42" + + # Loading without callbacks still works; the sandbox just has none. + bare = eryx.SandboxFactory.load(path) + result = bare.create_sandbox().execute("print(list_callbacks())") + assert result.stdout.strip() == "[]" + + +class TestFactoryCallbacksWithSetupCode: + def test_setup_code_and_callbacks_together(self): + factory = eryx.SandboxFactory( + setup_code="base = 100", + callbacks=[{"name": "add", "fn": add, "description": "Adds two numbers"}], + ) + sandbox = factory.create_sandbox() + result = sandbox.execute("print((await add(a=base, b=1))['sum'])") + assert result.stdout.strip() == "101" diff --git a/crates/eryx-runtime/src/preinit.rs b/crates/eryx-runtime/src/preinit.rs index 140fb1aa..f632068d 100644 --- a/crates/eryx-runtime/src/preinit.rs +++ b/crates/eryx-runtime/src/preinit.rs @@ -23,21 +23,21 @@ //! # Example //! //! ```rust,ignore -//! use eryx_runtime::preinit::pre_initialize; +//! use eryx_runtime::preinit::{PreInitOptions, pre_initialize_with_options}; //! //! // Pre-initialize with native extensions -//! let preinit_component = pre_initialize( -//! &python_stdlib_path, -//! Some(&site_packages_path), -//! &["numpy", "pandas"], // Modules to import during pre-init -//! &native_extensions, -//! Some("import numpy as np; arr = np.zeros(10)"), // Optional setup code +//! let preinit_component = pre_initialize_with_options( +//! PreInitOptions::new(&python_stdlib_path) +//! .site_packages(&site_packages_path) +//! .imports(["numpy", "pandas"]) // Modules to import during pre-init +//! .extensions(native_extensions) +//! .setup_code("import numpy as np; arr = np.zeros(10)"), // Optional setup code //! ).await?; //! ``` use anyhow::{Result, anyhow}; use std::collections::HashSet; -use std::path::Path; +use std::path::{Path, PathBuf}; use tempfile::TempDir; use wasmtime::{ Config, Engine, Store, @@ -48,6 +48,92 @@ use wasmtime_wizer::{WasmtimeWizerComponent, Wizer}; use crate::linker::{NativeExtension, link_with_extensions}; +/// A callback the sandbox will offer at runtime, as the guest sees it. +/// +/// Passing the declarations of the callbacks a sandbox will register to +/// [`PreInitOptions::callbacks`] installs their Python wrappers during +/// pre-initialization, so a fresh instance whose host registers the same set +/// (compared by name, description and schema) skips the per-instance setup +/// entirely. The set is matched irrespective of order. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CallbackDeclaration { + /// Unique name of the callback (e.g. `"http.get"`). + pub name: String, + /// Human-readable description. + pub description: String, + /// JSON Schema for the callback's arguments, serialized. + pub parameters_schema_json: String, +} + +/// Everything [`pre_initialize_with_options`] needs. +#[derive(Debug, Clone)] +pub struct PreInitOptions { + python_stdlib: PathBuf, + site_packages: Option, + imports: Vec, + extensions: Vec, + setup_code: Option, + callbacks: Vec, +} + +impl PreInitOptions { + /// Options for pre-initializing with the Python standard library at + /// `python_stdlib` and nothing else. + pub fn new(python_stdlib: impl Into) -> Self { + Self { + python_stdlib: python_stdlib.into(), + site_packages: None, + imports: Vec::new(), + extensions: Vec::new(), + setup_code: None, + callbacks: Vec::new(), + } + } + + /// Mount `path` as `/site-packages` during pre-initialization. + #[must_use] + pub fn site_packages(mut self, path: impl Into) -> Self { + self.site_packages = Some(path.into()); + self + } + + /// Modules to import during pre-init (e.g. `["numpy", "pandas"]`). + #[must_use] + pub fn imports(mut self, imports: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.imports = imports.into_iter().map(Into::into).collect(); + self + } + + /// Native extensions to link into the component. + #[must_use] + pub fn extensions(mut self, extensions: Vec) -> Self { + self.extensions = extensions; + self + } + + /// Python code to run after the imports, baked into the snapshot. Use + /// this to pre-create objects (e.g. a Jinja2 `SandboxedEnvironment`) so + /// every sandbox starts with them in copy-on-write memory. + #[must_use] + pub fn setup_code(mut self, code: impl Into) -> Self { + self.setup_code = Some(code.into()); + self + } + + /// Callbacks the sandboxes will register at runtime; see + /// [`CallbackDeclaration`]. Setup code cannot invoke them (there is no + /// host to answer during pre-initialization). + #[must_use] + pub fn callbacks(mut self, callbacks: Vec) -> Self { + self.callbacks = callbacks; + self + } +} + /// Context for the pre-initialization runtime. struct PreInitCtx { wasi: WasiCtx, @@ -55,6 +141,8 @@ struct PreInitCtx { /// Temp directory for dummy files - must be kept alive during pre-init #[allow(dead_code)] temp_dir: Option, + /// What the `list-callbacks` import answers, sorted by name. + callbacks: Vec, } impl std::fmt::Debug for PreInitCtx { @@ -103,10 +191,46 @@ pub async fn pre_initialize( extensions: &[NativeExtension], setup_code: Option<&str>, ) -> Result> { - let imports: Vec = imports.iter().map(|s| (*s).to_string()).collect(); + let mut options = PreInitOptions::new(python_stdlib) + .imports(imports.iter().copied()) + .extensions(extensions.to_vec()); + if let Some(path) = site_packages { + options = options.site_packages(path); + } + if let Some(code) = setup_code { + options = options.setup_code(code); + } + pre_initialize_with_options(options).await +} + +/// Pre-initialize a Python component according to `options`. +/// +/// Links the component with the native extensions, runs the Python +/// interpreter's initialization, imports the requested modules, runs the setup +/// code, installs the declared callbacks' wrappers, and captures the resulting +/// memory state into the returned component. +/// +/// # Errors +/// +/// Returns an error if pre-initialization fails (e.g., Python init error, +/// import failure, or setup code exception). +pub async fn pre_initialize_with_options(options: PreInitOptions) -> Result> { + let PreInitOptions { + python_stdlib, + site_packages, + imports, + extensions, + setup_code, + mut callbacks, + } = options; + let python_stdlib = python_stdlib.as_path(); + let site_packages = site_packages.as_deref(); + // The guest compares the host's declarations with the installed set as a + // serialized list, so both sides present them in the same (name) order. + callbacks.sort_by(|a, b| a.name.cmp(&b.name)); // Link the component with real WASI adapter. - let original_component = link_with_extensions(extensions) + let original_component = link_with_extensions(&extensions) .map_err(|e| anyhow!("Failed to link component with extensions: {}", e))?; // Phase 1: Instrument the component (synchronous). @@ -166,12 +290,14 @@ pub async fn pre_initialize( let wasi = wasi_builder.build(); + let has_callbacks = !callbacks.is_empty(); let mut store = Store::new( &engine, PreInitCtx { wasi, table, temp_dir, + callbacks, }, ); @@ -194,10 +320,17 @@ pub async fn pre_initialize( // If setup code is provided, execute it after imports so its state // (variables, objects, etc.) gets captured in the Wizer snapshot. - if let Some(code) = setup_code { + if let Some(code) = &setup_code { call_execute_code(&mut store, &instance, code, "setup code").await?; } + // The guest installs the declared callbacks' wrappers on its first + // execute(). If nothing above ran one, run a no-op so the installation + // still lands in the snapshot. + if has_callbacks && imports.is_empty() && setup_code.is_none() { + call_execute_code(&mut store, &instance, "pass", "callback installation").await?; + } + // CRITICAL: Call finalize-preinit to reset WASI state AFTER all imports. // This clears file handles from the WASI adapter and wasi-libc so they // don't get captured in the memory snapshot. Without this, restored @@ -490,14 +623,34 @@ fn add_sandbox_stubs(linker: &mut Linker) -> Result<()> { )?; // list-callbacks: func() -> list + // + // Answers with the declarations from `PreInitOptions::callbacks` so the + // guest installs their wrappers into the snapshot; empty by default. linker.root().func_new( "list-callbacks", - |_ctx: wasmtime::StoreContextMut<'_, PreInitCtx>, + |ctx: wasmtime::StoreContextMut<'_, PreInitCtx>, _func_ty: wasmtime::component::types::ComponentFunc, _params: &[Val], results: &mut [Val]| { - // Return empty list - results[0] = Val::List(vec![]); + let declared = ctx + .data() + .callbacks + .iter() + .map(|cb| { + Val::Record(vec![ + ("name".to_string(), Val::String(cb.name.clone())), + ( + "description".to_string(), + Val::String(cb.description.clone()), + ), + ( + "parameters-schema-json".to_string(), + Val::String(cb.parameters_schema_json.clone()), + ), + ]) + }) + .collect(); + results[0] = Val::List(declared); Ok(()) }, )?; diff --git a/crates/eryx/benches/execution.rs b/crates/eryx/benches/execution.rs index cf18492c..32b2b75d 100644 --- a/crates/eryx/benches/execution.rs +++ b/crates/eryx/benches/execution.rs @@ -124,12 +124,17 @@ impl TypedCallback for WorkCallback { // Helpers // ============================================================================ +/// Whether to collect traces, honouring `ERYX_PROFILE_TRACE`. +/// +/// Trace collection (`sys.settrace`) is on by default; `ERYX_PROFILE_TRACE=0` +/// turns it off to measure without per-event trace overhead. +fn collect_trace() -> bool { + !std::env::var("ERYX_PROFILE_TRACE").is_ok_and(|v| v.trim() == "0") +} + fn create_sandbox() -> Sandbox { - // Trace collection (sys.settrace) is on by default; `ERYX_PROFILE_TRACE=0` - // turns it off to measure without per-event trace overhead. - let collect_trace = !std::env::var("ERYX_PROFILE_TRACE").is_ok_and(|v| v.trim() == "0"); Sandbox::embedded() - .with_trace_collection(collect_trace) + .with_trace_collection(collect_trace()) .with_callback(NoopCallback) .with_callback(EchoCallback) .with_callback(WorkCallback) @@ -197,6 +202,61 @@ fn bench_stateless_execution(c: &mut Criterion) { group.finish(); } +/// Stateless execution from a snapshot that has the three benchmark callbacks +/// baked in (`PreInitOptions::callbacks`), so a fresh instance skips the +/// per-instance callback setup that dominates `stateless_execution/pass`. +/// +/// Needs the `preinit` feature (`--features embedded,preinit`); otherwise the +/// group is skipped. Pre-initialization itself takes ~20 s before the first +/// sample. +fn bench_stateless_execution_baked_callbacks(c: &mut Criterion) { + #[cfg(not(feature = "preinit"))] + { + let _ = c; + eprintln!("stateless_execution_baked: skipped (enable the `preinit` feature)"); + } + #[cfg(feature = "preinit")] + { + use eryx::preinit::{PreInitOptions, callback_declaration, pre_initialize_with_options}; + use std::path::PathBuf; + + let rt = tokio::runtime::Runtime::new().unwrap(); + let stdlib = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../eryx-wasm-runtime/tests/python-stdlib"); + let callbacks: [&dyn eryx::Callback; 3] = [&NoopCallback, &EchoCallback, &WorkCallback]; + let declarations = callbacks + .iter() + .map(|cb| callback_declaration(*cb)) + .collect(); + let preinit = rt + .block_on(pre_initialize_with_options( + PreInitOptions::new(&stdlib).callbacks(declarations), + )) + .expect("pre-initialization should succeed"); + let precompiled = eryx::PythonExecutor::precompile(&preinit).expect("precompile"); + let artifact = eryx::PrecompiledArtifact::new(precompiled); + + // SAFETY: the artifact was produced by `PythonExecutor::precompile` just above. + let sandbox = unsafe { Sandbox::builder().with_precompiled_artifact(artifact) } + .with_python_stdlib(&stdlib) + .with_trace_collection(collect_trace()) + .with_callback(NoopCallback) + .with_callback(EchoCallback) + .with_callback(WorkCallback) + .build() + .expect("Failed to create sandbox"); + + let mut group = c.benchmark_group("stateless_execution_baked"); + group.sample_size(10); + group.measurement_time(Duration::from_secs(10)); + group.bench_function("pass", |b| { + b.to_async(&rt) + .iter(|| async { sandbox.execute("pass").await.expect("Execution failed") }); + }); + group.finish(); + } +} + // ============================================================================ // Session Execution Benchmarks (reused WASM instance) // ============================================================================ @@ -460,6 +520,7 @@ criterion_group!( bench_sandbox_creation, bench_session_creation, bench_stateless_execution, + bench_stateless_execution_baked_callbacks, bench_session_execution, bench_callback_overhead, bench_parallel_callbacks, diff --git a/crates/eryx/src/lib.rs b/crates/eryx/src/lib.rs index 90a7c5ed..0b5ca34b 100644 --- a/crates/eryx/src/lib.rs +++ b/crates/eryx/src/lib.rs @@ -80,7 +80,26 @@ mod wasm; #[cfg(feature = "preinit")] pub mod preinit { pub use eryx_runtime::linker::NativeExtension; - pub use eryx_runtime::preinit::{PreInitError, pre_initialize}; + pub use eryx_runtime::preinit::{ + CallbackDeclaration, PreInitError, PreInitOptions, pre_initialize, + pre_initialize_with_options, + }; + + /// The declaration the guest sees for `callback`, for + /// [`PreInitOptions::callbacks`]. + /// + /// Sandboxes created from a snapshot that baked this declaration and that + /// register a callback with the same name, description and schema skip the + /// per-instance callback setup. + #[must_use] + pub fn callback_declaration(callback: &dyn crate::Callback) -> CallbackDeclaration { + CallbackDeclaration { + name: callback.name().to_string(), + description: callback.description().to_string(), + parameters_schema_json: serde_json::to_string(&callback.parameters_schema()) + .unwrap_or_else(|_| "{}".to_string()), + } + } } pub use callback::{ diff --git a/crates/eryx/src/wasm.rs b/crates/eryx/src/wasm.rs index e5b36150..4c6d3f77 100644 --- a/crates/eryx/src/wasm.rs +++ b/crates/eryx/src/wasm.rs @@ -734,15 +734,23 @@ impl SandboxImportsWithStore for HasSelf { impl SandboxImports for ExecutorState { /// List all available callbacks for introspection. + /// + /// Sorted by name: the guest compares this list, serialized, with the set + /// whose wrappers it has installed (possibly baked into the pre-init + /// snapshot from `PreInitOptions::callbacks`, which sorts the same way), so + /// registration order must not affect whether the sets match. fn list_callbacks(&mut self) -> Vec { - self.callbacks + let mut infos: Vec = self + .callbacks .iter() .map(|cb| CallbackInfo { name: cb.name.clone(), description: cb.description.clone(), parameters_schema_json: cb.parameters_schema_json.clone(), }) - .collect() + .collect(); + infos.sort_by(|a, b| a.name.cmp(&b.name)); + infos } /// Report a trace event to the host. diff --git a/crates/eryx/tests/preinit.rs b/crates/eryx/tests/preinit.rs index eb117e88..2cd06108 100644 --- a/crates/eryx/tests/preinit.rs +++ b/crates/eryx/tests/preinit.rs @@ -11,6 +11,7 @@ use eryx::Sandbox; use eryx::preinit::pre_initialize; use std::collections::HashMap; +use std::future::Future; use std::path::{Path, PathBuf}; use std::sync::Arc; use tokio::sync::{Mutex, OnceCell}; @@ -508,3 +509,137 @@ async fn preinit_setup_code_error_is_reported() { "error should mention 'setup code': {err}" ); } + +// ============================================================================= +// Baked Callback Tests +// ============================================================================= + +/// A callback that returns its arguments. +struct Echo; + +impl eryx::Callback for Echo { + fn name(&self) -> &str { + "echo" + } + fn description(&self) -> &str { + "Returns its arguments" + } + fn parameters_schema(&self) -> eryx::Schema { + eryx::empty_schema() + } + fn invoke( + &self, + args: serde_json::Value, + ) -> std::pin::Pin< + Box> + Send + '_>, + > { + Box::pin(async move { Ok(args) }) + } +} + +/// A callback that answers "pong". +struct Ping; + +impl eryx::Callback for Ping { + fn name(&self) -> &str { + "ping" + } + fn description(&self) -> &str { + "Answers pong" + } + fn parameters_schema(&self) -> eryx::Schema { + eryx::empty_schema() + } + fn invoke( + &self, + _args: serde_json::Value, + ) -> std::pin::Pin< + Box> + Send + '_>, + > { + Box::pin(async move { Ok(serde_json::json!("pong")) }) + } +} + +/// Pre-initialize with the given callbacks' declarations baked in. +async fn preinit_with_callbacks(stdlib: &Path, callbacks: &[&dyn eryx::Callback]) -> Vec { + use eryx::preinit::{PreInitOptions, callback_declaration, pre_initialize_with_options}; + + let declarations = callbacks + .iter() + .map(|cb| callback_declaration(*cb)) + .collect(); + pre_initialize_with_options(PreInitOptions::new(stdlib).callbacks(declarations)) + .await + .expect("pre-initialization with callbacks should succeed") +} + +/// Callbacks baked into the snapshot are callable from a sandbox registering +/// the same set, and a sandbox registering a different set gets that set. +#[tokio::test] +async fn preinit_baked_callbacks() { + let stdlib = get_stdlib_path(); + let preinit_bytes = preinit_with_callbacks(&stdlib, &[&Echo]).await; + + // Same set as baked: the wrapper is already installed and works. + let sandbox = Sandbox::builder() + .with_wasm_bytes(preinit_bytes.clone()) + .with_python_stdlib(&stdlib) + .with_callback(Echo) + .build() + .expect("sandbox creation should succeed"); + let result = sandbox + .execute( + "print(sorted(c['name'] for c in list_callbacks()))\nprint((await echo(data=7))['data'])", + ) + .await + .expect("execution should succeed"); + assert_eq!(result.stdout.trim(), "['echo']\n7"); + + // A different set is installed as before, replacing the baked one. + let sandbox = Sandbox::builder() + .with_wasm_bytes(preinit_bytes.clone()) + .with_python_stdlib(&stdlib) + .with_callback(Ping) + .build() + .expect("sandbox creation should succeed"); + let result = sandbox + .execute("print(sorted(c['name'] for c in list_callbacks()), await ping())") + .await + .expect("execution should succeed"); + assert_eq!(result.stdout.trim(), "['ping'] pong"); + + // No callbacks registered: the guest reflects the host, not the snapshot. + let sandbox = Sandbox::builder() + .with_wasm_bytes(preinit_bytes) + .with_python_stdlib(&stdlib) + .build() + .expect("sandbox creation should succeed"); + let result = sandbox + .execute("print(list_callbacks())") + .await + .expect("execution should succeed"); + assert_eq!(result.stdout.trim(), "[]"); +} + +/// Declaration order does not matter: the snapshot and the host both present +/// callbacks sorted by name. +#[tokio::test] +async fn preinit_baked_callbacks_ignore_registration_order() { + let stdlib = get_stdlib_path(); + let preinit_bytes = preinit_with_callbacks(&stdlib, &[&Ping, &Echo]).await; + + let sandbox = Sandbox::builder() + .with_wasm_bytes(preinit_bytes) + .with_python_stdlib(&stdlib) + .with_callback(Echo) + .with_callback(Ping) + .build() + .expect("sandbox creation should succeed"); + let result = sandbox + .execute( + "print([c['name'] for c in list_callbacks()], await ping(), (await echo(x=1))['x'])", + ) + .await + .expect("execution should succeed"); + assert_eq!(result.stdout.trim(), "['echo', 'ping'] pong 1"); +}