rill_runtime/
server.rs

1use std::sync::Arc;
2
3use rill_runtime_protocol::{RUNTIME_API_VERSION, RuntimeRequest, RuntimeResponse};
4use serde::Deserialize;
5use serde_json::Value;
6
7use crate::package::LoadedModelPack;
8
9pub const LINEAR_REGRESSION_CAPABILITY: &str = "rillml.linearRegression.predict";
10
11/// Consumers can implement this trait to add business-specific invocation logic.
12pub trait InvokeHandler: Send + Sync + std::fmt::Debug {
13    fn invoke(&self, capability: &str, input: &Value) -> Result<Value, String>;
14}
15
16/// Business-neutral linear-regression handler used by the distributed runtime binary.
17#[derive(Debug, Clone)]
18pub struct LinearRegressionInvokeHandler {
19    weights: Vec<f64>,
20    intercept: f64,
21}
22
23#[derive(Debug, Deserialize)]
24#[serde(rename_all = "camelCase", deny_unknown_fields)]
25struct LinearRegressionModel {
26    kind: String,
27    weights: Vec<f64>,
28    intercept: f64,
29}
30
31#[derive(Debug, Deserialize)]
32#[serde(deny_unknown_fields)]
33struct LinearRegressionInput {
34    features: Vec<f64>,
35}
36
37impl LinearRegressionInvokeHandler {
38    pub fn from_pack(pack: &LoadedModelPack) -> Result<Self, String> {
39        if pack.manifest.capabilities.as_slice() != [LINEAR_REGRESSION_CAPABILITY] {
40            return Err(format!(
41                "standalone runtime requires exactly the {LINEAR_REGRESSION_CAPABILITY} capability"
42            ));
43        }
44        let model: LinearRegressionModel = serde_json::from_value(pack.model.clone())
45            .map_err(|error| format!("invalid linear-regression model: {error}"))?;
46        if model.kind != "linearRegression" {
47            return Err("unsupported built-in model kind".into());
48        }
49        if model.weights.is_empty() || model.weights.len() > 65_536 {
50            return Err("linear-regression weights must contain 1..=65536 values".into());
51        }
52        if !model.intercept.is_finite() || model.weights.iter().any(|value| !value.is_finite()) {
53            return Err("linear-regression model values must be finite".into());
54        }
55        Ok(Self {
56            weights: model.weights,
57            intercept: model.intercept,
58        })
59    }
60}
61
62impl InvokeHandler for LinearRegressionInvokeHandler {
63    fn invoke(&self, capability: &str, input: &Value) -> Result<Value, String> {
64        if capability != LINEAR_REGRESSION_CAPABILITY {
65            return Err("unsupported capability".into());
66        }
67        let input: LinearRegressionInput = serde_json::from_value(input.clone())
68            .map_err(|error| format!("invalid linear-regression input: {error}"))?;
69        if input.features.len() != self.weights.len() {
70            return Err(format!(
71                "expected {} features, received {}",
72                self.weights.len(),
73                input.features.len()
74            ));
75        }
76        if input.features.iter().any(|value| !value.is_finite()) {
77            return Err("linear-regression input values must be finite".into());
78        }
79        let prediction = self
80            .weights
81            .iter()
82            .zip(&input.features)
83            .try_fold(self.intercept, |sum, (weight, feature)| {
84                let next = sum + weight * feature;
85                next.is_finite().then_some(next)
86            })
87            .ok_or_else(|| "linear-regression prediction overflowed".to_string())?;
88        Ok(serde_json::json!({ "prediction": prediction }))
89    }
90}
91
92#[derive(Debug, Clone)]
93pub struct RuntimeEngine {
94    pack: LoadedModelPack,
95    invoke_handler: Option<Arc<dyn InvokeHandler>>,
96}
97
98impl RuntimeEngine {
99    pub fn new(pack: LoadedModelPack) -> Self {
100        Self {
101            pack,
102            invoke_handler: None,
103        }
104    }
105
106    pub fn with_invoke_handler(mut self, handler: Arc<dyn InvokeHandler>) -> Self {
107        self.invoke_handler = Some(handler);
108        self
109    }
110
111    pub fn handle(&self, request: RuntimeRequest) -> RuntimeResponse {
112        let request_id = request.request_id().to_string();
113        if request_id.is_empty() || request_id.len() > 128 {
114            return self.error(request_id, "invalidRequestId", "invalid request id", false);
115        }
116        if request.api_version() != RUNTIME_API_VERSION {
117            return self.error(
118                request_id,
119                "incompatibleApiVersion",
120                "runtime API version is not supported",
121                false,
122            );
123        }
124
125        match request {
126            RuntimeRequest::Handshake {
127                request_id,
128                client_name,
129                client_version,
130                ..
131            } => {
132                if client_name.is_empty()
133                    || client_name.len() > 96
134                    || client_version.is_empty()
135                    || client_version.len() > 48
136                {
137                    return self.error(
138                        request_id,
139                        "invalidClientIdentity",
140                        "invalid client identity",
141                        false,
142                    );
143                }
144                RuntimeResponse::Handshake {
145                    request_id,
146                    api_version: RUNTIME_API_VERSION,
147                    runtime_version: env!("CARGO_PKG_VERSION").into(),
148                    model_pack_id: self.pack.manifest.id.clone(),
149                    model_pack_version: self.pack.manifest.version.clone(),
150                    capabilities: self.pack.manifest.capabilities.clone(),
151                }
152            }
153            RuntimeRequest::Health { request_id, .. } => RuntimeResponse::Health {
154                request_id,
155                api_version: RUNTIME_API_VERSION,
156                healthy: true,
157                model_pack_id: self.pack.manifest.id.clone(),
158                model_pack_version: self.pack.manifest.version.clone(),
159            },
160            RuntimeRequest::Invoke {
161                request_id,
162                capability,
163                input,
164                ..
165            } => {
166                if !self
167                    .pack
168                    .manifest
169                    .capabilities
170                    .iter()
171                    .any(|declared| declared == &capability)
172                {
173                    return self.error(
174                        request_id,
175                        "unsupportedCapability",
176                        "capability is not declared by the loaded model pack",
177                        false,
178                    );
179                }
180                match &self.invoke_handler {
181                    Some(handler) => match handler.invoke(&capability, &input) {
182                        Ok(output) => RuntimeResponse::Result {
183                            request_id,
184                            api_version: RUNTIME_API_VERSION,
185                            output,
186                        },
187                        Err(message) => self.error(request_id, "invokeFailed", &message, false),
188                    },
189                    None => self.error(
190                        request_id,
191                        "noInvokeHandler",
192                        "no invoke handler registered",
193                        false,
194                    ),
195                }
196            }
197        }
198    }
199
200    fn error(
201        &self,
202        request_id: String,
203        code: &str,
204        message: &str,
205        retryable: bool,
206    ) -> RuntimeResponse {
207        RuntimeResponse::Error {
208            request_id,
209            api_version: RUNTIME_API_VERSION,
210            code: code.into(),
211            message: message.into(),
212            retryable,
213        }
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use rill_runtime_protocol::{MODEL_PACK_FORMAT_VERSION, ModelPackManifest};
220
221    use super::*;
222
223    fn engine() -> RuntimeEngine {
224        RuntimeEngine::new(LoadedModelPack {
225            manifest: ModelPackManifest {
226                format_version: MODEL_PACK_FORMAT_VERSION,
227                id: "rillml.example.default".into(),
228                version: "0.5.0".into(),
229                runtime_api_version: RUNTIME_API_VERSION,
230                min_runtime_version: "0.5.0".into(),
231                publisher_key_id: "test".into(),
232                capabilities: vec!["rillml.example".into()],
233            },
234            model: serde_json::json!({}),
235        })
236    }
237
238    #[test]
239    fn handshake_reports_loaded_pack() {
240        let response = engine().handle(RuntimeRequest::Handshake {
241            request_id: "hello".into(),
242            api_version: RUNTIME_API_VERSION,
243            client_name: "example-host".into(),
244            client_version: "0.9.0".into(),
245        });
246        assert!(matches!(
247            response,
248            RuntimeResponse::Handshake { model_pack_id, .. }
249                if model_pack_id == "rillml.example.default"
250        ));
251    }
252
253    #[test]
254    fn incompatible_api_is_a_typed_error() {
255        let response = engine().handle(RuntimeRequest::Health {
256            request_id: "health".into(),
257            api_version: RUNTIME_API_VERSION + 1,
258        });
259        assert!(matches!(
260            response,
261            RuntimeResponse::Error { code, .. } if code == "incompatibleApiVersion"
262        ));
263    }
264
265    #[test]
266    fn invoke_without_handler_returns_no_invoke_handler_error() {
267        let response = engine().handle(RuntimeRequest::Invoke {
268            request_id: "invoke-1".into(),
269            api_version: RUNTIME_API_VERSION,
270            capability: "rillml.example".into(),
271            input: serde_json::json!({}),
272        });
273        assert!(matches!(
274            response,
275            RuntimeResponse::Error { code, .. } if code == "noInvokeHandler"
276        ));
277    }
278
279    #[test]
280    fn invoke_rejects_capability_not_declared_by_signed_manifest() {
281        let response = engine().handle(RuntimeRequest::Invoke {
282            request_id: "invoke-undeclared".into(),
283            api_version: RUNTIME_API_VERSION,
284            capability: "undeclared.capability".into(),
285            input: serde_json::json!({}),
286        });
287        assert!(matches!(
288            response,
289            RuntimeResponse::Error { code, .. } if code == "unsupportedCapability"
290        ));
291    }
292
293    #[test]
294    fn linear_regression_handler_validates_and_predicts() {
295        let pack = LoadedModelPack {
296            manifest: ModelPackManifest {
297                format_version: MODEL_PACK_FORMAT_VERSION,
298                id: "rillml.example.default".into(),
299                version: "0.5.1".into(),
300                runtime_api_version: RUNTIME_API_VERSION,
301                min_runtime_version: "0.5.1".into(),
302                publisher_key_id: "test".into(),
303                capabilities: vec![LINEAR_REGRESSION_CAPABILITY.into()],
304            },
305            model: serde_json::json!({
306                "kind": "linearRegression",
307                "weights": [0.5, -0.25],
308                "intercept": 1.0
309            }),
310        };
311        let handler = LinearRegressionInvokeHandler::from_pack(&pack).unwrap();
312        let engine = RuntimeEngine::new(pack).with_invoke_handler(Arc::new(handler));
313        let response = engine.handle(RuntimeRequest::Invoke {
314            request_id: "invoke-linear".into(),
315            api_version: RUNTIME_API_VERSION,
316            capability: LINEAR_REGRESSION_CAPABILITY.into(),
317            input: serde_json::json!({"features": [4.0, 2.0]}),
318        });
319        assert!(matches!(
320            response,
321            RuntimeResponse::Result { output, .. } if output["prediction"] == 2.5
322        ));
323    }
324
325    #[test]
326    fn linear_regression_handler_rejects_invalid_model() {
327        let mut pack = LoadedModelPack {
328            manifest: ModelPackManifest {
329                format_version: MODEL_PACK_FORMAT_VERSION,
330                id: "rillml.example.default".into(),
331                version: "0.5.1".into(),
332                runtime_api_version: RUNTIME_API_VERSION,
333                min_runtime_version: "0.5.1".into(),
334                publisher_key_id: "test".into(),
335                capabilities: vec![LINEAR_REGRESSION_CAPABILITY.into()],
336            },
337            model: serde_json::json!({
338                "kind": "linearRegression",
339                "weights": [],
340                "intercept": 0.0
341            }),
342        };
343        assert!(LinearRegressionInvokeHandler::from_pack(&pack).is_err());
344
345        pack.model = serde_json::json!({
346            "kind": "unknown",
347            "weights": [1.0],
348            "intercept": 0.0
349        });
350        assert!(LinearRegressionInvokeHandler::from_pack(&pack).is_err());
351    }
352}
rill_runtime/server.rs

rill_runtime/
server.rs
