Annoyingly high idle CPU usage

[SeanPedersen]

Hey,
I built a basic server with two CLIP ONNX endpoints and I noticed that the server process has very high idle CPU usage (my system gets noticably slown down). Any idea how to fix / improve this?

Cargo.toml

[package]
name = "rust-onnx"
version = "0.1.0"
edition = "2021"

[dependencies]
ort = "=2.0.0-rc.9"
tokenizers = { version = "0.21.0", features = ["http"] }
rocket = { version = "0.5.0", features = ["json"] }
serde = { version = "1.0", features = ["derive"] }
tokio = { version = "1", features = ["full"] }
image = "0.25.5"
ndarray = "0.16.1"

main.rs

mod model;
mod server;

#[rocket::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    println!("Starting ONNX model server...");

    // Paths to models and tokenizer name
    let text_model_path = "text_model.onnx";
    let vision_model_path = "vision_model.onnx";
    let tokenizer_name = "jinaai/jina-clip-v1";

    // Run the server with both models
    server::run_server(text_model_path, vision_model_path, tokenizer_name).await
}

server.rs

use crate::model::{Embedding, Model};
use rocket::http::Status;
use rocket::{serde::json::Json, State};

#[derive(Debug, serde::Deserialize)]
pub struct TextRequest {
    text: String,
}

#[derive(Debug, serde::Deserialize)]
pub struct ImageRequest {
    image_path: String,
}

#[rocket::post("/embed_text", format = "json", data = "<request>")]
pub async fn get_text_embedding(
    model: &State<Model>,
    request: Json<TextRequest>,
) -> Result<Json<Embedding>, Status> {
    match model.get_text_embedding(&request.text) {
        Ok(embedding) => Ok(Json(embedding)),
        Err(_) => Err(Status::InternalServerError),
    }
}

#[rocket::post("/embed_image", format = "json", data = "<request>")]
pub async fn get_image_embedding(
    model: &State<Model>,
    request: Json<ImageRequest>,
) -> Result<Json<Embedding>, Status> {
    match model.get_image_embedding(&request.image_path) {
        Ok(embedding) => Ok(Json(embedding)),
        Err(_) => Err(Status::InternalServerError),
    }
}

#[rocket::get("/health")]
pub fn health_check() -> &'static str {
    "OK"
}

pub async fn run_server(
    text_model_path: &str,
    vision_model_path: &str,
    tokenizer_name: &str,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    // Load both models
    let model = Model::new(text_model_path, vision_model_path, tokenizer_name)?;

    // Configure and launch Rocket server
    let figment = rocket::Config::figment()
        .merge(("port", 8001))
        .merge(("address", "0.0.0.0"))
        .merge(("workers", 2)); // Limit worker threads

    rocket::custom(figment)
        .mount(
            "/",
            rocket::routes![get_text_embedding, get_image_embedding, health_check],
        )
        .manage(model)
        .launch()
        .await?;

    Ok(())
}

model.rs

use image::DynamicImage;
use ndarray::Array4;
use ort::session::{builder::GraphOptimizationLevel, Session};
use serde::Serialize;
use std::path::Path;
use std::sync::Arc;
use tokenizers::Tokenizer;

#[derive(Debug, Serialize)]
pub struct Embedding {
    pub values: Vec<f32>,
    pub shape: Vec<usize>,
}

pub struct Model {
    text_session: Session,
    vision_session: Session,
    tokenizer: Tokenizer,
}

impl Model {
    pub fn new(
        text_model_path: &str,
        vision_model_path: &str,
        tokenizer_name: &str,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        // Load the text ONNX model
        let text_session = Session::builder()?
            .with_optimization_level(GraphOptimizationLevel::Level3)?
            .with_intra_threads(2)?
            .with_inter_threads(2)?
            .commit_from_file(text_model_path)?;

        // Load the vision ONNX model
        let vision_session = Session::builder()?
            .with_optimization_level(GraphOptimizationLevel::Level3)?
            .with_intra_threads(2)?
            .with_inter_threads(2)?
            .commit_from_file(vision_model_path)?;

        // Load the tokenizer
        let tokenizer = Tokenizer::from_pretrained(tokenizer_name, None)?;

        Ok(Model {
            text_session,
            vision_session,
            tokenizer,
        })
    }

