Edge Inference: Real-Time Personalization with WASM and ONNX Runtime

# model.py
import torch
import torch.nn as nn

class RankingModel(nn.Module):
    def __init__(self, embedding_dim: int):
        super(RankingModel, self).__init__()
        # Example: User embedding dim = 32, Item embedding dim = 32 -> input_dim = 64
        input_dim = embedding_dim * 2
        self.layer_stack = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )

    def forward(self, user_embedding, item_embedding):
        # Concatenate user and item embeddings
        x = torch.cat((user_embedding, item_embedding), dim=1)
        return self.layer_stack(x)

# --- Training would happen here --- 
# For this example, we'll just instantiate and export a dummy model.

def export_model_to_onnx():
    EMBEDDING_DIM = 32
    BATCH_SIZE = 1 # We'll use dynamic axes to handle variable batch sizes
    
    model = RankingModel(embedding_dim=EMBEDDING_DIM)
    model.eval() # Set model to evaluation mode

    # Create dummy inputs that match the expected dimensions
    dummy_user_embedding = torch.randn(BATCH_SIZE, EMBEDDING_DIM, requires_grad=False)
    dummy_item_embedding = torch.randn(BATCH_SIZE, EMBEDDING_DIM, requires_grad=False)

    # Define input and output names
    input_names = ["user_embedding", "item_embedding"]
    output_names = ["score"]

    print("Exporting model to ONNX...")
    torch.onnx.export(
        model,
        (dummy_user_embedding, dummy_item_embedding),
        "ranker.onnx",
        input_names=input_names,
        output_names=output_names,
        opset_version=14,
        dynamic_axes={
            'user_embedding': {0: 'batch_size'},
            'item_embedding': {0: 'batch_size'},
            'score': {0: 'batch_size'}
        },
        verbose=True
    )
    print("Model exported to ranker.onnx")

if __name__ == "__main__":
    export_model_to_onnx()

#!/bin/bash
set -e

# Clone the ONNX Runtime repository
# We pin to a specific version for reproducibility
ORT_VERSION="1.14.1"

git clone --single-branch --branch v${ORT_VERSION} --depth 1 https://github.com/microsoft/onnxruntime.git
cd onnxruntime

# Run the build script
# This script is provided by the ONNX Runtime team
# Key flags:
# --build_wasm: Enables the WebAssembly build
# --disable_wasm_exception_catching: Reduces binary size. We'll handle errors in JS.
# --disable_rtti: Disables Run-Time Type Information for smaller binary.
# --minimal_build: Only include ops required by our model.
# --config Release: Standard release optimizations.
# --add_wasm_flags "-s MODULARIZE=1 -s EXPORT_ES6=1": Crucial for modern JS module compatibility

./build.sh --config Release --build_wasm --disable_wasm_exception_catching --disable_rtti --minimal_build --parallel --add_wasm_flags "-s MODULARIZE=1 -s EXPORT_ES6=1"

# The output will be in build/js/packages/onnxruntime-web/dist/
# We are interested in onnxruntime-core.wasm and the JS wrapper.

# Use the Emscripten SDK image which has all the C++ -> WASM toolchains
FROM emscripten/emsdk:3.1.25

WORKDIR /src

# Copy the build script into the container
COPY build_onnx_wasm.sh .
RUN chmod +x build_onnx_wasm.sh

# Run the build
RUN ./build_onnx_wasm.sh

# This is a multi-stage build. The final stage will just copy the artifacts.
FROM alpine:latest
WORKDIR /artifacts

# Copy the essential build artifacts from the builder stage
COPY --from=0 /src/onnxruntime/build/js/packages/onnxruntime-web/dist/ort-wasm.wasm .
COPY --from=0 /src/onnxruntime/build/js/packages/onnxruntime-web/dist/ort-wasm.min.js .

# The final artifacts will be in the /artifacts directory of this image

npm create cloudflare@latest edge-inference-worker -- --type=module-worker
cd edge-inference-worker
npm install onnxruntime-common # Provides type definitions and tensor helpers

name = "edge-inference-worker"
main = "src/index.ts"
compatibility_date = "2023-10-30"

[vars]
ENVIRONMENT = "development"

# Make the WASM module available to the worker
[[wasm_modules]]
name = "ort_wasm"
path = "vendor/ort-wasm.wasm"

# Make the ONNX model available as a binding
[[r2_buckets]]
binding = "MODEL_BUCKET"
bucket_name = "my-ml-models"

# For local dev, we can use a site binding to load the model
[site]
bucket = "."

// src/index.ts
import { Tensor, InferenceSession } from 'onnxruntime-common';
import ortInit from '../vendor/ort-wasm.min.js';

// Import the WASM module we defined in wrangler.toml
import ortWasmModule from 'ort_wasm';

// Define the environment bindings for TypeScript
export interface Env {
    // This binding gives us access to the .onnx file
    __STATIC_CONTENT: KVNamespace;
    // In production, you would use R2
    MODEL_BUCKET: R2Bucket;
}

// --- Singleton Pattern for Caching --- 
// These variables are defined in the global scope of the worker.
// They will persist between requests on a warm instance, avoiding
// the expensive re-initialization cost.
let ort: any;
let session: InferenceSession;
let modelArrayBuffer: ArrayBuffer;

