epwalsh / batched-fn

🦀 Rust server plugin for deploying deep learning models with batched prediction
https://crates.io/crates/batched-fn
Apache License 2.0
19 stars 2 forks source link

Awaiting inside the closure #14

Open RohanGautam opened 3 years ago

RohanGautam commented 3 years ago

Hello! thanks for this library 😄

I'm having trouble executing async tasks form within the closure. From what I can tell, the closure cannot be an async function, and the model inference that I'm doing needs to be await-ed from within a tokio runtime.

Here's what I'm trying to do :

batched_fn! {
    handler = |batch: Batch<Input>, model: &Model| -> Batch<Output> {
        let output = model.predict(batch.clone()).await; // note the await
        println!("Processed batch {:?} -> {:?}", batch, output);
        output
    };
    config = {
        max_batch_size: 4,
        max_delay: 50,
    };
    context = {
        model: Model::load(),
    };
};

Any way this could be possible? I'm willing to help if you can point me in the right direction!

epwalsh commented 3 years ago

Hi @RohanGautam, you are correct that the handler closure cannot be async. The batched_fn! macro was designed with the assumption that the handler would be blocking, and so handler is always ran in it's own thread.

I haven't tried this, but I think a valid work-around for your use-case would be to create a Tokio runtime in the context and use that runtime to execute your async predict function.