kinfer 0.4.2__cp311-cp311-macosx_11_0_arm64.whl → 0.5.1__cp311-cp311-macosx_11_0_arm64.whl

kinfer/export/__init__.py

@@ -0,0 +1,3 @@
+"""Defines the export API."""
+
+from .serialize import *

kinfer/export/jax.py

@@ -10,20 +10,27 @@ from jax._src.stages import Wrapped
 from jax.experimental import jax2tf
 from onnx.onnx_pb import ModelProto
 
-from kinfer.export.common import get_shape
+from kinfer.rust_bindings import PyInputType, PyModelMetadata
 
 logger = logging.getLogger(__name__)
 
 
 def export_fn(
     model: Wrapped,
+    metadata: PyModelMetadata,
     *,
-    num_joints: int | None = None,
-    num_commands: int | None = None,
-    carry_shape: tuple[int, ...] | None = None,
     opset: int = 13,
 ) -> ModelProto:
-    """Export a JAX function to ONNX."""
+    """Export a JAX function to ONNX.
+
+    Args:
+        model: The model to export.
+        metadata: The metadata for the model.
+        opset: The ONNX opset to use.
+
+    Returns:
+        The ONNX model as a `ModelProto`.
+    """
     if not isinstance(model, Wrapped):
         raise ValueError("Model must be a Wrapped function")
 