async function initializeOrt(env: Env) {
    if (ort && session) {
        // Already initialized on this warm worker
        return;
    }

    console.log('Performing cold start initialization...');
    const start = Date.now();

    // 1. Initialize the ONNX Runtime
    // The JS wrapper needs to be initialized with the WASM module.
    ort = await ortInit({ wasm: { module: ortWasmModule } });
    // Disable multithreading for serverless environments
    ort.env.wasm.numThreads = 1;
    console.log(`ORT initialized in ${Date.now() - start}ms`);

    // 2. Load the model data
    // In local dev with `[site]`, files are served via a KV namespace.
    // In production, this would be `env.MODEL_BUCKET.get('ranker.onnx')`.
    const modelResponse = await env.__STATIC_CONTENT.get('ranker.onnx', 'arrayBuffer');
    if (!modelResponse) {
        throw new Error('Model not found!');
    }
    modelArrayBuffer = modelResponse;
    console.log(`Model loaded in ${Date.now() - start}ms`);

    // 3. Create the inference session
    // This parses the model and prepares it for execution.
    // It's a key operation to cache.
    session = await InferenceSession.create(modelArrayBuffer, {
        executionProviders: ['wasm'],
        graphOptimizationLevel: 'all',
    });
    console.log(`Session created in ${Date.now() - start}ms. Total cold start: ${Date.now() - start}ms`);
}

export default {
    async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {
        try {
            // Initialization is deferred until the first request
            await initializeOrt(env);

            // --- Example Scenario: Re-ranking a list of items ---
            // Assume we get a user ID and a list of item IDs to rank
            const { userId, itemIds } = await request.json<{ userId: string; itemIds: string[] }>();

            // In a real app, you would fetch these embeddings from a low-latency key-value store
            // like Cloudflare KV or Redis.
            const userEmbedding = await fetchUserEmbedding(userId);
            const itemEmbeddings = await fetchItemEmbeddings(itemIds);

            const numItems = itemIds.length;

            // Prepare model inputs. The user embedding is repeated for each item.
            const userEmbeddingTiled = new Float32Array(numItems * 32);
            for (let i = 0; i < numItems; i++) {
                userEmbeddingTiled.set(userEmbedding, i * 32);
            }
            
            // The item embeddings are flattened into a single array.
            const itemEmbeddingsFlattened = new Float32Array(numItems * 32);
            itemEmbeddings.forEach((emb, i) => {
                itemEmbeddingsFlattened.set(emb, i * 32);
            });

            // Create ONNX Runtime Tensors
            const userTensor = new Tensor('float32', userEmbeddingTiled, [numItems, 32]);
            const itemTensor = new Tensor('float32', itemEmbeddingsFlattened, [numItems, 32]);

            const feeds = {
                user_embedding: userTensor,
                item_embedding: itemTensor,
            };

            // --- Run Inference --- 
            const inferenceStart = Date.now();
            const results = await session.run(feeds);
            const inferenceTime = Date.now() - inferenceStart;

            const scores = results.score.data as Float32Array;

            // Post-process: Combine items with their scores and sort
            const rankedItems = itemIds
                .map((id, i) => ({ id, score: scores[i] }))
                .sort((a, b) => b.score - a.score);

            return new Response(JSON.stringify({ rankedItems, inferenceTime }), {
                headers: { 'Content-Type': 'application/json' },
            });

        } catch (e: any) {
            console.error(e);
            return new Response(`Error during inference: ${e.message}`, { status: 500 });
        }
    },
};

// --- Mock data fetching functions ---
async function fetchUserEmbedding(userId: string): Promise<Float32Array> {
    // In production, fetch from KV/Upstash
    return new Float32Array(32).fill(0.5);
}

async function fetchItemEmbeddings(itemIds: string[]): Promise<Float32Array[]> {
    // In production, use kv.getMultiple()
    return itemIds.map(() => new Float32Array(32).map(() => Math.random()));
}

# quantize.py
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

model_fp32 = 'ranker.onnx'
model_quant = 'ranker.quant.onnx'

print("Starting quantization...")
quantize_dynamic(model_fp32,
                 model_quant,
                 weight_type=QuantType.QInt8)

print(f"Quantized model saved to {model_quant}")

// Simplified A/B testing logic within initializeOrt

// 1. Read config from a low-TTL KV key
const modelConfig = await env.CONFIG_KV.get('model_config', 'json');

// 2. Determine model version based on request
const useBeta = request.headers.get('X-Use-Beta-Model') === 'true';
const model_version = useBeta ? modelConfig.beta_version : modelConfig.default_version;
const model_key = `ranker_${model_version}.quant.onnx`;

// 3. Check if the correct session is already cached
// The singleton 'session' object can be a map of version -> InferenceSession
if (sessions[model_version]) {
    // Session for this version is already warm
    return;
}

// 4. Load the specific model version from R2 and create a new session
const modelBuffer = await env.MODEL_BUCKET.get(model_key).arrayBuffer();
sessions[model_version] = await InferenceSession.create(modelBuffer, ...);

// Inside the fetch handler
const timings = {};

timings.start = Date.now();
await initializeOrt(env);
timings.init_complete = Date.now();

// ... run inference ...
timings.inference_complete = Date.now();

const logData = {
    userId,
    numItems,
    cold_start: timings.init_complete - timings.start > 10, // A simple heuristic
    init_duration: timings.init_complete - timings.start,
    inference_duration: timings.inference_complete - timings.init_complete,
    // ... other metadata
};

// Use console.log with JSON for structured logging services
console.log(JSON.stringify(logData));