kinfer 0.5.4__cp312-cp312-manylinux_2_28_x86_64.whl

kinfer/rust/src/model.rs

@@ -0,0 +1,354 @@
+use crate::logger::StepLogger;
+use crate::types::{InputType, ModelMetadata};
+use async_trait::async_trait;
+use chrono;
+use flate2::read::GzDecoder;
+use ndarray::{Array, IxDyn};
+use ort::session::Session;
+use ort::value::Value;
+use std::collections::HashMap;
+use std::io::Read;
+use std::path::Path;
+use std::sync::Arc;
+use tar::Archive;
+use tokio::fs::File;
+use tokio::io::AsyncReadExt;
+
+#[derive(Debug, thiserror::Error)]
+pub enum ModelError {
+    #[error("IO error: {0}")]
+    Io(#[from] std::io::Error),
+    #[error("Provider error: {0}")]
+    Provider(String),
+}
+
+#[async_trait]
+pub trait ModelProvider: Send + Sync {
+    async fn get_inputs(
+        &self,
+        input_types: &[InputType],
+        metadata: &ModelMetadata,
+    ) -> Result<HashMap<InputType, Array<f32, IxDyn>>, ModelError>;
+
+    async fn take_action(
+        &self,
+        action: Array<f32, IxDyn>,
+        metadata: &ModelMetadata,
+    ) -> Result<(), ModelError>;
+}
+
+pub struct ModelRunner {
+    init_session: Session,
+    step_session: Session,
+    metadata: ModelMetadata,
+    provider: Arc<dyn ModelProvider>,
+    logger: Option<Arc<StepLogger>>,
+}
+
+impl ModelRunner {
+    pub async fn new<P: AsRef<Path>>(
+        model_path: P,
+        input_provider: Arc<dyn ModelProvider>,
+    ) -> Result<Self, Box<dyn std::error::Error>> {
+        let mut file = File::open(model_path).await?;
+
+        // Read entire file into memory
+        let mut buffer = Vec::new();
+        file.read_to_end(&mut buffer).await?;
+
+        // Decompress and read the tar archive from memory
+        let gz = GzDecoder::new(&buffer[..]);
+        let mut archive = Archive::new(gz);
+
+        // Extract and validate joint names
+        let mut metadata: Option<String> = None;
+        let mut init_fn: Option<Vec<u8>> = None;
+        let mut step_fn: Option<Vec<u8>> = None;
+
+        for entry in archive.entries()? {
+            let mut entry = entry?;
+            let path = entry.path()?;
+            let path_str = path.to_string_lossy();
+
+            match path_str.as_ref() {
+                "metadata.json" => {
+                    let mut contents = String::new();
+                    entry.read_to_string(&mut contents)?;
+                    metadata = Some(contents);
+                }
+                "init_fn.onnx" => {
+                    let size = entry.size() as usize;
+                    let mut contents = vec![0u8; size];
+                    entry.read_exact(&mut contents)?;
+                    assert_eq!(contents.len(), entry.size() as usize);
+                    init_fn = Some(contents);
+                }
+                "step_fn.onnx" => {
+                    let size = entry.size() as usize;
+                    let mut contents = vec![0u8; size];
+                    entry.read_exact(&mut contents)?;
+                    assert_eq!(contents.len(), entry.size() as usize);
+                    step_fn = Some(contents);
+                }
+                _ => return Err("Unknown entry".into()),
+            }
+        }
+
+        // Reads the files.
+        let metadata = ModelMetadata::model_validate_json(
+            metadata.ok_or("metadata.json not found in archive")?,
+        )?;
+        let init_session = Session::builder()?
+            .commit_from_memory(&init_fn.ok_or("init_fn.onnx not found in archive")?)?;
+        let step_session = Session::builder()?
+            .commit_from_memory(&step_fn.ok_or("step_fn.onnx not found in archive")?)?;
+
+        // Validate init_fn has no inputs and one output
+        if !init_session.inputs.is_empty() {
+            return Err("init_fn should not have any inputs".into());
+        }
+        if init_session.outputs.len() != 1 {
+            return Err("init_fn should have exactly one output".into());
+        }
+
+        // Get carry shape from init_fn output
+        let carry_shape = init_session.outputs[0]
+            .output_type
+            .tensor_dimensions()
+            .ok_or("Missing tensor type")?
+            .to_vec();
+
+        // 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)?;
+            }
+
+            // Use uuid if found, otherwise timestamp
+            let log_name = std::env::var("KINFER_LOG_UUID").unwrap_or_else(|_| {
+                chrono::Utc::now().format("%Y-%m-%d_%H-%M-%S").to_string()
+            });
+
+            let log_file_path = log_dir_path.join(format!("{}.ndjson", log_name));
+
+            Some(StepLogger::new(log_file_path).map(Arc::new)?)
+        } else {
+            None
+        };
+
+        Ok(Self {
+            init_session,
+            step_session,
+            metadata,
+            provider: input_provider,
+            logger,
+        })
+    }
+
+    fn validate_step_fn(
+        session: &Session,
+        metadata: &ModelMetadata,
+        carry_shape: &[i64],
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        // Validate inputs
+        for input in &session.inputs {
+            let dims = input.input_type.tensor_dimensions().ok_or(format!(
+                "Input {} is not a tensor with known dimensions",
+                input.name
+            ))?;
+
+            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());
+            }
+        }
+
+        // Validate outputs
+        if session.outputs.len() != 2 {
+            return Err("Step function must have exactly 2 outputs".into());
+        }
+
+        let output_shape = session.outputs[0]
+            .output_type
+            .tensor_dimensions()
+            .ok_or("Missing tensor type")?;
+        let num_joints = metadata.joint_names.len();
+        if *output_shape != vec![num_joints as i64] {
+            return Err(format!(
+                "Expected output shape [{num_joints}], got {:?}",
+                output_shape
+            )
+            .into());
+        }
+
+        let infered_carry_shape = session.outputs[1]
+            .output_type
+            .tensor_dimensions()
+            .ok_or("Missing tensor type")?;
+        if *infered_carry_shape != *carry_shape {
+            return Err(format!(
+                "Expected carry shape {:?}, got {:?}",
+                carry_shape, infered_carry_shape
+            )
+            .into());
+        }
+
+        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)?;
+        let output_tensor = outputs[0].try_extract_tensor::<f32>()?;
+        Ok(output_tensor.view().to_owned())
+    }
+
+    pub async fn step(
+        &self,
+        carry: Array<f32, IxDyn>,
+    ) -> Result<(Array<f32, IxDyn>, Array<f32, IxDyn>), Box<dyn std::error::Error>> {
+        // Gets the model input names.
+        let input_names: Vec<String> = self
+            .step_session
+            .inputs
+            .iter()
+            .map(|i| i.name.clone())
+            .collect();
+
+        // Calls the relevant getter methods in parallel.
+        let mut input_types = Vec::new();
+        let mut inputs = HashMap::new();
+        for name in &input_names {
+            match name.as_str() {
+                "joint_angles" => {
+                    input_types.push(InputType::JointAngles);
+                }
+                "joint_angular_velocities" => {
+                    input_types.push(InputType::JointAngularVelocities);
+                }
+                "initial_heading" => {
+                    input_types.push(InputType::InitialHeading);
+                }
+                "quaternion" => {
+                    input_types.push(InputType::Quaternion);
+                }
+                "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());
+                }
+                _ => return Err(format!("Unknown input name: {}", name).into()),
+            }
+        }
+
+        // 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_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));
+        }
+
+        // Run the model
+        let outputs = self.step_session.run(input_values)?;
+        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 initial_heading_opt = inputs
+                .get(&InputType::InitialHeading)
+                .map(|a| a.as_slice().unwrap());
+            let quaternion_opt = inputs
+                .get(&InputType::Quaternion)
+                .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,
+                initial_heading_opt,
+                quaternion_opt,
+                projected_g_opt,
+                accel_opt,
+                gyro_opt,
+                command_opt,
+                output_opt,
+            );
+        }
+
+        Ok((
+            output_tensor.view().to_owned(),
+            carry_tensor.view().to_owned(),
+        ))
+    }
+
+    pub async fn take_action(
+        &self,
+        action: Array<f32, IxDyn>,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        self.provider.take_action(action, &self.metadata).await?;
+        Ok(())
+    }
+}

