Enabling wasi-nn for Preview2 Apps

[gustavojardim]

Summary of Changes

These changes enable running WebAssembly modules using the wasi-nn API on Kubernetes with the Wasmtime runtime.

1. wasi-nn API Integration in Wasmtime Shim

• Added wasmtime-wasi-nn as an optional dependency and feature in Cargo.toml.
• Conditional compilation for wasi-nn and PyTorch backend.
• Created and initialized WasiNnCtx in the WASI context when enabled.
• Registered the wasi-nn interface in the component linker.

2. WASI Context Update

• Updated context types and builder usage to support the new wasi-nn context.
• Adjusted trait implementations to use unified WasiView and resource table.

3. Dependency and Feature Updates

• Modified Cargo.toml to add wasmtime-wasi-nn and a pytorch feature.
• Updated feature flags and dependency lists for optional wasi-nn and PyTorch support.

[Copilot]

Pull request overview

This pull request enables wasi-nn (WebAssembly System Interface for Neural Networks) support for Preview2 applications in the Wasmtime runtime, with optional PyTorch backend integration. The changes include a major wasmtime version upgrade from 33.0.2 to 40.0.1 and associated API updates.

Changes:

• Upgrades wasmtime and related dependencies from v33 to v40
• Adds optional wasmtime-wasi-nn dependency with PyTorch backend feature
• Updates WASI context structure to include WasiNnCtx for neural network operations
• Modifies trait implementations to work with unified WasiView and resource table patterns

Reviewed changes

Copilot reviewed 5 out of 6 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
rust-toolchain.toml	Changes Rust toolchain from pinned version 1.85.0 to "stable"
Cargo.toml	Updates wasmtime dependencies to v40.0.1 and adds wasmtime-wasi-nn
crates/containerd-shim-wasmtime/Cargo.toml	Adds pytorch feature flag dependent on wasmtime-wasi-nn/pytorch
crates/containerd-shim-wasmtime/src/instance.rs	Integrates wasi-nn context into WasiPreview2Ctx, updates imports for wasmtime v40 API changes, modifies trait implementations
crates/containerd-shim-wasmtime/src/http_proxy.rs	Adds wasi-nn context initialization for HTTP proxy request handling
Cargo.lock	Updates all transitive dependencies for wasmtime v40 ecosystem

Comments suppressed due to low confidence (3)

crates/containerd-shim-wasmtime/src/instance.rs:139

• This line contains trailing whitespace, which is inconsistent with the project's code style and can cause unnecessary diff noise in version control.
  crates/containerd-shim-wasmtime/src/http_proxy.rs:185
• The wasi-nn context initialization logic is duplicated between instance.rs (lines 111-119) and http_proxy.rs (lines 177-185). This duplication makes the code harder to maintain and increases the risk of inconsistencies.

Consider extracting this into a helper function that can be reused in both locations, for example:

#[cfg(feature = "wasmtime-wasi-nn")]
fn create_wasi_nn_ctx() -> WasiNnCtx {
    let backends = vec![
        #[cfg(feature = "pytorch")]
        Backend::from(wasmtime_wasi_nn::backend::pytorch::PytorchBackend::default()),
    ];
    let registry = Registry::from(InMemoryRegistry::new());
    WasiNnCtx::new(backends, registry)
}

            wasi_nn: {
                // Create wasi-nn with backends
                let backends = vec![
                    #[cfg(feature = "pytorch")]
                    wasmtime_wasi_nn::Backend::from(wasmtime_wasi_nn::backend::pytorch::PytorchBackend::default()),
                ];
                let registry = wasmtime_wasi_nn::Registry::from(wasmtime_wasi_nn::InMemoryRegistry::new());
                wasmtime_wasi_nn::wit::WasiNnCtx::new(backends, registry)
            },

crates/containerd-shim-wasmtime/src/instance.rs:106

• This line contains trailing whitespace, which is inconsistent with the project's code style and can cause unnecessary diff noise in version control.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

This line contains trailing whitespace, which is inconsistent with the project's code style and can cause unnecessary diff noise in version control.

This issue also appears in the following locations of the same file:

• line 139
• line 106

[Copilot]

The "pytorch" feature depends on "wasmtime-wasi-nn/pytorch" but "wasmtime-wasi-nn" itself is optional. This creates a situation where enabling the "pytorch" feature alone won't work because the base dependency is optional.

The features should be structured as:

[features]
default = []
wasmtime-wasi-nn = ["dep:wasmtime-wasi-nn"]
pytorch = ["wasmtime-wasi-nn", "wasmtime-wasi-nn/pytorch"]

This ensures that enabling "pytorch" automatically enables the base "wasmtime-wasi-nn" dependency.

Suggested change

	pytorch = ["wasmtime-wasi-nn/pytorch"]
	wasmtime-wasi-nn = ["dep:wasmtime-wasi-nn"]
	pytorch = ["wasmtime-wasi-nn", "wasmtime-wasi-nn/pytorch"]

[Copilot]

The Rust toolchain channel was changed from a pinned version "1.85.0" to "stable". This makes the build less reproducible as "stable" is a moving target that changes every 6 weeks with new Rust releases. This could lead to unexpected build failures or behavioral changes when the stable version updates.

Consider keeping a pinned version for reproducibility, or at least document why this change was necessary for wasi-nn support.

Suggested change

	channel="stable"
	# Pin Rust for reproducible builds. 1.85.0 is known to work with wasi-nn;
	# update this version intentionally if newer Rust is required.
	channel = "1.85.0"

[Copilot]

The backends vector will be empty when the "pytorch" feature is not enabled, as the only element in the vec! is conditionally compiled. An empty backends vector may not be the intended behavior for wasi-nn initialization.

Consider either:

1. Making "wasmtime-wasi-nn" feature depend on at least one backend being enabled
2. Adding a compile-time check to ensure at least one backend is available
3. Adding a default backend that's always available
4. Documenting that wasi-nn can run with zero backends if that's intended

[Copilot]

Same issue as in instance.rs: the backends vector will be empty when the "pytorch" feature is not enabled, which may not be the intended behavior. This code is duplicated from instance.rs and should be kept consistent with the solution applied there.

This issue also appears in the following locations of the same file:

• line 177

[Copilot]

The code is missing a WasiNnView trait implementation for WasiPreview2Ctx. While WasiNnView::new() is being called in the linker setup, wasmtime-wasi-nn typically requires the host struct to implement the WasiNnView trait directly for the generated bindings to work correctly.

Add the missing trait implementation:

#[cfg(feature = "wasmtime-wasi-nn")]
impl wasmtime_wasi_nn::wit::WasiNnView for WasiPreview2Ctx {
    fn ctx(&mut self) -> &mut WasiNnCtx {
        &mut self.wasi_nn
    }
    
    fn table(&mut self) -> &mut ResourceTable {
        &mut self.resource_table
    }
}

Suggested change

	#[cfg(feature = "wasmtime-wasi-nn")]
	impl wasmtime_wasi_nn::wit::WasiNnView for WasiPreview2Ctx {
	fn ctx(&mut self) -> &mut WasiNnCtx {
	&mut self.wasi_nn
	}
	fn table(&mut self) -> &mut ResourceTable {
	&mut self.resource_table
	}
	}

[Mossaka]

Thanks for contributing! There are some CI failures and let me address them and then I will come back to review this PR.