furnace-server 0.11.0

Fast, unified, Rust-native, Axum-compatible inference endpoints for Burn models and LLMs: KServe / Triton and OpenAI APIs from one pure-Rust server
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! The [`BurnModel`] implementations the crate ships: [`Model::from_fn`]'s wrapper and the
//! file-based MLP. The dummies for tests and benchmarks are in `testing`.
//!
//! [`Model::from_fn`]: super::Model::from_fn

use burn::tensor::{Device, Tensor};

use super::BurnModel;
use crate::backend::BackendKind;
use crate::burn_model::BurnModelContainer;
use crate::error::Result;
use crate::tensor::{HostTensor, IoSpec, take_inputs};

/// [`Model::from_fn`]: a forward function from one float tensor to one
pub(super) struct FnModel<F, const DI: usize, const DO: usize> {
    pub(super) name: String,
    pub(super) backend: BackendKind,
    pub(super) device: Device,
    pub(super) input: IoSpec,
    pub(super) output: IoSpec,
    pub(super) forward: F,
}

impl<F, const DI: usize, const DO: usize> std::fmt::Debug for FnModel<F, DI, DO> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FnModel")
            .field("name", &self.name)
            .field("backend", &self.backend)
            .finish_non_exhaustive()
    }
}

impl<F, const DI: usize, const DO: usize> BurnModel for FnModel<F, DI, DO>
where
    F: Fn(Tensor<DI>) -> Tensor<DO> + Send + Sync,
{
    fn name(&self) -> &str {
        &self.name
    }

    fn backend(&self) -> BackendKind {
        self.backend
    }

    fn inputs(&self) -> Vec<IoSpec> {
        vec![self.input.clone()]
    }

    fn outputs(&self) -> Vec<IoSpec> {
        vec![self.output.clone()]
    }

    fn run(&self, inputs: Vec<HostTensor>) -> Result<Vec<HostTensor>> {
        let [input] = take_inputs(inputs)?;
        let output = (self.forward)(input.into_float::<DI>(&self.device)?);
        Ok(vec![HostTensor::from_float(output)?])
    }
}

/// One `f32` tensor in, one out, for the models below
pub(super) fn single_f32_run(
    inputs: Vec<HostTensor>,
    device: &Device,
    forward: impl FnOnce(Tensor<2>) -> Result<Tensor<2>>,
    output_shape: &[usize],
) -> Result<Vec<HostTensor>> {
    let [input] = take_inputs(inputs)?;
    let mut shape = vec![input.batch()];
    shape.extend_from_slice(output_shape);
    let output = forward(input.into_float_flat(device)?)?;
    Ok(vec![HostTensor {
        shape,
        values: HostTensor::from_float(output)?.values,
    }])
}

/// File-based MLP model (`.bpk` weights + `.json` metadata)
#[derive(Debug)]
pub struct MlpModel {
    container: BurnModelContainer,
    input_shape: Vec<usize>,
    output_shape: Vec<usize>,
}

impl MlpModel {
    /// Serve a loaded MLP container, with the shapes its metadata declares
    pub fn new(container: BurnModelContainer) -> Self {
        let input_shape = container.input_shape();
        let output_shape = container.output_shape();
        Self {
            container,
            input_shape,
            output_shape,
        }
    }
}

impl BurnModel for MlpModel {
    fn name(&self) -> &str {
        &self.container.metadata.name
    }

    fn backend(&self) -> BackendKind {
        self.container.backend
    }

    fn inputs(&self) -> Vec<IoSpec> {
        vec![IoSpec::fp32("input", &self.input_shape)]
    }

    fn outputs(&self) -> Vec<IoSpec> {
        vec![IoSpec::fp32("output", &self.output_shape)]
    }

    fn run(&self, inputs: Vec<HostTensor>) -> Result<Vec<HostTensor>> {
        single_f32_run(
            inputs,
            self.container.device(),
            |x| Ok(self.container.predict(x)),
            &self.output_shape,
        )
    }
}