kinfer/rust/src/runtime.rs

@@ -0,0 +1,107 @@
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use tokio::runtime::Runtime;
+
+use crate::model::{ModelError, ModelRunner};
+use crate::types::InputType;
+use std::time::Duration;
+use tokio::time::interval;
+
+pub struct ModelRuntime {
+    model_runner: Arc<ModelRunner>,
+    dt: Duration,
+    slowdown_factor: i32,
+    magnitude_factor: f32,
+    running: Arc<AtomicBool>,
+    runtime: Option<Runtime>,
+}
+
+impl ModelRuntime {
+    pub fn new(model_runner: Arc<ModelRunner>, dt: u64) -> Self {
+        Self {
+            model_runner,
+            dt: Duration::from_millis(dt),
+            slowdown_factor: 1,
+            magnitude_factor: 1.0,
+            running: Arc::new(AtomicBool::new(false)),
+            runtime: None,
+        }
+    }
+
+    pub fn set_slowdown_factor(&mut self, slowdown_factor: i32) {
+        assert!(slowdown_factor >= 1);
+        self.slowdown_factor = slowdown_factor;
+    }
+
+    pub fn set_magnitude_factor(&mut self, magnitude_factor: f32) {
+        assert!(magnitude_factor >= 0.0);
+        assert!(magnitude_factor <= 1.0);
+        self.magnitude_factor = magnitude_factor;
+    }
+
+    pub fn start(&mut self) -> Result<(), ModelError> {
+        if self.running.load(Ordering::Relaxed) {
+            return Ok(());
+        }
+
+        let running = self.running.clone();
+        let model_runner = self.model_runner.clone();
+        let dt = self.dt;
+        let slowdown_factor = self.slowdown_factor;
+        let magnitude_factor = self.magnitude_factor;
+
+        let runtime = Runtime::new()?;
+        running.store(true, Ordering::Relaxed);
+
+        runtime.spawn(async move {
+            let mut carry = model_runner
+                .init()
+                .await
+                .map_err(|e| ModelError::Provider(e.to_string()))?;
+
+            let model_inputs = model_runner
+                .get_inputs(&[InputType::JointAngles])
+                .await
+                .map_err(|e| ModelError::Provider(e.to_string()))?;
+            let mut joint_positions = model_inputs[&InputType::JointAngles].clone();
+
+            // Wait for the first tick, since it happens immediately.
+            let mut interval = interval(dt);
+            interval.tick().await;
+
+            while running.load(Ordering::Relaxed) {
+                let (output, next_carry) = model_runner
+                    .step(carry)
+                    .await
+                    .map_err(|e| ModelError::Provider(e.to_string()))?;
+                carry = next_carry;
+
+                for i in 1..(slowdown_factor + 1) {
+                    if !running.load(Ordering::Relaxed) {
+                        break;
+                    }
+                    let t = i as f32 / slowdown_factor as f32;
+                    let interp_joint_positions = &joint_positions * (1.0 - t) + &output * t;
+                    model_runner
+                        .take_action(interp_joint_positions * magnitude_factor)
+                        .await
+                        .map_err(|e| ModelError::Provider(e.to_string()))?;
+                    interval.tick().await;
+                }
+
+                joint_positions = output;
+            }
+            Ok::<(), ModelError>(())
+        });
+
+        self.runtime = Some(runtime);
+        Ok(())
+    }
+
+    pub fn stop(&mut self) {
+        self.running.store(false, Ordering::Relaxed);
+        if let Some(runtime) = self.runtime.take() {
+            runtime.shutdown_background();
+        }
+    }
+}

