kinfer 0.5.1__cp312-cp312-macosx_11_0_arm64.whl → 0.5.4__cp312-cp312-macosx_11_0_arm64.whl

kinfer/rust/Cargo.toml

@@ -24,6 +24,7 @@ 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"

kinfer/rust/src/logger.rs

@@ -16,6 +16,8 @@ struct NdjsonStep<'a> {
     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]>,
@@ -88,6 +90,8 @@ impl StepLogger {
         &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]>,
@@ -99,6 +103,8 @@ impl StepLogger {
             t_us: Self::now_us() as u64,
             joint_angles,
             joint_vels,
+            initial_heading,
+            quaternion,
             projected_g,
             accel,
             gyro,

kinfer/rust/src/model.rs

@@ -129,9 +129,12 @@ impl ModelRunner {
                 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));
+            // 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 {
@@ -241,6 +244,12 @@ impl ModelRunner {
                 "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);
                 }
@@ -296,6 +305,12 @@ impl ModelRunner {
             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());
@@ -313,6 +328,8 @@ impl ModelRunner {
             lg.log_step(
                 joint_angles_opt,
                 joint_vels_opt,
+                initial_heading_opt,
+                quaternion_opt,
                 projected_g_opt,
                 accel_opt,
                 gyro_opt,

kinfer/rust/src/types.rs

@@ -22,6 +22,8 @@ impl ModelMetadata {
 pub enum InputType {
     JointAngles,
     JointAngularVelocities,
+    InitialHeading,
+    Quaternion,
     ProjectedGravity,
     Accelerometer,
     Gyroscope,
@@ -35,6 +37,8 @@ impl InputType {
         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",
@@ -48,6 +52,8 @@ impl InputType {
         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],
@@ -61,6 +67,8 @@ impl InputType {
         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),
@@ -75,6 +83,8 @@ impl InputType {
         vec![
             "joint_angles",
             "joint_angular_velocities",
+            "initial_heading",
+            "quaternion",
             "projected_gravity",
             "accelerometer",
             "gyroscope",

kinfer/rust_bindings/src/lib.rs

@@ -17,6 +17,19 @@ use std::sync::Mutex;
 
 type StepResult = (Py<PyArrayDyn<f32>>, Py<PyArrayDyn<f32>>);
 
+// Custom error type for Send/Sync compatibility
+#[derive(Debug)]
+struct SendError(String);
+
+unsafe impl Send for SendError {}
+unsafe impl Sync for SendError {}
+
+impl std::fmt::Display for SendError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
 #[pyfunction]
 #[gen_stub_pyfunction]
 fn get_version() -> String {
@@ -204,9 +217,15 @@ impl ModelProviderABC {
         metadata: PyModelMetadata,
     ) -> PyResult<()> {
         let n = action.len()?;
-        assert_eq!(metadata.joint_names.len(), n); // TODO: this is wrong
+        if metadata.joint_names.len() != n {
+            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+                "Expected {} joints, got {} action elements",
+                metadata.joint_names.len(),
+                n
+            )));
+        }
         Err(PyNotImplementedError::new_err(format!(
-            "Must override take_action with {} action",
+            "Must override take_action with {} action elements",
             n
         )))
     }
@@ -286,6 +305,7 @@ impl ModelProvider for PyModelProvider {
 #[derive(Clone)]
 struct PyModelRunner {
     runner: Arc<ModelRunner>,
+    runtime: Arc<tokio::runtime::Runtime>,
 }
 
 #[gen_stub_pymethods]
@@ -295,7 +315,11 @@ impl PyModelRunner {
     fn __new__(model_path: String, provider: Py<ModelProviderABC>) -> PyResult<Self> {
         let input_provider = Arc::new(PyModelProvider::__new__(provider));
 
-        let runner = tokio::runtime::Runtime::new()?.block_on(async {
+        // Create a single runtime to be reused for all operations
+        let runtime = Arc::new(tokio::runtime::Runtime::new()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?);
+
+        let runner = runtime.block_on(async {
             ModelRunner::new(model_path, input_provider)
                 .await
                 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
@@ -303,17 +327,27 @@ impl PyModelRunner {
 
         Ok(Self {
             runner: Arc::new(runner),
+            runtime,
         })
     }
 
+    // Reuse runtime and release GIL
     fn init(&self) -> PyResult<Py<PyArrayDyn<f32>>> {
         let runner = self.runner.clone();
-        let result = tokio::runtime::Runtime::new()?.block_on(async {
-            runner
-                .init()
-                .await
-                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
-        })?;
+        let runtime = self.runtime.clone();
+
+        let result = Python::with_gil(|py| {
+            // Release GIL during async operation
+            py.allow_threads(|| {
+                runtime.block_on(async {
+                    runner
+                        .init()
+                        .await
+                        .map_err(|e| SendError(e.to_string()))
+                })
+            })
+        })
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.0))?;
 
         Python::with_gil(|py| {
             let array = numpy::PyArray::from_array(py, &result);
@@ -321,20 +355,31 @@ impl PyModelRunner {
         })
     }
 
+    // Reuse runtime and release GIL
     fn step(&self, carry: Py<PyArrayDyn<f32>>) -> PyResult<StepResult> {
         let runner = self.runner.clone();
+        let runtime = self.runtime.clone();
+        
+        // Extract the carry array from Python with GIL
         let carry_array = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
             let carry_array = carry.bind(py);
             Ok(carry_array.to_owned_array())
         })?;
 
-        let result = tokio::runtime::Runtime::new()?.block_on(async {
-            runner
-                .step(carry_array)
-                .await
-                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
-        })?;
+        // Release GIL during computation
+        let result = Python::with_gil(|py| {
+            py.allow_threads(|| {
+                runtime.block_on(async {
+                    runner
+                        .step(carry_array)
+                        .await
+                        .map_err(|e| SendError(e.to_string()))
+                })
+            })
+        })
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.0))?;
 
+        // Reacquire the GIL to convert results back to Python objects
         Python::with_gil(|py| {
             let (output, carry) = result;
             let output_array = numpy::PyArray::from_array(py, &output);
@@ -343,19 +388,29 @@ impl PyModelRunner {
         })
     }
 
+    // Reuse runtime and release GIL
     fn take_action(&self, action: Py<PyArrayDyn<f32>>) -> PyResult<()> {
         let runner = self.runner.clone();
+        let runtime = self.runtime.clone();
+
+        // Extract action data with GIL
         let action_array = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
             let action_array = action.bind(py);
             Ok(action_array.to_owned_array())
         })?;
-
-        tokio::runtime::Runtime::new()?.block_on(async {
-            runner
-                .take_action(action_array)
-                .await
-                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
-        })?;
+        
+        // Release GIL during computation
+        Python::with_gil(|py| {
+            py.allow_threads(|| {
+                runtime.block_on(async {
+                    runner
+                        .take_action(action_array)
+                        .await
+                        .map_err(|e| SendError(e.to_string()))
+                })
+            })
+        })
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.0))?;
 
         Ok(())
     }
@@ -428,4 +483,4 @@ fn rust_bindings(m: &Bound<PyModule>) -> PyResult<()> {
     Ok(())
 }
 
-define_stub_info_gatherer!(stub_info);
+define_stub_info_gatherer!(stub_info);

{kinfer-0.5.1.dist-info → kinfer-0.5.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kinfer
-Version: 0.5.1
+Version: 0.5.4
 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
@@ -56,3 +56,8 @@ Dynamic: summary
 This package is designed to support running real-time robotics models.
 
 For more information, see the documentation [here](https://docs.kscale.dev/docs/k-infer).
+
+To enable logging, set your log path in the environment variable `KINFER_LOG_PATH`. For example,
+```
+export KINFER_LOG_PATH=/home/dpsh/kinfer-logs
+```

{kinfer-0.5.1.dist-info → kinfer-0.5.4.dist-info}/RECORD

@@ -2,13 +2,13 @@ kinfer/requirements.txt,sha256=e8RzCNVZGJ9Jb874Eom0psUCOqp0wGf8LJZ86ysto0w,127
 kinfer/__init__.py,sha256=i5da6ND827Cgn8PFKzDCmEBk14ptQasLa_9fdof4Y9c,398
 kinfer/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kinfer/rust_bindings.pyi,sha256=lPGHMDsydr4Ue-CNHRjQtJWmq5sE5daVN8PbxJgOZY4,2048
-kinfer/rust_bindings.cpython-312-darwin.so,sha256=9a3I3Y7iqwQ8BNH8AjN-ZgM6_ETJG4WhAbSUWdBDceg,2069584
-kinfer/rust/Cargo.toml,sha256=lShTLTyccAHem8ywWOSJjbmVaADTfgUdwxwK70PpnkY,665
+kinfer/rust_bindings.cpython-312-darwin.so,sha256=5Thj3t4xz9-85MJ8ckTmb9yZ6C7_NIg5OQCgEuy2vbk,2069520
+kinfer/rust/Cargo.toml,sha256=NoWaLluRJHV6mOCeVXk8q8Pb4028Z3ZPqLspBMHAPTQ,703
 kinfer/rust/src/runtime.rs,sha256=7BRZmyMoV6FvYs2_XnLk0DFrxt7qgz35loCUJRSsuxM,3448
-kinfer/rust/src/logger.rs,sha256=iDb8hiRjxnH8l-Qzth5XcLk2ICnnX1JlBxyOzaZ4t64,3919
-kinfer/rust/src/types.rs,sha256=KNIM6sHvX1gnQZc74uRTgmzbgS4FZjxll5jcuNGIDc8,2768
+kinfer/rust/src/logger.rs,sha256=5Uvv60zp6HV4zmaD9bA1vENDtE19uYrRzA-xRY26bl8,4124
+kinfer/rust/src/types.rs,sha256=wNLlXXmGmcE7uF4tKuZ7ChRy9-PI-JWq2iuh7diQS_g,3187
 kinfer/rust/src/lib.rs,sha256=25LuAdONp9c3lI0_g2kAf73yrtzbYc9GtJjm1zl9IwU,120
-kinfer/rust/src/model.rs,sha256=RQR0y7u9lPoSBNPnBWYesfxDkBSBBJpkyiH6qqCeQRI,11504
+kinfer/rust/src/model.rs,sha256=WluovJlgKSrAlAOkKlhhiVnD-OhN9E6xznZRPn2LFFI,12176
 kinfer/scripts/plot_ndjson.py,sha256=eg-wwuDLRDcuVdr1Jc5W4yeJG5Bf--UjaK6kk9XoC4w,6179
 kinfer/export/serialize.py,sha256=QjsvTn7cIRroZXXONXQUrEd1tR69oLrszojuKoK-Xuk,3208
 kinfer/export/__init__.py,sha256=6a8cClA2H3AAjAhY6ucmxEnfCRZ30FDn2L7SArCu7nY,56
@@ -17,10 +17,10 @@ kinfer/export/pytorch.py,sha256=nRo9TKClRhPcLJQDR3oHNCbTYjf_3gLBxGfzZBScPxc,1303
 kinfer/rust_bindings/Cargo.toml,sha256=i1RGB9VNd9Q4FJ6gGwjZJQYo8DBBvpVWf3GJ95EfVgM,637
 kinfer/rust_bindings/pyproject.toml,sha256=jLcJuHCnQRh9HWR_R7a9qLHwj6LMBgnHyeKK_DruO1Y,135
 kinfer/rust_bindings/rust_bindings.pyi,sha256=lPGHMDsydr4Ue-CNHRjQtJWmq5sE5daVN8PbxJgOZY4,2048
-kinfer/rust_bindings/src/lib.rs,sha256=Uqi76Na_fGiecMeW5mSSmtYdz2wU8clLxYawbYomNOc,12662
+kinfer/rust_bindings/src/lib.rs,sha256=f5TBqP6tYmb3vXQiisCsxAPB37_mtoaZSFvc_FlOeKY,14473
 kinfer/rust_bindings/src/bin/stub_gen.rs,sha256=hhoVGnaSfazbSfj5a4x6mPicGPOgWQAfsDmiPej0B6Y,133
-kinfer-0.5.1.dist-info/RECORD,,
-kinfer-0.5.1.dist-info/WHEEL,sha256=V1loQ6TpxABu1APUg0MoTRBOzSKT5xVc3skizX-ovCU,136
-kinfer-0.5.1.dist-info/top_level.txt,sha256=6mY_t3PYr3Dm0dpqMk80uSnArbvGfCFkxOh1QWtgDEo,7
-kinfer-0.5.1.dist-info/METADATA,sha256=9S6I6udXfzxWQGS0l-XibfwGPY2qrC7koyLhvVievNY,1831
-kinfer-0.5.1.dist-info/licenses/LICENSE,sha256=Qw-Z0XTwS-diSW91e_jLeBPX9zZbAatOJTBLdPHPaC0,1069
+kinfer-0.5.4.dist-info/RECORD,,
+kinfer-0.5.4.dist-info/WHEEL,sha256=V1loQ6TpxABu1APUg0MoTRBOzSKT5xVc3skizX-ovCU,136
+kinfer-0.5.4.dist-info/top_level.txt,sha256=6mY_t3PYr3Dm0dpqMk80uSnArbvGfCFkxOh1QWtgDEo,7
+kinfer-0.5.4.dist-info/METADATA,sha256=Da03N1PNxRcIrOj0ub_N19j9v8n0siXLjxW5ZCtsiOQ,1983
+kinfer-0.5.4.dist-info/licenses/LICENSE,sha256=Qw-Z0XTwS-diSW91e_jLeBPX9zZbAatOJTBLdPHPaC0,1069