Skip to main content

bb_runtime/contracts/
data_source.rs

1//! `bb::DataSource` — Contract trait for data loaders.
2//!
3//! Each method takes the engine's `&mut RuntimeResourceRef<'_>` ctx
4//! plus a [`CompletionHandle`] AND returns [`ContractResponse`]. See
5//! [`crate::contracts::index`] for the sync (Now) vs async (Later)
6//! semantics.
7//!
8//! `ctx` exposes `ctx.dependency::<T>("<slot>")` for the source to
9//! reach any concrete declared in `#[depends(...)]` (e.g. a Backend
10//! to materialize the batch tensor on-device).
11//!
12//! ## Associated type: `Sample`
13//!
14//! `Sample` is the element type for both the batch tensor and the
15//! optional labels tensor returned by [`DataSource::next_batch`].
16//! Implement as `[f32]` for flat f32 sample batches. The second
17//! slot (`Box<Self::Sample>`) holds labels; unsupervised sources
18//! return a zero-length boxed slice.
19
20use crate::completion::{CompletionHandle, ContractResponse};
21use crate::runtime::RuntimeResourceRef;
22
23/// User-facing Contract trait for a data source / data loader.
24pub trait DataSource: Send + Sync {
25    /// Sample storage type. One associated type covers both the batch
26    /// tensor and the optional labels tensor. Implement as `[f32]` for
27    /// flat f32 sample batches.
28    type Sample: ?Sized + bb_ir::types::Storage;
29    /// Library-maker-defined error type.
30    type Error: std::error::Error + std::fmt::Display + Send + Sync + 'static;
31
32    /// Fetch the next batch. Returns `(batch, labels)` as boxed
33    /// `Self::Sample` slices; the second slot is zero-length for
34    /// unsupervised sources.
35    // ?Sized + Storage DST tuple `(Box<Self::Sample>, Box<Self::Sample>)` in CompletionHandle's generic param trips the type-complexity lint; unavoidable for DST-bound associated type.
36    #[allow(clippy::type_complexity)]
37    fn next_batch(
38        &mut self,
39        ctx: &mut RuntimeResourceRef<'_>,
40        completion: CompletionHandle<(Box<Self::Sample>, Box<Self::Sample>), Self::Error>,
41    ) -> ContractResponse<(Box<Self::Sample>, Box<Self::Sample>), Self::Error>;
42
43    /// Reset to the beginning of the data stream.
44    fn reset(
45        &mut self,
46        ctx: &mut RuntimeResourceRef<'_>,
47        completion: CompletionHandle<(), Self::Error>,
48    ) -> ContractResponse<(), Self::Error>;
49
50    /// One-shot notification signalling the source's data has
51    /// finished loading (e.g. dataset download complete).
52    fn on_data_loaded(
53        &mut self,
54        ctx: &mut RuntimeResourceRef<'_>,
55        completion: CompletionHandle<(), Self::Error>,
56    ) -> ContractResponse<(), Self::Error>;
57}