kinfer/rust/src/types.rs

@@ -0,0 +1,96 @@
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct ModelMetadata {
+    pub joint_names: Vec<String>,
+    pub num_commands: Option<usize>,
+    pub carry_size: Vec<usize>,
+}
+
+impl ModelMetadata {
+    pub fn model_validate_json(json: String) -> Result<Self, Box<dyn std::error::Error>> {
+        Ok(serde_json::from_str(&json)?)
+    }
+
+    pub fn to_json(&self) -> Result<String, Box<dyn std::error::Error>> {
+        Ok(serde_json::to_string(self)?)
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
+pub enum InputType {
+    JointAngles,
+    JointAngularVelocities,
+    InitialHeading,
+    Quaternion,
+    ProjectedGravity,
+    Accelerometer,
+    Gyroscope,
+    Command,
+    Time,
+    Carry,
+}
+
+impl InputType {
+    pub fn get_name(&self) -> &str {
+        match self {
+            InputType::JointAngles => "joint_angles",
+            InputType::JointAngularVelocities => "joint_angular_velocities",
+            InputType::InitialHeading => "initial_heading",
+            InputType::Quaternion => "quaternion",
+            InputType::ProjectedGravity => "projected_gravity",
+            InputType::Accelerometer => "accelerometer",
+            InputType::Gyroscope => "gyroscope",
+            InputType::Command => "command",
+            InputType::Time => "time",
+            InputType::Carry => "carry",
+        }
+    }
+
+    pub fn get_shape(&self, metadata: &ModelMetadata) -> Vec<usize> {
+        match self {
+            InputType::JointAngles => vec![metadata.joint_names.len()],
+            InputType::JointAngularVelocities => vec![metadata.joint_names.len()],
+            InputType::InitialHeading => vec![1],
+            InputType::Quaternion => vec![4],
+            InputType::ProjectedGravity => vec![3],
+            InputType::Accelerometer => vec![3],
+            InputType::Gyroscope => vec![3],
+            InputType::Command => vec![metadata.num_commands.unwrap_or(0)],
+            InputType::Time => vec![1],
+            InputType::Carry => metadata.carry_size.clone(),
+        }
+    }
+
+    pub fn from_name(name: &str) -> Result<Self, Box<dyn std::error::Error>> {
+        match name {
+            "joint_angles" => Ok(InputType::JointAngles),
+            "joint_angular_velocities" => Ok(InputType::JointAngularVelocities),
+            "initial_heading" => Ok(InputType::InitialHeading),
+            "quaternion" => Ok(InputType::Quaternion),
+            "projected_gravity" => Ok(InputType::ProjectedGravity),
+            "accelerometer" => Ok(InputType::Accelerometer),
+            "gyroscope" => Ok(InputType::Gyroscope),
+            "command" => Ok(InputType::Command),
+            "time" => Ok(InputType::Time),
+            "carry" => Ok(InputType::Carry),
+            _ => Err(format!("Unknown input type: {}", name).into()),
+        }
+    }
+
+    pub fn get_names() -> Vec<&'static str> {
+        vec![
+            "joint_angles",
+            "joint_angular_velocities",
+            "initial_heading",
+            "quaternion",
+            "projected_gravity",
+            "accelerometer",
+            "gyroscope",
+            "command",
+            "time",
+            "carry",
+        ]
+    }
+}

kinfer/rust_bindings/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+
+name = "rust_bindings"
+version.workspace = true
+edition.workspace = true
+description.workspace = true
+authors.workspace = true
+repository.workspace = true
+license.workspace = true
+readme.workspace = true
+
+[lib]
+
+name = "rust_bindings"
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+
+async-trait = "0.1"
+kinfer = { path = "../rust" }
+ndarray = "0.16.1"
+numpy = ">= 0.24.0"
+pyo3 = { version = ">= 0.24.0", features = ["extension-module"] }
+pyo3-async-runtimes = { version = ">= 0.24.0", features = ["attributes", "tokio-runtime"] }
+pyo3-stub-gen = ">= 0.6.0"
+tokio = { version = "1", features = ["rt-multi-thread", "macros"] }

kinfer/rust_bindings/pyproject.toml

@@ -0,0 +1,7 @@
+[build-system]
+requires = ["maturin>=1.1,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "rust_bindings"
+requires-python = ">=3.11"

kinfer/rust_bindings/rust_bindings.pyi

@@ -0,0 +1,46 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401
+
+import builtins
+import numpy
+import numpy.typing
+import typing
+
+class ModelProviderABC:
+    def __new__(cls) -> ModelProviderABC: ...
+    def get_inputs(self, input_types:typing.Sequence[builtins.str], metadata:PyModelMetadata) -> builtins.dict[builtins.str, numpy.typing.NDArray[numpy.float32]]: ...
+    def take_action(self, action:numpy.typing.NDArray[numpy.float32], metadata:PyModelMetadata) -> None: ...
+
+class PyInputType:
+    def __new__(cls, input_type:builtins.str) -> PyInputType: ...
+    def get_name(self) -> builtins.str: ...
+    def get_shape(self, metadata:PyModelMetadata) -> builtins.list[builtins.int]: ...
+    def __repr__(self) -> builtins.str: ...
+    def __eq__(self, other:typing.Any) -> builtins.bool: ...
+
+class PyModelMetadata:
+    def __new__(self, joint_names:typing.Sequence[builtins.str], num_commands:typing.Optional[builtins.int], carry_size:typing.Sequence[builtins.int]) -> PyModelMetadata: ...
+    def to_json(self) -> builtins.str: ...
+    def __repr__(self) -> builtins.str: ...
+    def __eq__(self, other:typing.Any) -> builtins.bool: ...
+
+class PyModelProvider:
+    ...
+
+class PyModelRunner:
+    def __new__(cls, model_path:builtins.str, provider:ModelProviderABC) -> PyModelRunner: ...
+    def init(self) -> numpy.typing.NDArray[numpy.float32]: ...
+    def step(self, carry:numpy.typing.NDArray[numpy.float32]) -> tuple[numpy.typing.NDArray[numpy.float32], numpy.typing.NDArray[numpy.float32]]: ...
+    def take_action(self, action:numpy.typing.NDArray[numpy.float32]) -> None: ...
+
+class PyModelRuntime:
+    def __new__(cls, model_runner:PyModelRunner, dt:builtins.int) -> PyModelRuntime: ...
+    def set_slowdown_factor(self, slowdown_factor:builtins.int) -> None: ...
+    def set_magnitude_factor(self, magnitude_factor:builtins.float) -> None: ...
+    def start(self) -> None: ...
+    def stop(self) -> None: ...
+
+def get_version() -> builtins.str: ...
+
+def metadata_from_json(json:builtins.str) -> PyModelMetadata: ...
+

kinfer/rust_bindings/src/bin/stub_gen.rs

@@ -0,0 +1,7 @@
+use pyo3_stub_gen::Result;
+
+fn main() -> Result<()> {
+    let stub = rust_bindings::stub_info()?;
+    stub.generate()?;
+    Ok(())
+}