    pub fn get_text_embedding(
        &self,
        text: &str,
    ) -> Result<Embedding, Box<dyn std::error::Error + Send + Sync>> {
        // Tokenize input
        let input = self.tokenizer.encode(text, true)?;

        // Prepare input for ONNX inference
        let tokens = Arc::new(
            input
                .get_ids()
                .iter()
                .map(|i| *i as i64)
                .collect::<Vec<_>>()
                .into_boxed_slice(),
        );

        // Create input tensor
        let input_shape = vec![1, tokens.len() as i64];
        let model_input = (input_shape, Arc::clone(&tokens));

        // Run inference with the text model
        let outputs = self
            .text_session
            .run(ort::inputs!["input_ids" => model_input]?)?;

        // Extract embedding
        let text_embeds = outputs.get("text_embeds").unwrap();

        let mut embeds: Vec<f32> = text_embeds
            .try_extract_tensor::<f32>()
            .unwrap()
            .as_slice()
            .unwrap()
            .to_vec();

        // Normalize the embedding vector
        self.normalize_vector(&mut embeds);

        let shape = text_embeds.shape()?;

        Ok(Embedding {
            values: embeds,
            shape: shape.iter().map(|&d| d as usize).collect(),
        })
    }

    pub fn get_image_embedding(
        &self,
        image_path: &str,
    ) -> Result<Embedding, Box<dyn std::error::Error + Send + Sync>> {
        // Load the image
        let img = self.load_image(image_path)?;

        // Preprocess the image (using standard 224x224 dimensions for vision models)
        let input_tensor = self.preprocess_image(img, 224, 224)?;

        // Convert to flat vector and prepare for ONNX input
        let input_vec: Vec<f32> = input_tensor.iter().cloned().collect();
        let input_data = Arc::new(input_vec.into_boxed_slice());

        // Shape is [batch_size, channels, height, width]
        let input_shape = vec![1, 3, 224, 224];
        let model_input = (input_shape, Arc::clone(&input_data));

        // Run inference with the vision model
        let outputs = self
            .vision_session
            .run(ort::inputs!["pixel_values" => model_input]?)?;

        // Extract the embedding tensor from outputs
        let image_embeds = outputs.get("image_embeds").unwrap();

        let mut embeds: Vec<f32> = image_embeds
            .try_extract_tensor::<f32>()
            .unwrap()
            .as_slice()
            .unwrap()
            .to_vec();

        // Normalize the embedding vector
        self.normalize_vector(&mut embeds);

        let shape = image_embeds.shape()?;

        Ok(Embedding {
            values: embeds,
            shape: shape.iter().map(|&d| d as usize).collect(),
        })
    }

    // Helper method to normalize a vector to unit length (L2 normalization)
    fn normalize_vector(&self, vector: &mut Vec<f32>) {
        // Calculate the L2 norm (Euclidean length) of the vector
        let norm: f32 = vector.iter().map(|&x| x * x).sum::<f32>().sqrt();

        // Avoid division by zero
        if norm > 1e-10 {
            // Divide each component by the norm to normalize
            for val in vector.iter_mut() {
                *val /= norm;
            }
        }
    }

    fn load_image<P: AsRef<Path>>(
        &self,
        path: P,
    ) -> Result<DynamicImage, Box<dyn std::error::Error + Send + Sync>> {
        let img = image::open(path)?;
        Ok(img)
    }

    fn preprocess_image(
        &self,
        img: DynamicImage,
        width: u32,
        height: u32,
    ) -> Result<Array4<f32>, Box<dyn std::error::Error + Send + Sync>> {
        // Resize the image to the required dimensions
        let img = img.resize_exact(width, height, image::imageops::FilterType::Triangle);

        // Create a buffer to store the normalized pixel values
        let mut input_tensor = Array4::<f32>::zeros((1, 3, height as usize, width as usize));

        // Extract RGB values and normalize them
        let img_rgb = img.to_rgb8();
        for y in 0..height {
            for x in 0..width {
                let pixel = img_rgb.get_pixel(x, y);

                // Normalize to [0, 1] range and then apply specific normalization
                input_tensor[[0, 0, y as usize, x as usize]] =
                    (pixel[0] as f32 / 255.0 - 0.485) / 0.229;
                input_tensor[[0, 1, y as usize, x as usize]] =
                    (pixel[1] as f32 / 255.0 - 0.456) / 0.224;
                input_tensor[[0, 2, y as usize, x as usize]] =
                    (pixel[2] as f32 / 255.0 - 0.406) / 0.225;
            }
        }

        Ok(input_tensor)
    }
}

[decahedron1]

Are you certain the CPU usage is coming from ONNX Runtime? I believe its worker threads do occasionally spin but usually for a very short time before putting the thread to sleep (iirc that's a bit faster for high throughput scenarios). I don't get any idle CPU usage when I modify the yolov8 example to put the main thread to sleep after running the model, so I suspect ort isn't the cause of your issue here.

(btw, you should be using Session::run_async so you don't block the async runtime)

[SeanPedersen]

Thanks for the pointer, yes it seems sth else caused the high idle usage.