adding RecordingScheduler to capture SchedulerOutput and ModelRunnerO…

…utput

Signed-off-by: Ryan Olson <[email protected]>

.gitignore

@@ -91,3 +91,11 @@ generated-values.yaml
 .build/
 **/.devcontainer/.env
 TensorRT-LLM
+
+
+# START Ruler Generated Files
+/.cursor/instructions.md
+/.cursor/instructions.md.bak
+/CLAUDE.md
+/CLAUDE.md.bak
+# END Ruler Generated Files

lib/bindings/python/rust/llm/block_manager/vllm.rs

@@ -31,6 +31,7 @@ use crate::to_pyerr;
 mod block_list;
 mod connector;
 mod request;
+// mod scheduler;  // TODO: Fix PyO3 bindings
 mod slot;
 
 pub use block_list::{BlockListType, BlockState, BlockStates, KvbmBlockList};
@@ -53,6 +54,15 @@ fn _vllm_integration(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // TODO: use TRTLLM own integration module
     m.add_class::<connector::trtllm_worker::PyTrtllmKvConnectorWorker>()?;
     m.add_class::<connector::trtllm_leader::PyTrtllmKvConnectorLeader>()?;
+
+    // Add scheduler recorder and conversion functions
+    // TODO: Fix PyO3 bindings for these types
+    // m.add_class::<scheduler::recorder_bindings::PySchedulerRecorder>()?;
+    // m.add_function(wrap_pyfunction!(scheduler::scheduler_types::convert_scheduler_output, m)?)?;
+    // m.add_function(wrap_pyfunction!(scheduler::scheduler_types::convert_model_runner_output, m)?)?;
+    // m.add_function(wrap_pyfunction!(scheduler::scheduler_types::convert_engine_core_outputs, m)?)?;
+    // m.add_function(wrap_pyfunction!(scheduler::recorder_bindings::load_scheduler_trace, m)?)?;
+
     Ok(())
 }
 

lib/bindings/python/rust/llm/block_manager/vllm/scheduler/mod.rs

@@ -0,0 +1,7 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+pub mod recorder_bindings;
+pub mod scheduler_types;
+
+use pyo3::prelude::*;

lib/bindings/python/rust/llm/block_manager/vllm/scheduler/recorder_bindings.rs

@@ -0,0 +1,83 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Python bindings for the SchedulerRecorder
+
+use dynamo_llm::integrations::vllm::recorder::SchedulerRecorder as RustRecorder;
+use dynamo_llm::integrations::vllm::types::*;
+use pyo3::prelude::*;
+use std::path::PathBuf;
+
+/// Python-accessible SchedulerRecorder
+#[pyclass(name = "SchedulerRecorder")]
+pub struct PySchedulerRecorder {
+    inner: RustRecorder,
+}
+
+#[pymethods]
+impl PySchedulerRecorder {
+    /// Create a new SchedulerRecorder
+    #[new]
+    #[pyo3(signature = (model, vllm_version))]
+    fn new(model: String, vllm_version: String) -> Self {
+        Self {
+            inner: RustRecorder::new(model, vllm_version),
+        }
+    }
+
+    /// Record a scheduler output (already converted to Rust)
+    fn record_schedule_output(&mut self, output: SchedulerOutput) -> PyResult<()> {
+        self.inner.record_schedule_output(output);
+        Ok(())
+    }
+
+    /// Record a model runner output (already converted to Rust)
+    fn record_model_runner_output(&mut self, output: ModelRunnerOutput) -> PyResult<()> {
+        self.inner.record_model_runner_output(output);
+        Ok(())
+    }
+
+    /// Record engine core outputs (already converted to Rust)
+    fn record_engine_core_outputs(&mut self, outputs: EngineCoreOutputs) -> PyResult<()> {
+        self.inner.record_engine_core_outputs(outputs);
+        Ok(())
+    }
+
+    /// Move to the next iteration
+    fn next_iteration(&mut self) -> PyResult<()> {
+        self.inner.next_iteration();
+        Ok(())
+    }
+
+    /// Get the current iteration number
+    fn current_iteration(&self) -> u64 {
+        self.inner.current_iteration()
+    }
+
+    /// Save the recording to a JSON file
+    fn save_to_file(&mut self, path: String) -> PyResult<()> {
+        let path = PathBuf::from(path);
+        self.inner
+            .save_to_file(&path)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("{}", e)))
+    }
+
+    /// Clear all recordings
+    fn clear(&mut self) -> PyResult<()> {
+        self.inner.clear();
+        Ok(())
+    }
+
+    /// Get the number of recorded iterations
+    fn num_iterations(&self) -> usize {
+        self.inner.get_trace().iterations.len()
+    }
+}
+
+/// Load a recording from a JSON file
+#[pyfunction]
+pub fn load_scheduler_trace(path: String) -> PyResult<SchedulerTrace> {
+    let path = PathBuf::from(path);
+    RustRecorder::load_from_file(&path)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("{}", e)))
+}