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}