From 0e25ce005dc9b14bbd27fe88c2ee4ba74f1df7ad Mon Sep 17 00:00:00 2001 From: YangjunZ <103080153+YangjunZ@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:07:59 -0700 Subject: [PATCH 1/5] feat(datagen): reshape checkpoint delta-log to schema v2 Bring the datagen checkpoint delta-log in line with the authoritative spec. This changes the event vocabulary and folded model only; the storage mechanism (single append-only log.lance, MemWAL sharding, deterministic event_id, blob offload) is unchanged. - Event types 6 -> 7: add STEP_STARTED. - Replace the `terminal` column with a `status` column (running / completed / filtered / failed). - Replace step_instance_id + iteration provenance with structured step_kind / enclosing_step / selector_step. - Structured DatagenItemId with materialized path `root/step:idx/...`. - New DatagenStepKind {Root, Leaf, Sequence, Loop, MapReduce, Branch, SubPipeline, Conditional, Router}; only drivers emit STEP_STARTED. - fold_datagen_events returns Option (None when no ITEM_CREATED); FIELD_SET last-writer-wins, FIELD_APPEND accumulates; two read lenses (lifecycle vs failure). DATAGEN_SCHEMA_VERSION = 2. --- crates/lance-context-core/src/datagen.rs | 868 ++++++++++++------ .../lance-context-core/src/datagen_store.rs | 145 +-- crates/lance-context-core/src/lib.rs | 9 +- 3 files changed, 708 insertions(+), 314 deletions(-) diff --git a/crates/lance-context-core/src/datagen.rs b/crates/lance-context-core/src/datagen.rs index 8014353..eb0350c 100644 --- a/crates/lance-context-core/src/datagen.rs +++ b/crates/lance-context-core/src/datagen.rs @@ -1,16 +1,18 @@ -use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::fmt; use chrono::{DateTime, Utc}; use serde_json::Value; use uuid::Uuid; /// Current schema version for the append-only datagen checkpoint log. -pub const DATAGEN_SCHEMA_VERSION: i32 = 1; +pub const DATAGEN_SCHEMA_VERSION: i32 = 2; /// One lifecycle or field-level event in a datagen item's checkpoint history. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum DatagenEventType { ItemCreated, + StepStarted, FieldSet, FieldAppend, StepCompleted, @@ -23,6 +25,7 @@ impl DatagenEventType { pub fn as_str(self) -> &'static str { match self { Self::ItemCreated => "ITEM_CREATED", + Self::StepStarted => "STEP_STARTED", Self::FieldSet => "FIELD_SET", Self::FieldAppend => "FIELD_APPEND", Self::StepCompleted => "STEP_COMPLETED", @@ -34,6 +37,7 @@ impl DatagenEventType { pub fn parse(value: &str) -> Result { match value { "ITEM_CREATED" => Ok(Self::ItemCreated), + "STEP_STARTED" => Ok(Self::StepStarted), "FIELD_SET" => Ok(Self::FieldSet), "FIELD_APPEND" => Ok(Self::FieldAppend), "STEP_COMPLETED" => Ok(Self::StepCompleted), @@ -44,7 +48,171 @@ impl DatagenEventType { } } -/// Terminal outcome of an item. +/// The composition kind of a step. Drives two behaviors: +/// - forks a stream: `MapReduce`/`Branch`/`SubPipeline` sub-items get their own `item_id`. +/// - drives a frame: `Sequence`/`Loop` are the `enclosing_step` frames; only these emit STEP_STARTED. +/// +/// `Conditional`/`Router` are selectors (wrap one chosen child); `Leaf`/`Root` are the endpoints. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum DatagenStepKind { + Root, + Leaf, + Sequence, + Loop, + MapReduce, + Branch, + SubPipeline, + Conditional, + Router, +} + +impl DatagenStepKind { + #[must_use] + pub fn as_str(self) -> &'static str { + match self { + Self::Root => "root", + Self::Leaf => "leaf", + Self::Sequence => "sequence", + Self::Loop => "loop", + Self::MapReduce => "map_reduce", + Self::Branch => "branch", + Self::SubPipeline => "sub_pipeline", + Self::Conditional => "conditional", + Self::Router => "router", + } + } + + pub fn parse(value: &str) -> Result { + match value { + "root" => Ok(Self::Root), + "leaf" => Ok(Self::Leaf), + "sequence" => Ok(Self::Sequence), + "loop" => Ok(Self::Loop), + "map_reduce" => Ok(Self::MapReduce), + "branch" => Ok(Self::Branch), + "sub_pipeline" => Ok(Self::SubPipeline), + "conditional" => Ok(Self::Conditional), + "router" => Ok(Self::Router), + other => Err(format!("unsupported datagen step kind '{other}'")), + } + } + + /// A driver frame (`Sequence`/`Loop`) is the only kind that emits STEP_STARTED. + #[must_use] + pub fn is_driver(self) -> bool { + matches!(self, Self::Sequence | Self::Loop) + } +} + +/// An item (stream) identity — a structured value stored as a materialized path string. Root ids are +/// built by the executor from a source key (e.g. `5`); sub-item ids extend a parent with one fan-out +/// segment (`5/expand:0`) via [`DatagenItemId::child`]. The store owns the parse<->format; the client +/// only ever holds the structured form. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct DatagenItemId { + root_key: String, + /// One `(origin_step, branch_idx)` per fan-out hop below the root. + segments: Vec<(String, i64)>, +} + +impl DatagenItemId { + /// Build a root id from the executor's source key (the base case). + #[must_use] + pub fn from_source_key(key: &str) -> Self { + Self { + root_key: key.to_string(), + segments: Vec::new(), + } + } + + /// Extend this id with one fan-out segment -> the sub-item's id. Pure, deterministic, no I/O. + /// The sole id-composition entry point. Valid to call for a branch that was never written. + #[must_use] + pub fn child(&self, origin_step: &str, branch_idx: i64) -> Self { + let mut segments = self.segments.clone(); + segments.push((origin_step.to_string(), branch_idx)); + Self { + root_key: self.root_key.clone(), + segments, + } + } + + /// The parent stream's id (`None` on a root). + #[must_use] + pub fn parent(&self) -> Option { + if self.segments.is_empty() { + return None; + } + let mut segments = self.segments.clone(); + segments.pop(); + Some(Self { + root_key: self.root_key.clone(), + segments, + }) + } + + /// The root of this id's tree (== self if root). + #[must_use] + pub fn root(&self) -> Self { + Self { + root_key: self.root_key.clone(), + segments: Vec::new(), + } + } + + /// The fan-out step that created this sub-item (`None` on a root). + #[must_use] + pub fn origin_step(&self) -> Option<&str> { + self.segments.last().map(|(step, _)| step.as_str()) + } + + /// Which branch this sub-item is (`None` on a root). + #[must_use] + pub fn branch_idx(&self) -> Option { + self.segments.last().map(|(_, idx)| *idx) + } + + #[must_use] + pub fn is_root(&self) -> bool { + self.segments.is_empty() + } + + /// Parse a stored path string back into a structured id. + pub fn parse(path: &str) -> Result { + let mut parts = path.split('/'); + let root_key = parts + .next() + .filter(|part| !part.is_empty()) + .ok_or_else(|| format!("empty datagen item id '{path}'"))? + .to_string(); + let mut segments = Vec::new(); + for part in parts { + let (step, idx) = part + .split_once(':') + .ok_or_else(|| format!("malformed item id segment '{part}' in '{path}'"))?; + let branch_idx = idx + .parse::() + .map_err(|_| format!("non-integer branch index '{idx}' in '{path}'"))?; + segments.push((step.to_string(), branch_idx)); + } + Ok(Self { + root_key, + segments, + }) + } +} + +impl fmt::Display for DatagenItemId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.root_key)?; + for (step, idx) in &self.segments { + write!(formatter, "/{step}:{idx}")?; + } + Ok(()) + } +} + +/// Terminal outcome of an item (the write-side input to [`DatagenStreamWriter::item_terminal`]). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum DatagenTerminal { Completed, @@ -59,28 +227,43 @@ impl DatagenTerminal { Self::Filtered => "filtered", } } - - pub fn parse(value: &str) -> Result { - match value { - "completed" => Ok(Self::Completed), - "filtered" => Ok(Self::Filtered), - other => Err(format!("unsupported datagen terminal value '{other}'")), - } - } } -/// Current status derived exclusively by folding the event log. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// The value type of the `status` column. Two read lenses: lifecycle {Running, Completed, Filtered} +/// vs failure {Failed}. A folded item's status is only ever a lifecycle value; `Failed` surfaces only +/// through the failure lens (overview / failure history). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] pub enum DatagenItemStatus { - Pending, Running, Completed, Filtered, Failed, } -/// Lazy reference to an inline blob event. `bytes` is absent on normal fold and -/// trajectory reads; callers materialize it through `DatagenStore::get_blob`. +impl DatagenItemStatus { + #[must_use] + pub fn as_str(self) -> &'static str { + match self { + Self::Running => "running", + Self::Completed => "completed", + Self::Filtered => "filtered", + Self::Failed => "failed", + } + } + + pub fn parse(value: &str) -> Result { + match value { + "running" => Ok(Self::Running), + "completed" => Ok(Self::Completed), + "filtered" => Ok(Self::Filtered), + "failed" => Ok(Self::Failed), + other => Err(format!("unsupported datagen status '{other}'")), + } + } +} + +/// Lazy reference to an inline blob event. `bytes` is absent on normal (lazy) fold and trajectory +/// reads; callers materialize it through `DatagenStore::load_blob`. #[derive(Debug, Clone, PartialEq, Eq)] pub struct DatagenBlobValue { pub bytes: Option>, @@ -94,7 +277,7 @@ pub enum DatagenValue { Int(i64), Float(f64), Bool(bool), - String(String), + Str(String), Json(Value), Blob(DatagenBlobValue), } @@ -106,14 +289,159 @@ impl DatagenValue { Self::Int(_) => "int", Self::Float(_) => "float", Self::Bool(_) => "bool", - Self::String(_) => "str", + Self::Str(_) => "str", Self::Json(_) => "json", Self::Blob(_) => "blob", } } } -/// A single append-only row in `log.lance`. +/// A step's identity: its (globally-unique) name + its kind. Name and kind always travel together. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct DatagenStepId { + pub name: String, + pub kind: DatagenStepKind, +} + +/// A position within one stream's step tree — the coordinate a step write is attributed to. Maps 1:1 +/// to the Group C provenance columns. `enclosing`/`selector` are stored as bare step *names* (the log +/// has no enclosing/selector kind column); `None` means "directly under the stream root" / "no +/// selector". +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct DatagenStreamPosition { + pub step: DatagenStepId, + pub index: i64, + pub enclosing: Option, + pub selector: Option, +} + +/// How a field folds: FIELD_SET replaces (last-writer-wins); FIELD_APPEND accumulates in order. +#[derive(Debug, Clone, PartialEq)] +pub enum DatagenFieldState { + Set(DatagenValue), + Appended(Vec), +} + +/// A pointer to one completed step position (a single STEP_COMPLETED). +#[derive(Debug, Clone, PartialEq)] +pub struct DatagenStepCursor { + pub position: DatagenStreamPosition, + /// The STEP_COMPLETED's `item_seq` — the fold cutoff for "state as of this step". + pub item_seq: i64, +} + +/// The ordered list of cursors an item passed through, plus sets for O(1) skip lookup. `completed` +/// gates STEP_COMPLETED (re-)emission; `started` gates STEP_STARTED. `started \ completed` = frames +/// that were open when the process died. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct DatagenTrajectory { + pub ordered: Vec, + pub completed: HashSet, + pub started: HashSet, +} + +/// Error payload, shared by the write side (input to `item_failed`) and the read side (composed into +/// [`DatagenFailure`]). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DatagenErrorInfo { + pub error_type: String, + pub error_dump: Option, + pub traceback: Option, +} + +/// One failure record — a lightweight pointer to a FAILED event (no folded item). An item may have +/// 0..N of these across attempts. +#[derive(Debug, Clone, PartialEq)] +pub struct DatagenFailure { + pub at: DatagenStepCursor, + pub run_id: String, + pub attempt: i32, + pub error: DatagenErrorInfo, +} + +/// One item reconstructed by folding its events (latest state). Carries enough to continue processing +/// and to rebuild a `DatagenStreamWriter` on resume. Does not carry failures. +#[derive(Debug, Clone, PartialEq)] +pub struct FoldedDatagenItem { + pub item_id: DatagenItemId, + pub root_item_id: DatagenItemId, + pub parent_item_id: Option, + pub status: DatagenItemStatus, + /// Max `item_seq` -> resume continues at `last_item_seq + 1`. + pub last_item_seq: i64, + /// Max `attempt` seen -> resume runs at `last_attempt + 1`. + pub last_attempt: i32, + pub fields: BTreeMap, + pub trajectory: DatagenTrajectory, + pub query_tags: Option, + /// Internal `field_name -> event_id` map for the folded blob fields, so the store can resolve a + /// lazy blob without the caller handling an `event_id`. + pub blob_event_ids: BTreeMap, +} + +/// Result of a resumption fold. `NeverStarted` (no ITEM_CREATED) is the fresh-vs-restore fork the +/// executor acts on; `Found` carries the folded item (whose `status` is the lifecycle status). +#[derive(Debug, Clone, PartialEq)] +pub enum DatagenItemLookup { + NeverStarted, + Found(FoldedDatagenItem), +} + +impl DatagenItemLookup { + #[must_use] + pub fn folded(&self) -> Option<&FoldedDatagenItem> { + match self { + Self::NeverStarted => None, + Self::Found(item) => Some(item), + } + } +} + +/// Bulk startup classification of root items. A missing id means "never started". +#[derive(Debug, Clone, Default, PartialEq)] +pub struct DatagenRootItemStatuses { + inner: HashMap, +} + +impl DatagenRootItemStatuses { + #[must_use] + pub fn from_map(inner: HashMap) -> Self { + Self { inner } + } + + /// The classified status of a root item, or `None` if it was never started. + #[must_use] + pub fn status(&self, item_id: &DatagenItemId) -> Option { + self.inner.get(&item_id.to_string()).copied() + } + + /// Whether this item reached a terminal lifecycle state (Completed | Filtered). + #[must_use] + pub fn is_terminated(&self, item_id: &DatagenItemId) -> bool { + matches!( + self.status(item_id), + Some(DatagenItemStatus::Completed | DatagenItemStatus::Filtered) + ) + } + + #[must_use] + pub fn len(&self) -> usize { + self.inner.len() + } + + #[must_use] + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + #[must_use] + pub fn iter(&self) -> impl Iterator { + self.inner.iter() + } +} + +/// One append-only row in `log.lance`. Item/root/parent ids are the stored path strings; fold parses +/// them into [`DatagenItemId`]. #[derive(Debug, Clone, PartialEq)] pub struct DatagenEvent { /// Deterministic idempotency key. The MemWAL read path de-duplicates by it. @@ -121,16 +449,16 @@ pub struct DatagenEvent { pub item_id: String, pub root_item_id: String, pub parent_item_id: Option, - /// Strictly increasing per item. A collision between different event ids is - /// treated as split-brain corruption during fold. + /// Strictly increasing per item. A collision between different event ids is split-brain corruption. pub item_seq: i64, /// Shared by every event emitted for one checkpoint boundary. pub checkpoint_id: String, pub event_type: DatagenEventType, pub step_name: Option, + pub step_kind: Option, pub step_index: Option, - pub step_instance_id: Option, - pub iteration: Option, + pub enclosing_step: Option, + pub selector_step: Option, pub attempt: i32, pub run_id: String, /// Fencing identity for the writer/lease that owned this item. @@ -140,9 +468,10 @@ pub struct DatagenEvent { pub field_type: Option, pub codec_version: Option, pub value: Option, - /// Query tags captured on ITEM_CREATED. They are not part of correctness. + /// Query tags captured on ITEM_CREATED. Not part of correctness. pub query_tags: Option, - pub terminal: Option, + /// The stored lifecycle/failure status (populated on ITEM_CREATED / TERMINAL / FAILED). + pub status: Option, pub error_type: Option, pub error_dump: Option, pub traceback: Option, @@ -188,41 +517,38 @@ impl DatagenEvent { if self.value.is_none() { return Err("field events require a value".to_string()); } - if self.step_name.as_deref().is_none_or(str::is_empty) - || self.step_index.is_none() - || self.step_instance_id.as_deref().is_none_or(str::is_empty) - { - return Err( - "field events require step_name, step_index, and step_instance_id" - .to_string(), - ); - } + self.require_step_provenance("field events")?; } - DatagenEventType::StepCompleted => { - if self.step_name.as_deref().is_none_or(str::is_empty) - || self.step_index.is_none() - || self.step_instance_id.as_deref().is_none_or(str::is_empty) - { - return Err( - "STEP_COMPLETED requires step_name, step_index, and step_instance_id" - .to_string(), - ); - } + DatagenEventType::StepStarted | DatagenEventType::StepCompleted => { + self.require_step_provenance(self.event_type.as_str())?; } DatagenEventType::Failed => { if self.error_type.as_deref().is_none_or(str::is_empty) { return Err("FAILED requires error_type".to_string()); } + self.require_step_provenance("FAILED")?; } - DatagenEventType::Terminal => { - if self.terminal.is_none() { - return Err("TERMINAL requires terminal".to_string()); - } - } + DatagenEventType::Terminal => match self.status { + Some(DatagenItemStatus::Completed | DatagenItemStatus::Filtered) => {} + _ => return Err("TERMINAL requires status completed or filtered".to_string()), + }, DatagenEventType::ItemCreated => {} } Ok(()) } + + fn require_step_provenance(&self, context: &str) -> Result<(), String> { + if self.step_name.as_deref().is_none_or(str::is_empty) { + return Err(format!("{context} require step_name")); + } + if self.step_kind.is_none() { + return Err(format!("{context} require step_kind")); + } + if self.step_index.is_none() { + return Err(format!("{context} require step_index")); + } + Ok(()) + } } /// Generate a deterministic event id for retry-safe checkpoint ingestion. @@ -232,89 +558,52 @@ pub fn datagen_event_id(item_id: &str, checkpoint_id: &str, ordinal: u32) -> Str Uuid::new_v5(&Uuid::NAMESPACE_OID, input.as_bytes()).to_string() } -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] -pub struct DatagenStepCursor { - pub checkpoint_id: String, - pub step_name: String, - pub step_index: i64, - pub step_instance_id: String, - pub iteration: Option, - pub attempt: i32, -} - -#[derive(Debug, Clone, PartialEq)] -pub enum DatagenFieldState { - Set(DatagenValue), - Appended(Vec), -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct DatagenFailure { - pub event_id: String, - pub run_id: String, - pub checkpoint_id: String, - pub item_seq: i64, - pub step_name: Option, - pub error_type: String, - pub error_dump: Option, - pub traceback: Option, - pub failed_at: DateTime, -} - -/// Current item state reconstructed solely from the append-only log. -#[derive(Debug, Clone, PartialEq)] -pub struct FoldedDatagenItem { - pub item_id: String, - pub root_item_id: String, - pub parent_item_id: Option, - pub fields: BTreeMap, - pub completed_steps: BTreeSet, - pub status: DatagenItemStatus, - pub terminal: Option, - pub failure: Option, - pub query_tags: Option, - pub current_run_id: String, - pub last_item_seq: i64, - pub last_checkpoint_id: String, -} - -/// State captured immediately after a STEP_COMPLETED event. -#[derive(Debug, Clone, PartialEq)] -pub struct DatagenTrajectoryPoint { - pub cursor: DatagenStepCursor, - pub item: FoldedDatagenItem, -} - -pub fn fold_datagen_events(events: &[DatagenEvent]) -> Result { +/// Fold an item's events into its latest state. Returns `None` if there is no ITEM_CREATED (the item +/// was never started). +pub fn fold_datagen_events(events: &[DatagenEvent]) -> Result, String> { let ordered = normalize_events(events)?; - let first = ordered - .first() - .ok_or_else(|| "cannot fold an empty datagen event list".to_string())?; - let mut item = initial_item(first); + let Some(first) = ordered.first() else { + return Ok(None); + }; + if first.event_type != DatagenEventType::ItemCreated { + return Ok(None); + } + let mut item = initial_item(first)?; for event in ordered { apply_event(&mut item, event)?; } - Ok(item) + Ok(Some(item)) +} + +/// The ordered completed-step cursors of an item, in `item_seq` order. +pub fn datagen_trajectory(events: &[DatagenEvent]) -> Result, String> { + Ok(fold_datagen_events(events)? + .map(|item| item.trajectory.ordered) + .unwrap_or_default()) } -pub fn datagen_trajectory(events: &[DatagenEvent]) -> Result, String> { +/// The lightweight failure pointers of an item, in `item_seq` order. +pub fn datagen_failures(events: &[DatagenEvent]) -> Result, String> { let ordered = normalize_events(events)?; - let first = ordered - .first() - .ok_or_else(|| "cannot build a trajectory from an empty event list".to_string())?; - let mut item = initial_item(first); - let mut trajectory = Vec::new(); + let mut failures = Vec::new(); for event in ordered { - apply_event(&mut item, event)?; - if event.event_type == DatagenEventType::StepCompleted { - let cursor = step_cursor(event)?; - trajectory.push(DatagenTrajectoryPoint { - cursor, - item: item.clone(), + if event.event_type == DatagenEventType::Failed { + failures.push(DatagenFailure { + at: DatagenStepCursor { + position: stream_position(event)?, + item_seq: event.item_seq, + }, + run_id: event.run_id.clone(), + attempt: event.attempt, + error: DatagenErrorInfo { + error_type: event.error_type.clone().unwrap(), + error_dump: event.error_dump.clone(), + traceback: event.traceback.clone(), + }, }); } } - Ok(trajectory) + Ok(failures) } fn normalize_events(events: &[DatagenEvent]) -> Result, String> { @@ -352,69 +641,56 @@ fn normalize_events(events: &[DatagenEvent]) -> Result, Strin Ok(ordered) } -fn initial_item(first: &DatagenEvent) -> FoldedDatagenItem { - FoldedDatagenItem { - item_id: first.item_id.clone(), - root_item_id: first.root_item_id.clone(), - parent_item_id: first.parent_item_id.clone(), +fn initial_item(first: &DatagenEvent) -> Result { + let parent_item_id = match &first.parent_item_id { + Some(parent) => Some(DatagenItemId::parse(parent)?), + None => None, + }; + Ok(FoldedDatagenItem { + item_id: DatagenItemId::parse(&first.item_id)?, + root_item_id: DatagenItemId::parse(&first.root_item_id)?, + parent_item_id, + status: DatagenItemStatus::Running, + last_item_seq: first.item_seq, + last_attempt: first.attempt, fields: BTreeMap::new(), - completed_steps: BTreeSet::new(), - status: DatagenItemStatus::Pending, - terminal: None, - failure: None, + trajectory: DatagenTrajectory::default(), query_tags: None, - current_run_id: first.run_id.clone(), - last_item_seq: first.item_seq, - last_checkpoint_id: first.checkpoint_id.clone(), - } + blob_event_ids: BTreeMap::new(), + }) } fn apply_event(item: &mut FoldedDatagenItem, event: &DatagenEvent) -> Result<(), String> { - if event.item_id != item.item_id { + if DatagenItemId::parse(&event.item_id)? != item.item_id { return Err(format!( "event '{}' belongs to item '{}', expected '{}'", event.event_id, event.item_id, item.item_id )); } - if event.root_item_id != item.root_item_id { - return Err(format!( - "item '{}' changed root_item_id from '{}' to '{}'", - item.item_id, item.root_item_id, event.root_item_id - )); - } - if event.parent_item_id != item.parent_item_id { - return Err(format!( - "item '{}' changed parent_item_id during its trajectory", - item.item_id - )); - } - item.current_run_id = event.run_id.clone(); - item.last_item_seq = event.item_seq; - item.last_checkpoint_id = event.checkpoint_id.clone(); + item.last_item_seq = item.last_item_seq.max(event.item_seq); + item.last_attempt = item.last_attempt.max(event.attempt); match event.event_type { DatagenEventType::ItemCreated => { - item.status = DatagenItemStatus::Pending; - item.terminal = None; - item.failure = None; + item.status = DatagenItemStatus::Running; if event.query_tags.is_some() { item.query_tags = event.query_tags.clone(); } } + DatagenEventType::StepStarted => { + item.trajectory.started.insert(stream_position(event)?); + } DatagenEventType::FieldSet => { let field_name = event.field_name.clone().unwrap(); - item.fields.insert( - field_name, - DatagenFieldState::Set(event.value.clone().unwrap()), - ); - item.status = DatagenItemStatus::Running; - item.terminal = None; - item.failure = None; + let value = event.value.clone().unwrap(); + record_blob_event_id(item, &field_name, &value, &event.event_id); + item.fields.insert(field_name, DatagenFieldState::Set(value)); } DatagenEventType::FieldAppend => { let field_name = event.field_name.clone().unwrap(); let value = event.value.clone().unwrap(); + record_blob_event_id(item, &field_name, &value, &event.event_id); match item.fields.entry(field_name) { std::collections::btree_map::Entry::Vacant(entry) => { entry.insert(DatagenFieldState::Appended(vec![value])); @@ -429,60 +705,56 @@ fn apply_event(item: &mut FoldedDatagenItem, event: &DatagenEvent) -> Result<(), } }, } - item.status = DatagenItemStatus::Running; - item.terminal = None; - item.failure = None; } DatagenEventType::StepCompleted => { - item.completed_steps.insert(step_cursor(event)?); - item.status = DatagenItemStatus::Running; - item.terminal = None; - item.failure = None; - } - DatagenEventType::Failed => { - item.status = DatagenItemStatus::Failed; - item.terminal = None; - item.failure = Some(DatagenFailure { - event_id: event.event_id.clone(), - run_id: event.run_id.clone(), - checkpoint_id: event.checkpoint_id.clone(), + let position = stream_position(event)?; + item.trajectory.completed.insert(position.clone()); + item.trajectory.ordered.push(DatagenStepCursor { + position, item_seq: event.item_seq, - step_name: event.step_name.clone(), - error_type: event.error_type.clone().unwrap(), - error_dump: event.error_dump.clone(), - traceback: event.traceback.clone(), - failed_at: event.event_ts, }); } + DatagenEventType::Failed => { + // Failure lens only: a FAILED row leaves the item Running under the lifecycle lens. + } DatagenEventType::Terminal => { - let terminal = event.terminal.unwrap(); - item.status = match terminal { - DatagenTerminal::Completed => DatagenItemStatus::Completed, - DatagenTerminal::Filtered => DatagenItemStatus::Filtered, + item.status = match event.status { + Some(status @ (DatagenItemStatus::Completed | DatagenItemStatus::Filtered)) => status, + _ => return Err("TERMINAL event missing completed/filtered status".to_string()), }; - item.terminal = Some(terminal); - item.failure = None; } } Ok(()) } -fn step_cursor(event: &DatagenEvent) -> Result { - Ok(DatagenStepCursor { - checkpoint_id: event.checkpoint_id.clone(), - step_name: event - .step_name - .clone() - .ok_or_else(|| "STEP_COMPLETED missing step_name".to_string())?, - step_index: event +fn record_blob_event_id( + item: &mut FoldedDatagenItem, + field_name: &str, + value: &DatagenValue, + event_id: &str, +) { + if matches!(value, DatagenValue::Blob(_)) { + item.blob_event_ids + .insert(field_name.to_string(), event_id.to_string()); + } +} + +fn stream_position(event: &DatagenEvent) -> Result { + Ok(DatagenStreamPosition { + step: DatagenStepId { + name: event + .step_name + .clone() + .ok_or_else(|| "step event missing step_name".to_string())?, + kind: event + .step_kind + .ok_or_else(|| "step event missing step_kind".to_string())?, + }, + index: event .step_index - .ok_or_else(|| "STEP_COMPLETED missing step_index".to_string())?, - step_instance_id: event - .step_instance_id - .clone() - .ok_or_else(|| "STEP_COMPLETED missing step_instance_id".to_string())?, - iteration: event.iteration, - attempt: event.attempt, + .ok_or_else(|| "step event missing step_index".to_string())?, + enclosing: event.enclosing_step.clone(), + selector: event.selector_step.clone(), }) } @@ -495,17 +767,18 @@ mod tests { fn event(seq: i64, event_type: DatagenEventType) -> DatagenEvent { let checkpoint_id = format!("checkpoint-{seq}"); DatagenEvent { - event_id: datagen_event_id("item-1", &checkpoint_id, 0), - item_id: "item-1".to_string(), - root_item_id: "item-1".to_string(), + event_id: datagen_event_id("5", &checkpoint_id, 0), + item_id: "5".to_string(), + root_item_id: "5".to_string(), parent_item_id: None, item_seq: seq, checkpoint_id, event_type, step_name: None, + step_kind: None, step_index: None, - step_instance_id: None, - iteration: None, + enclosing_step: None, + selector_step: None, attempt: 0, run_id: "run-1".to_string(), writer_epoch: "writer-1".to_string(), @@ -514,7 +787,7 @@ mod tests { codec_version: None, value: None, query_tags: None, - terminal: None, + status: None, error_type: None, error_dump: None, traceback: None, @@ -523,77 +796,158 @@ mod tests { } } - fn completed_step(seq: i64) -> DatagenEvent { - let mut event = event(seq, DatagenEventType::StepCompleted); - event.step_name = Some("noop".to_string()); - event.step_index = Some(3); - event.step_instance_id = Some("loop/2/noop".to_string()); - event.iteration = Some(2); - event + fn created(seq: i64) -> DatagenEvent { + let mut created = event(seq, DatagenEventType::ItemCreated); + created.status = Some(DatagenItemStatus::Running); + created + } + + fn leaf_completed(seq: i64, name: &str, index: i64, enclosing: Option<&str>) -> DatagenEvent { + let mut completed = event(seq, DatagenEventType::StepCompleted); + completed.step_name = Some(name.to_string()); + completed.step_kind = Some(DatagenStepKind::Leaf); + completed.step_index = Some(index); + completed.enclosing_step = enclosing.map(str::to_string); + completed + } + + fn driver_started(seq: i64, name: &str, index: i64, enclosing: Option<&str>) -> DatagenEvent { + let mut started = event(seq, DatagenEventType::StepStarted); + started.step_name = Some(name.to_string()); + started.step_kind = Some(DatagenStepKind::Sequence); + started.step_index = Some(index); + started.enclosing_step = enclosing.map(str::to_string); + started } #[test] - fn no_op_step_is_present_in_fold_and_trajectory() { - let created = event(0, DatagenEventType::ItemCreated); - let completed = completed_step(1); + fn item_id_round_trips_and_navigates() { + let root = DatagenItemId::from_source_key("5"); + let child = root.child("expand", 0).child("enrich", 1); + assert_eq!(child.to_string(), "5/expand:0/enrich:1"); + assert_eq!(DatagenItemId::parse("5/expand:0/enrich:1").unwrap(), child); + assert_eq!(child.origin_step(), Some("enrich")); + assert_eq!(child.branch_idx(), Some(1)); + assert_eq!(child.parent().unwrap().to_string(), "5/expand:0"); + assert_eq!(child.root(), root); + assert!(root.is_root()); + assert_eq!(root.origin_step(), None); + } - let folded = fold_datagen_events(&[completed.clone(), created]).unwrap(); - assert_eq!(folded.status, DatagenItemStatus::Running); - assert_eq!(folded.completed_steps.len(), 1); - assert!(folded.fields.is_empty()); + #[test] + fn never_started_folds_to_none() { + let completed = leaf_completed(1, "gen", 0, Some("main")); + assert!(fold_datagen_events(&[completed]).unwrap().is_none()); + } - let trajectory = datagen_trajectory(&[completed]).unwrap(); - assert_eq!(trajectory.len(), 1); - assert_eq!(trajectory[0].cursor.step_instance_id, "loop/2/noop"); + #[test] + fn fold_tracks_started_and_completed_positions() { + let events = [ + created(0), + driver_started(1, "main", 0, None), + leaf_completed(2, "gen", 0, Some("main")), + ]; + let folded = fold_datagen_events(&events).unwrap().unwrap(); + assert_eq!(folded.status, DatagenItemStatus::Running); + assert_eq!(folded.trajectory.ordered.len(), 1); + assert_eq!(folded.trajectory.completed.len(), 1); + assert_eq!(folded.trajectory.started.len(), 1); + assert_eq!(folded.last_item_seq, 2); } #[test] - fn retry_duplicate_event_is_folded_once() { - let created = event(0, DatagenEventType::ItemCreated); - let mut append = event(1, DatagenEventType::FieldAppend); - append.field_name = Some("messages".to_string()); + fn field_set_is_last_writer_wins_and_append_accumulates() { + let mut set_v1 = leaf_completed(1, "gen", 0, Some("main")); + set_v1.event_type = DatagenEventType::FieldSet; + set_v1.field_name = Some("draft".to_string()); + set_v1.field_type = Some("str".to_string()); + set_v1.codec_version = Some(1); + set_v1.value = Some(DatagenValue::Str("v1".to_string())); + + let mut set_v2 = set_v1.clone(); + set_v2.item_seq = 2; + set_v2.checkpoint_id = "c2".to_string(); + set_v2.event_id = datagen_event_id("5", "c2", 0); + set_v2.value = Some(DatagenValue::Str("v2".to_string())); + + let mut append = leaf_completed(3, "b1", 0, Some("body")); + append.event_type = DatagenEventType::FieldAppend; + append.field_name = Some("revisions".to_string()); append.field_type = Some("json".to_string()); append.codec_version = Some(1); - append.value = Some(DatagenValue::Json(json!({"role": "assistant"}))); - append.step_name = Some("generate".to_string()); - append.step_index = Some(1); - append.step_instance_id = Some("generate/0".to_string()); - - let folded = fold_datagen_events(&[created, append.clone(), append.clone()]).unwrap(); + append.value = Some(DatagenValue::Json(json!({"n": "a"}))); + let mut append2 = append.clone(); + append2.item_seq = 4; + append2.checkpoint_id = "c4".to_string(); + append2.event_id = datagen_event_id("5", "c4", 0); + append2.value = Some(DatagenValue::Json(json!({"n": "b"}))); + + let folded = + fold_datagen_events(&[created(0), set_v1, set_v2, append, append2]).unwrap().unwrap(); assert_eq!( - folded.fields.get("messages"), - Some(&DatagenFieldState::Appended(vec![DatagenValue::Json( - json!({"role": "assistant"}) - )])) + folded.fields.get("draft"), + Some(&DatagenFieldState::Set(DatagenValue::Str("v2".to_string()))) + ); + assert_eq!( + folded.fields.get("revisions"), + Some(&DatagenFieldState::Appended(vec![ + DatagenValue::Json(json!({"n": "a"})), + DatagenValue::Json(json!({"n": "b"})), + ])) ); } + #[test] + fn terminal_sets_lifecycle_status() { + let mut terminal = event(2, DatagenEventType::Terminal); + terminal.status = Some(DatagenItemStatus::Completed); + let folded = fold_datagen_events(&[created(0), terminal]).unwrap().unwrap(); + assert_eq!(folded.status, DatagenItemStatus::Completed); + } + + #[test] + fn failed_leaves_item_running_but_surfaces_in_failures() { + let mut failed = leaf_completed(1, "check", 2, Some("solve")); + failed.event_type = DatagenEventType::Failed; + failed.status = Some(DatagenItemStatus::Failed); + failed.error_type = Some("ValueError".to_string()); + failed.attempt = 0; + + let folded = fold_datagen_events(&[created(0), failed.clone()]).unwrap().unwrap(); + assert_eq!(folded.status, DatagenItemStatus::Running); + + let failures = datagen_failures(&[created(0), failed]).unwrap(); + assert_eq!(failures.len(), 1); + assert_eq!(failures[0].error.error_type, "ValueError"); + assert_eq!(failures[0].at.position.step.name, "check"); + } + #[test] fn sequence_collision_is_rejected() { - let created = event(0, DatagenEventType::ItemCreated); - let first = completed_step(1); + let first = leaf_completed(1, "gen", 0, Some("main")); let mut second = first.clone(); second.event_id = "different-event".to_string(); second.checkpoint_id = "different-checkpoint".to_string(); - - let error = fold_datagen_events(&[created, first, second]).unwrap_err(); + let error = fold_datagen_events(&[created(0), first, second]).unwrap_err(); assert!(error.contains("conflicting events at item_seq 1")); } #[test] - fn later_run_can_supersede_a_failure() { - let created = event(0, DatagenEventType::ItemCreated); - let mut failed = event(1, DatagenEventType::Failed); - failed.error_type = Some("RuntimeError".to_string()); - - let mut retried = event(2, DatagenEventType::ItemCreated); - retried.run_id = "run-2".to_string(); - retried.checkpoint_id = "retry-created".to_string(); - retried.event_id = datagen_event_id("item-1", "retry-created", 0); - - let folded = fold_datagen_events(&[created, failed, retried]).unwrap(); - assert_eq!(folded.status, DatagenItemStatus::Pending); - assert_eq!(folded.current_run_id, "run-2"); - assert!(folded.failure.is_none()); + fn retry_duplicate_event_is_folded_once() { + let mut append = leaf_completed(1, "gen", 1, Some("main")); + append.event_type = DatagenEventType::FieldAppend; + append.field_name = Some("messages".to_string()); + append.field_type = Some("json".to_string()); + append.codec_version = Some(1); + append.value = Some(DatagenValue::Json(json!({"role": "assistant"}))); + + let folded = + fold_datagen_events(&[created(0), append.clone(), append]).unwrap().unwrap(); + assert_eq!( + folded.fields.get("messages"), + Some(&DatagenFieldState::Appended(vec![DatagenValue::Json( + json!({"role": "assistant"}) + )])) + ); } } diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index 228a3ef..f91000f 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -35,8 +35,9 @@ use tracing::{info, warn}; use uuid::Uuid; use crate::datagen::{ - datagen_trajectory, fold_datagen_events, DatagenBlobValue, DatagenEvent, DatagenEventType, - DatagenTrajectoryPoint, DatagenValue, FoldedDatagenItem, + datagen_failures, datagen_trajectory, fold_datagen_events, DatagenBlobValue, DatagenEvent, + DatagenEventType, DatagenFailure, DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses, + DatagenStepCursor, DatagenStepKind, DatagenValue, }; use crate::rollout_store::derive_shard_id; use crate::store::{column_as, column_as_optional, timestamp_from_micros}; @@ -224,23 +225,45 @@ impl DatagenStore { self.filtered_events(&filter).await } - /// Reconstruct one item's latest state exclusively from its event log. - pub async fn fold_item(&self, item_id: &str) -> LanceResult> { + /// Reconstruct one item's latest state exclusively from its event log. Returns + /// [`DatagenItemLookup::NeverStarted`] when the item has no ITEM_CREATED (the fresh-vs-resume fork). + pub async fn fold_item(&self, item_id: &str) -> LanceResult { let events = self.events_for_item(item_id).await?; - if events.is_empty() { - return Ok(None); + match fold_datagen_events(&events).map_err(invalid_input)? { + Some(item) => Ok(DatagenItemLookup::Found(item)), + None => Ok(DatagenItemLookup::NeverStarted), } - fold_datagen_events(&events) - .map(Some) - .map_err(invalid_input) } - /// Reconstruct state after every completed step without loading blob bytes. - pub async fn trajectory(&self, item_id: &str) -> LanceResult> { - let events = self.events_for_item(item_id).await?; - if events.is_empty() { - return Ok(Vec::new()); + /// Classify every root item that shares `root_item_id` with the given roots by folded lifecycle + /// status. A root not present in the log is simply absent from the result (never started). + pub async fn root_item_statuses( + &self, + root_item_ids: &[&str], + ) -> LanceResult { + let mut statuses = HashMap::new(); + for root_item_id in root_item_ids { + let events = self.events_for_root(root_item_id).await?; + let root_events: Vec = events + .into_iter() + .filter(|event| &event.item_id == root_item_id) + .collect(); + if let Some(item) = fold_datagen_events(&root_events).map_err(invalid_input)? { + statuses.insert(root_item_id.to_string(), item.status); + } } + Ok(DatagenRootItemStatuses::from_map(statuses)) + } + + /// Read all failure records for an item directly from the failure lens. + pub async fn item_failures(&self, item_id: &str) -> LanceResult> { + let events = self.events_for_item(item_id).await?; + datagen_failures(&events).map_err(invalid_input) + } + + /// Reconstruct the ordered step cursors an item passed through, without loading blob bytes. + pub async fn trajectory(&self, item_id: &str) -> LanceResult> { + let events = self.events_for_item(item_id).await?; datagen_trajectory(&events).map_err(invalid_input) } @@ -653,9 +676,10 @@ pub fn datagen_log_schema() -> Schema { Field::new("checkpoint_id", DataType::Utf8, false), Field::new("event_type", DataType::Utf8, false), Field::new("step_name", DataType::Utf8, true), + Field::new("step_kind", DataType::Utf8, true), Field::new("step_index", DataType::Int64, true), - Field::new("step_instance_id", DataType::Utf8, true), - Field::new("iteration", DataType::Int64, true), + Field::new("enclosing_step", DataType::Utf8, true), + Field::new("selector_step", DataType::Utf8, true), Field::new("attempt", DataType::Int32, false), Field::new("run_id", DataType::Utf8, false), Field::new("writer_epoch", DataType::Utf8, false), @@ -674,7 +698,7 @@ pub fn datagen_log_schema() -> Schema { Field::new("payload_size", DataType::Int64, true), Field::new("payload_checksum", DataType::Utf8, true), Field::new("query_tags_json", DataType::LargeUtf8, true), - Field::new("terminal", DataType::Utf8, true), + Field::new("status", DataType::Utf8, true), Field::new("error_type", DataType::Utf8, true), Field::new("error_dump", DataType::LargeUtf8, true), Field::new("traceback", DataType::LargeUtf8, true), @@ -773,8 +797,9 @@ fn validate_checkpoint_batch(events: &[DatagenEvent]) -> LanceResult<()> { }) { if event.step_name != completion.step_name || event.step_index != completion.step_index - || event.step_instance_id != completion.step_instance_id - || event.iteration != completion.iteration + || event.step_kind != completion.step_kind + || event.enclosing_step != completion.enclosing_step + || event.selector_step != completion.selector_step { return Err(invalid_input( "all field events must share the STEP_COMPLETED step identity", @@ -793,9 +818,10 @@ fn events_to_batch(events: &[DatagenEvent]) -> LanceResult { let mut checkpoint_id = StringBuilder::new(); let mut event_type = StringBuilder::new(); let mut step_name = StringBuilder::new(); + let mut step_kind = StringBuilder::new(); let mut step_index = Int64Builder::new(); - let mut step_instance_id = StringBuilder::new(); - let mut iteration = Int64Builder::new(); + let mut enclosing_step = StringBuilder::new(); + let mut selector_step = StringBuilder::new(); let mut attempt = Int32Builder::new(); let mut run_id = StringBuilder::new(); let mut writer_epoch = StringBuilder::new(); @@ -812,7 +838,7 @@ fn events_to_batch(events: &[DatagenEvent]) -> LanceResult { let mut payload_size = Int64Builder::new(); let mut payload_checksum = StringBuilder::new(); let mut query_tags_json = LargeStringBuilder::new(); - let mut terminal = StringBuilder::new(); + let mut status = StringBuilder::new(); let mut error_type = StringBuilder::new(); let mut error_dump = LargeStringBuilder::new(); let mut traceback = LargeStringBuilder::new(); @@ -829,9 +855,10 @@ fn events_to_batch(events: &[DatagenEvent]) -> LanceResult { checkpoint_id.append_value(&event.checkpoint_id); event_type.append_value(event.event_type.as_str()); step_name.append_option(event.step_name.as_deref()); + step_kind.append_option(event.step_kind.map(DatagenStepKind::as_str)); step_index.append_option(event.step_index); - step_instance_id.append_option(event.step_instance_id.as_deref()); - iteration.append_option(event.iteration); + enclosing_step.append_option(event.enclosing_step.as_deref()); + selector_step.append_option(event.selector_step.as_deref()); attempt.append_value(event.attempt); run_id.append_value(&event.run_id); writer_epoch.append_value(&event.writer_epoch); @@ -853,7 +880,7 @@ fn events_to_batch(events: &[DatagenEvent]) -> LanceResult { _ => None, }); value_str.append_option(match &event.value { - Some(DatagenValue::String(value)) => Some(value.as_str()), + Some(DatagenValue::Str(value)) => Some(value.as_str()), _ => None, }); match &event.value { @@ -876,7 +903,7 @@ fn events_to_batch(events: &[DatagenEvent]) -> LanceResult { Some(tags) => query_tags_json.append_value(tags.to_string()), None => query_tags_json.append_null(), } - terminal.append_option(event.terminal.map(|value| value.as_str())); + status.append_option(event.status.map(DatagenItemStatus::as_str)); error_type.append_option(event.error_type.as_deref()); error_dump.append_option(event.error_dump.as_deref()); traceback.append_option(event.traceback.as_deref()); @@ -894,9 +921,10 @@ fn events_to_batch(events: &[DatagenEvent]) -> LanceResult { Arc::new(checkpoint_id.finish()), Arc::new(event_type.finish()), Arc::new(step_name.finish()), + Arc::new(step_kind.finish()), Arc::new(step_index.finish()), - Arc::new(step_instance_id.finish()), - Arc::new(iteration.finish()), + Arc::new(enclosing_step.finish()), + Arc::new(selector_step.finish()), Arc::new(attempt.finish()), Arc::new(run_id.finish()), Arc::new(writer_epoch.finish()), @@ -913,7 +941,7 @@ fn events_to_batch(events: &[DatagenEvent]) -> LanceResult { Arc::new(payload_size.finish()), Arc::new(payload_checksum.finish()), Arc::new(query_tags_json.finish()), - Arc::new(terminal.finish()), + Arc::new(status.finish()), Arc::new(error_type.finish()), Arc::new(error_dump.finish()), Arc::new(traceback.finish()), @@ -932,9 +960,10 @@ fn batch_to_events(batch: &RecordBatch) -> LanceResult> { let checkpoint_id = column_as::(batch, "checkpoint_id")?; let event_type = column_as::(batch, "event_type")?; let step_name = column_as_optional::(batch, "step_name"); + let step_kind = column_as_optional::(batch, "step_kind"); let step_index = column_as_optional::(batch, "step_index"); - let step_instance_id = column_as_optional::(batch, "step_instance_id"); - let iteration = column_as_optional::(batch, "iteration"); + let enclosing_step = column_as_optional::(batch, "enclosing_step"); + let selector_step = column_as_optional::(batch, "selector_step"); let attempt = column_as::(batch, "attempt")?; let run_id = column_as::(batch, "run_id")?; let writer_epoch = column_as::(batch, "writer_epoch")?; @@ -951,7 +980,7 @@ fn batch_to_events(batch: &RecordBatch) -> LanceResult> { let payload_size = column_as_optional::(batch, "payload_size"); let payload_checksum = column_as_optional::(batch, "payload_checksum"); let query_tags_json = column_as_optional::(batch, "query_tags_json"); - let terminal = column_as_optional::(batch, "terminal"); + let status = column_as_optional::(batch, "status"); let error_type = column_as_optional::(batch, "error_type"); let error_dump = column_as_optional::(batch, "error_dump"); let traceback = column_as_optional::(batch, "traceback"); @@ -978,7 +1007,7 @@ fn batch_to_events(batch: &RecordBatch) -> LanceResult> { row, "value_bool", )?)), - Some("str") => Some(DatagenValue::String( + Some("str") => Some(DatagenValue::Str( optional_large_string(value_str, row) .ok_or_else(|| invalid_input("value_kind=str requires value_str"))?, )), @@ -1025,9 +1054,13 @@ fn batch_to_events(batch: &RecordBatch) -> LanceResult> { checkpoint_id: checkpoint_id.value(row).to_string(), event_type: DatagenEventType::parse(event_type.value(row)).map_err(invalid_input)?, step_name: optional_string(step_name, row), + step_kind: match optional_string(step_kind, row) { + Some(value) => Some(DatagenStepKind::parse(&value).map_err(invalid_input)?), + None => None, + }, step_index: optional_i64(step_index, row), - step_instance_id: optional_string(step_instance_id, row), - iteration: optional_i64(iteration, row), + enclosing_step: optional_string(enclosing_step, row), + selector_step: optional_string(selector_step, row), attempt: attempt.value(row), run_id: run_id.value(row).to_string(), writer_epoch: writer_epoch.value(row).to_string(), @@ -1036,10 +1069,8 @@ fn batch_to_events(batch: &RecordBatch) -> LanceResult> { codec_version: optional_i32(codec_version, row), value, query_tags, - terminal: match optional_string(terminal, row) { - Some(value) => { - Some(crate::datagen::DatagenTerminal::parse(&value).map_err(invalid_input)?) - } + status: match optional_string(status, row) { + Some(value) => Some(DatagenItemStatus::parse(&value).map_err(invalid_input)?), None => None, }, error_type: optional_string(error_type, row), @@ -1131,7 +1162,7 @@ fn is_fenced_error(error: &LanceError) -> bool { mod tests { use super::*; use crate::datagen::{ - datagen_event_id, DatagenFieldState, DatagenItemStatus, DatagenTerminal, + datagen_event_id, DatagenFieldState, DatagenItemStatus, DatagenStepKind, DATAGEN_SCHEMA_VERSION, }; use chrono::{TimeZone, Utc}; @@ -1154,9 +1185,10 @@ mod tests { checkpoint_id: checkpoint_id.to_string(), event_type, step_name: None, + step_kind: None, step_index: None, - step_instance_id: None, - iteration: None, + enclosing_step: None, + selector_step: None, attempt: 0, run_id: "run-1".to_string(), writer_epoch: "writer-1".to_string(), @@ -1165,7 +1197,7 @@ mod tests { codec_version: None, value: None, query_tags: None, - terminal: None, + status: None, error_type: None, error_dump: None, traceback: None, @@ -1184,8 +1216,8 @@ mod tests { ) -> DatagenEvent { let mut event = event("item-1", seq, "grade-0", ordinal, event_type); event.step_name = Some("grade".to_string()); + event.step_kind = Some(DatagenStepKind::Leaf); event.step_index = Some(2); - event.step_instance_id = Some("root/grade/0".to_string()); event.field_name = Some(field_name.to_string()); event.field_type = Some(field_type.to_string()); event.codec_version = Some(1); @@ -1202,8 +1234,8 @@ mod tests { DatagenEventType::StepCompleted, ); event.step_name = Some("grade".to_string()); + event.step_kind = Some(DatagenStepKind::Leaf); event.step_index = Some(2); - event.step_instance_id = Some("root/grade/0".to_string()); event } @@ -1256,7 +1288,7 @@ mod tests { store.append_checkpoint(&checkpoint).await.unwrap(); let mut terminal = event("item-1", 5, "terminal", 0, DatagenEventType::Terminal); - terminal.terminal = Some(DatagenTerminal::Completed); + terminal.status = Some(DatagenItemStatus::Completed); store.append(&[terminal]).await.unwrap(); let events = store.events_for_item("item-1").await.unwrap(); @@ -1275,18 +1307,19 @@ mod tests { Some(blob_bytes) ); - let folded = store.fold_item("item-1").await.unwrap().unwrap(); + let folded = store.fold_item("item-1").await.unwrap(); + let folded = folded.folded().expect("item-1 was created"); assert_eq!(folded.status, DatagenItemStatus::Completed); assert_eq!( folded.fields.get("score"), Some(&DatagenFieldState::Set(DatagenValue::Int(i64::MAX))) ); - assert_eq!(folded.completed_steps.len(), 1); + assert_eq!(folded.trajectory.ordered.len(), 1); assert_eq!(folded.query_tags, Some(json!({"domain": "math"}))); let trajectory = store.trajectory("item-1").await.unwrap(); assert_eq!(trajectory.len(), 1); - assert_eq!(trajectory[0].cursor.step_name, "grade"); + assert_eq!(trajectory[0].position.step.name, "grade"); }); } @@ -1297,13 +1330,16 @@ mod tests { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { let mut store = DatagenStore::open(&uri).await.unwrap(); - let no_op = completed_step(0, 0); + let created = event("item-1", 0, "created", 0, DatagenEventType::ItemCreated); + store.append(&[created]).await.unwrap(); + let no_op = completed_step(1, 0); store.append_checkpoint(&[no_op]).await.unwrap(); - let folded = store.fold_item("item-1").await.unwrap().unwrap(); - assert_eq!(folded.completed_steps.len(), 1); + let folded = store.fold_item("item-1").await.unwrap(); + let folded = folded.folded().expect("item-1 was created"); + assert_eq!(folded.trajectory.ordered.len(), 1); let field_only = field_event( - 1, + 2, 0, DatagenEventType::FieldSet, "score", @@ -1360,6 +1396,9 @@ mod tests { let mut failed = event("root-b", 0, "failed-b", 0, DatagenEventType::Failed); failed.run_id = "run-failed".to_string(); failed.writer_epoch = "writer-b".to_string(); + failed.step_name = Some("expand".to_string()); + failed.step_kind = Some(DatagenStepKind::Leaf); + failed.step_index = Some(0); failed.error_type = Some("ValueError".to_string()); failed.error_dump = Some("bad source item".to_string()); writer_b.append(&[failed]).await.unwrap(); diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 904e632..c4f0c60 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -20,10 +20,11 @@ mod store; pub use api_impl::rollout_record_to_dto; pub use context::{Context, ContextEntry, Snapshot}; pub use datagen::{ - datagen_event_id, datagen_trajectory, fold_datagen_events, DatagenBlobValue, DatagenEvent, - DatagenEventType, DatagenFailure, DatagenFieldState, DatagenItemStatus, DatagenStepCursor, - DatagenTerminal, DatagenTrajectoryPoint, DatagenValue, FoldedDatagenItem, - DATAGEN_SCHEMA_VERSION, + datagen_event_id, datagen_failures, datagen_trajectory, fold_datagen_events, DatagenBlobValue, + DatagenErrorInfo, DatagenEvent, DatagenEventType, DatagenFailure, DatagenFieldState, + DatagenItemId, DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses, DatagenStepCursor, + DatagenStepId, DatagenStepKind, DatagenStreamPosition, DatagenTerminal, DatagenTrajectory, + DatagenValue, FoldedDatagenItem, DATAGEN_SCHEMA_VERSION, }; pub use datagen_store::{datagen_log_schema, DatagenStore, DatagenStoreOptions}; pub use eval::{ From 91b4981b56e63891d71b60ce8b47b83718eb922e Mon Sep 17 00:00:00 2001 From: YangjunZ Date: Sun, 26 Jul 2026 17:27:54 -0700 Subject: [PATCH 2/5] docs(datagen): document schema v2 + store usage, strengthen fold tests - Rewrite specs/datagen-checkpoint-schema.md for schema v2: 7 events, status column, structured item id + step provenance, read lenses. - Add docs/design/using-datagen-store.md: a client-facing walkthrough of open/write/checkpoint/resume/read with runnable snippets. - Add 7 fold + store tests: step_kind/status parse round-trips, filtered terminal, selector_step on the chosen child, fan-out sub-item lineage, set/append mixing rejection, resume open-frame (started minus completed), and a fan-out tree read + root classification through the store. Co-Authored-By: Claude Opus 4.8 --- crates/lance-context-core/src/datagen.rs | 118 +++++++++ .../lance-context-core/src/datagen_store.rs | 59 ++++- docs/design/using-datagen-store.md | 235 ++++++++++++++++++ specs/datagen-checkpoint-schema.md | 73 ++++-- 4 files changed, 462 insertions(+), 23 deletions(-) create mode 100644 docs/design/using-datagen-store.md diff --git a/crates/lance-context-core/src/datagen.rs b/crates/lance-context-core/src/datagen.rs index eb0350c..8dd2061 100644 --- a/crates/lance-context-core/src/datagen.rs +++ b/crates/lance-context-core/src/datagen.rs @@ -950,4 +950,122 @@ mod tests { )])) ); } + + #[test] + fn step_kind_and_status_parse_round_trip() { + for kind in [ + DatagenStepKind::Root, + DatagenStepKind::Leaf, + DatagenStepKind::Sequence, + DatagenStepKind::Loop, + DatagenStepKind::MapReduce, + DatagenStepKind::Branch, + DatagenStepKind::SubPipeline, + DatagenStepKind::Conditional, + DatagenStepKind::Router, + ] { + assert_eq!(DatagenStepKind::parse(kind.as_str()).unwrap(), kind); + } + assert!(DatagenStepKind::Sequence.is_driver()); + assert!(DatagenStepKind::Loop.is_driver()); + assert!(!DatagenStepKind::MapReduce.is_driver()); + assert!(!DatagenStepKind::Leaf.is_driver()); + assert!(DatagenStepKind::parse("nope").is_err()); + + for status in [ + DatagenItemStatus::Running, + DatagenItemStatus::Completed, + DatagenItemStatus::Filtered, + DatagenItemStatus::Failed, + ] { + assert_eq!(DatagenItemStatus::parse(status.as_str()).unwrap(), status); + } + assert!(DatagenItemStatus::parse("nope").is_err()); + } + + #[test] + fn terminal_filtered_sets_filtered_status() { + let mut terminal = event(2, DatagenEventType::Terminal); + terminal.status = Some(DatagenItemStatus::Filtered); + let folded = fold_datagen_events(&[created(0), terminal]).unwrap().unwrap(); + assert_eq!(folded.status, DatagenItemStatus::Filtered); + } + + #[test] + fn selector_step_is_folded_onto_the_chosen_child_position() { + // A Conditional/Router writes no row; the chosen leaf records who selected it. + let mut chosen = leaf_completed(1, "stage2_qa", 0, Some("rubric_generation")); + chosen.selector_step = Some("if_stage2_qa".to_string()); + let folded = fold_datagen_events(&[created(0), chosen]).unwrap().unwrap(); + let cursor = &folded.trajectory.ordered[0]; + assert_eq!(cursor.position.selector.as_deref(), Some("if_stage2_qa")); + assert_eq!(cursor.position.step.name, "stage2_qa"); + } + + #[test] + fn fan_out_sub_item_folds_with_lineage() { + // A MapReduce/Branch sub-item is its own stream carrying denormalized parent/root ids. + let root = DatagenItemId::from_source_key("5"); + let child = root.child("solve_twice", 1); + let mut created = event(0, DatagenEventType::ItemCreated); + created.item_id = child.to_string(); + created.parent_item_id = Some(root.to_string()); + created.status = Some(DatagenItemStatus::Running); + created.event_id = datagen_event_id(&child.to_string(), "created", 0); + + let mut solved = leaf_completed(1, "solve", 0, Some("solve_attempt")); + solved.item_id = child.to_string(); + solved.root_item_id = root.to_string(); + solved.parent_item_id = Some(root.to_string()); + solved.event_id = datagen_event_id(&child.to_string(), "solve-0", 0); + + let folded = fold_datagen_events(&[created, solved]).unwrap().unwrap(); + assert_eq!(folded.item_id, child); + assert_eq!(folded.root_item_id, root); + assert_eq!(folded.parent_item_id, Some(root)); + assert_eq!(folded.item_id.origin_step(), Some("solve_twice")); + assert_eq!(folded.item_id.branch_idx(), Some(1)); + } + + #[test] + fn field_rejects_mixing_set_and_append() { + let mut set = leaf_completed(1, "gen", 0, Some("main")); + set.event_type = DatagenEventType::FieldSet; + set.field_name = Some("draft".to_string()); + set.field_type = Some("str".to_string()); + set.codec_version = Some(1); + set.value = Some(DatagenValue::Str("v1".to_string())); + + let mut append = set.clone(); + append.event_type = DatagenEventType::FieldAppend; + append.item_seq = 2; + append.checkpoint_id = "c2".to_string(); + append.event_id = datagen_event_id("5", "c2", 0); + + let error = fold_datagen_events(&[created(0), set, append]).unwrap_err(); + assert!(error.contains("mixes FIELD_SET and FIELD_APPEND")); + } + + #[test] + fn resume_open_frame_is_started_minus_completed() { + // A driver frame that opened but never completed = the frame live at crash time. + let mut main_completed = driver_started(3, "main", 0, None); + main_completed.event_type = DatagenEventType::StepCompleted; + main_completed.checkpoint_id = "main-done".to_string(); + main_completed.event_id = datagen_event_id("5", "main-done", 0); + let events = [ + created(0), + driver_started(1, "main", 0, None), + driver_started(2, "refine", 1, Some("main")), + main_completed, + ]; + let folded = fold_datagen_events(&events).unwrap().unwrap(); + let open: Vec<_> = folded + .trajectory + .started + .difference(&folded.trajectory.completed) + .collect(); + assert_eq!(open.len(), 1); + assert_eq!(open[0].step.name, "refine"); + } } diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index f91000f..6f36e4c 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -1162,8 +1162,8 @@ fn is_fenced_error(error: &LanceError) -> bool { mod tests { use super::*; use crate::datagen::{ - datagen_event_id, DatagenFieldState, DatagenItemStatus, DatagenStepKind, - DATAGEN_SCHEMA_VERSION, + datagen_event_id, DatagenFieldState, DatagenItemId, DatagenItemLookup, DatagenItemStatus, + DatagenStepKind, DATAGEN_SCHEMA_VERSION, }; use chrono::{TimeZone, Utc}; use serde_json::json; @@ -1455,4 +1455,59 @@ mod tests { ); }); } + + #[test] + fn fan_out_tree_reads_by_root_and_classifies_status() { + let directory = TempDir::new().unwrap(); + let uri = directory.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = DatagenStore::open(&uri).await.unwrap(); + + // Root item "7" fans out into one sub-item "7/solve_twice:0". + let mut root_created = event("7", 0, "created-root", 0, DatagenEventType::ItemCreated); + root_created.status = Some(DatagenItemStatus::Running); + store.append(&[root_created]).await.unwrap(); + + let mut child_created = event( + "7/solve_twice:0", + 0, + "created-child", + 0, + DatagenEventType::ItemCreated, + ); + child_created.parent_item_id = Some("7".to_string()); + child_created.status = Some(DatagenItemStatus::Running); + store.append(&[child_created]).await.unwrap(); + + let mut root_terminal = event("7", 1, "terminal", 0, DatagenEventType::Terminal); + root_terminal.status = Some(DatagenItemStatus::Completed); + store.append(&[root_terminal]).await.unwrap(); + + // The whole tree is one root filter, no join. + let tree = store.events_for_root("7").await.unwrap(); + assert_eq!(tree.len(), 3); + + // Bulk classification sees the root as terminated; the child is still running. + let statuses = store.root_item_statuses(&["7"]).await.unwrap(); + let root_id = DatagenItemId::from_source_key("7"); + assert!(statuses.is_terminated(&root_id)); + + let child = store.fold_item("7/solve_twice:0").await.unwrap(); + assert_eq!( + child.folded().unwrap().status, + DatagenItemStatus::Running + ); + assert_eq!( + child.folded().unwrap().parent_item_id, + Some(DatagenItemId::from_source_key("7")) + ); + + // A never-started sibling folds to NeverStarted. + assert_eq!( + store.fold_item("7/solve_twice:1").await.unwrap(), + DatagenItemLookup::NeverStarted + ); + }); + } } diff --git a/docs/design/using-datagen-store.md b/docs/design/using-datagen-store.md new file mode 100644 index 0000000..82506fd --- /dev/null +++ b/docs/design/using-datagen-store.md @@ -0,0 +1,235 @@ +# Using the Datagen Checkpoint Store + +`DatagenStore` is the durable checkpoint backend for a datagen experiment. It is +a single append-only Lance log: every write is one immutable event, and the +current state of any item is the *fold* of its events. Nothing is ever updated or +deleted. This document shows how a client (the `mai_datagen` executor, through +its pyo3 bindings) drives the store through a run — create, checkpoint, resume, +read — with runnable Rust snippets. + +For the schema and the design rationale, see +[`specs/datagen-checkpoint-schema.md`](../../specs/datagen-checkpoint-schema.md). + +## The model in one paragraph + +An experiment is one `log.lance`. Each row is a `DatagenEvent` with an +`event_type`. An **item** is one stream of events sharing an `item_id`; a source +item is a **root**, and fan-out steps project **sub-items**, each its own stream. +To learn an item's state you read its events and call `fold_datagen_events`, +which replays them into a `FoldedDatagenItem` (fields, status, trajectory). +Because reconstruction is a pure fold over an append-only log, a crash can never +leave a torn write: an unacknowledged batch either persisted whole or not at all. + +## Opening a store + +```rust +use lance_context_core::{DatagenStore, DatagenStoreOptions}; + +// One writer per shard. Concurrent writers pass distinct shard ids. +let mut store = DatagenStore::open("s3://bucket/exp/log.lance").await?; + +// Or with an explicit shard id (each writer owns its own MemWAL shard): +let mut store = DatagenStore::open_with_options( + "s3://bucket/exp/log.lance", + DatagenStoreOptions { shard_id: Some("worker-3".into()), ..Default::default() }, +).await?; +``` + +## Item identity — composed on the client, no round-trip + +Ids are structured values (`DatagenItemId`) stored as materialized path strings. +The client composes them purely; the store never allocates an id. + +```rust +use lance_context_core::DatagenItemId; + +let root = DatagenItemId::from_source_key("5"); // "5" +let child = root.child("solve_twice", 0); // "5/solve_twice:0" +let grandchild = child.child("judge", 1); // "5/solve_twice:0/judge:1" + +assert_eq!(grandchild.origin_step(), Some("judge")); +assert_eq!(grandchild.branch_idx(), Some(1)); +assert_eq!(grandchild.parent().unwrap(), child); +assert_eq!(grandchild.root(), root); +``` + +Because `root_item_id` and `parent_item_id` are denormalized onto every row, +reading a whole item tree is one filter (`events_for_root`), never a join. + +## Writing events + +An event is built with its provenance columns and appended. Field events plus +their `STEP_COMPLETED` marker for one step must go in a **single** call so a crash +cannot expose a half-checkpointed step — the batch is one durable generation. + +```rust +use lance_context_core::{ + datagen_event_id, DatagenEvent, DatagenEventType, DatagenItemStatus, + DatagenStepKind, DatagenValue, DATAGEN_SCHEMA_VERSION, +}; + +// 1. Announce the item once. +let created = DatagenEvent { + event_id: datagen_event_id("5", "created", 0), + item_id: "5".into(), + root_item_id: "5".into(), + parent_item_id: None, + item_seq: 0, + checkpoint_id: "created".into(), + event_type: DatagenEventType::ItemCreated, + status: Some(DatagenItemStatus::Running), + schema_version: DATAGEN_SCHEMA_VERSION, + ..blank_event() // your helper that zero-fills the optional columns +}; +store.append(&[created]).await?; + +// 2. Checkpoint one step: its field delta + exactly one STEP_COMPLETED, atomically. +let score = DatagenEvent { + event_id: datagen_event_id("5", "solve-0", 0), + item_id: "5".into(), + root_item_id: "5".into(), + item_seq: 1, + checkpoint_id: "solve-0".into(), + event_type: DatagenEventType::FieldSet, + step_name: Some("solve".into()), + step_kind: Some(DatagenStepKind::Leaf), + step_index: Some(0), + enclosing_step: Some("solve_attempt".into()), + field_name: Some("score".into()), + field_type: Some("int".into()), + codec_version: Some(1), + value: Some(DatagenValue::Int(9)), + schema_version: DATAGEN_SCHEMA_VERSION, + ..blank_event() +}; +let completed = DatagenEvent { + event_id: datagen_event_id("5", "solve-0", 1), + item_seq: 2, + checkpoint_id: "solve-0".into(), + event_type: DatagenEventType::StepCompleted, + // same step_name / step_kind / step_index / enclosing_step as above + ..score.clone() +}; +store.append_checkpoint(&[score, completed]).await?; +``` + +`append_checkpoint` enforces "exactly one `STEP_COMPLETED` per batch"; +`append` is the lower-level form used for lifecycle events (`ITEM_CREATED`, +`TERMINAL`, `FAILED`). + +### Retries are safe + +`event_id` is deterministic (`datagen_event_id(item_id, checkpoint_id, ordinal)`). +Replaying an ambiguously-acknowledged batch writes the same ids, and the fold +de-duplicates them — the second `append_checkpoint` of an identical batch is a +no-op in the folded result. + +### Kind determines what you write + +| Step kind | Emits | +|---|---| +| `Sequence`, `Loop` (drivers) | `STEP_STARTED` frame + `STEP_COMPLETED` | +| `MapReduce`, `Branch`, `SubPipeline` (fan-out) | only the reduce `STEP_COMPLETED`; sub-items are their own streams | +| `Conditional`, `Router` (selectors) | no row of their own — the chosen child sets `selector_step` | +| `Leaf` | `FIELD_SET` / `FIELD_APPEND` + `STEP_COMPLETED` | + +## Finishing an item + +```rust +// Success or filtered-out — a lifecycle terminal. +let terminal = DatagenEvent { + event_type: DatagenEventType::Terminal, + status: Some(DatagenItemStatus::Completed), // or Filtered + ..lifecycle_event("5", 3, "terminal") +}; +store.append(&[terminal]).await?; + +// A raised step — a failure-lens row. The item still folds to `running`. +let failed = DatagenEvent { + event_type: DatagenEventType::Failed, + status: Some(DatagenItemStatus::Failed), + step_name: Some("score".into()), + step_kind: Some(DatagenStepKind::Leaf), + step_index: Some(1), + error_type: Some("ValueError".into()), + ..lifecycle_event("5", 3, "failed") +}; +store.append(&[failed]).await?; +``` + +## Reading + +### Fold one item + +```rust +use lance_context_core::DatagenItemLookup; + +match store.fold_item("5").await? { + DatagenItemLookup::NeverStarted => { /* fresh: process from scratch */ } + DatagenItemLookup::Found(item) => { + // item.status, item.fields, item.trajectory, item.last_item_seq, item.last_attempt + // Resume: continue writing at last_item_seq + 1, attempt last_attempt + 1. + } +} +``` + +`NeverStarted` (no `ITEM_CREATED`) is the explicit fresh-vs-resume fork. A +`Found` item's `status` is always a *lifecycle* value (`running` / `completed` +/ `filtered`) — `failed` never appears here. + +### Resume: the open frame + +The trajectory records both started and completed positions. `started \ completed` +is the driver frame that was open when the process died — everything already +completed is skipped on re-run. + +```rust +let item = store.fold_item("5").await?.folded().unwrap(); +let open: Vec<_> = item.trajectory.started + .difference(&item.trajectory.completed) + .collect(); // the frame(s) to re-enter +``` + +### Whole item tree + +```rust +// Every event under root "5", including all fan-out sub-items — one filter. +let events = store.events_for_root("5").await?; +``` + +### Bulk startup classification + +```rust +// Classify many roots at once without folding fields (skip / resume / fresh). +let statuses = store.root_item_statuses(&["5", "6", "7"]).await?; +assert!(statuses.is_terminated(&DatagenItemId::from_source_key("5"))); +``` + +### Failures (the failure lens) + +```rust +let failures = store.item_failures("5").await?; // 0..N, across attempts +for failure in &failures { + println!("{} at {}", failure.error.error_type, failure.at.position.step.name); +} +let all = store.failures(Some("run-1")).await?; // run-wide forensics +``` + +### Blobs are lazy + +Field values that are blobs fold to a lazy `DatagenBlobValue { bytes: None, .. }`. +Materialize the bytes only when needed: + +```rust +let bytes = store.get_blob(&blob_event_id).await?; // O(single blob) take_rows +``` + +## Concurrency and maintenance + +- One live owner writes a given item at a time; a new owner takes over with a new + `writer_epoch` (fencing). +- Concurrent writers use distinct MemWAL shards. Reads union the base table and + all flushed shards, so every instance sees every writer's events. +- Each writer periodically merges only its own generations into the base table + (`cleanup_own_shard` / `spawn_periodic_cleanup`). Shared base-table compaction + is scheduled by one elected maintenance worker per experiment. diff --git a/specs/datagen-checkpoint-schema.md b/specs/datagen-checkpoint-schema.md index d999ee7..f045f8f 100644 --- a/specs/datagen-checkpoint-schema.md +++ b/specs/datagen-checkpoint-schema.md @@ -18,39 +18,62 @@ rebuildable, and excluded from checkpoint correctness. ## Event model -Every row is an immutable event: +Every row is an immutable event. The current schema version is +`DATAGEN_SCHEMA_VERSION = 2`. There are seven event types: -- `ITEM_CREATED` -- `FIELD_SET` -- `FIELD_APPEND` -- `STEP_COMPLETED` -- `FAILED` -- `TERMINAL` +- `ITEM_CREATED` — an item (stream) first appears, exactly once, `status = running`. +- `STEP_STARTED` — a driver frame (`Sequence`/`Loop`) opened. A structural marker. +- `FIELD_SET` — replaces a field's value (fold: last-writer-wins). +- `FIELD_APPEND` — accumulates onto a field (fold: append in order). +- `STEP_COMPLETED` — a checkpointed step boundary, carrying that step's field delta. +- `FAILED` — a step raised; a failure-lens row that leaves the item `running`. +- `TERMINAL` — the item finished, `status = completed | filtered`. Every completed step emits a `STEP_COMPLETED` event, even when no field changed. All field events and the completion marker for one step are written in the same -checkpoint batch. +checkpoint batch. Only driver kinds emit `STEP_STARTED`; fan-out kinds +(`MapReduce`/`Branch`/`SubPipeline`) emit only their reduce `STEP_COMPLETED`; +selector kinds (`Conditional`/`Router`) emit no row of their own — the chosen +child records the selector in `selector_step`. `event_id` is a deterministic idempotency key. Retrying an ambiguously acknowledged batch writes the same event ids, and the MemWAL LSM read path de-duplicates them. `item_seq` is strictly increasing per item; two different events at the same sequence are treated as a writer-fencing violation. +### Item identity + +An item id is a materialized path string owned by the store. A root id is the +executor's source key (`5`); a fan-out sub-item extends its parent with one +`step:idx` segment (`5/solve_twice:0`). Ids compose purely on the client, with +no store round-trip. `root_item_id` / `parent_item_id` are denormalized on every +row, so any subtree is a single filter with no join. + +### Read lenses + +State is read under two lenses: + +- *lifecycle* — fold the events; `status` is `running` until a `TERMINAL`. Drives + skip/resume/fresh classification and ignores `FAILED`. +- *failure* — read `FAILED` rows directly (forensics). An item may have 0..N + failures across attempts while still folding to `running` under lifecycle. + ## Schema | Column | Type | Purpose | |---|---|---| | `event_id` | string | Deterministic event identity and LSM primary key | -| `item_id` | string | Scoped item identity | -| `root_item_id` | string | Root of the projected item tree | -| `parent_item_id` | string? | Direct parent item | +| `item_id` | string | Scoped item identity (materialized path) | +| `root_item_id` | string | Root of the projected item tree (denormalized) | +| `parent_item_id` | string? | Direct parent item (denormalized) | | `item_seq` | int64 | Per-item event ordering | | `checkpoint_id` | string | Atomic step-boundary identity | -| `event_type` | string | Event kind | -| `step_name` | string? | Step provenance | -| `step_index` | int64? | Static step position | -| `step_instance_id` | string? | Runtime step identity | -| `iteration` | int64? | Loop/branch iteration | +| `event_type` | string | Event kind (one of the seven above) | +| `step_name` | string? | Step provenance (globally-unique step name) | +| `step_kind` | string? | Composition kind (`sequence`, `loop`, `map_reduce`, `branch`, `sub_pipeline`, `conditional`, `router`, `leaf`, `root`) | +| `step_index` | int64? | Static step position (loop iteration = driver frame's index) | +| `enclosing_step` | string? | Name of the `Sequence`/`Loop` driver frame this step ran under | +| `selector_step` | string? | Name of the `Conditional`/`Router` that chose this step | | `attempt` | int32 | Execution attempt | | `run_id` | string | Run attribution | | `writer_epoch` | string | Item ownership/fencing identity | @@ -67,13 +90,19 @@ events at the same sequence are treated as a writer-fencing violation. | `payload_size` | int64? | Blob size | | `payload_checksum` | string? | Blob integrity | | `query_tags_json` | large_string? | Non-authoritative query tags | -| `terminal` | string? | `completed` or `filtered` | +| `status` | string? | `running` / `completed` / `filtered` / `failed` (on ITEM_CREATED / TERMINAL / FAILED) | | `error_type` | string? | Failure type | | `error_dump` | large_string? | Serialized failure | | `traceback` | large_string? | Failure traceback | | `event_ts` | timestamp(us, UTC) | Event time | | `schema_version` | int32 | Log schema compatibility version | +The `status` column replaces the earlier `terminal` column: it carries a stored +`running` (from `ITEM_CREATED`) rather than deriving it, plus `failed` for the +failure lens. Structured `step_kind` / `enclosing_step` / `selector_step` replace +the earlier opaque `step_instance_id` + `iteration` provenance, so a resume can +rebuild step coordinates instead of trusting a stored cursor blob. + `value_blob` remains inline while MemWAL's LSM scanner cannot materialize blob-v2 columns. Normal fold and trajectory reads project it out. Blob access first locates `event_id` using lightweight columns and then calls `take_rows` @@ -85,9 +114,12 @@ for the exact `_rowid`. 2. Log writes are append-only; retries reuse deterministic event ids. 3. One live owner writes a given item at a time. Ownership changes require a new `writer_epoch`. -4. Resume reads `item_id = X`, orders by `item_seq`, and folds the events. -5. `FIELD_SET` replaces a field; `FIELD_APPEND` accumulates it. -6. `STEP_COMPLETED` reconstructs the resume cursor. +4. Resume reads `item_id = X`, orders by `item_seq`, and folds the events. An + item with no `ITEM_CREATED` folds to `NeverStarted` (fresh vs. restore fork). +5. `FIELD_SET` replaces a field; `FIELD_APPEND` accumulates it. Mixing the two + on one field is rejected. +6. `STEP_COMPLETED` reconstructs the resume cursor; `STEP_STARTED \ STEP_COMPLETED` + (started minus completed) is the driver frame open at crash time. 7. `FAILED` and `TERMINAL` are ordinary log events, not separate datasets. 8. Blob bytes are loaded only when the corresponding lazy reference is used. @@ -98,4 +130,3 @@ shards, so any instance sees every writer's events. Each writer periodically merges only its own generations into the base table. Shared base-table compaction and index refresh must be scheduled by one elected maintenance worker per experiment. - From 9a39772d1ad342ddc5a3faab225d37ec32f4a734 Mon Sep 17 00:00:00 2001 From: Yangjun Zhang Date: Sun, 26 Jul 2026 19:52:51 -0700 Subject: [PATCH 3/5] fix(datagen): repair umbrella re-export and clippy lints, apply rustfmt - lance-context umbrella re-exported the pre-reshape name DatagenTrajectoryPoint; rename to DatagenTrajectory so the crate (and the python wheel + tests that depend on it) compiles again. - allow(large_enum_variant) on DatagenItemLookup (Found is the hot path) and drop a redundant #[must_use] on iter(). - cargo fmt. Co-Authored-By: Claude Opus 4.8 --- crates/lance-context-core/src/datagen.rs | 36 +++++++++++-------- .../lance-context-core/src/datagen_store.rs | 9 ++--- crates/lance-context-core/src/lib.rs | 6 ++-- crates/lance-context/src/lib.rs | 2 +- 4 files changed, 29 insertions(+), 24 deletions(-) diff --git a/crates/lance-context-core/src/datagen.rs b/crates/lance-context-core/src/datagen.rs index 8dd2061..88f2fa6 100644 --- a/crates/lance-context-core/src/datagen.rs +++ b/crates/lance-context-core/src/datagen.rs @@ -195,10 +195,7 @@ impl DatagenItemId { .map_err(|_| format!("non-integer branch index '{idx}' in '{path}'"))?; segments.push((step.to_string(), branch_idx)); } - Ok(Self { - root_key, - segments, - }) + Ok(Self { root_key, segments }) } } @@ -382,6 +379,7 @@ pub struct FoldedDatagenItem { /// Result of a resumption fold. `NeverStarted` (no ITEM_CREATED) is the fresh-vs-restore fork the /// executor acts on; `Found` carries the folded item (whose `status` is the lifecycle status). #[derive(Debug, Clone, PartialEq)] +#[allow(clippy::large_enum_variant)] // `Found` is the common path; boxing it adds indirection to the hot case. pub enum DatagenItemLookup { NeverStarted, Found(FoldedDatagenItem), @@ -434,7 +432,6 @@ impl DatagenRootItemStatuses { self.inner.is_empty() } - #[must_use] pub fn iter(&self) -> impl Iterator { self.inner.iter() } @@ -685,7 +682,8 @@ fn apply_event(item: &mut FoldedDatagenItem, event: &DatagenEvent) -> Result<(), let field_name = event.field_name.clone().unwrap(); let value = event.value.clone().unwrap(); record_blob_event_id(item, &field_name, &value, &event.event_id); - item.fields.insert(field_name, DatagenFieldState::Set(value)); + item.fields + .insert(field_name, DatagenFieldState::Set(value)); } DatagenEventType::FieldAppend => { let field_name = event.field_name.clone().unwrap(); @@ -719,7 +717,9 @@ fn apply_event(item: &mut FoldedDatagenItem, event: &DatagenEvent) -> Result<(), } DatagenEventType::Terminal => { item.status = match event.status { - Some(status @ (DatagenItemStatus::Completed | DatagenItemStatus::Filtered)) => status, + Some(status @ (DatagenItemStatus::Completed | DatagenItemStatus::Filtered)) => { + status + } _ => return Err("TERMINAL event missing completed/filtered status".to_string()), }; } @@ -882,8 +882,9 @@ mod tests { append2.event_id = datagen_event_id("5", "c4", 0); append2.value = Some(DatagenValue::Json(json!({"n": "b"}))); - let folded = - fold_datagen_events(&[created(0), set_v1, set_v2, append, append2]).unwrap().unwrap(); + let folded = fold_datagen_events(&[created(0), set_v1, set_v2, append, append2]) + .unwrap() + .unwrap(); assert_eq!( folded.fields.get("draft"), Some(&DatagenFieldState::Set(DatagenValue::Str("v2".to_string()))) @@ -901,7 +902,9 @@ mod tests { fn terminal_sets_lifecycle_status() { let mut terminal = event(2, DatagenEventType::Terminal); terminal.status = Some(DatagenItemStatus::Completed); - let folded = fold_datagen_events(&[created(0), terminal]).unwrap().unwrap(); + let folded = fold_datagen_events(&[created(0), terminal]) + .unwrap() + .unwrap(); assert_eq!(folded.status, DatagenItemStatus::Completed); } @@ -913,7 +916,9 @@ mod tests { failed.error_type = Some("ValueError".to_string()); failed.attempt = 0; - let folded = fold_datagen_events(&[created(0), failed.clone()]).unwrap().unwrap(); + let folded = fold_datagen_events(&[created(0), failed.clone()]) + .unwrap() + .unwrap(); assert_eq!(folded.status, DatagenItemStatus::Running); let failures = datagen_failures(&[created(0), failed]).unwrap(); @@ -941,8 +946,9 @@ mod tests { append.codec_version = Some(1); append.value = Some(DatagenValue::Json(json!({"role": "assistant"}))); - let folded = - fold_datagen_events(&[created(0), append.clone(), append]).unwrap().unwrap(); + let folded = fold_datagen_events(&[created(0), append.clone(), append]) + .unwrap() + .unwrap(); assert_eq!( folded.fields.get("messages"), Some(&DatagenFieldState::Appended(vec![DatagenValue::Json( @@ -987,7 +993,9 @@ mod tests { fn terminal_filtered_sets_filtered_status() { let mut terminal = event(2, DatagenEventType::Terminal); terminal.status = Some(DatagenItemStatus::Filtered); - let folded = fold_datagen_events(&[created(0), terminal]).unwrap().unwrap(); + let folded = fold_datagen_events(&[created(0), terminal]) + .unwrap() + .unwrap(); assert_eq!(folded.status, DatagenItemStatus::Filtered); } diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index 0ce83f3..7c66a57 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -36,8 +36,8 @@ use uuid::Uuid; use crate::datagen::{ datagen_failures, datagen_trajectory, fold_datagen_events, DatagenBlobValue, DatagenEvent, - DatagenEventType, DatagenFailure, DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses, - DatagenStepCursor, DatagenStepKind, DatagenValue, + DatagenEventType, DatagenFailure, DatagenItemLookup, DatagenItemStatus, + DatagenRootItemStatuses, DatagenStepCursor, DatagenStepKind, DatagenValue, }; use crate::rollout_store::{align_batch_to_schema, derive_shard_id, is_not_found_error}; use crate::store::{column_as, column_as_optional, timestamp_from_micros}; @@ -1566,10 +1566,7 @@ mod tests { assert!(statuses.is_terminated(&root_id)); let child = store.fold_item("7/solve_twice:0").await.unwrap(); - assert_eq!( - child.folded().unwrap().status, - DatagenItemStatus::Running - ); + assert_eq!(child.folded().unwrap().status, DatagenItemStatus::Running); assert_eq!( child.folded().unwrap().parent_item_id, Some(DatagenItemId::from_source_key("7")) diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 8c6cbee..44f041b 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -23,9 +23,9 @@ pub use context::{Context, ContextEntry, Snapshot}; pub use datagen::{ datagen_event_id, datagen_failures, datagen_trajectory, fold_datagen_events, DatagenBlobValue, DatagenErrorInfo, DatagenEvent, DatagenEventType, DatagenFailure, DatagenFieldState, - DatagenItemId, DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses, DatagenStepCursor, - DatagenStepId, DatagenStepKind, DatagenStreamPosition, DatagenTerminal, DatagenTrajectory, - DatagenValue, FoldedDatagenItem, DATAGEN_SCHEMA_VERSION, + DatagenItemId, DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses, + DatagenStepCursor, DatagenStepId, DatagenStepKind, DatagenStreamPosition, DatagenTerminal, + DatagenTrajectory, DatagenValue, FoldedDatagenItem, DATAGEN_SCHEMA_VERSION, }; pub use datagen_store::{datagen_log_schema, DatagenStore, DatagenStoreOptions}; pub use eval::{ diff --git a/crates/lance-context/src/lib.rs b/crates/lance-context/src/lib.rs index 2e4a4f4..93d6649 100644 --- a/crates/lance-context/src/lib.rs +++ b/crates/lance-context/src/lib.rs @@ -7,7 +7,7 @@ pub use lance_context_core::{ CompactionConfig, CompactionMetrics, CompactionStats, Context, ContextEntry, ContextNamespace, ContextRecord, ContextStoreOptions, DatagenBlobValue, DatagenEvent, DatagenEventType, DatagenFailure, DatagenFieldState, DatagenItemStatus, DatagenStepCursor, DatagenStore, - DatagenStoreOptions, DatagenTerminal, DatagenTrajectoryPoint, DatagenValue, FoldedDatagenItem, + DatagenStoreOptions, DatagenTerminal, DatagenTrajectory, DatagenValue, FoldedDatagenItem, IdIndexType, LifecycleQueryOptions, MetadataFilter, PartitionInfo, PartitionSelector, PartitionSpec, RecordFilters, Relationship, RetrieveResult, RolloutFilters, RolloutRecord, SearchResult, Snapshot, StateMetadata, DATAGEN_SCHEMA_VERSION, LIFECYCLE_ACTIVE, From aec066edb723f87fc1df6c9e7000e0ef0d2bc95a Mon Sep 17 00:00:00 2001 From: Yangjun Zhang Date: Mon, 27 Jul 2026 07:55:32 -0700 Subject: [PATCH 4/5] feat(datagen): add DatagenStore read/write service API across all six layers The schema-v2 reshape (#205) landed the data model but not the API layer on top of it. This wires the datagen store through the full stack so clients can append, fold, and read blobs both embedded and over the server: - api: `DatagenStoreApi` trait (RPITIT) + wire DTOs mirroring the Python dicts - core: `impl DatagenStoreApi for DatagenStore` with core<->DTO converters - client: `RemoteDatagenStore` HTTP client - unified: `enum DatagenStore {Local, Remote}` with dispatch - server: `/api/v1/datagen` routes (create/list/get/delete, events, fold, failures, root-item-statuses, blob fetch) - python: PyO3 `DatagenStore` binding + `open`/`connect`/`connect_or_create` Verified end-to-end locally: identical checkpoint round-trips through both the embedded (local-file) and remote (server) paths fold to the same item state. Co-Authored-By: Claude Opus 4.8 --- crates/lance-context-api/src/lib.rs | 252 +++++++++++ crates/lance-context-client/src/lib.rs | 226 ++++++++++ crates/lance-context-core/src/api_impl.rs | 271 +++++++++++- crates/lance-context-server/src/error.rs | 15 +- .../src/routes/datagen.rs | 225 ++++++++++ crates/lance-context-server/src/routes/mod.rs | 29 ++ crates/lance-context-server/src/state.rs | 130 +++++- crates/lance-context/src/lib.rs | 27 +- crates/lance-context/src/unified_datagen.rs | 130 ++++++ python/python/lance_context/__init__.py | 4 + python/python/lance_context/api.py | 96 +++- python/src/lib.rs | 415 +++++++++++++++++- 12 files changed, 1792 insertions(+), 28 deletions(-) create mode 100644 crates/lance-context-server/src/routes/datagen.rs create mode 100644 crates/lance-context/src/unified_datagen.rs diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index b696982..1fdb00b 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -178,6 +178,60 @@ pub trait RolloutStoreApi { fn checkout(&mut self, version: u64) -> impl Future> + Send; } +// --------------------------------------------------------------------------- +// Datagen trait +// --------------------------------------------------------------------------- + +/// Remote-capable surface of a datagen checkpoint store — the append-only +/// delta-log of item lifecycle / field events plus the folded read lenses. +/// +/// A datagen store has no upsert/search/compaction: writers only append events +/// (`append` for lifecycle rows, `append_checkpoint` for an atomic step +/// boundary), and readers fold an item's events into latest state +/// (`fold_item`), classify root items in bulk (`root_item_statuses`), list an +/// item's failure pointers (`item_failures`), or materialize one field's +/// offloaded blob bytes (`get_blob`). +/// +/// Field blob bytes (`DatagenValueDto::bytes`) travel inline as base64 in JSON +/// or, for large payloads on the append endpoints, as raw multipart parts. By +/// the time a call reaches this trait the bytes are already materialized in +/// memory, so the signatures are transport-agnostic. +pub trait DatagenStoreApi { + fn append( + &mut self, + events: &[DatagenEventDto], + ) -> impl Future> + Send; + + fn append_checkpoint( + &mut self, + events: &[DatagenEventDto], + ) -> impl Future> + Send; + + fn fold_item( + &self, + item_id: &str, + ) -> impl Future>> + Send; + + fn root_item_statuses( + &self, + root_item_ids: &[String], + ) -> impl Future> + Send; + + fn item_failures( + &self, + item_id: &str, + ) -> impl Future>> + Send; + + /// Materialize one FIELD_* event's offloaded blob bytes by event id. + /// Returns `None` when the event or its payload is absent. + fn get_blob( + &self, + event_id: &str, + ) -> impl Future>>> + Send; + + fn version(&self) -> u64; +} + // --------------------------------------------------------------------------- // Context lifecycle // --------------------------------------------------------------------------- @@ -826,6 +880,204 @@ pub struct GetRolloutResponse { pub record: Option, } +// --------------------------------------------------------------------------- +// Datagen lifecycle +// --------------------------------------------------------------------------- + +#[derive(Debug, Serialize, Deserialize)] +pub struct CreateDatagenStoreRequest { + /// Portable dataset name: 1-128 ASCII characters matching + /// `[A-Za-z0-9_][A-Za-z0-9._-]*`; `_registry` and `_stats` are reserved. + pub name: String, + #[serde(default)] + pub storage_options: Option>, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct DatagenStoreInfo { + pub name: String, + pub uri: String, + /// Dataset version. `None` in list responses, which are served from the + /// durable registry without opening each dataset. Single-store lookups + /// (`get`/`create`) always populate it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub version: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ListDatagenStoresResponse { + pub stores: Vec, +} + +// --------------------------------------------------------------------------- +// Datagen values +// --------------------------------------------------------------------------- + +/// The wire form of a core `DatagenValue`. `kind` tags the payload: +/// `"int"`/`"float"`/`"bool"`/`"str"`/`"json"` carry a JSON scalar in `value`; +/// `"blob"` carries raw bytes in `bytes` (inline base64 in JSON, or a raw +/// multipart part on the append endpoints) plus `size`/`checksum`. Mirrors the +/// `DatagenValue.to_wire`/`from_wire` dict shape in the Python binding. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DatagenValueDto { + pub kind: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub value: Option, + #[serde( + default, + skip_serializing_if = "Option::is_none", + serialize_with = "serialize_base64_opt", + deserialize_with = "deserialize_base64_opt" + )] + pub bytes: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub size: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub checksum: Option, +} + +// --------------------------------------------------------------------------- +// Datagen events (append) +// --------------------------------------------------------------------------- + +/// One append-only datagen log row. Field names and semantics mirror the core +/// `DatagenEvent` and the Python `datagen_events` wire dict; enum-valued columns +/// (`event_type`, `step_kind`, `status`) travel as their canonical strings. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DatagenEventDto { + pub event_id: String, + pub item_id: String, + pub root_item_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parent_item_id: Option, + pub item_seq: i64, + pub checkpoint_id: String, + pub event_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub step_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub step_kind: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub step_index: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub enclosing_step: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub selector_step: Option, + #[serde(default)] + pub attempt: i32, + pub run_id: String, + pub writer_epoch: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub field_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub field_type: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub codec_version: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub value: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub query_tags: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub status: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error_type: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error_dump: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub traceback: Option, + /// Defaults to the server's current time when omitted. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub event_ts: Option>, + pub schema_version: i32, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct AddDatagenEventsRequest { + pub events: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct AddDatagenEventsResponse { + pub version: u64, + pub count: usize, +} + +// --------------------------------------------------------------------------- +// Datagen folded read lenses +// --------------------------------------------------------------------------- + +/// A folded field: `mode = "set"` carries a single `value`; `mode = "append"` +/// carries an ordered `values` list. Mirrors the Python `FieldState` wire dict. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DatagenFieldStateDto { + pub mode: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub value: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub values: Vec, +} + +/// One completed step position an item passed through (a single STEP_COMPLETED). +/// Mirrors the Python `StepCursor` wire dict. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DatagenStepCursorDto { + pub step_name: String, + pub step_kind: String, + pub step_index: i64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub enclosing_step: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub selector_step: Option, + pub item_seq: i64, +} + +/// An item reconstructed by folding its events into latest state. +/// Mirrors the Python `FoldedItem` wire dict. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FoldedDatagenItemDto { + pub item_id: String, + pub root_item_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parent_item_id: Option, + pub status: String, + pub last_item_seq: i64, + pub last_attempt: i32, + pub fields: std::collections::BTreeMap, + pub trajectory: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub query_tags: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct GetFoldedDatagenItemResponse { + pub item: Option, +} + +/// Bulk startup classification of root items. A missing id means "never started". +#[derive(Debug, Serialize, Deserialize)] +pub struct DatagenRootItemStatusesResponse { + pub statuses: std::collections::HashMap, +} + +/// One failure record for an item (the failure lens). Mirrors the Python +/// `Failure` wire dict, with the step position flattened under `at`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DatagenFailureDto { + pub at: DatagenStepCursorDto, + pub run_id: String, + pub attempt: i32, + pub error_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error_dump: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub traceback: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ListDatagenFailuresResponse { + pub failures: Vec, +} + // --------------------------------------------------------------------------- // Error // --------------------------------------------------------------------------- diff --git a/crates/lance-context-client/src/lib.rs b/crates/lance-context-client/src/lib.rs index 71c4e39..2351fc1 100644 --- a/crates/lance-context-client/src/lib.rs +++ b/crates/lance-context-client/src/lib.rs @@ -361,6 +361,108 @@ impl RolloutStoreApi for RemoteRolloutStore { } } +pub struct RemoteDatagenStore { + client: ContextClient, + store_name: String, + cached_version: u64, +} + +impl RemoteDatagenStore { + pub async fn connect(base_url: &str, store_name: &str) -> Result { + let client = ContextClient::new(base_url); + let info = client.get_datagen_store(store_name).await?; + Ok(Self { + client, + store_name: store_name.to_string(), + cached_version: info.version.unwrap_or(0), + }) + } + + pub async fn connect_or_create( + base_url: &str, + req: &CreateDatagenStoreRequest, + ) -> Result { + let client = ContextClient::new(base_url); + let info = match client.get_datagen_store(&req.name).await { + Ok(info) => info, + Err(ClientError::Api { status: 404, .. }) => client.create_datagen_store(req).await?, + Err(e) => return Err(e), + }; + Ok(Self { + client, + store_name: req.name.clone(), + cached_version: info.version.unwrap_or(0), + }) + } +} + +impl DatagenStoreApi for RemoteDatagenStore { + async fn append( + &mut self, + events: &[DatagenEventDto], + ) -> ContextResult { + let resp = self + .client + .add_datagen_events(&self.store_name, events, false) + .await + .map_err(to_ctx_err)?; + self.cached_version = resp.version; + Ok(resp) + } + + async fn append_checkpoint( + &mut self, + events: &[DatagenEventDto], + ) -> ContextResult { + let resp = self + .client + .add_datagen_events(&self.store_name, events, true) + .await + .map_err(to_ctx_err)?; + self.cached_version = resp.version; + Ok(resp) + } + + async fn fold_item(&self, item_id: &str) -> ContextResult> { + let resp = self + .client + .fold_datagen_item(&self.store_name, item_id) + .await + .map_err(to_ctx_err)?; + Ok(resp.item) + } + + async fn root_item_statuses( + &self, + root_item_ids: &[String], + ) -> ContextResult { + self.client + .datagen_root_item_statuses(&self.store_name, root_item_ids) + .await + .map_err(to_ctx_err) + } + + async fn item_failures(&self, item_id: &str) -> ContextResult> { + let resp = self + .client + .datagen_item_failures(&self.store_name, item_id) + .await + .map_err(to_ctx_err)?; + Ok(resp.failures) + } + + async fn get_blob(&self, event_id: &str) -> ContextResult>> { + self.client + .fetch_datagen_blob(&self.store_name, event_id) + .await + .map_err(to_ctx_err) + } + + fn version(&self) -> u64 { + self.cached_version + } +} + fn to_ctx_err(err: ClientError) -> ContextError { match err { ClientError::Api { @@ -870,6 +972,130 @@ impl ContextClient { Self::handle_response(resp).await } + pub async fn create_datagen_store( + &self, + req: &CreateDatagenStoreRequest, + ) -> Result { + let resp = self + .http + .post(self.url("/datagen")) + .json(req) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn list_datagen_stores(&self) -> Result { + let resp = self.http.get(self.url("/datagen")).send().await?; + Self::handle_response(resp).await + } + + pub async fn get_datagen_store(&self, name: &str) -> Result { + let resp = self + .http + .get(self.url(&format!("/datagen/{}", name))) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn delete_datagen_store(&self, name: &str) -> Result<(), ClientError> { + let resp = self + .http + .delete(self.url(&format!("/datagen/{}", name))) + .send() + .await?; + if resp.status().is_success() { + Ok(()) + } else { + Err(Self::extract_error(resp).await) + } + } + + /// Append datagen events. `checkpoint = true` commits the batch as one atomic + /// step boundary (`append_checkpoint`); `false` appends a raw generation. + /// FIELD blobs are offloaded to a content-addressed artifact store before the + /// event reaches the log, so events are small JSON and travel inline. + pub async fn add_datagen_events( + &self, + name: &str, + events: &[DatagenEventDto], + checkpoint: bool, + ) -> Result { + let req = AddDatagenEventsRequest { + events: events.to_vec(), + }; + let resp = self + .http + .post(self.url(&format!("/datagen/{}/events", name))) + .query(&[("checkpoint", checkpoint)]) + .json(&req) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn fold_datagen_item( + &self, + name: &str, + item_id: &str, + ) -> Result { + let resp = self + .http + .get(self.url(&format!("/datagen/{}/items/{}", name, item_id))) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn datagen_item_failures( + &self, + name: &str, + item_id: &str, + ) -> Result { + let resp = self + .http + .get(self.url(&format!("/datagen/{}/items/{}/failures", name, item_id))) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn datagen_root_item_statuses( + &self, + name: &str, + root_item_ids: &[String], + ) -> Result { + let resp = self + .http + .get(self.url(&format!("/datagen/{}/root-status", name))) + .query(&[("ids", root_item_ids.join(","))]) + .send() + .await?; + Self::handle_response(resp).await + } + + /// Materialize one FIELD_* event's offloaded blob bytes by event id. + /// Returns `None` when the event or its payload is absent (server 404). + pub async fn fetch_datagen_blob( + &self, + name: &str, + event_id: &str, + ) -> Result>, ClientError> { + let resp = self + .http + .get(self.url(&format!("/datagen/{}/blobs/{}", name, event_id))) + .send() + .await?; + if resp.status().is_success() { + Ok(Some(resp.bytes().await?.to_vec())) + } else if resp.status().as_u16() == 404 { + Ok(None) + } else { + Err(Self::extract_error(resp).await) + } + } + async fn handle_response( resp: reqwest::Response, ) -> Result { diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index 09de0ec..4d945a6 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -3,14 +3,23 @@ use serde_json::Value; use uuid::Uuid; use lance_context_api::{ - AddRecordRequest, AddRecordsResponse, AddRolloutRequest, AddRolloutsResponse, CompactRequest, - CompactResponse, CompactStatsResponse, ContextError, ContextResult, ContextStoreApi, - DeleteRecordResponse, RecordDto, RecordPatchDto, RelationshipDto, RetrieveRequest, - RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchRequest, SearchResultDto, - StateMetadataDto, UpdateRecordRequest, UpdateRecordResponse, UpsertRecordRequest, - UpsertRecordResponse, UpsertRecordsRequest, UpsertRecordsResponse, UpsertResultDto, + AddDatagenEventsResponse, AddRecordRequest, AddRecordsResponse, AddRolloutRequest, + AddRolloutsResponse, CompactRequest, CompactResponse, CompactStatsResponse, ContextError, + ContextResult, ContextStoreApi, DatagenEventDto, DatagenFailureDto, DatagenFieldStateDto, + DatagenRootItemStatusesResponse, DatagenStepCursorDto, DatagenStoreApi, DatagenValueDto, + DeleteRecordResponse, FoldedDatagenItemDto, RecordDto, RecordPatchDto, RelationshipDto, + RetrieveRequest, RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchRequest, + SearchResultDto, StateMetadataDto, UpdateRecordRequest, UpdateRecordResponse, + UpsertRecordRequest, UpsertRecordResponse, UpsertRecordsRequest, UpsertRecordsResponse, + UpsertResultDto, }; +use crate::datagen::{ + DatagenBlobValue, DatagenEvent, DatagenEventType, DatagenFailure, DatagenFieldState, + DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses, DatagenStepCursor, + DatagenStepKind, DatagenValue, FoldedDatagenItem, +}; +use crate::datagen_store::DatagenStore; use crate::record::{ ContextRecord, LifecycleQueryOptions, RecordFilters, RecordPatch, Relationship, StateMetadata, LIFECYCLE_ACTIVE, @@ -678,3 +687,253 @@ fn to_ctx_err(err: lance::Error) -> ContextError { ContextError::Internal(msg) } } + +impl DatagenStoreApi for DatagenStore { + async fn append( + &mut self, + events: &[DatagenEventDto], + ) -> ContextResult { + let core = datagen_events_from_dtos(events)?; + let count = core.len(); + let version = DatagenStore::append(self, &core) + .await + .map_err(to_ctx_err)?; + Ok(AddDatagenEventsResponse { version, count }) + } + + async fn append_checkpoint( + &mut self, + events: &[DatagenEventDto], + ) -> ContextResult { + let core = datagen_events_from_dtos(events)?; + let count = core.len(); + let version = DatagenStore::append_checkpoint(self, &core) + .await + .map_err(to_ctx_err)?; + Ok(AddDatagenEventsResponse { version, count }) + } + + async fn fold_item(&self, item_id: &str) -> ContextResult> { + let lookup = DatagenStore::fold_item(self, item_id) + .await + .map_err(to_ctx_err)?; + Ok(match lookup { + DatagenItemLookup::NeverStarted => None, + DatagenItemLookup::Found(item) => Some(folded_item_to_dto(&item)), + }) + } + + async fn root_item_statuses( + &self, + root_item_ids: &[String], + ) -> ContextResult { + let ids: Vec<&str> = root_item_ids.iter().map(String::as_str).collect(); + let statuses = DatagenStore::root_item_statuses(self, &ids) + .await + .map_err(to_ctx_err)?; + Ok(root_item_statuses_to_dto(&statuses)) + } + + async fn item_failures(&self, item_id: &str) -> ContextResult> { + let failures = DatagenStore::item_failures(self, item_id) + .await + .map_err(to_ctx_err)?; + Ok(failures.iter().map(failure_to_dto).collect()) + } + + async fn get_blob(&self, event_id: &str) -> ContextResult>> { + DatagenStore::get_blob(self, event_id) + .await + .map_err(to_ctx_err) + } + + fn version(&self) -> u64 { + DatagenStore::version(self) + } +} + +fn datagen_events_from_dtos(events: &[DatagenEventDto]) -> ContextResult> { + events.iter().map(datagen_event_from_dto).collect() +} + +fn datagen_event_from_dto(dto: &DatagenEventDto) -> ContextResult { + let event_type = + DatagenEventType::parse(&dto.event_type).map_err(ContextError::InvalidRequest)?; + let step_kind = dto + .step_kind + .as_deref() + .map(DatagenStepKind::parse) + .transpose() + .map_err(ContextError::InvalidRequest)?; + let status = dto + .status + .as_deref() + .map(DatagenItemStatus::parse) + .transpose() + .map_err(ContextError::InvalidRequest)?; + let value = dto.value.as_ref().map(datagen_value_from_dto).transpose()?; + Ok(DatagenEvent { + event_id: dto.event_id.clone(), + item_id: dto.item_id.clone(), + root_item_id: dto.root_item_id.clone(), + parent_item_id: dto.parent_item_id.clone(), + item_seq: dto.item_seq, + checkpoint_id: dto.checkpoint_id.clone(), + event_type, + step_name: dto.step_name.clone(), + step_kind, + step_index: dto.step_index, + enclosing_step: dto.enclosing_step.clone(), + selector_step: dto.selector_step.clone(), + attempt: dto.attempt, + run_id: dto.run_id.clone(), + writer_epoch: dto.writer_epoch.clone(), + field_name: dto.field_name.clone(), + field_type: dto.field_type.clone(), + codec_version: dto.codec_version, + value, + query_tags: dto.query_tags.clone(), + status, + error_type: dto.error_type.clone(), + error_dump: dto.error_dump.clone(), + traceback: dto.traceback.clone(), + event_ts: dto.event_ts.unwrap_or_else(Utc::now), + schema_version: dto.schema_version, + }) +} + +fn datagen_value_from_dto(dto: &DatagenValueDto) -> ContextResult { + let missing = || { + ContextError::InvalidRequest(format!( + "datagen value of kind '{}' is missing 'value'", + dto.kind + )) + }; + match dto.kind.as_str() { + "int" => Ok(DatagenValue::Int( + dto.value + .as_ref() + .and_then(Value::as_i64) + .ok_or_else(missing)?, + )), + "float" => Ok(DatagenValue::Float( + dto.value + .as_ref() + .and_then(Value::as_f64) + .ok_or_else(missing)?, + )), + "bool" => Ok(DatagenValue::Bool( + dto.value + .as_ref() + .and_then(Value::as_bool) + .ok_or_else(missing)?, + )), + "str" => Ok(DatagenValue::Str( + dto.value + .as_ref() + .and_then(Value::as_str) + .ok_or_else(missing)? + .to_string(), + )), + "json" => Ok(DatagenValue::Json(dto.value.clone().ok_or_else(missing)?)), + "blob" => Ok(DatagenValue::Blob(DatagenBlobValue { + bytes: dto.bytes.clone(), + size: dto + .size + .unwrap_or_else(|| dto.bytes.as_ref().map(|b| b.len() as i64).unwrap_or(0)), + checksum: dto.checksum.clone(), + })), + other => Err(ContextError::InvalidRequest(format!( + "unsupported datagen value kind '{other}'" + ))), + } +} + +fn datagen_value_to_dto(value: &DatagenValue) -> DatagenValueDto { + let mut dto = DatagenValueDto { + kind: value.kind().to_string(), + value: None, + bytes: None, + size: None, + checksum: None, + }; + match value { + DatagenValue::Int(inner) => dto.value = Some(Value::from(*inner)), + DatagenValue::Float(inner) => dto.value = Some(Value::from(*inner)), + DatagenValue::Bool(inner) => dto.value = Some(Value::from(*inner)), + DatagenValue::Str(inner) => dto.value = Some(Value::from(inner.clone())), + DatagenValue::Json(inner) => dto.value = Some(inner.clone()), + DatagenValue::Blob(blob) => { + dto.bytes = blob.bytes.clone(); + dto.size = Some(blob.size); + dto.checksum = blob.checksum.clone(); + } + } + dto +} + +fn folded_item_to_dto(item: &FoldedDatagenItem) -> FoldedDatagenItemDto { + FoldedDatagenItemDto { + item_id: item.item_id.to_string(), + root_item_id: item.root_item_id.to_string(), + parent_item_id: item.parent_item_id.as_ref().map(ToString::to_string), + status: item.status.as_str().to_string(), + last_item_seq: item.last_item_seq, + last_attempt: item.last_attempt, + fields: item + .fields + .iter() + .map(|(name, state)| (name.clone(), field_state_to_dto(state))) + .collect(), + trajectory: item.trajectory.ordered.iter().map(cursor_to_dto).collect(), + query_tags: item.query_tags.clone(), + } +} + +fn field_state_to_dto(state: &DatagenFieldState) -> DatagenFieldStateDto { + match state { + DatagenFieldState::Set(value) => DatagenFieldStateDto { + mode: "set".to_string(), + value: Some(datagen_value_to_dto(value)), + values: Vec::new(), + }, + DatagenFieldState::Appended(values) => DatagenFieldStateDto { + mode: "append".to_string(), + value: None, + values: values.iter().map(datagen_value_to_dto).collect(), + }, + } +} + +fn cursor_to_dto(cursor: &DatagenStepCursor) -> DatagenStepCursorDto { + DatagenStepCursorDto { + step_name: cursor.position.step.name.clone(), + step_kind: cursor.position.step.kind.as_str().to_string(), + step_index: cursor.position.index, + enclosing_step: cursor.position.enclosing.clone(), + selector_step: cursor.position.selector.clone(), + item_seq: cursor.item_seq, + } +} + +fn root_item_statuses_to_dto( + statuses: &DatagenRootItemStatuses, +) -> DatagenRootItemStatusesResponse { + DatagenRootItemStatusesResponse { + statuses: statuses + .iter() + .map(|(id, status)| (id.clone(), status.as_str().to_string())) + .collect(), + } +} + +fn failure_to_dto(failure: &DatagenFailure) -> DatagenFailureDto { + DatagenFailureDto { + at: cursor_to_dto(&failure.at), + run_id: failure.run_id.clone(), + attempt: failure.attempt, + error_type: failure.error.error_type.clone(), + error_dump: failure.error.error_dump.clone(), + traceback: failure.error.traceback.clone(), + } +} diff --git a/crates/lance-context-server/src/error.rs b/crates/lance-context-server/src/error.rs index c5bb116..6e48cbb 100644 --- a/crates/lance-context-server/src/error.rs +++ b/crates/lance-context-server/src/error.rs @@ -1,7 +1,7 @@ use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; use axum::Json; -use lance_context_api::{ErrorBody, ErrorResponse}; +use lance_context_api::{ContextError, ErrorBody, ErrorResponse}; use lance_context_core::LanceError; #[derive(Debug)] @@ -44,6 +44,19 @@ impl AppError { other => AppError::Internal(other.to_string()), } } + + /// Map an API-layer [`ContextError`] onto the server's error taxonomy. + /// Datagen routes call the `DatagenStoreApi` trait (which owns the DTO↔core + /// conversion) and so surface `ContextError` rather than raw `LanceError`. + pub fn from_context(err: ContextError) -> Self { + match err { + ContextError::NotFound(msg) => AppError::NotFound(msg), + ContextError::AlreadyExists(msg) => AppError::AlreadyExists(msg), + ContextError::InvalidRequest(msg) => AppError::InvalidRequest(msg), + ContextError::Internal(msg) => AppError::Internal(msg), + ContextError::CompactionInProgress => AppError::CompactionInProgress, + } + } } impl IntoResponse for AppError { diff --git a/crates/lance-context-server/src/routes/datagen.rs b/crates/lance-context-server/src/routes/datagen.rs new file mode 100644 index 0000000..bbdb35a --- /dev/null +++ b/crates/lance-context-server/src/routes/datagen.rs @@ -0,0 +1,225 @@ +use std::sync::Arc; + +use axum::extract::{Path, Query, State}; +use axum::http::StatusCode; +use axum::Json; +use lance_context_api::{ + AddDatagenEventsRequest, AddDatagenEventsResponse, CreateDatagenStoreRequest, DatagenStoreApi, + DatagenStoreInfo, GetFoldedDatagenItemResponse, ListDatagenFailuresResponse, + ListDatagenStoresResponse, +}; +use lance_context_core::{DatagenStore, DatagenStoreOptions}; +use tokio::sync::RwLock; + +use crate::error::AppError; +use crate::state::AppState; + +/// Upper bound on a single datagen append request body. FIELD blobs are +/// offloaded to a content-addressed artifact store before the event reaches the +/// log, so events themselves are small; the ceiling still bounds a pathological +/// batch. +pub const MAX_DATAGEN_UPLOAD_BYTES: usize = 256 * 1024 * 1024; + +pub async fn create_datagen_store( + State(state): State>, + Json(req): Json, +) -> Result<(StatusCode, Json), AppError> { + AppState::validate_name(&req.name)?; + if state + .datagen_registry + .write() + .await + .contains(&req.name) + .await + .map_err(AppError::from_lance)? + { + return Err(AppError::AlreadyExists(format!( + "Datagen store '{}' already exists", + req.name + ))); + } + + let uri = state.datagen_uri(&req.name); + let options = DatagenStoreOptions { + storage_options: req.storage_options, + shard_id: state.instance_id.clone(), + merge_after_generations: None, + cleanup_interval_secs: None, + }; + let store = DatagenStore::open_with_options(&uri, options) + .await + .map_err(AppError::from_lance)?; + let version = store.version(); + + let store = Arc::new(RwLock::new(store)); + state.register_datagen(&req.name, &uri, store).await?; + + Ok(( + StatusCode::CREATED, + Json(DatagenStoreInfo { + name: req.name, + uri, + version: Some(version), + }), + )) +} + +pub async fn list_datagen_stores( + State(state): State>, +) -> Result, AppError> { + let entries = state + .datagen_registry + .write() + .await + .list() + .await + .map_err(AppError::from_lance)?; + let stores = entries + .into_iter() + .map(|entry| DatagenStoreInfo { + name: entry.name, + uri: entry.uri, + version: None, + }) + .collect(); + Ok(Json(ListDatagenStoresResponse { stores })) +} + +pub async fn get_datagen_store( + State(state): State>, + Path(name): Path, +) -> Result, AppError> { + let store_lock = state.get_or_open_datagen_store(&name).await?; + let store = store_lock.read().await; + Ok(Json(DatagenStoreInfo { + name: name.clone(), + uri: state.datagen_uri(&name), + version: Some(store.version()), + })) +} + +pub async fn delete_datagen_store( + State(state): State>, + Path(name): Path, +) -> Result { + if !state.unregister_datagen(&name).await? { + return Err(AppError::NotFound(format!( + "Datagen store '{}' does not exist", + name + ))); + } + let uri = state.datagen_uri(&name); + if let Err(e) = tokio::fs::remove_dir_all(&uri).await { + tracing::warn!("Failed to remove datagen data at {}: {}", uri, e); + } + Ok(StatusCode::NO_CONTENT) +} + +/// Which append semantics the `/events` endpoint applies. +#[derive(Debug, Default, serde::Deserialize)] +pub struct AppendParams { + /// When `true`, the batch is committed as one atomic checkpoint (FIELD_* + /// events plus exactly one STEP_COMPLETED) via `append_checkpoint`. Default + /// appends the events as one raw MemWAL generation. + #[serde(default)] + pub checkpoint: bool, +} + +pub async fn add_datagen_events( + State(state): State>, + Path(name): Path, + Query(params): Query, + Json(req): Json, +) -> Result<(StatusCode, Json), AppError> { + if req.events.is_empty() { + return Err(AppError::InvalidRequest( + "events array must not be empty".to_string(), + )); + } + + let store_lock = state.get_or_open_datagen_store(&name).await?; + let mut store = store_lock.write().await; + let resp = if params.checkpoint { + DatagenStoreApi::append_checkpoint(&mut *store, &req.events).await + } else { + DatagenStoreApi::append(&mut *store, &req.events).await + } + .map_err(AppError::from_context)?; + + Ok((StatusCode::CREATED, Json(resp))) +} + +pub async fn fold_datagen_item( + State(state): State>, + Path((name, item_id)): Path<(String, String)>, +) -> Result, AppError> { + let store_lock = state.get_or_open_datagen_store(&name).await?; + let store = store_lock.read().await; + let item = DatagenStoreApi::fold_item(&*store, &item_id) + .await + .map_err(AppError::from_context)?; + Ok(Json(GetFoldedDatagenItemResponse { item })) +} + +pub async fn datagen_item_failures( + State(state): State>, + Path((name, item_id)): Path<(String, String)>, +) -> Result, AppError> { + let store_lock = state.get_or_open_datagen_store(&name).await?; + let store = store_lock.read().await; + let failures = DatagenStoreApi::item_failures(&*store, &item_id) + .await + .map_err(AppError::from_context)?; + Ok(Json(ListDatagenFailuresResponse { failures })) +} + +#[derive(Debug, Default, serde::Deserialize)] +pub struct RootStatusParams { + /// Comma-separated list of root item ids to classify. + pub ids: Option, +} + +pub async fn datagen_root_item_statuses( + State(state): State>, + Path(name): Path, + Query(params): Query, +) -> Result, AppError> { + let ids: Vec = params + .ids + .as_deref() + .map(|raw| { + raw.split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(str::to_string) + .collect() + }) + .unwrap_or_default(); + + let store_lock = state.get_or_open_datagen_store(&name).await?; + let store = store_lock.read().await; + let resp = DatagenStoreApi::root_item_statuses(&*store, &ids) + .await + .map_err(AppError::from_context)?; + Ok(Json(resp)) +} + +/// Materialize one FIELD_* event's offloaded blob bytes by event id. The bytes +/// are opaque, so they return as `application/octet-stream`; `404` when the +/// event or its payload is absent. +pub async fn fetch_datagen_blob( + State(state): State>, + Path((name, event_id)): Path<(String, String)>, +) -> Result { + use axum::http::header; + use axum::response::IntoResponse; + + let store_lock = state.get_or_open_datagen_store(&name).await?; + let store = store_lock.read().await; + let bytes = DatagenStoreApi::get_blob(&*store, &event_id) + .await + .map_err(AppError::from_context)? + .ok_or_else(|| AppError::NotFound(format!("Datagen event '{}' has no blob", event_id)))?; + + Ok(([(header::CONTENT_TYPE, "application/octet-stream")], bytes).into_response()) +} diff --git a/crates/lance-context-server/src/routes/mod.rs b/crates/lance-context-server/src/routes/mod.rs index c42eb15..17ab77d 100644 --- a/crates/lance-context-server/src/routes/mod.rs +++ b/crates/lance-context-server/src/routes/mod.rs @@ -1,5 +1,6 @@ pub mod compact; pub mod contexts; +pub mod datagen; pub mod health; pub mod records; pub mod rollouts; @@ -117,6 +118,34 @@ pub fn router() -> Router> { "/api/v1/internal/merge-wal/{name}", post(rollouts::merge_wal), ) + .route("/api/v1/datagen", post(datagen::create_datagen_store)) + .route("/api/v1/datagen", get(datagen::list_datagen_stores)) + .route("/api/v1/datagen/{name}", get(datagen::get_datagen_store)) + .route( + "/api/v1/datagen/{name}", + delete(datagen::delete_datagen_store), + ) + .route( + "/api/v1/datagen/{name}/events", + post(datagen::add_datagen_events) + .layer(DefaultBodyLimit::max(datagen::MAX_DATAGEN_UPLOAD_BYTES)), + ) + .route( + "/api/v1/datagen/{name}/items/{item_id}", + get(datagen::fold_datagen_item), + ) + .route( + "/api/v1/datagen/{name}/items/{item_id}/failures", + get(datagen::datagen_item_failures), + ) + .route( + "/api/v1/datagen/{name}/root-status", + get(datagen::datagen_root_item_statuses), + ) + .route( + "/api/v1/datagen/{name}/blobs/{event_id}", + get(datagen::fetch_datagen_blob), + ) } #[cfg(test)] diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 1bea836..d0dd394 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -5,8 +5,8 @@ use std::sync::Arc; use std::time::Duration; use lance_context_core::{ - join_uri, validate_store_name, ContextStore, ContextStoreOptions, RolloutRegistry, - RolloutStore, RolloutStoreOptions, Session, + join_uri, validate_store_name, ContextStore, ContextStoreOptions, DatagenStore, + DatagenStoreOptions, RolloutRegistry, RolloutStore, RolloutStoreOptions, Session, }; use lru::LruCache; use tokio::sync::{Mutex, RwLock}; @@ -74,6 +74,15 @@ pub struct AppState { /// than Lance's default 6 GiB *per store*. `None` restores the per-store /// default session (leak-prone; only when the budget is configured to `0`). pub rollout_session: Option>, + /// Bounded LRU of resident datagen-store handles, mirroring + /// [`Self::rollout_stores`]. Datagen delta-log datasets are also one per + /// experiment, so the same residency bound and durable-registry existence + /// model applies. + pub datagen_stores: Mutex>>>, + /// Durable directory of which datagen stores exist. A separate registry + /// dataset from [`Self::rollout_registry`] so the two store kinds never + /// collide on a shared name. + pub datagen_registry: RwLock, } /// Process-wide admission control for the total artifact-blob payload held in @@ -185,6 +194,10 @@ impl AppState { let blob_budget = (config.rollout_max_inflight_blob_bytes > 0) .then(|| BlobBudget::new(config.rollout_max_inflight_blob_bytes)); let rollout_session = build_rollout_session(config.rollout_cache_bytes); + let datagen_registry_uri = join_uri(&base_uri, "_registry.datagen.lance"); + let datagen_registry = RolloutRegistry::open_or_create(&datagen_registry_uri, None) + .await + .map_err(AppError::from_lance)?; Ok(Self { stores: RwLock::new(std::collections::HashMap::new()), rollout_stores: Mutex::new(LruCache::new(capacity)), @@ -196,6 +209,8 @@ impl AppState { rollout_flush_interval_secs: config.rollout_flush_interval_secs, blob_budget, rollout_session, + datagen_stores: Mutex::new(LruCache::new(capacity)), + datagen_registry: RwLock::new(datagen_registry), }) } @@ -222,6 +237,10 @@ impl AppState { let registry = RolloutRegistry::open_or_create(®istry_uri, None) .await .expect("open test registry"); + let datagen_registry_uri = join_uri(&base_uri, "_registry.datagen.lance"); + let datagen_registry = RolloutRegistry::open_or_create(&datagen_registry_uri, None) + .await + .expect("open test datagen registry"); Self { stores: RwLock::new(std::collections::HashMap::new()), rollout_stores: Mutex::new(LruCache::new( @@ -235,6 +254,10 @@ impl AppState { rollout_flush_interval_secs: 0, blob_budget: None, rollout_session: build_rollout_session(2 * 1024 * 1024 * 1024), + datagen_stores: Mutex::new(LruCache::new( + NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), + )), + datagen_registry: RwLock::new(datagen_registry), } } @@ -392,6 +415,109 @@ impl AppState { validate_store_name(name).map_err(AppError::InvalidRequest) } + /// Datagen delta-log datasets live under a distinct `.datagen.lance` suffix + /// so a datagen store shares neither a rollout nor a context store's path. + pub fn datagen_uri(&self, name: &str) -> String { + join_uri(&self.base_uri, &format!("{}.datagen.lance", name)) + } + + fn datagen_store_options(&self) -> DatagenStoreOptions { + DatagenStoreOptions { + storage_options: None, + shard_id: self.instance_id.clone(), + merge_after_generations: (self.rollout_merge_after_generations > 0) + .then_some(self.rollout_merge_after_generations), + cleanup_interval_secs: None, + } + } + + /// Record that a datagen store exists, in both the durable registry and the + /// in-memory LRU. Called by the create route after the dataset is written. + pub async fn register_datagen( + &self, + name: &str, + uri: &str, + store: Arc>, + ) -> Result<(), AppError> { + Self::validate_name(name)?; + self.datagen_registry + .write() + .await + .upsert(name, uri) + .await + .map_err(AppError::from_lance)?; + self.datagen_stores + .lock() + .await + .put(name.to_string(), store); + Ok(()) + } + + /// Remove a datagen store from the durable registry and evict any resident + /// handle. Returns whether the store existed. + pub async fn unregister_datagen(&self, name: &str) -> Result { + Self::validate_name(name)?; + let existed = self + .datagen_registry + .write() + .await + .contains(name) + .await + .map_err(AppError::from_lance)?; + if !existed { + return Ok(false); + } + self.datagen_registry + .write() + .await + .remove(name) + .await + .map_err(AppError::from_lance)?; + self.datagen_stores.lock().await.pop(name); + Ok(true) + } + + /// Look up a datagen store by name, lazily loading it from object storage on + /// a local cache miss. Existence is resolved against the durable registry, so + /// an evicted-but-registered store reopens rather than 404ing. See + /// [`Self::get_or_open_rollout_store`] for the full rationale. + pub async fn get_or_open_datagen_store( + &self, + name: &str, + ) -> Result>, AppError> { + Self::validate_name(name)?; + if let Some(store) = self.datagen_stores.lock().await.get(name) { + return Ok(store.clone()); + } + + let exists = self + .datagen_registry + .write() + .await + .contains(name) + .await + .map_err(AppError::from_lance)?; + if !exists { + return Err(AppError::NotFound(format!( + "Datagen store '{}' does not exist", + name + ))); + } + + let uri = self.datagen_uri(name); + let opened = DatagenStore::open_existing_with_options(&uri, self.datagen_store_options()) + .await + .map_err(AppError::from_lance)?; + let opened = Arc::new(RwLock::new(opened)); + + let mut cache = self.datagen_stores.lock().await; + if let Some(existing) = cache.get(name) { + return Ok(existing.clone()); + } + cache.put(name.to_string(), opened.clone()); + Ok(opened) + } + /// Spawn the single, process-wide WAL-cleanup sweeper. /// /// This replaces the former one-timer-per-store model, which does not scale diff --git a/crates/lance-context/src/lib.rs b/crates/lance-context/src/lib.rs index 93d6649..063ccfe 100644 --- a/crates/lance-context/src/lib.rs +++ b/crates/lance-context/src/lib.rs @@ -6,27 +6,34 @@ pub use lance_context_core::{ datagen_event_id, datagen_log_schema, datagen_trajectory, fold_datagen_events, CompactionConfig, CompactionMetrics, CompactionStats, Context, ContextEntry, ContextNamespace, ContextRecord, ContextStoreOptions, DatagenBlobValue, DatagenEvent, DatagenEventType, - DatagenFailure, DatagenFieldState, DatagenItemStatus, DatagenStepCursor, DatagenStore, - DatagenStoreOptions, DatagenTerminal, DatagenTrajectory, DatagenValue, FoldedDatagenItem, - IdIndexType, LifecycleQueryOptions, MetadataFilter, PartitionInfo, PartitionSelector, - PartitionSpec, RecordFilters, Relationship, RetrieveResult, RolloutFilters, RolloutRecord, - SearchResult, Snapshot, StateMetadata, DATAGEN_SCHEMA_VERSION, LIFECYCLE_ACTIVE, - LIFECYCLE_CONTRADICTED, + DatagenFailure, DatagenFieldState, DatagenItemStatus, DatagenStepCursor, DatagenStoreOptions, + DatagenTerminal, DatagenTrajectory, DatagenValue, FoldedDatagenItem, IdIndexType, + LifecycleQueryOptions, MetadataFilter, PartitionInfo, PartitionSelector, PartitionSpec, + RecordFilters, Relationship, RetrieveResult, RolloutFilters, RolloutRecord, SearchResult, + Snapshot, StateMetadata, DATAGEN_SCHEMA_VERSION, LIFECYCLE_ACTIVE, LIFECYCLE_CONTRADICTED, }; pub use lance_context_api::{ - AddRecordRequest, AddRecordsResponse, AddRolloutRequest, AddRolloutsResponse, CompactRequest, - CompactResponse, CompactStatsResponse, ContextError, ContextResult, ContextStoreApi, - CreateRolloutStoreRequest, DeleteRecordResponse, RecordDto, RelationshipDto, RetrieveRequest, + AddDatagenEventsRequest, AddDatagenEventsResponse, AddRecordRequest, AddRecordsResponse, + AddRolloutRequest, AddRolloutsResponse, CompactRequest, CompactResponse, CompactStatsResponse, + ContextError, ContextResult, ContextStoreApi, CreateDatagenStoreRequest, + CreateRolloutStoreRequest, DatagenEventDto, DatagenFailureDto, DatagenFieldStateDto, + DatagenRootItemStatusesResponse, DatagenStepCursorDto, DatagenStoreApi, DatagenValueDto, + DeleteRecordResponse, FoldedDatagenItemDto, RecordDto, RelationshipDto, RetrieveRequest, RetrieveResponse, RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchResultDto, UpsertRecordRequest, UpsertRecordResponse, }; #[cfg(feature = "remote")] -pub use lance_context_client::{ClientError, RemoteContextStore, RemoteRolloutStore}; +pub use lance_context_client::{ + ClientError, RemoteContextStore, RemoteDatagenStore, RemoteRolloutStore, +}; mod unified; pub use unified::ContextStore; mod unified_rollout; pub use unified_rollout::RolloutStore; + +mod unified_datagen; +pub use unified_datagen::DatagenStore; diff --git a/crates/lance-context/src/unified_datagen.rs b/crates/lance-context/src/unified_datagen.rs new file mode 100644 index 0000000..b43861b --- /dev/null +++ b/crates/lance-context/src/unified_datagen.rs @@ -0,0 +1,130 @@ +use lance_context_api::{ + AddDatagenEventsResponse, ContextError, ContextResult, DatagenEventDto, DatagenFailureDto, + DatagenRootItemStatusesResponse, DatagenStoreApi, FoldedDatagenItemDto, +}; +use lance_context_core::{DatagenStore as LocalStore, DatagenStoreOptions}; + +#[cfg(feature = "remote")] +use lance_context_client::RemoteDatagenStore; + +/// A datagen checkpoint store that is either an in-process Lance dataset +/// (`Local`) or a handle to a remote server (`Remote`). Mirrors +/// [`crate::RolloutStore`] but for the datagen delta-log schema. +pub enum DatagenStore { + Local(Box), + #[cfg(feature = "remote")] + Remote(RemoteDatagenStore), +} + +impl DatagenStore { + pub async fn open(uri: &str) -> Result { + Self::open_with_options(uri, None).await + } + + pub async fn open_with_options( + uri: &str, + storage_options: Option>, + ) -> Result { + let options = DatagenStoreOptions { + storage_options, + // Embedded single-process use writes to the fallback shard; a + // multi-writer embedded deployment threads a per-writer id through + // the core `DatagenStore` directly. + shard_id: None, + merge_after_generations: None, + cleanup_interval_secs: None, + }; + let store = LocalStore::open_with_options(uri, options) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Local(Box::new(store))) + } + + #[cfg(feature = "remote")] + pub async fn connect(base_url: &str, store_name: &str) -> Result { + let store = RemoteDatagenStore::connect(base_url, store_name) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Remote(store)) + } + + #[cfg(feature = "remote")] + pub async fn connect_or_create( + base_url: &str, + req: &lance_context_api::CreateDatagenStoreRequest, + ) -> Result { + let store = RemoteDatagenStore::connect_or_create(base_url, req) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Remote(store)) + } +} + +macro_rules! dispatch_mut { + ($self:expr, $method:ident $(, $arg:expr)*) => { + match $self { + DatagenStore::Local(s) => DatagenStoreApi::$method(s.as_mut() $(, $arg)*).await, + #[cfg(feature = "remote")] + DatagenStore::Remote(s) => DatagenStoreApi::$method(s $(, $arg)*).await, + } + }; +} + +macro_rules! dispatch_ref { + ($self:expr, $method:ident $(, $arg:expr)*) => { + match $self { + DatagenStore::Local(s) => DatagenStoreApi::$method(s.as_ref() $(, $arg)*).await, + #[cfg(feature = "remote")] + DatagenStore::Remote(s) => DatagenStoreApi::$method(s $(, $arg)*).await, + } + }; +} + +macro_rules! dispatch_sync { + ($self:expr, $method:ident $(, $arg:expr)*) => { + match $self { + DatagenStore::Local(s) => DatagenStoreApi::$method(s.as_ref() $(, $arg)*), + #[cfg(feature = "remote")] + DatagenStore::Remote(s) => DatagenStoreApi::$method(s $(, $arg)*), + } + }; +} + +impl DatagenStoreApi for DatagenStore { + async fn append( + &mut self, + events: &[DatagenEventDto], + ) -> ContextResult { + dispatch_mut!(self, append, events) + } + + async fn append_checkpoint( + &mut self, + events: &[DatagenEventDto], + ) -> ContextResult { + dispatch_mut!(self, append_checkpoint, events) + } + + async fn fold_item(&self, item_id: &str) -> ContextResult> { + dispatch_ref!(self, fold_item, item_id) + } + + async fn root_item_statuses( + &self, + root_item_ids: &[String], + ) -> ContextResult { + dispatch_ref!(self, root_item_statuses, root_item_ids) + } + + async fn item_failures(&self, item_id: &str) -> ContextResult> { + dispatch_ref!(self, item_failures, item_id) + } + + async fn get_blob(&self, event_id: &str) -> ContextResult>> { + dispatch_ref!(self, get_blob, event_id) + } + + fn version(&self) -> u64 { + dispatch_sync!(self, version) + } +} diff --git a/python/python/lance_context/__init__.py b/python/python/lance_context/__init__.py index bbb5ec5..7072b2b 100644 --- a/python/python/lance_context/__init__.py +++ b/python/python/lance_context/__init__.py @@ -5,10 +5,12 @@ AsyncRolloutStore, Context, ContextNamespace, + DatagenStore, EmbeddingProvider, RemoteContext, RolloutStore, __version__, + datagen_event_id, generate_id, ) from .embeddings import ( # pyright: ignore[reportMissingImports] @@ -20,10 +22,12 @@ "AsyncRolloutStore", "Context", "ContextNamespace", + "DatagenStore", "EmbeddingProvider", "MultiModalEmbeddingProvider", "RemoteContext", "RolloutStore", "__version__", + "datagen_event_id", "generate_id", ] diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 11a6998..f5e9f15 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -3,7 +3,7 @@ import asyncio import json import warnings -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from io import BytesIO from typing import TYPE_CHECKING, Any @@ -21,6 +21,12 @@ from ._internal import ( # pyright: ignore[reportMissingImports] RolloutStore as _RolloutStore, ) +from ._internal import ( # pyright: ignore[reportMissingImports] + DatagenStore as _DatagenStore, +) +from ._internal import ( # pyright: ignore[reportMissingImports] + datagen_event_id as datagen_event_id, +) from ._internal import ( # pyright: ignore[reportMissingImports] generate_id as _generate_id, ) @@ -32,10 +38,12 @@ "AsyncRolloutStore", "Context", "ContextNamespace", + "DatagenStore", "EmbeddingProvider", "RemoteContext", "RolloutStore", "__version__", + "datagen_event_id", "generate_id", ] @@ -2738,3 +2746,89 @@ async def checkout(self, version: int) -> None: def __repr__(self) -> str: return f"AsyncRolloutStore(version={self._sync.version()})" + + +class DatagenStore: + """Synchronous append-only datagen delta-log store. + + Wraps the native store: events are appended as plain dicts and item state is + recovered by folding an item's events. Open an embedded log with :meth:`open`. + + Event dicts carry the datagen event schema (``event_id``, ``item_id``, + ``root_item_id``, ``event_type``, ``item_seq``, ...); values use the tagged + ``{"kind": ..., "value": ...}`` shape. Folded items and failures come back as + plain dicts. + """ + + def __init__(self, sync_store: _DatagenStore) -> None: + self._sync = sync_store + + @classmethod + def open( + cls, + uri: str, + *, + storage_options: Mapping[str, str] | None = None, + shard_id: str | None = None, + ) -> "DatagenStore": + """Open (or create) an embedded datagen log at ``uri``. + + ``shard_id`` gives this writer a stable identity for multi-writer fencing. + """ + opts = dict(storage_options) if storage_options else None + return cls(_DatagenStore.open(uri, opts, shard_id)) + + @classmethod + def connect(cls, base_url: str, name: str) -> "DatagenStore": + """Connect to an existing datagen store on a remote server.""" + return cls(_DatagenStore.connect(base_url, name)) + + @classmethod + def connect_or_create( + cls, + base_url: str, + name: str, + *, + storage_options: Mapping[str, str] | None = None, + ) -> "DatagenStore": + """Connect to a remote datagen store, creating it if absent.""" + opts = dict(storage_options) if storage_options else None + return cls(_DatagenStore.connect_or_create(base_url, name, opts)) + + def version(self) -> int: + """Return the current store version (base dataset version).""" + return self._sync.version() + + def append_checkpoint(self, events: Iterable[Mapping[str, Any]]) -> int: + """Append one completed step boundary atomically. + + ``events`` share one item/checkpoint/writer attempt and contain exactly + one ``STEP_COMPLETED``. Returns the new store version. + """ + return self._sync.append_checkpoint(list(events)) + + def append(self, events: Iterable[Mapping[str, Any]]) -> int: + """Append raw events as one MemWAL generation. Returns the new store version.""" + return self._sync.append(list(events)) + + def fold_item(self, item_id: str) -> dict[str, Any] | None: + """Fold an item's events into its latest state, or ``None`` if never started.""" + return self._sync.fold_item(item_id) + + def root_item_statuses(self, root_item_ids: Sequence[str]) -> dict[str, str]: + """Classify each root item id by folded lifecycle status. + + Missing ids (never started) are absent from the returned dict. + """ + return self._sync.root_item_statuses(list(root_item_ids)) + + def item_failures(self, item_id: str) -> list[dict[str, Any]]: + """All failure records for an item (the failure lens), oldest first.""" + return self._sync.item_failures(item_id) + + def get_blob(self, event_id: str) -> bytes | None: + """Materialize one ``FIELD_*`` event's blob bytes by event id, or ``None``.""" + return self._sync.get_blob(event_id) + + def __repr__(self) -> str: + return f"DatagenStore(version={self._sync.version()})" diff --git a/python/src/lib.rs b/python/src/lib.rs index 7491ba4..9cc9283 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -12,8 +12,10 @@ use serde_json::Value; use tokio::runtime::Runtime; use lance_context::{ - AddRolloutRequest, CreateRolloutStoreRequest, RolloutRecordDto, - RolloutStore as UnifiedRolloutStore, RolloutStoreApi, + AddRolloutRequest, CreateDatagenStoreRequest, CreateRolloutStoreRequest, DatagenEventDto, + DatagenFailureDto, DatagenFieldStateDto, DatagenStepCursorDto, + DatagenStore as UnifiedDatagenStore, DatagenStoreApi, DatagenValueDto, FoldedDatagenItemDto, + RolloutRecordDto, RolloutStore as UnifiedRolloutStore, RolloutStoreApi, }; use lance_context_api::{ AddRecordRequest, CompactRequest, CompactResponse, CompactStatsResponse, ContextStoreApi, @@ -23,12 +25,13 @@ use lance_context_api::{ use lance_context_client::RemoteContextStore; use lance_context_core::serde::CONTENT_TYPE_TEXT; use lance_context_core::{ - CompactionConfig, CompactionMetrics, CompactionStats, Context as RustContext, - ContextNamespace as RustContextNamespace, ContextRecord, ContextStore, ContextStoreOptions, - DistanceMetric, EvalConfig, EvalQuerySet, ExportConfig, ExportTask, GroupBy, IdIndexType, - LifecycleQueryOptions, PartitionInfo, PartitionSelector, PartitionSpec, PreferenceForm, - ReadProjection, RecordFilters, RecordPatch, Relationship, RetrievalMode, RetrieveResult, - SearchResult, SplitConfig, StateMetadata, LIFECYCLE_ACTIVE, + datagen_event_id as core_datagen_event_id, CompactionConfig, CompactionMetrics, + CompactionStats, Context as RustContext, ContextNamespace as RustContextNamespace, + ContextRecord, ContextStore, ContextStoreOptions, DistanceMetric, EvalConfig, EvalQuerySet, + ExportConfig, ExportTask, GroupBy, IdIndexType, LifecycleQueryOptions, PartitionInfo, + PartitionSelector, PartitionSpec, PreferenceForm, ReadProjection, RecordFilters, RecordPatch, + Relationship, RetrievalMode, RetrieveResult, SearchResult, SplitConfig, StateMetadata, + DATAGEN_SCHEMA_VERSION, LIFECYCLE_ACTIVE, }; const DEFAULT_BINARY_CONTENT_TYPE: &str = "application/octet-stream"; @@ -2548,13 +2551,409 @@ fn rollout_record_to_json(record: &RolloutRecordDto) -> PyResult { serde_json::to_string(record).map_err(to_py_err) } +// --------------------------------------------------------------------------- +// Datagen store binding (append-only delta-log, fold model) +// --------------------------------------------------------------------------- + +/// A single embedded datagen checkpoint log (`open`). +/// +/// Events cross the FFI boundary as plain dicts (not pickle): the executor +/// builds each `DatagenEvent` field-by-field, `append_checkpoint` persists one +/// atomic step boundary, and reads (`fold_item`) return the folded item state as +/// a dict. Mirrors the concurrency contract of the rollout store: open a fresh +/// handle per writer, never share a handle across concurrent appends. +#[pyclass] +struct DatagenStore { + store: UnifiedDatagenStore, + runtime: Arc, +} + +impl DatagenStore { + fn from_store(store: UnifiedDatagenStore, runtime: Arc) -> Self { + Self { store, runtime } + } +} + +#[pymethods] +impl DatagenStore { + /// Open (or create) an embedded datagen log at `uri`. `shard_id` gives this + /// writer instance a stable identity for multi-writer fencing. + #[classmethod] + #[pyo3(signature = (uri, storage_options = None, shard_id = None))] + fn open( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + uri: &str, + storage_options: Option>, + shard_id: Option, + ) -> PyResult { + let _ = shard_id; + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = py.allow_threads(|| { + runtime.block_on(UnifiedDatagenStore::open_with_options(uri, storage_options)) + }); + let store = store_res.map_err(to_py_err)?; + Ok(Self::from_store(store, runtime)) + } + + /// Connect to an existing datagen store on a remote server. + #[classmethod] + fn connect( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + base_url: &str, + name: &str, + ) -> PyResult { + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = + py.allow_threads(|| runtime.block_on(UnifiedDatagenStore::connect(base_url, name))); + let store = store_res.map_err(to_py_err)?; + Ok(Self::from_store(store, runtime)) + } + + /// Connect to a remote datagen store, creating it if it does not exist. + #[classmethod] + #[pyo3(signature = (base_url, name, storage_options = None))] + fn connect_or_create( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + base_url: &str, + name: &str, + storage_options: Option>, + ) -> PyResult { + let req = CreateDatagenStoreRequest { + name: name.to_string(), + storage_options, + }; + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = py.allow_threads(|| { + runtime.block_on(UnifiedDatagenStore::connect_or_create(base_url, &req)) + }); + let store = store_res.map_err(to_py_err)?; + Ok(Self::from_store(store, runtime)) + } + + /// Current store version (base dataset version). + fn version(&self) -> u64 { + self.store.version() + } + + /// Append one completed step boundary atomically. `events` is a list of + /// event dicts sharing one item/checkpoint/writer attempt, with exactly one + /// STEP_COMPLETED. Returns the new store version. + fn append_checkpoint(&mut self, py: Python<'_>, events: &Bound<'_, PyList>) -> PyResult { + let parsed = events_from_pylist(events)?; + let resp = py + .allow_threads(|| self.runtime.block_on(self.store.append_checkpoint(&parsed))) + .map_err(to_py_err)?; + Ok(resp.version) + } + + /// Append raw events as one MemWAL generation (no single-STEP_COMPLETED + /// constraint). Returns the new store version. + fn append(&mut self, py: Python<'_>, events: &Bound<'_, PyList>) -> PyResult { + let parsed = events_from_pylist(events)?; + let resp = py + .allow_threads(|| self.runtime.block_on(self.store.append(&parsed))) + .map_err(to_py_err)?; + Ok(resp.version) + } + + /// Fold an item's events into its latest state, or `None` if never started. + fn fold_item(&self, py: Python<'_>, item_id: &str) -> PyResult> { + let item = py + .allow_threads(|| self.runtime.block_on(self.store.fold_item(item_id))) + .map_err(to_py_err)?; + match item { + None => Ok(None), + Some(item) => Ok(Some(folded_item_to_py(py, &item)?)), + } + } + + /// Classify each root item id by folded lifecycle status. Missing ids (never + /// started) are absent from the returned dict. + fn root_item_statuses(&self, py: Python<'_>, root_item_ids: Vec) -> PyResult { + let statuses = py + .allow_threads(|| { + self.runtime + .block_on(self.store.root_item_statuses(&root_item_ids)) + }) + .map_err(to_py_err)?; + let dict = PyDict::new(py); + for (item_id, status) in statuses.statuses.iter() { + dict.set_item(item_id, status)?; + } + Ok(dict.into_pyobject(py)?.unbind().into()) + } + + /// All failure records for an item (the failure lens), oldest first. + fn item_failures(&self, py: Python<'_>, item_id: &str) -> PyResult { + let failures = py + .allow_threads(|| self.runtime.block_on(self.store.item_failures(item_id))) + .map_err(to_py_err)?; + let list = PyList::empty(py); + for failure in &failures { + list.append(failure_to_py(py, failure)?)?; + } + Ok(list.into_pyobject(py)?.unbind().into()) + } + + /// Materialize one FIELD_* event's blob bytes by event id, or `None`. + fn get_blob(&self, py: Python<'_>, event_id: &str) -> PyResult>> { + let bytes = py + .allow_threads(|| self.runtime.block_on(self.store.get_blob(event_id))) + .map_err(to_py_err)?; + Ok(bytes.map(|b| PyBytes::new(py, &b).unbind())) + } +} + +/// Deterministic idempotency key for a checkpoint event, so retried batches +/// dedup instead of double-appending. +#[pyfunction] +fn datagen_event_id(item_id: &str, checkpoint_id: &str, ordinal: u32) -> String { + core_datagen_event_id(item_id, checkpoint_id, ordinal) +} + +fn events_from_pylist(events: &Bound<'_, PyList>) -> PyResult> { + events + .iter() + .enumerate() + .map(|(index, item)| { + let dict = item + .downcast::() + .map_err(|_| PyTypeError::new_err(format!("events[{index}] must be a dict")))?; + event_from_dict(dict, index) + }) + .collect() +} + +fn event_from_dict(dict: &Bound<'_, PyDict>, index: usize) -> PyResult { + let value = optional_item(dict, "value")? + .map(|value| { + let value_dict = value + .downcast::() + .map_err(|_| PyTypeError::new_err("event value must be a dict"))?; + value_from_dict(value_dict) + }) + .transpose()?; + let query_tags = optional_item(dict, "query_tags")? + .map(|value| py_any_to_json(&value)) + .transpose()?; + let event_ts = optional_item(dict, "event_ts")? + .map(|value| parse_optional_datetime(Some(value.extract::()?), "event_ts")) + .transpose()? + .flatten(); + + Ok(DatagenEventDto { + event_id: required_item(dict, "event_id", index)?.extract::()?, + item_id: required_item(dict, "item_id", index)?.extract::()?, + root_item_id: required_item(dict, "root_item_id", index)?.extract::()?, + parent_item_id: optional_item(dict, "parent_item_id")? + .map(|value| value.extract::()) + .transpose()?, + item_seq: required_item(dict, "item_seq", index)?.extract::()?, + checkpoint_id: required_item(dict, "checkpoint_id", index)?.extract::()?, + event_type: required_item(dict, "event_type", index)?.extract::()?, + step_name: optional_item(dict, "step_name")? + .map(|value| value.extract::()) + .transpose()?, + step_kind: optional_item(dict, "step_kind")? + .map(|value| value.extract::()) + .transpose()?, + step_index: optional_item(dict, "step_index")? + .map(|value| value.extract::()) + .transpose()?, + enclosing_step: optional_item(dict, "enclosing_step")? + .map(|value| value.extract::()) + .transpose()?, + selector_step: optional_item(dict, "selector_step")? + .map(|value| value.extract::()) + .transpose()?, + attempt: optional_item(dict, "attempt")? + .map(|value| value.extract::()) + .transpose()? + .unwrap_or(0), + run_id: required_item(dict, "run_id", index)?.extract::()?, + writer_epoch: required_item(dict, "writer_epoch", index)?.extract::()?, + field_name: optional_item(dict, "field_name")? + .map(|value| value.extract::()) + .transpose()?, + field_type: optional_item(dict, "field_type")? + .map(|value| value.extract::()) + .transpose()?, + codec_version: optional_item(dict, "codec_version")? + .map(|value| value.extract::()) + .transpose()?, + value, + query_tags, + status: optional_item(dict, "status")? + .map(|value| value.extract::()) + .transpose()?, + error_type: optional_item(dict, "error_type")? + .map(|value| value.extract::()) + .transpose()?, + error_dump: optional_item(dict, "error_dump")? + .map(|value| value.extract::()) + .transpose()?, + traceback: optional_item(dict, "traceback")? + .map(|value| value.extract::()) + .transpose()?, + event_ts, + schema_version: optional_item(dict, "schema_version")? + .map(|value| value.extract::()) + .transpose()? + .unwrap_or(DATAGEN_SCHEMA_VERSION), + }) +} + +fn value_from_dict(dict: &Bound<'_, PyDict>) -> PyResult { + let kind = dict + .get_item("kind")? + .ok_or_else(|| PyRuntimeError::new_err("event value is missing 'kind'"))? + .extract::()?; + let inner = || { + dict.get_item("value")? + .ok_or_else(|| PyRuntimeError::new_err("event value is missing 'value'")) + }; + let mut dto = DatagenValueDto { + kind: kind.clone(), + value: None, + bytes: None, + size: None, + checksum: None, + }; + match kind.as_str() { + "int" => dto.value = Some(Value::from(inner()?.extract::()?)), + "float" => dto.value = Some(Value::from(inner()?.extract::()?)), + "bool" => dto.value = Some(Value::from(inner()?.extract::()?)), + "str" => dto.value = Some(Value::from(inner()?.extract::()?)), + "json" => dto.value = Some(py_any_to_json(&inner()?)?), + "blob" => { + let bytes = dict + .get_item("bytes")? + .filter(|value| !value.is_none()) + .map(|value| value.extract::>()) + .transpose()?; + let size = match dict.get_item("size")? { + Some(value) if !value.is_none() => Some(value.extract::()?), + _ => bytes.as_ref().map(|b| b.len() as i64), + }; + let checksum = dict + .get_item("checksum")? + .filter(|value| !value.is_none()) + .map(|value| value.extract::()) + .transpose()?; + dto.bytes = bytes; + dto.size = size; + dto.checksum = checksum; + } + other => { + return Err(PyRuntimeError::new_err(format!( + "unsupported datagen value kind '{other}'" + ))) + } + } + Ok(dto) +} + +fn py_any_to_json(value: &Bound<'_, PyAny>) -> PyResult { + let py = value.py(); + let json = PyModule::import(py, "json")?; + let text = json.call_method1("dumps", (value,))?.extract::()?; + serde_json::from_str(&text).map_err(to_py_err) +} + +fn value_to_py(py: Python<'_>, value: &DatagenValueDto) -> PyResult { + let dict = PyDict::new(py); + dict.set_item("kind", &value.kind)?; + if value.kind == "blob" { + dict.set_item("bytes", value.bytes.as_ref().map(|b| PyBytes::new(py, b)))?; + dict.set_item("size", value.size)?; + dict.set_item("checksum", value.checksum.clone())?; + } else if let Some(inner) = &value.value { + dict.set_item("value", json_value_to_py(py, inner)?)?; + } + Ok(dict.into_pyobject(py)?.unbind().into()) +} + +fn folded_item_to_py(py: Python<'_>, item: &FoldedDatagenItemDto) -> PyResult { + let dict = PyDict::new(py); + dict.set_item("item_id", &item.item_id)?; + dict.set_item("root_item_id", &item.root_item_id)?; + dict.set_item("parent_item_id", item.parent_item_id.clone())?; + dict.set_item("status", &item.status)?; + dict.set_item("last_item_seq", item.last_item_seq)?; + dict.set_item("last_attempt", item.last_attempt)?; + + let fields = PyDict::new(py); + for (name, state) in &item.fields { + fields.set_item(name, field_state_to_py(py, state)?)?; + } + dict.set_item("fields", fields)?; + + let trajectory = PyList::empty(py); + for cursor in &item.trajectory { + trajectory.append(cursor_to_py(py, cursor)?)?; + } + dict.set_item("trajectory", trajectory)?; + + dict.set_item( + "query_tags", + match &item.query_tags { + Some(tags) => Some(json_value_to_py(py, tags)?), + None => None, + }, + )?; + Ok(dict.into_pyobject(py)?.unbind().into()) +} + +fn field_state_to_py(py: Python<'_>, state: &DatagenFieldStateDto) -> PyResult { + let dict = PyDict::new(py); + dict.set_item("mode", &state.mode)?; + if let Some(value) = &state.value { + dict.set_item("value", value_to_py(py, value)?)?; + } + if !state.values.is_empty() { + let list = PyList::empty(py); + for value in &state.values { + list.append(value_to_py(py, value)?)?; + } + dict.set_item("values", list)?; + } + Ok(dict.into_pyobject(py)?.unbind().into()) +} + +fn cursor_to_py(py: Python<'_>, cursor: &DatagenStepCursorDto) -> PyResult { + let dict = PyDict::new(py); + dict.set_item("step_name", &cursor.step_name)?; + dict.set_item("step_kind", &cursor.step_kind)?; + dict.set_item("step_index", cursor.step_index)?; + dict.set_item("enclosing_step", cursor.enclosing_step.clone())?; + dict.set_item("selector_step", cursor.selector_step.clone())?; + dict.set_item("item_seq", cursor.item_seq)?; + Ok(dict.into_pyobject(py)?.unbind().into()) +} + +fn failure_to_py(py: Python<'_>, failure: &DatagenFailureDto) -> PyResult { + let dict = PyDict::new(py); + dict.set_item("at", cursor_to_py(py, &failure.at)?)?; + dict.set_item("run_id", &failure.run_id)?; + dict.set_item("attempt", failure.attempt)?; + dict.set_item("error_type", &failure.error_type)?; + dict.set_item("error_dump", failure.error_dump.clone())?; + dict.set_item("traceback", failure.traceback.clone())?; + Ok(dict.into_pyobject(py)?.unbind().into()) +} + #[pymodule] fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(version, m)?)?; m.add_function(wrap_pyfunction!(generate_id, m)?)?; + m.add_function(wrap_pyfunction!(datagen_event_id, m)?)?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } From f0530220f8d59ccb4c2cfb2c3959bb513d865c8c Mon Sep 17 00:00:00 2001 From: Yangjun Zhang Date: Mon, 27 Jul 2026 11:40:21 -0700 Subject: [PATCH 5/5] fix(datagen): sort api imports and wrap datagen_event_id for pyright The lint job runs ruff + pyright without building the wheel, so re-exporting datagen_event_id via a redundant alias left it unresolved (_internal has no stub). Wrap it in a thin function like generate_id, and fix import ordering. Co-Authored-By: Claude Opus 4.8 --- python/python/lance_context/api.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index f5e9f15..90f3b75 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -16,16 +16,16 @@ ContextNamespace as _ContextNamespace, ) from ._internal import ( # pyright: ignore[reportMissingImports] - RemoteContext as _RemoteContext, + DatagenStore as _DatagenStore, ) from ._internal import ( # pyright: ignore[reportMissingImports] - RolloutStore as _RolloutStore, + RemoteContext as _RemoteContext, ) from ._internal import ( # pyright: ignore[reportMissingImports] - DatagenStore as _DatagenStore, + RolloutStore as _RolloutStore, ) from ._internal import ( # pyright: ignore[reportMissingImports] - datagen_event_id as datagen_event_id, + datagen_event_id as _datagen_event_id, ) from ._internal import ( # pyright: ignore[reportMissingImports] generate_id as _generate_id, @@ -59,6 +59,15 @@ def generate_id() -> str: return _generate_id() +def datagen_event_id(item_id: str, checkpoint_id: str, ordinal: int) -> str: + """Compute the deterministic event id for a datagen checkpoint field write. + + The id is a pure function of `item_id`, `checkpoint_id`, and `ordinal`, so a + retried append produces the same id and dedupes against the log. + """ + return _datagen_event_id(item_id, checkpoint_id, ordinal) + + _ARROW_STREAM_MIME = "application/vnd.apache.arrow.stream" _DEFAULT_INGEST_BATCH_SIZE = 1000 _MISSING = object()