Skip to content

Commit

Permalink
Refactor AICI to use WebAssembly Component Model
Browse files Browse the repository at this point in the history
This is a significant change to the AICI Runtime (host) and AICI Controller
(guest) to use WASI components. As part of this change, a significant amount of
unsafe code is removed and the protocol is simplified to remove the need for
"BLOB" types and side-channels.

The protocol is documented in `wit/controller.wit`, and a minimal WASI runtime
is provided to AI CI controllers.

Some notes:
* The AICI runtime now longer directly reads and writes to the guest's memory.
  Instead, the guest provides a `Runner` resource (using WebAssembly Component
  terminology), which exposes the low-level protocol to the host as a
  constructor and trait with methods.
* The Blob protocols are removed entirely, replaced by the `Runner` resource.
  This and other side-channels for communicating with the runtime, e.g. allowed
  tokens (logit biases) outside of `MidProcessResult`, are removed.
* The (Variable) Storage and Tokenizer protocols are separate WebAssembly
  Components, which can be versioned independently of the runtime.
* Types are changed to be consistent with the WebAssembly interface, e.g.:
  `SeqId` is used in far more places to avoid casts.
  • Loading branch information
AaronFriel committed Aug 27, 2024
1 parent eec6e52 commit c9dc62d
Show file tree
Hide file tree
Showing 74 changed files with 2,208 additions and 2,028 deletions.
2 changes: 1 addition & 1 deletion .devcontainer/Dockerfile-prod-vllm
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ FROM rust:1.75.0-bookworm AS aicirt

WORKDIR /workspace

RUN rustup target add wasm32-wasi
RUN rustup target add wasm32-wasip2
RUN curl -fsSL https://deb.nodesource.com/setup_18.x | bash -
RUN apt-get install -y nodejs

Expand Down
2 changes: 1 addition & 1 deletion .devcontainer/common.dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ ENV RUSTUP_HOME=/usr/local/rustup \

RUN curl https://sh.rustup.rs -sSf | sh -s -- \
-y --no-modify-path --profile minimal --default-toolchain $RUST_VERSION
RUN rustup target add wasm32-wasi
RUN rustup target add wasm32-wasip2
RUN rustup component add rustfmt

# run as root please; note that settings in devcontainer.json are also needed...
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/aicirt-release.yml
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ jobs:
- uses: actions/checkout@v3
with:
submodules: true
- run: rustup target add wasm32-wasi
- run: rustup target add wasm32-wasip2
- uses: hendrikmuhs/[email protected]
- uses: Swatinem/rust-cache@v2
with:
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/aicirt.yml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ jobs:
- uses: actions/checkout@v3
with:
submodules: true
- run: rustup target add wasm32-wasi
- run: rustup target add wasm32-wasip2
- uses: hendrikmuhs/[email protected]
- uses: Swatinem/rust-cache@v2
with:
Expand Down
Loading

0 comments on commit c9dc62d

Please sign in to comment.