Add pre_step & ctx.truncate to Elodin Sim (#400)

* Add elodin-py simulation pre-step

* Add pre/post step truncate() to clear Simulation DB

* chore: Make Clippy happy.

* doc: Update release.md.

* style: Reformat.

---------

Co-authored-by: Shane Celis <shane.celis@gmail.com>
Co-authored-by: Shane Celis <scelis@singularityus.com>

Author: x46085

docs/internal/release.md

@@ -39,7 +39,9 @@ Doing a pre-release should be quick-and-dirty on-demand operation without any of
 1. Run `just tag v$VERSION` on the updated `main` branch. 
 2. Run `just promote v$VERSION` to promote the artifacts from GitHub to S3 and PyPi. 
 ```sh
-export VERSION=0.15.5-alpha.0
+export PYPI_TOKEN=pypi-...; # You'll need this token to publish the pypi package.
+export VERSION=0.16.0-alpha.1
+# Start release. Wait 40 minutes. Then publish pypi package.
 just tag v$VERSION origin/main && sleep $((40 * 60)) && just promote v$VERSION
 ``` 
 3. Once released, we want to bump the version of alpha so that what was just released and what's on the 'main' branch will not be easily confused. If the version released was `0.1.2-alpha.0` then the next version would be `0.1.2-alpha.1`. 

examples/betaflight-sitl/main.py

@@ -184,7 +184,7 @@ def __post_init__(self):
 _rc_channels_buffer = np.full(MAX_RC_CHANNELS, 1500, dtype=np.uint16)
 
 
-def sitl_post_step(tick: int, ctx: el.PostStepContext):
+def sitl_post_step(tick: int, ctx: el.StepContext):
     """
     Post-step callback for lockstep SITL synchronization.
 
@@ -229,6 +229,7 @@ def sitl_post_step(tick: int, ctx: el.PostStepContext):
                 pass  # Expected during initial warmup
         print(f"[SITL] Warmup complete ({warmup_count} responses at {config.pid_rate:.0f}Hz)")
         print("[SITL] Bridge ready")
+        ctx.truncate()  # Clears all data, resets tick to 0
 
     if start_time[0] is None:
         start_time[0] = time.time()
@@ -305,7 +306,7 @@ def sitl_post_step(tick: int, ctx: el.PostStepContext):
         s.max_motor = max(s.max_motor, np.max(s.motors))
 
         # Write motor commands back to Elodin-DB for physics simulation
-        # This uses the PostStepContext for direct DB access (no TCP overhead)
+        # This uses the StepContext for direct DB access (no TCP overhead)
         ctx.write_component("drone.motor_command", s.motors)
     except TimeoutError:
         pass  # Timeouts expected during bootgrace

libs/db/src/append_log.rs

@@ -175,4 +175,14 @@ impl<E: IntoBytes + Immutable> AppendLog<E> {
         drop(guard);
         Ok(end - size_of::<Header<E>>())
     }
+
+    /// Truncate the append log, clearing all data while preserving the header.
+    ///
+    /// This resets the committed_len back to just the header size, effectively
+    /// removing all stored data without deallocating the memory-mapped file.
+    pub fn truncate(&self) {
+        let _guard = self.write_lock.lock().unwrap();
+        self.committed_len()
+            .store(size_of::<Header<E>>() as u64, Ordering::SeqCst);
+    }
 }

libs/db/src/lib.rs

@@ -394,6 +394,18 @@ impl DB {
         Ok(())
     }
 
+    /// Truncate all component data and message logs, clearing all data while preserving schemas and metadata.
+    ///
+    /// This effectively resets the database to an empty state, ready for fresh data.
+    /// The vtable generation is incremented to signal that clients should refresh their views.
+    pub fn truncate(&self) {
+        self.with_state(|state| {
+            state.truncate_all();
+        });
+        self.last_updated.store(Timestamp(i64::MIN));
+        self.vtable_gen.fetch_add(1, atomic::Ordering::SeqCst);
+    }
+
     pub fn copy_native(&self, target_db_path: impl AsRef<Path>) -> Result<PathBuf, Error> {
         let final_db_dir = target_db_path.as_ref().to_path_buf();
         let parent_dir = final_db_dir
@@ -751,6 +763,16 @@ impl State {
         msg_log.set_metadata(metadata)?;
         Ok(())
     }
+
+    /// Truncate all components and message logs, clearing all data while preserving schemas and metadata.
+    pub fn truncate_all(&self) {
+        for component in self.components.values() {
+            component.truncate();
+        }
+        for msg_log in self.msg_logs.values() {
+            msg_log.truncate();
+        }
+    }
 }
 
 #[derive(Clone, Debug, PartialEq)]
@@ -954,6 +976,11 @@ impl Component {
     fn sync_all(&self) -> Result<(), Error> {
         self.time_series.sync_all()
     }
+
+    /// Truncate the component, clearing all time-series data while preserving the schema.
+    pub fn truncate(&self) {
+        self.time_series.truncate();
+    }
 }
 
 struct DBSink<'a> {

libs/db/src/msg_log.rs

@@ -136,6 +136,15 @@ impl MsgLog {
         Ok(())
     }
 
+    /// Truncate the message log, clearing all messages while preserving metadata.
+    ///
+    /// This resets all underlying append logs, effectively removing all stored
+    /// messages without deallocating the underlying files.
+    pub fn truncate(&self) {
+        self.timestamps.truncate();
+        self.bufs.truncate();
+    }
+
     pub fn set_metadata(&mut self, metadata: MsgMetadata) -> Result<(), Error> {
         let metadata = self.metadata.insert(metadata);
         let metadata_path = self.path.join("metadata");
@@ -165,6 +174,11 @@ impl BufLog {
         Ok(())
     }
 
+    fn truncate(&self) {
+        self.offsets.truncate();
+        self.data_log.truncate();
+    }
+
     pub fn get_msg(&self, index: usize) -> Option<&[u8]> {
         let buf = self.bufs().get(index)?;
         let data = match buf.len as usize {

libs/db/src/time_series.rs

@@ -142,6 +142,15 @@ impl TimeSeries {
         Ok(())
     }
 
+    /// Truncate the time series, clearing all data while preserving the schema.
+    ///
+    /// This resets both the index and data append logs, effectively removing
+    /// all stored time-series data without deallocating the underlying files.
+    pub fn truncate(&self) {
+        self.index.truncate();
+        self.data.truncate();
+    }
+
     pub fn push_buf(&self, timestamp: Timestamp, buf: &[u8]) -> Result<(), Error> {
         let len = self.index.len() as usize;
 

libs/nox-ecs/src/impeller2_server.rs

@@ -5,7 +5,10 @@ use impeller2_wkt::{ComponentMetadata, EntityMetadata};
 use nox_ecs::Error;
 use std::{
     collections::HashSet,
-    sync::{Arc, atomic},
+    sync::{
+        Arc,
+        atomic::{self, AtomicU64, Ordering},
+    },
     time::{Duration, Instant},
 };
 use stellarator::struc_con::{Joinable, Thread};
@@ -25,14 +28,15 @@ impl Server {
 
     pub async fn run(self) -> Result<(), Error> {
         tracing::info!("running server");
-        self.run_with_cancellation(|| false, |_, _, _| {}, false)
+        self.run_with_cancellation(|| false, |_, _, _, _| {}, |_, _, _, _| {}, false)
             .await
     }
 
     pub async fn run_with_cancellation(
         self,
         is_cancelled: impl Fn() -> bool + 'static,
-        post_step: impl Fn(u64, &Arc<DB>, Timestamp) + 'static,
+        pre_step: impl Fn(u64, &Arc<DB>, &Arc<AtomicU64>, Timestamp) + 'static,
+        post_step: impl Fn(u64, &Arc<DB>, &Arc<AtomicU64>, Timestamp) + 'static,
         interactive: bool,
     ) -> Result<(), Error> {
         tracing::info!("running server with cancellation");
@@ -41,6 +45,8 @@ impl Server {
         let start_time = Timestamp::now();
         init_db(&db, &mut world.world, start_time)?;
         let tick_db = db.clone();
+        // Shared tick counter that can be reset by StepContext::truncate()
+        let tick_counter = Arc::new(AtomicU64::new(0));
         let stream: Thread<Option<Result<(), Error>>> =
             stellarator::struc_con::stellar(move || async move {
                 let mut handles = vec![];
@@ -51,8 +57,10 @@ impl Server {
             });
         let tick = stellarator::spawn(tick(
             tick_db,
+            tick_counter,
             world,
             is_cancelled,
+            pre_step,
             post_step,
             start_time,
             interactive,
@@ -167,7 +175,10 @@ pub fn copy_db_to_world(state: &State, world: &mut WorldExec<Compiled>) {
             let Some(component) = state.get_component(pair_id) else {
                 continue;
             };
-            let (_, head) = component.time_series.latest().unwrap();
+            // After truncate(), time series may be empty - skip if no data
+            let Some((_, head)) = component.time_series.latest() else {
+                continue;
+            };
 
             // Check if the value has changed
             let current_value = &column.buffer[offset..offset + size];
@@ -262,16 +273,18 @@ pub fn commit_world_head(
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn tick(
     db: Arc<DB>,
+    tick_counter: Arc<AtomicU64>,
     mut world: WorldExec<Compiled>,
     is_cancelled: impl Fn() -> bool + 'static,
-    post_step: impl Fn(u64, &Arc<DB>, Timestamp) + 'static,
-    mut timestamp: Timestamp,
+    pre_step: impl Fn(u64, &Arc<DB>, &Arc<AtomicU64>, Timestamp) + 'static,
+    post_step: impl Fn(u64, &Arc<DB>, &Arc<AtomicU64>, Timestamp) + 'static,
+    start_timestamp: Timestamp,
     interactive: bool,
 ) {
     // XXX This is what world.run ultimately calls.
-    let mut tick = 0;
     let external_controls: HashSet<ComponentId> = external_controls(&world).collect();
     let wait_for_write: Vec<ComponentId> = wait_for_write(&world).collect();
     let wait_for_write_pair_ids: Vec<PairId> = get_pair_ids(&world, &wait_for_write).unwrap();
@@ -281,15 +294,22 @@ async fn tick(
         .metadata
         .run_time_step
         .map(|time_step| time_step.0);
+    let time_step = world.world.sim_time_step().0;
     while db.recording_cell.wait().await {
         let start = Instant::now();
+        // Read tick from shared counter (can be reset by StepContext::truncate())
+        let tick = tick_counter.load(Ordering::SeqCst);
+        // Calculate timestamp based on current tick
+        let timestamp = start_timestamp + time_step * (tick as u32);
         if tick >= world.world.max_tick() {
             db.recording_cell.set_playing(false);
             world.world.metadata.max_tick = u64::MAX;
             if !interactive {
                 return;
             }
         }
+        // Python pre_step func runs (before copy_db_to_world so writes are picked up).
+        pre_step(tick, &db, &tick_counter, timestamp);
         db.with_state(|state| copy_db_to_world(state, &mut world));
         // JAX runs.
         if let Err(err) = world.run() {
@@ -314,16 +334,15 @@ async fn tick(
         if is_cancelled() {
             return;
         }
-        // Python func runs.
-        post_step(tick, &db, timestamp);
+        // Python post_step func runs.
+        post_step(tick, &db, &tick_counter, timestamp);
         // We only wait if there is a run_time_step set and it's >= the time elapsed.
         if let Some(run_time_step) = run_time_step.as_ref()
             && let Some(sleep_time) = run_time_step.checked_sub(start.elapsed())
         {
             stellarator::sleep(sleep_time).await;
         }
-        tick += 1;
-        timestamp += world.world.sim_time_step().0;
+        tick_counter.fetch_add(1, Ordering::SeqCst);
     }
 }
 

libs/nox-py/python/elodin/__init__.py

@@ -403,6 +403,7 @@ def run(
         max_ticks: Optional[int] = None,
         optimize: bool = False,
         is_canceled: Optional[callable] = None,
+        pre_step: Optional[callable] = None,
         post_step: Optional[callable] = None,
         db_path: Optional[str] = None,
         interactive: bool = True,
@@ -421,6 +422,7 @@ def run(
             max_ticks,
             optimize,
             is_canceled,
+            pre_step,
             post_step,
             db_path,
             interactive,

libs/nox-py/python/elodin/elodin.pyi

@@ -22,11 +22,11 @@ class Integrator:
     Rk4: Integrator
     SemiImplicit: Integrator
 
-class PostStepContext:
-    """Context object passed to post_step callbacks, providing direct DB read/write access.
+class StepContext:
+    """Context object passed to pre_step and post_step callbacks, providing direct DB read/write access.
 
     This enables SITL workflows to read sensor data and write component data (like motor
-    commands from Betaflight) back to the database within the same process.
+    commands from Betaflight) directly to the database within the same process.
     """
     @property
     def tick(self) -> int:
@@ -87,6 +87,15 @@ class PostStepContext:
             ValueError: If any write data size doesn't match the component schema
         """
         ...
+    def truncate(self) -> None:
+        """Truncate all component data and message logs in the database, resetting tick to 0.
+
+        This clears all stored time-series data while preserving component schemas and metadata.
+        The simulation tick will be reset to 0, effectively starting fresh.
+
+        Use this to control the freshness of the database and ensure reliable data from a known tick.
+        """
+        ...
 
 class ComponentType:
     def __init__(self, ty: PrimitiveType, shape: Tuple[int, ...]): ...
@@ -120,7 +129,8 @@ class WorldBuilder:
         max_ticks: Optional[int] = None,
         optimize: bool = False,
         is_canceled: Optional[Callable[[], bool]] = None,
-        post_step: Optional[Callable[[int, PostStepContext], None]] = None,
+        pre_step: Optional[Callable[[int, StepContext], None]] = None,
+        post_step: Optional[Callable[[int, StepContext], None]] = None,
         db_path: Optional[str] = None,
         interactive: bool = True,
     ): ...

libs/nox-py/src/lib.rs

@@ -21,10 +21,10 @@ mod error;
 mod exec;
 mod graph;
 mod linalg;
-mod post_step_context;
 mod query;
 mod s10;
 mod spatial;
+mod step_context;
 mod system;
 mod ukf;
 mod world_builder;
@@ -36,9 +36,9 @@ pub use error::*;
 pub use exec::*;
 pub use graph::*;
 pub use linalg::*;
-pub use post_step_context::*;
 pub use query::*;
 pub use spatial::*;
+pub use step_context::*;
 pub use system::*;
 pub use world_builder::*;
 
@@ -143,7 +143,7 @@ pub fn elodin(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<QueryMetadata>()?;
     m.add_class::<SystemBuilder>()?;
     m.add_class::<System>()?;
-    m.add_class::<PostStepContext>()?;
+    m.add_class::<StepContext>()?;
     m.add_function(wrap_pyfunction!(six_dof, m)?)?;
     m.add_function(wrap_pyfunction!(skew, m)?)?;
     m.add_function(wrap_pyfunction!(_get_cache_dir, m)?)?;