@@ -33,12 +40,7 @@ def export_fn(
     # Gets the dummy input tensors for exporting the model.
     tf_args = []
     for name in input_names:
-        shape = get_shape(
-            name,
-            num_joints=num_joints,
-            num_commands=num_commands,
-            carry_shape=carry_shape,
-        )
+        shape = PyInputType(name).get_shape(metadata)
         tf_args.append(tf.TensorSpec(shape, tf.float32, name=name))
 
     finalised_fn = finalise_fn(model)

kinfer/export/pytorch.py

@@ -12,23 +12,18 @@ import torch
 from onnx.onnx_pb import ModelProto
 from torch._C import FunctionSchema
 
-from kinfer.export.common import get_shape
+from kinfer.rust_bindings import PyInputType, PyModelMetadata
 
 
 def export_fn(
     model: torch.jit.ScriptFunction,
-    *,
-    num_joints: int | None = None,
-    num_commands: int | None = None,
-    carry_shape: tuple[int, ...] | None = None,
+    metadata: PyModelMetadata,
 ) -> ModelProto:
     """Exports a PyTorch function to ONNX.
 
     Args:
         model: The model to export.
-        num_joints: The number of joints in the model.
-        num_commands: The number of commands in the model.
-        carry_shape: The shape of the carry tensor.
+        metadata: The metadata for the model.
 
     Returns:
         The ONNX model as a `ModelProto`.
@@ -42,12 +37,7 @@ def export_fn(
     # Gets the dummy input tensors for exporting the model.
     args = []
     for name in input_names:
-        shape = get_shape(
-            name,
-            num_joints=num_joints,
-            num_commands=num_commands,
-            carry_shape=carry_shape,
-        )
+        shape = PyInputType(name).get_shape(metadata)
         args.append(torch.zeros(shape))
 
     buffer = io.BytesIO()

kinfer/export/serialize.py

@@ -4,34 +4,30 @@ __all__ = [
     "pack",
 ]
 
-
 import io
+import logging
 import tarfile
 
 from onnx.onnx_pb import ModelProto
 
-from kinfer.common.types import Metadata
-from kinfer.export.common import get_shape
+from kinfer.rust_bindings import PyInputType, PyModelMetadata
+
+logger = logging.getLogger(__name__)
 
 
 def pack(
     init_fn: ModelProto,
     step_fn: ModelProto,
-    joint_names: list[str],
-    num_commands: int | None = None,
-    carry_shape: tuple[int, ...] | None = None,
+    metadata: PyModelMetadata,
 ) -> bytes:
     """Packs the initialization function and step function into a directory.
 
     Args:
         init_fn: The initialization function.
         step_fn: The step function.
-        joint_names: The list of joint names, in the order that the model
-            expects them to be provided.
-        num_commands: The number of commands in the model.
-        carry_shape: The shape of the carry tensor.
+        metadata: The metadata for the model.
     """
-    num_joints = len(joint_names)
+    num_joints = len(metadata.joint_names)  # type: ignore[attr-defined]
 
     # Checks the `init` function.
     if len(init_fn.graph.input) > 0:
@@ -40,20 +36,21 @@ def pack(
         raise ValueError(f"`init` function should have exactly 1 output! Got {len(init_fn.graph.output)}")
     init_carry = init_fn.graph.output[0]
     init_carry_shape = tuple(dim.dim_value for dim in init_carry.type.tensor_type.shape.dim)
-    if carry_shape is not None and init_carry_shape != carry_shape:
-        raise ValueError(f"Expected carry shape {carry_shape} for output `{init_carry.name}`, got {init_carry_shape}")
+
+    if metadata.carry_size != init_carry_shape:  # type: ignore[attr-defined]
+        logger.warning(
+            "Updating carry size from %s to %s to match the `init` function",
+            metadata.carry_size,  # type: ignore[attr-defined]
+            init_carry_shape,
+        )
+        metadata.carry_size = init_carry_shape  # type: ignore[attr-defined]
 
     # Checks the `step` function.
     for step_input in step_fn.graph.input:
         step_input_type = step_input.type.tensor_type
         shape = tuple(dim.dim_value for dim in step_input_type.shape.dim)
-        expected_shape = get_shape(
-            step_input.name,
-            num_joints=num_joints,
-            num_commands=num_commands,
-            carry_shape=carry_shape,
-        )
-        if shape != expected_shape:
+        expected_shape = PyInputType(step_input.name).get_shape(metadata)
+        if shape != tuple(expected_shape):
             raise ValueError(f"Expected shape {expected_shape} for input `{step_input.name}`, got {shape}")
 
     if len(step_fn.graph.output) != 2:
@@ -69,12 +66,6 @@ def pack(
     if output_carry_shape != init_carry_shape:
         raise ValueError(f"Expected carry shape {init_carry_shape} for output carry, got {output_carry_shape}")
 
-    # Builds the metadata object.
-    metadata = Metadata(
-        joint_names=joint_names,
-        num_commands=num_commands,
-    )
-
     buffer = io.BytesIO()
 
     with tarfile.open(fileobj=buffer, mode="w:gz") as tar:
@@ -86,7 +77,7 @@ def pack(
 
         add_file_bytes("init_fn.onnx", init_fn.SerializeToString())
         add_file_bytes("step_fn.onnx", step_fn.SerializeToString())
-        add_file_bytes("metadata.json", metadata.model_dump_json().encode("utf-8"))
+        add_file_bytes("metadata.json", metadata.to_json().encode("utf-8"))
 
     buffer.seek(0)
 

kinfer/requirements.txt

@@ -6,3 +6,8 @@ onnxruntime==1.20.0
 
 # Serialization
 pydantic
+
+# Plotting
+matplotlib
+numpy
+pathlib

kinfer/rust/Cargo.toml

@@ -17,8 +17,11 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 
 async-trait = "0.1"
+chrono = "0.4.41"
+crossbeam-channel = "0.5.15"
 flate2 = "1.0"
 futures-util = "0.3.30"
+log = "0.4"
 ndarray = "0.16.1"
 ort = { version = "2.0.0-rc.9", features = [ "load-dynamic" ] }
 serde = { version = "1.0", features = ["derive"] }

kinfer/rust/src/lib.rs

@@ -1,5 +1,8 @@
+pub mod logger;
 pub mod model;
 pub mod runtime;
+pub mod types;
 
 pub use model::*;
 pub use runtime::*;
+pub use types::*;

kinfer/rust/src/logger.rs

@@ -0,0 +1,135 @@
+use std::{
+    fs::OpenOptions,
+    io::{BufWriter, Write},
+    path::{Path, PathBuf},
+    sync::atomic::{AtomicU64, Ordering},
+    thread,
+};
+
+use crossbeam_channel::{bounded, Sender};
+use log::{info, warn};
+use serde::Serialize;
+
+#[derive(Serialize)]
+struct NdjsonStep<'a> {
+    step_id: u64,
+    t_us: u64,
+    joint_angles: Option<&'a [f32]>,
+    joint_vels: Option<&'a [f32]>,
+    projected_g: Option<&'a [f32]>,
+    accel: Option<&'a [f32]>,
+    gyro: Option<&'a [f32]>,
+    command: Option<&'a [f32]>,
+    output: Option<&'a [f32]>,
+}
+
+// Channel capacity for non-blocking logging.
+// ~1000 entires at 50Hz is 20 seconds of buffering.
+// Warns if messages are dropped due to full buffer.
+const CHANNEL_CAP: usize = 1024;
+
+// Flush buffered writes every 100 log entries.
+// At 50Hz control frequency, this flushes every 2 seconds.
+const FLUSH_EVERY: u64 = 100;
+
+pub struct StepLogger {
+    tx: Option<Sender<Vec<u8>>>,
+    worker: Option<thread::JoinHandle<()>>,
+    next_id: AtomicU64,
+}
+
+impl StepLogger {
+    pub fn new(path: impl AsRef<Path>) -> std::io::Result<Self> {
+        let path: PathBuf = path.as_ref().into();
+        if let Some(parent) = path.parent() {
+            std::fs::create_dir_all(parent)?;
+        }
+        info!("kinfer: logging to NDJSON: {}", path.display());
+
+        // I/O objects created here, but moved into the worker thread.
+        let file = OpenOptions::new().create(true).append(true).open(&path)?;
+        let mut bw = BufWriter::new(file);
+
+        // Bounded channel -> back-pressure capped at CHANNEL_CAP lines
+        let (tx, rx) = bounded::<Vec<u8>>(CHANNEL_CAP);
+
+        let worker = thread::spawn(move || {
+            let mut line_ctr: u64 = 0;
+            for msg in rx {
+                // drains until all senders dropped
+                let _ = bw.write_all(&msg);
+                line_ctr += 1;
+                if line_ctr % FLUSH_EVERY == 0 {
+                    let _ = bw.flush();
+                }
+            }
+            // Final flush on graceful shutdown
+            let _ = bw.flush();
+        });
+
+        Ok(Self {
+            tx: Some(tx),
+            worker: Some(worker),
+            next_id: AtomicU64::new(0),
+        })
+    }
+
+    #[inline]
+    fn now_us() -> u128 {
+        std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_micros()
+    }
+
+    /// Non-blocking; drops a line if the channel is full.
+    #[allow(clippy::too_many_arguments)]
+    pub fn log_step(
+        &self,
+        joint_angles: Option<&[f32]>,
+        joint_vels: Option<&[f32]>,
+        projected_g: Option<&[f32]>,
+        accel: Option<&[f32]>,
+        gyro: Option<&[f32]>,
+        command: Option<&[f32]>,
+        output: Option<&[f32]>,
+    ) {
+        let record = NdjsonStep {
+            step_id: self.next_id.fetch_add(1, Ordering::Relaxed),
+            t_us: Self::now_us() as u64,
+            joint_angles,
+            joint_vels,
+            projected_g,
+            accel,
+            gyro,
+            command,
+            output,
+        };
+
+        // Serialise directly into a Vec<u8>; then push newline and send.
+        if let Ok(mut line) = serde_json::to_vec(&record) {
+            line.push(b'\n');
+            if let Some(tx) = &self.tx {
+                if tx.try_send(line).is_err() {
+                    warn!(
+                        "kinfer: logging buffer full, dropped message (step_id: {})",
+                        record.step_id
+                    );
+                }
+            }
+        }
+    }
+}
+
+/// Ensure the worker drains and flushes before program exit.
+impl Drop for StepLogger {
+    fn drop(&mut self) {
+        if let Some(tx) = self.tx.take() {
+            drop(tx); // Drop sender to close channel
+        }
+        // Wait for worker to finish
+        if let Some(worker) = self.worker.take() {
+            let _ = worker.join();
+        }
+    }
+}

kinfer/rust/src/model.rs

@@ -1,10 +1,11 @@
+use crate::logger::StepLogger;
+use crate::types::{InputType, ModelMetadata};
 use async_trait::async_trait;
+use chrono;
 use flate2::read::GzDecoder;
-use futures_util::future;
 use ndarray::{Array, IxDyn};
 use ort::session::Session;
 use ort::value::Value;
-use serde::Deserialize;
 use std::collections::HashMap;
 use std::io::Read;
 use std::path::Path;
@@ -13,18 +14,6 @@ use tar::Archive;
 use tokio::fs::File;
 use tokio::io::AsyncReadExt;
 
-#[derive(Debug, Deserialize)]
-struct ModelMetadata {
-    joint_names: Vec<String>,
-    num_commands: Option<usize>,
-}
-
-impl ModelMetadata {
-    fn model_validate_json(json: String) -> Result<Self, Box<dyn std::error::Error>> {
-        Ok(serde_json::from_str(&json)?)
-    }
-}
-
 #[derive(Debug, thiserror::Error)]
 pub enum ModelError {
     #[error("IO error: {0}")]
@@ -35,23 +24,16 @@ pub enum ModelError {
 
 #[async_trait]
 pub trait ModelProvider: Send + Sync {
-    async fn get_joint_angles(
-        &self,
-        joint_names: &[String],
-    ) -> Result<Array<f32, IxDyn>, ModelError>;
-    async fn get_joint_angular_velocities(
+    async fn get_inputs(
         &self,
-        joint_names: &[String],
-    ) -> Result<Array<f32, IxDyn>, ModelError>;
-    async fn get_projected_gravity(&self) -> Result<Array<f32, IxDyn>, ModelError>;
-    async fn get_accelerometer(&self) -> Result<Array<f32, IxDyn>, ModelError>;
-    async fn get_gyroscope(&self) -> Result<Array<f32, IxDyn>, ModelError>;
-    async fn get_command(&self) -> Result<Array<f32, IxDyn>, ModelError>;
-    async fn get_carry(&self, carry: Array<f32, IxDyn>) -> Result<Array<f32, IxDyn>, ModelError>;
+        input_types: &[InputType],
+        metadata: &ModelMetadata,
+    ) -> Result<HashMap<InputType, Array<f32, IxDyn>>, ModelError>;
+
     async fn take_action(
         &self,
-        joint_names: Vec<String>,
         action: Array<f32, IxDyn>,
+        metadata: &ModelMetadata,
     ) -> Result<(), ModelError>;
 }
 
@@ -60,6 +42,7 @@ pub struct ModelRunner {
     step_session: Session,
     metadata: ModelMetadata,
     provider: Arc<dyn ModelProvider>,
+    logger: Option<Arc<StepLogger>>,
 }
 
 impl ModelRunner {
@@ -138,11 +121,29 @@ impl ModelRunner {
         // Validate step_fn inputs and outputs
         Self::validate_step_fn(&step_session, &metadata, &carry_shape)?;
 
+        let logger = if let Ok(log_dir) = std::env::var("KINFER_LOG_PATH") {
+            let log_dir_path = std::path::Path::new(&log_dir);
+
+            // Create the directory if it doesn't exist
+            if !log_dir_path.exists() {
+                std::fs::create_dir_all(log_dir_path)?;
+            }
+
+            // Generate a timestamped filename
+            let timestamp = chrono::Utc::now().format("%Y-%m-%d_%H-%M-%S").to_string();
+            let log_file_path = log_dir_path.join(format!("{}.ndjson", timestamp));
+
+            Some(StepLogger::new(log_file_path).map(Arc::new)?)
+        } else {
+            None
+        };
+
         Ok(Self {
             init_session,
             step_session,
             metadata,
             provider: input_provider,
+            logger,
         })
     }
 
@@ -158,46 +159,15 @@ impl ModelRunner {
                 input.name
             ))?;
 
-            match input.name.as_str() {
-                "joint_angles" | "joint_angular_velocities" => {
-                    let num_joints = metadata.joint_names.len();
-                    if *dims != vec![num_joints as i64] {
-                        return Err(format!(
-                            "Expected shape [{num_joints}] for input `{}`, got {:?}",
-                            input.name, dims
-                        )
-                        .into());
-                    }
-                }
-                "projected_gravity" | "accelerometer" | "gyroscope" => {
-                    if *dims != vec![3] {
-                        return Err(format!(
-                            "Expected shape [3] for input `{}`, got {:?}",
-                            input.name, dims
-                        )
-                        .into());
-                    }
-                }
-                "command" => {
-                    let num_commands = metadata.num_commands.ok_or("num_commands is not set")?;
-                    if *dims != vec![num_commands as i64] {
-                        return Err(format!(
-                            "Expected shape [{num_commands}] for input `{}`, got {:?}",
-                            input.name, dims
-                        )
-                        .into());
-                    }
-                }
-                "carry" => {
-                    if dims != carry_shape {
-                        return Err(format!(
-                            "Expected shape {:?} for input `carry`, got {:?}",
-                            carry_shape, dims
-                        )
-                        .into());
-                    }
-                }
-                _ => return Err(format!("Unknown input name: {}", input.name).into()),
+            let input_type = InputType::from_name(&input.name)?;
+            let expected_shape = input_type.get_shape(metadata);
+            let expected_shape_i64: Vec<i64> = expected_shape.iter().map(|&x| x as i64).collect();
+            if *dims != expected_shape_i64 {
+                return Err(format!(
+                    "Expected input shape {:?}, got {:?}",
+                    expected_shape_i64, dims
+                )
+                .into());
             }
         }
 
@@ -234,6 +204,13 @@ impl ModelRunner {
         Ok(())
     }
 
+    pub async fn get_inputs(
+        &self,
+        input_types: &[InputType],
+    ) -> Result<HashMap<InputType, Array<f32, IxDyn>>, ModelError> {
+        self.provider.get_inputs(input_types, &self.metadata).await
+    }
+
     pub async fn init(&self) -> Result<Array<f32, IxDyn>, Box<dyn std::error::Error>> {
         let input_values: Vec<(&str, Value)> = Vec::new();
         let outputs = self.init_session.run(input_values)?;
@@ -254,36 +231,53 @@ impl ModelRunner {
             .collect();
 
         // Calls the relevant getter methods in parallel.
-        let mut futures = Vec::new();
+        let mut input_types = Vec::new();
+        let mut inputs = HashMap::new();
         for name in &input_names {
             match name.as_str() {
                 "joint_angles" => {
-                    futures.push(self.provider.get_joint_angles(&self.metadata.joint_names))
+                    input_types.push(InputType::JointAngles);
+                }
+                "joint_angular_velocities" => {
+                    input_types.push(InputType::JointAngularVelocities);
+                }
+                "projected_gravity" => {
+                    input_types.push(InputType::ProjectedGravity);
+                }
+                "accelerometer" => {
+                    input_types.push(InputType::Accelerometer);
+                }
+                "gyroscope" => {
+                    input_types.push(InputType::Gyroscope);
+                }
+                "command" => {
+                    input_types.push(InputType::Command);
+                }
+                "time" => {
+                    input_types.push(InputType::Time);
+                }
+                "carry" => {
+                    inputs.insert(InputType::Carry, carry.clone());
                 }
-                "joint_angular_velocities" => futures.push(
-                    self.provider
-                        .get_joint_angular_velocities(&self.metadata.joint_names),
-                ),
-                "projected_gravity" => futures.push(self.provider.get_projected_gravity()),
-                "accelerometer" => futures.push(self.provider.get_accelerometer()),
-                "gyroscope" => futures.push(self.provider.get_gyroscope()),
-                "command" => futures.push(self.provider.get_command()),
-                "carry" => futures.push(self.provider.get_carry(carry.clone())),
                 _ => return Err(format!("Unknown input name: {}", name).into()),
             }
         }
 
-        let results = future::try_join_all(futures).await?;
-        let mut inputs = HashMap::new();
-        for (name, value) in input_names.iter().zip(results) {
-            inputs.insert(name.clone(), value);
-        }
+        // Gets the input values.
+        let result = self
+            .provider
+            .get_inputs(&input_types, &self.metadata)
+            .await?;
+
+        // Adds the input values to the input map.
+        inputs.extend(result);
 
         // Convert inputs to ONNX values
         let mut input_values: Vec<(&str, Value)> = Vec::new();
         for input in &self.step_session.inputs {
+            let input_type = InputType::from_name(&input.name)?;
             let input_data = inputs
-                .get(&input.name)
+                .get(&input_type)
                 .ok_or_else(|| format!("Missing input: {}", input.name))?;
             let input_value = Value::from_array(input_data.view())?.into_dyn();
             input_values.push((input.name.as_str(), input_value));
@@ -294,6 +288,39 @@ impl ModelRunner {
         let output_tensor = outputs[0].try_extract_tensor::<f32>()?;
         let carry_tensor = outputs[1].try_extract_tensor::<f32>()?;
 
+        // Log the step if needed
+        if let Some(lg) = &self.logger {
+            let joint_angles_opt = inputs
+                .get(&InputType::JointAngles)
+                .map(|a| a.as_slice().unwrap());
+            let joint_vels_opt = inputs
+                .get(&InputType::JointAngularVelocities)
+                .map(|a| a.as_slice().unwrap());
+            let projected_g_opt = inputs
+                .get(&InputType::ProjectedGravity)
+                .map(|a| a.as_slice().unwrap());
+            let accel_opt = inputs
+                .get(&InputType::Accelerometer)
+                .map(|a| a.as_slice().unwrap());
+            let gyro_opt = inputs
+                .get(&InputType::Gyroscope)
+                .map(|a| a.as_slice().unwrap());
+            let command_opt = inputs
+                .get(&InputType::Command)
+                .map(|a| a.as_slice().unwrap());
+            let output_opt = Some(output_tensor.as_slice().unwrap());
+
+            lg.log_step(
+                joint_angles_opt,
+                joint_vels_opt,
+                projected_g_opt,
+                accel_opt,
+                gyro_opt,
+                command_opt,
+                output_opt,
+            );
+        }
+
         Ok((
             output_tensor.view().to_owned(),
             carry_tensor.view().to_owned(),
@@ -304,15 +331,7 @@ impl ModelRunner {
         &self,
         action: Array<f32, IxDyn>,
     ) -> Result<(), Box<dyn std::error::Error>> {
-        self.provider
-            .take_action(self.metadata.joint_names.clone(), action)
-            .await?;
+        self.provider.take_action(action, &self.metadata).await?;
         Ok(())
     }
-
-    pub async fn get_joint_angles(&self) -> Result<Array<f32, IxDyn>, Box<dyn std::error::Error>> {
-        let joint_names = &self.metadata.joint_names;
-        let joint_angles = self.provider.get_joint_angles(joint_names).await?;
-        Ok(joint_angles)
-    }
 }