kinfer 0.5.4__cp311-cp311-manylinux_2_28_x86_64.whl

kinfer/__init__.py

@@ -0,0 +1,16 @@
+"""Defines the kinfer API."""
+
+import os
+
+if "ORT_DYLIB_PATH" not in os.environ:
+    from pathlib import Path
+
+    import onnxruntime as ort
+
+    LIB_PATH = next((Path(ort.__file__).parent / "capi").glob("libonnxruntime.*"), None)
+    if LIB_PATH is not None:
+        os.environ["ORT_DYLIB_PATH"] = LIB_PATH.resolve().as_posix()
+
+from .rust_bindings import get_version
+
+__version__ = get_version()

kinfer/export/__init__.py

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

kinfer/export/jax.py

@@ -0,0 +1,55 @@
+"""Jax model export utilities."""
+
+import inspect
+import logging
+
+import tensorflow as tf
+import tf2onnx
+from equinox.internal._finalise_jaxpr import finalise_fn
+from jax._src.stages import Wrapped
+from jax.experimental import jax2tf
+from onnx.onnx_pb import ModelProto
+
+from kinfer.rust_bindings import PyInputType, PyModelMetadata
+
+logger = logging.getLogger(__name__)
+
+
+def export_fn(
+    model: Wrapped,
+    metadata: PyModelMetadata,
+    *,
+    opset: int = 13,
+) -> ModelProto:
+    """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")
+
+    params = inspect.signature(model).parameters
+    input_names = list(params.keys())
+
+    # Gets the dummy input tensors for exporting the model.
+    tf_args = []
+    for name in input_names:
+        shape = PyInputType(name).get_shape(metadata)
+        tf_args.append(tf.TensorSpec(shape, tf.float32, name=name))
+
+    finalised_fn = finalise_fn(model)
+    tf_fn = tf.function(jax2tf.convert(finalised_fn, enable_xla=False))
+
+    model_proto, _ = tf2onnx.convert.from_function(
+        tf_fn,
+        input_signature=tf_args,
+        opset=opset,
+        large_model=False,
+    )
+    return model_proto

kinfer/export/pytorch.py

@@ -0,0 +1,53 @@
+"""PyTorch model export utilities."""
+
+__all__ = [
+    "export_fn",
+]
+
+import io
+from typing import cast
+
+import onnx
+import torch
+from onnx.onnx_pb import ModelProto
+from torch._C import FunctionSchema
+
+from kinfer.rust_bindings import PyInputType, PyModelMetadata
+
+
+def export_fn(
+    model: torch.jit.ScriptFunction,
+    metadata: PyModelMetadata,
+) -> ModelProto:
+    """Exports a PyTorch function to ONNX.
+
+    Args:
+        model: The model to export.
+        metadata: The metadata for the model.
+
+    Returns:
+        The ONNX model as a `ModelProto`.
+    """
+    if not isinstance(model, torch.jit.ScriptFunction):
+        raise ValueError("Model must be a torch.jit.ScriptFunction")
+
+    schema = cast(FunctionSchema, model.schema)
+    input_names = [arg.name for arg in schema.arguments]
+
+    # Gets the dummy input tensors for exporting the model.
+    args = []
+    for name in input_names:
+        shape = PyInputType(name).get_shape(metadata)
+        args.append(torch.zeros(shape))
+
+    buffer = io.BytesIO()
+    torch.onnx.export(
+        model=model,
+        f=buffer,  # type: ignore[arg-type]
+        args=tuple(args),
+        input_names=input_names,
+        external_data=False,
+    )
+    buffer.seek(0)
+    model_bytes = buffer.read()
+    return onnx.load_from_string(model_bytes)

kinfer/export/serialize.py

@@ -0,0 +1,84 @@
+"""Functions for serializing and deserializing models."""
+
+__all__ = [
+    "pack",
+]
+
+import io
+import logging
+import tarfile
+
+from onnx.onnx_pb import ModelProto
+
+from kinfer.rust_bindings import PyInputType, PyModelMetadata
+
+logger = logging.getLogger(__name__)
+
+
+def pack(
+    init_fn: ModelProto,
+    step_fn: ModelProto,
+    metadata: PyModelMetadata,
+) -> bytes:
+    """Packs the initialization function and step function into a directory.
+
+    Args:
+        init_fn: The initialization function.
+        step_fn: The step function.
+        metadata: The metadata for the model.
+    """
+    num_joints = len(metadata.joint_names)  # type: ignore[attr-defined]
+
+    # Checks the `init` function.
+    if len(init_fn.graph.input) > 0:
+        raise ValueError(f"`init` function should not have any inputs! Got {len(init_fn.graph.input)}")
+    if len(init_fn.graph.output) != 1:
+        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 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 = 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:
+        raise ValueError(f"Step function must have exactly 2 outputs, got {len(step_fn.graph.output)}")
+
+    output_actions = step_fn.graph.output[0]
+    actions_shape = tuple(dim.dim_value for dim in output_actions.type.tensor_type.shape.dim)
+    if actions_shape != (num_joints,):
+        raise ValueError(f"Expected output shape {num_joints} for output `{output_actions.name}`, got {actions_shape}")
+
+    output_carry = step_fn.graph.output[1]
+    output_carry_shape = tuple(dim.dim_value for dim in output_carry.type.tensor_type.shape.dim)
+    if output_carry_shape != init_carry_shape:
+        raise ValueError(f"Expected carry shape {init_carry_shape} for output carry, got {output_carry_shape}")
+
+    buffer = io.BytesIO()
+
+    with tarfile.open(fileobj=buffer, mode="w:gz") as tar:
+
+        def add_file_bytes(name: str, data: bytes) -> None:  # noqa: ANN401
+            info = tarfile.TarInfo(name=name)
+            info.size = len(data)
+            tar.addfile(info, io.BytesIO(data))
+
+        add_file_bytes("init_fn.onnx", init_fn.SerializeToString())
+        add_file_bytes("step_fn.onnx", step_fn.SerializeToString())
+        add_file_bytes("metadata.json", metadata.to_json().encode("utf-8"))
+
+    buffer.seek(0)
+
+    return buffer.read()

kinfer/requirements.txt

@@ -0,0 +1,13 @@
+# requirements.txt
+
+# Machine Learning
+onnx
+onnxruntime==1.20.0
+
+# Serialization
+pydantic
+
+# Plotting
+matplotlib
+numpy
+pathlib

kinfer/rust/Cargo.toml

@@ -0,0 +1,36 @@
+[package]
+
+name = "kinfer"
+version.workspace = true
+edition.workspace = true
+description.workspace = true
+authors.workspace = true
+repository.workspace = true
+license.workspace = true
+readme.workspace = true
+
+[lib]
+
+name = "kinfer"
+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" ] }
+ort-sys = { version = "=2.0.0-rc.9" }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+tar = "0.4"
+thiserror = "1.0"
+tokio = { version = "1.0", features = ["full"] }
+
+[dev-dependencies]
+
+rand = "0.8"

kinfer/rust/src/lib.rs

@@ -0,0 +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,141 @@
+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]>,
+    initial_heading: Option<&'a [f32]>,
+    quaternion: 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]>,
+        initial_heading: Option<&[f32]>,
+        quaternion: 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,
+            initial_heading,
+            quaternion,
+            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();
+        }
+    }
+}