用 WasmEdge 和 YoMo 对实时数据流进行 AI 推理

YoMo 是一个用于辅助开发者方便构建分布式云系统（Geo-Distributed Cloud System）的编程框架。YoMo 的通讯层构建在 QUIC 协议之上，带来高速数据传输的同时，内置了 Streaming Serverless 的“流函数”，大幅提升了分布式云系统的开发体验。YoMo 构建的分布式云系统在近场算力和终端之间提供了超高速通讯机制，在 Metaverse、VR/AR、IoT 等领域有广泛的应用场景。

例如在 AR/VR、智能工厂里做实时 AI 推理的过程中，摄像头可以通过 YoMo 将实时的非结构化数据发送到近场 MEC （多访问边缘计算）设备中的计算节点，并自动执行托管的 AI 推理函数。当 AI 推理完成，YoMo将 AI 计算结果实时发送给端设备。

然而，YoMo 面临的挑战是在边缘计算节点中合并和管理由多个外部开发者编写的处理程序函数。这需要在不牺牲性能的情况下对这些函数进行 runtime 隔离。传统的软件容器解决方案，如 Docker，无法胜任这项任务，因为太重且速度太慢，无法处理实时任务。

WebAssembly 提供了一个轻量级高性能的软件容器。它非常适合作为 YoMo 数据处理 handler 函数的 runtime。

本文中，我们将向你展示如何为基于 Tensorflow 的图片识别创建 Rust 函数，将其编译为 WebAssembly，然后使用 YoMo 将其作为流数据 handler 运行。我们使用 WasmEdge 作为 WebAssembly runtime，因为与其它 WebAssembly runtime 相比，WasmEdge 提供了最佳性能和最高灵活度。WasmEdge 是唯一稳定支持 Tensorflow 的 WebAssembly 虚拟机。 YoMo 通过 WasmEdge 的 Golang API管理 WasmEdge VM 实例和容器内的 WebAssembly 字节码应用。

准备工作

显然，需要安装 Golang 。我们假设你已经安装。

同时，需要安装 YoMo CLI 应用程序。它安排和协调数据流和 handler 函数调用。

$ go install github.com/yomorun/cli/yomo@latest
$ yomo version
YoMo CLI version: v0.0.5

接下来，请安装 WasmEdge 及 Tensorflow 共享库。 WasmEdge 是领先的 WebAssembly runtime，由 CNCF 托管。 我们将使用它嵌入和运行来自 YoMo 的 WebAssembly 程序。

# Install WasmEdge
$ wget https://github.com/second-state/WasmEdge-go/releases/download/v0.8.1/install_wasmedge.sh
$ chmod +x ./install_wasmedge.sh
$ sudo ./install_wasmedge.sh /usr/local

# Install WasmEdge Tensorflow extension
$ wget https://github.com/second-state/WasmEdge-go/releases/download/v0.8.1/install_wasmedge_tensorflow_deps.sh
$ wget https://github.com/second-state/WasmEdge-go/releases/download/v0.8.1/install_wasmedge_tensorflow.sh
$ chmod +x ./install_wasmedge_tensorflow_deps.sh
$ chmod +x ./install_wasmedge_tensorflow.sh
$ sudo ./install_wasmedge_tensorflow_deps.sh /usr/local
$ sudo ./install_wasmedge_tensorflow.sh /usr/local

# Install WasmEdge Images extension
$ wget https://github.com/second-state/WasmEdge-go/releases/download/v0.8.1/install_wasmedge_image_deps.sh
$ wget https://github.com/second-state/WasmEdge-go/releases/download/v0.8.1/install_wasmedge_image.sh
$ chmod +x ./install_wasmedge_image_deps.sh
$ chmod +x ./install_wasmedge_image.sh
$ sudo ./install_wasmedge_image_deps.sh /usr/local
$ sudo ./install_wasmedge_image.sh /usr/local

最后， 因为我们的 demo WebAssembly 函数是用 Rust 编写的，因此你还需要安装 Rust 编译器和 rustwasmc工具链。

demo 的其它部分可以 fork 和 clone 源代码 repo。

$ git clone https://github.com/yomorun/yomo-wasmedge-tensorflow.git

图片分类函数

处理 YoMo 图片流的图片识别函数是用 Rust 编写的。 它用 WasmEdge Tensorflow API 来处理输入图片。

