kinfer 0.4.2__tar.gz → 0.5.1__tar.gz

{kinfer-0.4.2 → kinfer-0.5.1}/Cargo.toml

@@ -8,7 +8,7 @@ resolver = "2"
 
 [workspace.package]
 
-version = "0.4.2"
+version = "0.5.1"
 authors = ["Wesley Maa <wesley@kscale.dev>", "Benjamin Bolte <ben@kscale.dev>"]
 edition = "2021"
 description = "K-Scale Inference Library"

{kinfer-0.4.2 → kinfer-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kinfer
-Version: 0.4.2
+Version: 0.5.1
 Summary: Tool to make it easier to run a model on a real robot
 Home-page: https://github.com/kscalelabs/kinfer.git
 Author: K-Scale Labs
@@ -10,6 +10,9 @@ License-File: LICENSE
 Requires-Dist: onnx
 Requires-Dist: onnxruntime==1.20.0
 Requires-Dist: pydantic
+Requires-Dist: matplotlib
+Requires-Dist: numpy
+Requires-Dist: pathlib
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: darglint; extra == "dev"
@@ -21,7 +24,7 @@ Provides-Extra: pytorch
 Requires-Dist: torch; extra == "pytorch"
 Provides-Extra: jax
 Requires-Dist: tensorflow; extra == "jax"
-Requires-Dist: tf2onnx; extra == "jax"
+Requires-Dist: tf2onnx>=1.16.0; extra == "jax"
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: equinox; extra == "jax"
 Requires-Dist: numpy<2; extra == "jax"
@@ -34,7 +37,7 @@ Requires-Dist: ruff; extra == "all"
 Requires-Dist: types-tensorflow; extra == "all"
 Requires-Dist: torch; extra == "all"
 Requires-Dist: tensorflow; extra == "all"
-Requires-Dist: tf2onnx; extra == "all"
+Requires-Dist: tf2onnx>=1.16.0; extra == "all"
 Requires-Dist: jax; extra == "all"
 Requires-Dist: equinox; extra == "all"
 Requires-Dist: numpy<2; extra == "all"

kinfer-0.5.1/kinfer/export/__init__.py

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

{kinfer-0.4.2 → kinfer-0.5.1}/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-0.4.2 → kinfer-0.5.1}/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-0.4.2 → kinfer-0.5.1}/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-0.4.2 → kinfer-0.5.1}/kinfer/requirements.txt

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

{kinfer-0.4.2 → kinfer-0.5.1}/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-0.4.2 → kinfer-0.5.1}/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-0.5.1/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();
+        }
+    }
+}