lmflow/
kernel_runner.rs

1//! Single-kernel test harness shared by Rust, C++ and Python hosts.
2//!
3//! Unlike [`crate::Graph`], this harness does not parse YAML or create an executor.
4//! It drives one registered kernel directly through its contract and lifecycle.
5
6use std::collections::{BTreeMap, VecDeque};
7use std::ffi::c_void;
8use std::sync::Arc;
9
10use crate::config::GraphConfig;
11use crate::context::{Context, Options};
12use crate::kernel::{Contract, KernelInstance, PortTable};
13use crate::packet::Packet;
14use crate::runtime::GraphShared;
15use crate::status::{Error, Result};
16use crate::timestamp::Timestamp;
17
18pub struct KernelRunner {
19    kernel: KernelInstance,
20    contract: Contract,
21    context: Context,
22    opened: bool,
23    open_error: Option<Error>,
24    pending_inputs: Vec<Option<Packet>>,
25    outputs: Vec<VecDeque<Packet>>,
26}
27
28impl KernelRunner {
29    pub fn new(kernel_name: &str, input_ports: usize, output_ports: usize) -> Result<Self> {
30        let inputs = Arc::new(PortTable::build(
31            &(0..input_ports)
32                .map(|i| format!("in{i}"))
33                .collect::<Vec<_>>(),
34            "kernel runner input_ports",
35        )?);
36        let outputs = Arc::new(PortTable::build(
37            &(0..output_ports)
38                .map(|i| format!("out{i}"))
39                .collect::<Vec<_>>(),
40            "kernel runner output_ports",
41        )?);
42        let mut contract = Contract::new(inputs.clone(), outputs.clone());
43        unsafe {
44            KernelInstance::fill_contract(
45                kernel_name,
46                &mut contract as *mut Contract as *mut c_void,
47            )?;
48        }
49        if let Some(error) = contract.take_error() {
50            return Err(Error::InvalidArg(format!(
51                "kernel runner contract: {error}"
52            )));
53        }
54        let kernel = KernelInstance::create(kernel_name)?;
55        let shared = Arc::new(GraphShared::new(GraphConfig::default()));
56        let context = Context::new(
57            "kernel_runner".into(),
58            kernel_name.into(),
59            inputs,
60            outputs,
61            Arc::new(Options::new(serde_yaml::Value::Mapping(
62                serde_yaml::Mapping::new(),
63            ))),
64            Arc::new(BTreeMap::new()),
65            shared,
66        );
67        Ok(Self {
68            kernel,
69            contract,
70            context,
71            opened: false,
72            open_error: None,
73            pending_inputs: (0..input_ports).map(|_| None).collect(),
74            outputs: (0..output_ports).map(|_| VecDeque::new()).collect(),
75        })
76    }
77
78    pub fn contract(&self) -> &Contract {
79        &self.contract
80    }
81
82    pub fn set_options_json(&mut self, json: &str) -> Result<()> {
83        if self.opened {
84            return Err(Error::State(
85                "kernel runner options must be set before open".into(),
86            ));
87        }
88        let value: serde_yaml::Value = serde_json::from_str(json)
89            .map_err(|error| Error::InvalidArg(format!("kernel runner options: {error}")))?;
90        self.context.options = Arc::new(Options::new(value));
91        Ok(())
92    }
93
94    pub fn set_side_packet(&mut self, name: &str, packet: Packet) -> Result<()> {
95        if self.opened {
96            return Err(Error::State(
97                "kernel runner side packets must be set before open".into(),
98            ));
99        }
100        Arc::make_mut(&mut self.context.side_packets).insert(name.into(), packet);
101        Ok(())
102    }
103
104    pub fn open(&mut self) -> Result<()> {
105        if let Some(error) = &self.open_error {
106            return Err(error.clone());
107        }
108        if self.opened {
109            return Ok(());
110        }
111        for name in &self.contract.required_side_packets {
112            if !self.context.side_packets.contains_key(name) {
113                return Err(Error::InvalidArg(format!(
114                    "kernel runner missing required side packet `{name}`"
115                )));
116            }
117        }
118        self.context.reset();
119        self.context.input_ts = Timestamp::unstarted();
120        let status = unsafe { self.kernel.open(&mut self.context as *mut _ as *mut c_void) };
121        if status != 0 {
122            let error = self.context.take_error(status);
123            self.context.discard_staging();
124            self.open_error = Some(error.clone());
125            return Err(error);
126        }
127        // A successful callback still needs Close for cleanup if its outputs
128        // fail validation. Keep that obligation separate from usable Open state.
129        self.opened = true;
130        if let Err(error) = self.validate_outputs() {
131            self.context.discard_staging();
132            self.open_error = Some(error.clone());
133            return Err(error);
134        }
135        for (output, staged) in self.outputs.iter_mut().zip(&mut self.context.staging) {
136            output.extend(staged.drain(..));
137        }
138        Ok(())
139    }
140
141    pub fn process(
142        &mut self,
143        inputs: Vec<Option<Packet>>,
144        timestamp: Timestamp,
145    ) -> Result<Vec<Vec<Packet>>> {
146        self.open()?;
147        if inputs.len() != self.context.inputs.len() {
148            return Err(Error::InvalidArg(format!(
149                "kernel runner expected {} inputs, got {}",
150                self.context.inputs.len(),
151                inputs.len()
152            )));
153        }
154        self.context.reset();
155        for (port, packet) in inputs.iter().enumerate() {
156            if let Some(packet) = packet {
157                let expected = self.contract.input_types[port];
158                if expected != 0 && packet.type_id() != expected {
159                    return Err(Error::InvalidArg(format!(
160                        "kernel runner input {port} type mismatch: expected {}, got {}",
161                        expected,
162                        packet.type_id()
163                    )));
164                }
165            }
166        }
167        self.context.inputs = inputs
168            .into_iter()
169            .map(|packet| {
170                packet.map(|packet| {
171                    if packet.timestamp() == Timestamp::unset() {
172                        packet.at(timestamp)
173                    } else {
174                        packet
175                    }
176                })
177            })
178            .collect();
179        self.context.input_ts = timestamp;
180        let status = unsafe {
181            self.kernel
182                .process(&mut self.context as *mut _ as *mut c_void)
183        };
184        // Input borrows end with the callback, including error paths. Keeping
185        // them here pins consumed payloads and forces CoW on forwarded output.
186        self.context.clear_inputs();
187        if status != 0 {
188            let error = self.context.take_error(status);
189            self.context.discard_staging();
190            return Err(error);
191        }
192        if let Err(error) = self.validate_outputs() {
193            self.context.discard_staging();
194            return Err(error);
195        }
196        let staging: Vec<Vec<Packet>> = self
197            .context
198            .staging
199            .iter_mut()
200            .map(std::mem::take)
201            .collect();
202        for (port, packets) in staging.iter().enumerate() {
203            self.outputs[port].extend(packets.iter().cloned());
204        }
205        Ok(staging)
206    }
207
208    pub fn add_input(&mut self, port: usize, packet: Packet) -> Result<()> {
209        let slot = self
210            .pending_inputs
211            .get_mut(port)
212            .ok_or_else(|| Error::InvalidArg(format!("input port index {port} is out of range")))?;
213        if slot.is_some() {
214            return Err(Error::State(format!(
215                "kernel runner input port {port} already has a packet for the next process call"
216            )));
217        }
218        *slot = Some(packet);
219        Ok(())
220    }
221
222    pub fn process_pending(&mut self, timestamp: Timestamp) -> Result<Vec<Vec<Packet>>> {
223        let inputs = std::mem::take(&mut self.pending_inputs);
224        let result = self.process(inputs, timestamp);
225        self.pending_inputs = (0..self.context.inputs.len()).map(|_| None).collect();
226        result
227    }
228
229    pub fn try_output(&mut self, port: usize) -> Result<Option<Packet>> {
230        self.outputs
231            .get_mut(port)
232            .ok_or_else(|| Error::InvalidArg(format!("output port index {port} is out of range")))
233            .map(VecDeque::pop_front)
234    }
235
236    /// Close with the same `Done` activation timestamp used by Graph, collecting
237    /// any tail output carrying an explicit legal stream timestamp.
238    /// Pending, unprocessed inputs are discarded; completed output remains readable.
239    pub fn close(&mut self) -> Result<Vec<Vec<Packet>>> {
240        // Discard inputs for an invocation that never ran, including before
241        // Open and on Close errors. Preserve slots for the next session.
242        self.pending_inputs.fill(None);
243        if !self.opened {
244            return Ok((0..self.context.staging.len())
245                .map(|_| Vec::new())
246                .collect());
247        }
248        self.context.reset();
249        self.context.close_reason = if self.open_error.is_some() {
250            crate::runtime::CLOSE_ERROR
251        } else {
252            crate::runtime::CLOSE_NORMAL
253        };
254        let status = unsafe {
255            self.context.input_ts = Timestamp::done();
256            self.kernel
257                .close(&mut self.context as *mut _ as *mut c_void)
258        };
259        self.opened = false;
260        if let Some(error) = &self.open_error {
261            self.context.discard_staging();
262            return Err(error.clone());
263        }
264        if status != 0 {
265            let error = self.context.take_error(status);
266            self.context.discard_staging();
267            return Err(error);
268        }
269        if let Err(error) = self.validate_outputs() {
270            self.context.discard_staging();
271            return Err(error);
272        }
273        let staging: Vec<Vec<Packet>> = self
274            .context
275            .staging
276            .iter_mut()
277            .map(std::mem::take)
278            .collect();
279        for (port, packets) in staging.iter().enumerate() {
280            self.outputs[port].extend(packets.iter().cloned());
281        }
282        Ok(staging)
283    }
284
285    fn validate_outputs(&self) -> Result<()> {
286        for (port, packets) in self.context.staging.iter().enumerate() {
287            let expected = self.contract.output_types[port];
288            if expected != 0 {
289                if let Some(packet) = packets.iter().find(|packet| packet.type_id() != expected) {
290                    return Err(Error::Kernel(format!(
291                        "kernel runner output {port} type mismatch: expected {}, got {}",
292                        expected,
293                        packet.type_id()
294                    )));
295                }
296            }
297        }
298        Ok(())
299    }
300}
301
302impl Drop for KernelRunner {
303    fn drop(&mut self) {
304        let _ = self.close();
305    }
306}
lmflow/kernel_runner.rs

lmflow/
kernel_runner.rs