#[wasm_bindgen]
pub fn infer(image_data: &[u8]) -> String {
    // Load the TFLite model and its meta data (the text label for each recognized object number)
    let model_data: &[u8] = include_bytes!("lite-model_aiy_vision_classifier_food_V1_1.tflite");
    let labels = include_str!("aiy_food_V1_labelmap.txt");

    // Pre-process the image to a format that can be used by this model
    let flat_img = wasmedge_tensorflow_interface::load_jpg_image_to_rgb8(image_data, 192, 192);

    // Run the TFLite model using the WasmEdge Tensorflow API
    let mut session = wasmedge_tensorflow_interface::Session::new(&model_data, wasmedge_tensorflow_interface::ModelType::TensorFlowLite);
    session.add_input("input", &flat_img, &[1, 192, 192, 3])
           .run();
    let res_vec: Vec<u8> = session.get_output("MobilenetV1/Predictions/Softmax");

    // Find the object index in res_vec that has the greatest probability
    // Translate the probability into a confidence level
    // Translate the object index into a label from the model meta data food_name

    ret_str = format!(
        "It {} a <a href='https://www.google.com/search?q={}'>{}</a> in the picture",
        confidence, food_name, food_name
    );
    return ret_str;
}

你可以使用 rustwasmc 工具将该函数编译为 WebAssembly 字节码。

$ rustup default 1.50.0

$ cd flow/rust_mobilenet_food
$ rustwasmc  build `--enable-ext`
# The output WASM will be pkg/rust_mobilenet_food_lib_bg.wasm.

# Copy the wasm bytecode file to the flow/ directory
$ cp pkg/rust_mobilenet_food_lib_bg.wasm ../

与 YoMo 集成

在 YoMo 这端，我们使用 WasmEdge Golang API 来为图片识别函数启动和运行 WasmEdge 虚拟机。 app.go 文件在源代码项目中如下：

package main

... ...

var (
    vm      *wasmedge.VM
    vmConf  *wasmedge.Configure
    counter uint64
)

func main() {
    // Initialize WasmEdge's VM
    initVM()
    defer vm.Delete()
    defer vmConf.Delete()

    // Connect to Zipper service
    cli, err := client.NewServerless("image-recognition").Connect("localhost", 9000)
    if err != nil {
        log.Print("❌ Connect to zipper failure: ", err)
        return
    }

    defer cli.Close()
    cli.Pipe(Handler)
}

// Handler process the data in the stream
func Handler(rxStream rx.RxStream) rx.RxStream {
    stream := rxStream.
        Subscribe(ImageDataKey).
        OnObserve(decode).
        Encode(0x11)

    return stream
}

// decode Decode and perform image recognition
var decode = func(v []byte) (interface{}, error) {
    // get image binary
    p, _, _, err := y3.DecodePrimitivePacket(v)
    if err != nil {
        return nil, err
    }
    img := p.ToBytes()

    // recognize the image
    res, err := vm.ExecuteBindgen("infer", wasmedge.Bindgen_return_array, img)

    return hash, nil
}

... ...

// initVM initialize WasmEdge's VM
func initVM() {
    wasmedge.SetLogErrorLevel()
    vmConf = wasmedge.NewConfigure(wasmedge.WASI)
    vm = wasmedge.NewVMWithConfig(vmConf)

    var wasi = vm.GetImportObject(wasmedge.WASI)
    wasi.InitWasi(
        os.Args[1:],     /// The args
        os.Environ(),    /// The envs
        []string{".:."}, /// The mapping directories
        []string{},      /// The preopens will be empty
    )

    /// Register WasmEdge-tensorflow and WasmEdge-image
    var tfobj = wasmedge.NewTensorflowImportObject()
    var tfliteobj = wasmedge.NewTensorflowLiteImportObject()
    vm.RegisterImport(tfobj)
    vm.RegisterImport(tfliteobj)
    var imgobj = wasmedge.NewImageImportObject()
    vm.RegisterImport(imgobj)

    /// Instantiate wasm
    vm.LoadWasmFile("rust_mobilenet_food_lib_bg.wasm")
    vm.Validate()
    vm.Instantiate()
}

运行

最后，我们启动 YoMo 并查看整个数据处理管道的运行情况。从项目文件夹启动 YoMo CLI 应用程序。 yaml文件定义了 YoMo 应听取的端口以及触发传入数据的工作流 handler。请注意，流名称 image-recognition 与上面提到的数据 handler app.go 相匹配。

$ yomo serve -c ./zipper/workflow.yaml

通过运行上面提到的 app.go 程序启动 handler 程序。

$ cd flow
$ go run --tags "tensorflow image" app.go

通过给 YoMo 发送数据来启动模拟数据源。 视频是一系列图片帧。 app.go 中的 WasmEdge 函数将针对视频中的每个图片帧进行调用。

# Download a video file
$ wget -P source 'https://github.com/yomorun/yomo-wasmedge-tensorflow/releases/download/v0.1.0/hot-dog.mp4'

# Stream the video to YoMo
$ go run ./source/main.go ./source/hot-dog.mp4

你可以在控制台中看到 WasmEdge handler函数的输出。它将在视频的每个图片帧中检测到的对象的名称进行打印。

展望未来

本文讨论了如何使用 YoMo 框架中的 WasmEdge Tensorflow API 和 Golang SDK 来几近实时地处理图片流。

与 YoMo 合作，我们很快将在智能工厂的实际生产中部署 WasmEdge，用于各种装配线上的任务。 WasmEdge 是边缘计算的软件 runtime！