use super::*;
#[derive(Clone, Debug, Default)]
pub struct InMemoryPredictionStore {
blocks: Vec<PredictionBlock>,
}
impl InMemoryPredictionStore {
pub fn new() -> Self {
Self::default()
}
pub fn append(&mut self, block: PredictionBlock) -> Result<()> {
block.validate_content()?;
self.blocks.push(block);
Ok(())
}
pub fn blocks(&self) -> &[PredictionBlock] {
&self.blocks
}
pub fn find(
&self,
producer_node: Option<&NodeId>,
phase_partition: Option<&crate::oof::PredictionPartition>,
fold_id: Option<&FoldId>,
) -> Vec<&PredictionBlock> {
self.blocks
.iter()
.filter(|block| {
producer_node.is_none_or(|node_id| &block.producer_node == node_id)
&& phase_partition.is_none_or(|partition| &block.partition == partition)
&& fold_id.is_none_or(|requested| block.fold_id.as_ref() == Some(requested))
})
.collect()
}
}
#[derive(Clone, Debug, Default)]
pub struct InMemoryAggregatedPredictionStore {
blocks: Vec<AggregatedPredictionBlock>,
}
impl InMemoryAggregatedPredictionStore {
pub fn new() -> Self {
Self::default()
}
pub fn append(&mut self, block: AggregatedPredictionBlock) -> Result<()> {
block.validate_shape()?;
self.blocks.push(block);
Ok(())
}
pub fn blocks(&self) -> &[AggregatedPredictionBlock] {
&self.blocks
}
pub fn find(
&self,
producer_node: Option<&NodeId>,
phase_partition: Option<&PredictionPartition>,
fold_id: Option<&FoldId>,
prediction_level: Option<PredictionLevel>,
) -> Vec<&AggregatedPredictionBlock> {
self.blocks
.iter()
.filter(|block| {
producer_node.is_none_or(|node_id| &block.producer_node == node_id)
&& phase_partition.is_none_or(|partition| &block.partition == partition)
&& fold_id.is_none_or(|requested| block.fold_id.as_ref() == Some(requested))
&& prediction_level.is_none_or(|level| block.level == level)
})
.collect()
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct PredictionCacheMaterializationRequest {
pub run_id: RunId,
pub bundle_id: BundleId,
pub phase: Phase,
pub variant_id: Option<VariantId>,
pub requirement: BundlePredictionRequirement,
pub cache: BundlePredictionCacheRecord,
pub producer_controller_id: ControllerId,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct PredictionCacheMaterializationRecord {
pub run_id: RunId,
pub bundle_id: BundleId,
pub phase: Phase,
pub variant_id: Option<VariantId>,
pub requirement_key: String,
pub cache_id: String,
pub handle: HandleRef,
}
pub trait RuntimePredictionCacheStore {
fn load_blocks(&self, requirement_key: &str) -> Result<Vec<PredictionBlock>>;
fn load_aggregated_blocks(
&self,
requirement_key: &str,
) -> Result<Vec<AggregatedPredictionBlock>> {
Err(DagMlError::RuntimeValidation(format!(
"prediction cache store does not support aggregated requirement `{requirement_key}`"
)))
}
fn materialize(&self, request: &PredictionCacheMaterializationRequest) -> Result<HandleRef>;
}
pub const FILE_PREDICTION_CACHE_STORE_SCHEMA_VERSION: u32 = 1;
pub const FILE_PREDICTION_CACHE_MANIFEST_FILE: &str = "prediction_cache_manifest.json";
pub(crate) fn default_file_prediction_cache_store_schema_version() -> u32 {
FILE_PREDICTION_CACHE_STORE_SCHEMA_VERSION
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FilePredictionCacheEntry {
pub requirement_key: String,
pub cache_id: String,
pub file_name: String,
#[serde(default = "default_runtime_prediction_level")]
pub prediction_level: PredictionLevel,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub unit_ids: Vec<PredictionUnitId>,
pub block_count: usize,
pub row_count: usize,
pub content_fingerprint: String,
}
impl FilePredictionCacheEntry {
pub fn validate(&self) -> Result<()> {
validate_runtime_non_empty("requirement_key", &self.requirement_key)?;
validate_runtime_non_empty("cache_id", &self.cache_id)?;
validate_runtime_non_empty("file_name", &self.file_name)?;
validate_prediction_cache_file_name(&self.file_name)?;
if self.block_count == 0 {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache `{}` has zero block_count",
self.cache_id
)));
}
if self.row_count == 0 {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache `{}` has zero row_count",
self.cache_id
)));
}
if self.prediction_level != PredictionLevel::Sample && self.unit_ids.is_empty() {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache `{}` has no aggregated unit ids",
self.cache_id
)));
}
if self
.unit_ids
.iter()
.any(|unit_id| unit_id.level() != self.prediction_level)
{
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache `{}` has unit ids outside {:?}",
self.cache_id, self.prediction_level
)));
}
validate_runtime_fingerprint("prediction cache content", &self.content_fingerprint)
}
fn from_payload(payload: &crate::bundle::BundlePredictionCachePayload) -> Result<Self> {
Ok(Self {
requirement_key: payload.requirement_key.clone(),
cache_id: payload.cache_id.clone(),
file_name: prediction_cache_payload_file_name(payload)?,
prediction_level: payload.prediction_level,
unit_ids: payload
.aggregated_blocks
.iter()
.flat_map(|block| block.unit_ids.iter().cloned())
.collect(),
block_count: payload.block_count,
row_count: payload.row_count,
content_fingerprint: payload.content_fingerprint.clone(),
})
}
fn matches_record(&self, record: &BundlePredictionCacheRecord) -> bool {
self.requirement_key == record.requirement_key
&& self.cache_id == record.cache_id
&& self.prediction_level == record.prediction_level
&& self.unit_ids == record.unit_ids
&& self.block_count == record.block_count
&& self.row_count == record.row_count
&& self.content_fingerprint == record.content_fingerprint
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FilePredictionCacheManifest {
pub bundle_id: BundleId,
#[serde(default = "default_file_prediction_cache_store_schema_version")]
pub schema_version: u32,
#[serde(default)]
pub caches: Vec<FilePredictionCacheEntry>,
}
impl FilePredictionCacheManifest {
pub fn validate(&self) -> Result<()> {
if self.schema_version != FILE_PREDICTION_CACHE_STORE_SCHEMA_VERSION {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache manifest for bundle `{}` uses unsupported schema_version {}, expected {}",
self.bundle_id,
self.schema_version,
FILE_PREDICTION_CACHE_STORE_SCHEMA_VERSION
)));
}
let mut requirement_keys = BTreeSet::new();
let mut cache_ids = BTreeSet::new();
let mut file_names = BTreeSet::new();
for entry in &self.caches {
entry.validate()?;
if !requirement_keys.insert(entry.requirement_key.as_str()) {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache manifest for bundle `{}` has duplicate requirement `{}`",
self.bundle_id, entry.requirement_key
)));
}
if !cache_ids.insert(entry.cache_id.as_str()) {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache manifest for bundle `{}` has duplicate cache id `{}`",
self.bundle_id, entry.cache_id
)));
}
if !file_names.insert(entry.file_name.as_str()) {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache manifest for bundle `{}` has duplicate file `{}`",
self.bundle_id, entry.file_name
)));
}
}
Ok(())
}
pub fn validate_against_bundle(&self, bundle: &ExecutionBundle) -> Result<()> {
self.validate()?;
bundle.validate()?;
if self.bundle_id != bundle.bundle_id {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache manifest bundle `{}` does not match bundle `{}`",
self.bundle_id, bundle.bundle_id
)));
}
if self.caches.len() != bundle.prediction_caches.len() {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache manifest for bundle `{}` has {} cache(s) for {} bundle cache record(s)",
self.bundle_id,
self.caches.len(),
bundle.prediction_caches.len()
)));
}
let entries_by_requirement = self
.caches
.iter()
.map(|entry| (entry.requirement_key.as_str(), entry))
.collect::<BTreeMap<_, _>>();
for record in &bundle.prediction_caches {
let entry = entries_by_requirement
.get(record.requirement_key.as_str())
.ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"file prediction cache manifest for bundle `{}` is missing requirement `{}`",
self.bundle_id, record.requirement_key
))
})?;
if !entry.matches_record(record) {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache manifest entry `{}` does not match bundle cache record",
entry.cache_id
)));
}
}
Ok(())
}
}
#[derive(Clone, Debug)]
pub struct FilePredictionCacheStore {
root: PathBuf,
manifest: FilePredictionCacheManifest,
records_by_requirement: BTreeMap<String, BundlePredictionCacheRecord>,
materialization_records: RefCell<Vec<PredictionCacheMaterializationRecord>>,
}
impl FilePredictionCacheStore {
pub fn write_payload_set(
root: impl AsRef<Path>,
bundle: &ExecutionBundle,
payloads: &BundlePredictionCachePayloadSet,
) -> Result<FilePredictionCacheManifest> {
payloads.validate_against_bundle(bundle)?;
let root = root.as_ref();
fs::create_dir_all(root).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to create prediction cache store `{}`: {err}",
root.display()
))
})?;
let mut entries = Vec::new();
let records_by_requirement = bundle
.prediction_caches
.iter()
.map(|record| (record.requirement_key.as_str(), record))
.collect::<BTreeMap<_, _>>();
for payload in &payloads.caches {
let record = records_by_requirement
.get(payload.requirement_key.as_str())
.ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"prediction cache payload `{}` references unknown requirement `{}`",
payload.cache_id, payload.requirement_key
))
})?;
validate_prediction_cache_payload_matches_record(payload, record)?;
let entry = FilePredictionCacheEntry::from_payload(payload)?;
let payload_path = root.join(&entry.file_name);
write_runtime_json(&payload_path, payload, "prediction cache payload")?;
entries.push(entry);
}
entries.sort_by(|left, right| left.requirement_key.cmp(&right.requirement_key));
let manifest = FilePredictionCacheManifest {
bundle_id: bundle.bundle_id.clone(),
schema_version: FILE_PREDICTION_CACHE_STORE_SCHEMA_VERSION,
caches: entries,
};
manifest.validate_against_bundle(bundle)?;
write_runtime_json(
&root.join(FILE_PREDICTION_CACHE_MANIFEST_FILE),
&manifest,
"prediction cache manifest",
)?;
Ok(manifest)
}
pub fn open(root: impl Into<PathBuf>, bundle: &ExecutionBundle) -> Result<Self> {
bundle.validate()?;
let root = root.into();
let manifest: FilePredictionCacheManifest = read_runtime_json(
&root.join(FILE_PREDICTION_CACHE_MANIFEST_FILE),
"prediction cache manifest",
)?;
manifest.validate_against_bundle(bundle)?;
let records_by_requirement = bundle
.prediction_caches
.iter()
.cloned()
.map(|record| (record.requirement_key.clone(), record))
.collect::<BTreeMap<_, _>>();
Ok(Self {
root,
manifest,
records_by_requirement,
materialization_records: RefCell::new(Vec::new()),
})
}
pub fn manifest(&self) -> &FilePredictionCacheManifest {
&self.manifest
}
pub fn materialization_records(&self) -> Vec<PredictionCacheMaterializationRecord> {
self.materialization_records.borrow().clone()
}
fn payload_for_requirement(
&self,
requirement_key: &str,
) -> Result<crate::bundle::BundlePredictionCachePayload> {
let entry = self
.manifest
.caches
.iter()
.find(|entry| entry.requirement_key == requirement_key)
.ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"file prediction cache store is missing requirement `{requirement_key}`"
))
})?;
let record = self
.records_by_requirement
.get(requirement_key)
.ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"file prediction cache store has no bundle record for requirement `{requirement_key}`"
))
})?;
let payload: crate::bundle::BundlePredictionCachePayload = read_runtime_json(
&self.root.join(&entry.file_name),
"prediction cache payload",
)?;
validate_prediction_cache_payload_matches_record(&payload, record)?;
Ok(payload)
}
}
impl RuntimePredictionCacheStore for FilePredictionCacheStore {
fn load_blocks(&self, requirement_key: &str) -> Result<Vec<PredictionBlock>> {
let payload = self.payload_for_requirement(requirement_key)?;
if payload.prediction_level != PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache store requirement `{requirement_key}` contains {:?} predictions, not sample blocks",
payload.prediction_level
)));
}
Ok(payload.blocks)
}
fn load_aggregated_blocks(
&self,
requirement_key: &str,
) -> Result<Vec<AggregatedPredictionBlock>> {
let payload = self.payload_for_requirement(requirement_key)?;
if payload.prediction_level == PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache store requirement `{requirement_key}` contains sample predictions, not aggregated blocks"
)));
}
Ok(payload.aggregated_blocks)
}
fn materialize(&self, request: &PredictionCacheMaterializationRequest) -> Result<HandleRef> {
request.requirement.validate()?;
request.cache.validate()?;
let requirement_key = request.requirement.key();
let record = self
.records_by_requirement
.get(&requirement_key)
.ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"file prediction cache store is missing requirement `{requirement_key}`"
))
})?;
if record != &request.cache {
return Err(DagMlError::RuntimeValidation(format!(
"file prediction cache materialization request for `{requirement_key}` does not match bundle cache record"
)));
}
let payload = self.payload_for_requirement(&requirement_key)?;
validate_prediction_cache_payload_matches_record(&payload, record)?;
let fingerprint = stable_json_fingerprint(&(
&request.run_id,
&request.bundle_id,
request.phase,
&request.variant_id,
&request.cache.requirement_key,
&request.cache.cache_id,
request.cache.prediction_level,
&request.cache.content_fingerprint,
))?;
let handle = HandleRef {
handle: u64::from_str_radix(&fingerprint[..16], 16)
.expect("sha256 hex prefix should fit into u64"),
kind: HandleKind::Prediction,
owner_controller: request.producer_controller_id.clone(),
};
self.materialization_records
.borrow_mut()
.push(PredictionCacheMaterializationRecord {
run_id: request.run_id.clone(),
bundle_id: request.bundle_id.clone(),
phase: request.phase,
variant_id: request.variant_id.clone(),
requirement_key,
cache_id: request.cache.cache_id.clone(),
handle: handle.clone(),
});
Ok(handle)
}
}
pub(crate) fn prediction_cache_payload_file_name(
payload: &crate::bundle::BundlePredictionCachePayload,
) -> Result<String> {
let fingerprint = stable_json_fingerprint(&(
&payload.requirement_key,
&payload.cache_id,
payload.prediction_level,
&payload.content_fingerprint,
payload.block_count,
payload.row_count,
))?;
Ok(format!("prediction-cache-{}.json", &fingerprint[..16]))
}
pub(crate) fn validate_prediction_cache_file_name(file_name: &str) -> Result<()> {
if file_name == "." || file_name == ".." || file_name.contains('/') || file_name.contains('\\')
{
return Err(DagMlError::RuntimeValidation(format!(
"prediction cache file name `{file_name}` must be a plain file name"
)));
}
Ok(())
}
#[derive(Clone, Debug, PartialEq)]
pub struct ColumnarPredictionCacheBlock {
pub prediction_id: Option<String>,
pub producer_node: NodeId,
pub partition: PredictionPartition,
pub fold_id: Option<FoldId>,
pub prediction_level: PredictionLevel,
pub unit_ids: Vec<PredictionUnitId>,
pub sample_ids: Vec<SampleId>,
pub target_names: Vec<String>,
pub width: usize,
pub columns: Vec<Vec<f64>>,
}
impl ColumnarPredictionCacheBlock {
pub fn from_prediction_block(block: &PredictionBlock) -> Result<Self> {
let width = block.validate_shape()?;
let mut columns = vec![Vec::with_capacity(block.values.len()); width];
for row in &block.values {
for (column_idx, value) in row.iter().enumerate() {
columns[column_idx].push(*value);
}
}
Ok(Self {
prediction_id: block.prediction_id.clone(),
producer_node: block.producer_node.clone(),
partition: block.partition.clone(),
fold_id: block.fold_id.clone(),
prediction_level: PredictionLevel::Sample,
unit_ids: Vec::new(),
sample_ids: block.sample_ids.clone(),
target_names: block.target_names.clone(),
width,
columns,
})
}
pub fn from_aggregated_prediction_block(block: &AggregatedPredictionBlock) -> Result<Self> {
let width = block.validate_shape()?;
if block.level == PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"columnar aggregated prediction block for `{}` must use target/group level, got sample",
block.producer_node
)));
}
let mut columns = vec![Vec::with_capacity(block.values.len()); width];
for row in &block.values {
for (column_idx, value) in row.iter().enumerate() {
columns[column_idx].push(*value);
}
}
Ok(Self {
prediction_id: block.prediction_id.clone(),
producer_node: block.producer_node.clone(),
partition: block.partition.clone(),
fold_id: block.fold_id.clone(),
prediction_level: block.level,
unit_ids: block.unit_ids.clone(),
sample_ids: Vec::new(),
target_names: block.target_names.clone(),
width,
columns,
})
}
pub fn row_count(&self) -> usize {
match self.prediction_level {
PredictionLevel::Sample => self.sample_ids.len(),
PredictionLevel::Target | PredictionLevel::Group => self.unit_ids.len(),
PredictionLevel::Observation => 0,
}
}
pub fn value_count(&self) -> usize {
self.columns.iter().map(Vec::len).sum()
}
pub fn validate(&self) -> Result<()> {
match self.prediction_level {
PredictionLevel::Observation => {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction block for `{}` cannot store observation-level predictions",
self.producer_node
)));
}
PredictionLevel::Sample => {
if self.sample_ids.is_empty() {
return Err(DagMlError::RuntimeValidation(format!(
"columnar sample prediction block for `{}` has no sample ids",
self.producer_node
)));
}
if !self.unit_ids.is_empty() {
return Err(DagMlError::RuntimeValidation(format!(
"columnar sample prediction block for `{}` unexpectedly carries unit ids",
self.producer_node
)));
}
}
PredictionLevel::Target | PredictionLevel::Group => {
if !self.sample_ids.is_empty() {
return Err(DagMlError::RuntimeValidation(format!(
"columnar aggregated prediction block for `{}` unexpectedly carries sample ids",
self.producer_node
)));
}
if self.unit_ids.is_empty() {
return Err(DagMlError::RuntimeValidation(format!(
"columnar aggregated prediction block for `{}` has no unit ids",
self.producer_node
)));
}
if self
.unit_ids
.iter()
.any(|unit_id| unit_id.level() != self.prediction_level)
{
return Err(DagMlError::RuntimeValidation(format!(
"columnar aggregated prediction block for `{}` carries unit ids outside {:?}",
self.producer_node, self.prediction_level
)));
}
}
}
if self.width == 0 {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction block for `{}` has zero width",
self.producer_node
)));
}
if self.columns.len() != self.width {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction block for `{}` has {} column(s), expected {}",
self.producer_node,
self.columns.len(),
self.width
)));
}
for (column_idx, column) in self.columns.iter().enumerate() {
if column.len() != self.row_count() {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction block for `{}` column {} has {} value(s), expected {}",
self.producer_node,
column_idx,
column.len(),
self.row_count()
)));
}
}
if !self.target_names.is_empty() && self.target_names.len() != self.width {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction block for `{}` has {} target names for width {}",
self.producer_node,
self.target_names.len(),
self.width
)));
}
Ok(())
}
pub fn to_prediction_block(&self) -> Result<PredictionBlock> {
self.validate()?;
if self.prediction_level != PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction block for `{}` contains {:?} predictions, not sample predictions",
self.producer_node, self.prediction_level
)));
}
let values = (0..self.row_count())
.map(|row_idx| {
self.columns
.iter()
.map(|column| column[row_idx])
.collect::<Vec<_>>()
})
.collect();
let block = PredictionBlock {
prediction_id: self.prediction_id.clone(),
producer_node: self.producer_node.clone(),
partition: self.partition.clone(),
fold_id: self.fold_id.clone(),
sample_ids: self.sample_ids.clone(),
values,
target_names: self.target_names.clone(),
};
block.validate_shape()?;
Ok(block)
}
pub fn to_aggregated_prediction_block(&self) -> Result<AggregatedPredictionBlock> {
self.validate()?;
if self.prediction_level == PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction block for `{}` contains sample predictions, not aggregated predictions",
self.producer_node
)));
}
let values = (0..self.row_count())
.map(|row_idx| {
self.columns
.iter()
.map(|column| column[row_idx])
.collect::<Vec<_>>()
})
.collect();
let block = AggregatedPredictionBlock {
prediction_id: self.prediction_id.clone(),
producer_node: self.producer_node.clone(),
partition: self.partition.clone(),
fold_id: self.fold_id.clone(),
level: self.prediction_level,
unit_ids: self.unit_ids.clone(),
values,
target_names: self.target_names.clone(),
};
block.validate_shape()?;
Ok(block)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ColumnarPredictionCacheManifest {
pub requirement_key: String,
pub cache_id: String,
pub prediction_level: PredictionLevel,
pub block_count: usize,
pub row_count: usize,
pub prediction_width: usize,
pub value_count: usize,
pub estimated_value_bytes: usize,
pub content_fingerprint: String,
}
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct ColumnarPredictionCacheEntry {
cache: BundlePredictionCacheRecord,
blocks: Vec<ColumnarPredictionCacheBlock>,
}
impl ColumnarPredictionCacheEntry {
fn from_payload(
payload: BundlePredictionCachePayload,
cache: BundlePredictionCacheRecord,
) -> Result<Self> {
validate_prediction_cache_payload_matches_record(&payload, &cache)?;
let blocks = match payload.prediction_level {
PredictionLevel::Sample => payload
.blocks
.iter()
.map(ColumnarPredictionCacheBlock::from_prediction_block)
.collect::<Result<Vec<_>>>()?,
PredictionLevel::Target | PredictionLevel::Group => payload
.aggregated_blocks
.iter()
.map(ColumnarPredictionCacheBlock::from_aggregated_prediction_block)
.collect::<Result<Vec<_>>>()?,
PredictionLevel::Observation => {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache payload `{}` cannot use observation-level predictions",
payload.cache_id
)));
}
};
let entry = Self { cache, blocks };
entry.validate()?;
Ok(entry)
}
fn validate(&self) -> Result<()> {
self.cache.validate()?;
if self.blocks.len() != self.cache.block_count {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache `{}` has {} block(s), expected {}",
self.cache.cache_id,
self.blocks.len(),
self.cache.block_count
)));
}
let mut row_count = 0usize;
let mut value_count = 0usize;
for block in &self.blocks {
block.validate()?;
if block.prediction_level != self.cache.prediction_level {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache `{}` contains a {:?} block, expected {:?}",
self.cache.cache_id, block.prediction_level, self.cache.prediction_level
)));
}
if block.partition != self.cache.partition {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache `{}` contains a block from partition {:?}",
self.cache.cache_id, block.partition
)));
}
row_count += block.row_count();
value_count += block.value_count();
}
if row_count != self.cache.row_count {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache `{}` has {} row(s), expected {}",
self.cache.cache_id, row_count, self.cache.row_count
)));
}
let expected_values = self
.cache
.row_count
.checked_mul(self.cache.prediction_width)
.ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"columnar prediction cache `{}` value count overflow",
self.cache.cache_id
))
})?;
if value_count != expected_values {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache `{}` has {} value(s), expected {}",
self.cache.cache_id, value_count, expected_values
)));
}
Ok(())
}
fn to_blocks(&self) -> Result<Vec<PredictionBlock>> {
self.validate()?;
self.blocks
.iter()
.map(ColumnarPredictionCacheBlock::to_prediction_block)
.collect()
}
fn to_aggregated_blocks(&self) -> Result<Vec<AggregatedPredictionBlock>> {
self.validate()?;
self.blocks
.iter()
.map(ColumnarPredictionCacheBlock::to_aggregated_prediction_block)
.collect()
}
fn validate_against_cache_record(&self, cache: &BundlePredictionCacheRecord) -> Result<()> {
if &self.cache != cache {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache materialization request for `{}` does not match bundle cache record",
cache.requirement_key
)));
}
let (blocks, aggregated_blocks) = match self.cache.prediction_level {
PredictionLevel::Sample => (self.to_blocks()?, Vec::new()),
PredictionLevel::Target | PredictionLevel::Group => {
(Vec::new(), self.to_aggregated_blocks()?)
}
PredictionLevel::Observation => {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache `{}` cannot materialize observation-level predictions",
self.cache.cache_id
)));
}
};
let payload = BundlePredictionCachePayload {
requirement_key: self.cache.requirement_key.clone(),
cache_id: self.cache.cache_id.clone(),
format: self.cache.format.clone(),
partition: self.cache.partition.clone(),
prediction_level: self.cache.prediction_level,
block_count: self.cache.block_count,
row_count: self.cache.row_count,
content_fingerprint: self.cache.content_fingerprint.clone(),
blocks,
aggregated_blocks,
};
validate_prediction_cache_payload_matches_record(&payload, cache)
}
fn manifest(&self) -> ColumnarPredictionCacheManifest {
let value_count = self
.blocks
.iter()
.map(ColumnarPredictionCacheBlock::value_count)
.sum::<usize>();
ColumnarPredictionCacheManifest {
requirement_key: self.cache.requirement_key.clone(),
cache_id: self.cache.cache_id.clone(),
prediction_level: self.cache.prediction_level,
block_count: self.cache.block_count,
row_count: self.cache.row_count,
prediction_width: self.cache.prediction_width,
value_count,
estimated_value_bytes: value_count * std::mem::size_of::<f64>(),
content_fingerprint: self.cache.content_fingerprint.clone(),
}
}
}
#[derive(Clone, Debug, Default)]
pub struct ColumnarPredictionCacheStore {
entries: BTreeMap<String, ColumnarPredictionCacheEntry>,
materialization_records: RefCell<Vec<PredictionCacheMaterializationRecord>>,
}
impl ColumnarPredictionCacheStore {
pub fn from_payloads(
bundle: &ExecutionBundle,
payloads: BundlePredictionCachePayloadSet,
) -> Result<Self> {
payloads.validate_against_bundle(bundle)?;
let records_by_requirement = bundle
.prediction_caches
.iter()
.cloned()
.map(|cache| (cache.requirement_key.clone(), cache))
.collect::<BTreeMap<_, _>>();
let mut entries = BTreeMap::new();
for payload in payloads.caches {
let cache = records_by_requirement
.get(&payload.requirement_key)
.cloned()
.ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"columnar prediction cache payload `{}` references unknown requirement `{}`",
payload.cache_id, payload.requirement_key
))
})?;
let requirement_key = payload.requirement_key.clone();
let previous = entries.insert(
requirement_key,
ColumnarPredictionCacheEntry::from_payload(payload, cache)?,
);
debug_assert!(previous.is_none());
}
Ok(Self {
entries,
materialization_records: RefCell::new(Vec::new()),
})
}
pub fn entry_count(&self) -> usize {
self.entries.len()
}
pub fn manifests(&self) -> Vec<ColumnarPredictionCacheManifest> {
self.entries
.values()
.map(ColumnarPredictionCacheEntry::manifest)
.collect()
}
pub fn materialization_records(&self) -> Vec<PredictionCacheMaterializationRecord> {
self.materialization_records.borrow().clone()
}
}
impl RuntimePredictionCacheStore for ColumnarPredictionCacheStore {
fn load_blocks(&self, requirement_key: &str) -> Result<Vec<PredictionBlock>> {
let entry = self.entries.get(requirement_key).ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"columnar prediction cache store is missing requirement `{requirement_key}`"
))
})?;
if entry.cache.prediction_level != PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache store requirement `{requirement_key}` contains {:?} predictions, not sample blocks",
entry.cache.prediction_level
)));
}
entry.validate_against_cache_record(&entry.cache)?;
entry.to_blocks()
}
fn load_aggregated_blocks(
&self,
requirement_key: &str,
) -> Result<Vec<AggregatedPredictionBlock>> {
let entry = self.entries.get(requirement_key).ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"columnar prediction cache store is missing requirement `{requirement_key}`"
))
})?;
if entry.cache.prediction_level == PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache store requirement `{requirement_key}` contains sample predictions, not aggregated blocks"
)));
}
entry.validate_against_cache_record(&entry.cache)?;
entry.to_aggregated_blocks()
}
fn materialize(&self, request: &PredictionCacheMaterializationRequest) -> Result<HandleRef> {
request.requirement.validate()?;
request.cache.validate()?;
let requirement_key = request.requirement.key();
if requirement_key != request.cache.requirement_key {
return Err(DagMlError::RuntimeValidation(format!(
"columnar prediction cache materialization request for `{}` uses cache `{}` with mismatched requirement `{}`",
requirement_key, request.cache.cache_id, request.cache.requirement_key
)));
}
let entry = self.entries.get(&requirement_key).ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"columnar prediction cache store is missing requirement `{requirement_key}`"
))
})?;
entry.validate_against_cache_record(&request.cache)?;
let fingerprint = stable_json_fingerprint(&(
&request.run_id,
&request.bundle_id,
request.phase,
&request.variant_id,
&request.cache.requirement_key,
&request.cache.cache_id,
request.cache.prediction_level,
&request.cache.content_fingerprint,
))?;
let handle = HandleRef {
handle: u64::from_str_radix(&fingerprint[..16], 16)
.expect("sha256 hex prefix should fit into u64"),
kind: HandleKind::Prediction,
owner_controller: request.producer_controller_id.clone(),
};
self.materialization_records
.borrow_mut()
.push(PredictionCacheMaterializationRecord {
run_id: request.run_id.clone(),
bundle_id: request.bundle_id.clone(),
phase: request.phase,
variant_id: request.variant_id.clone(),
requirement_key,
cache_id: request.cache.cache_id.clone(),
handle: handle.clone(),
});
Ok(handle)
}
}
pub(crate) fn validate_runtime_non_empty(label: &str, value: &str) -> Result<()> {
if value.trim().is_empty() {
return Err(DagMlError::RuntimeValidation(format!("{label} is empty")));
}
Ok(())
}
pub(crate) fn validate_artifact_optional_text(
label: &str,
value: &Option<String>,
artifact_id: &ArtifactId,
) -> Result<()> {
let Some(value) = value else {
return Ok(());
};
if value.trim().is_empty() {
return Err(DagMlError::RuntimeValidation(format!(
"artifact `{artifact_id}` has empty {label}"
)));
}
if value.chars().any(char::is_control) {
return Err(DagMlError::RuntimeValidation(format!(
"artifact `{artifact_id}` has control characters in {label}"
)));
}
Ok(())
}
pub(crate) fn artifact_payload_path(root: &Path, artifact: &ArtifactRef) -> Result<PathBuf> {
artifact.validate_portable()?;
let uri = artifact
.uri
.as_deref()
.expect("portable artifact validation requires uri");
Ok(root.join(uri))
}
pub(crate) fn validate_artifact_payload_file(
root: &Path,
artifact: &ArtifactRef,
) -> Result<ArtifactPayloadMetadata> {
artifact.validate_portable()?;
let uri = artifact
.uri
.as_deref()
.expect("portable artifact validation requires uri")
.to_string();
let path = artifact_payload_path(root, artifact)?;
validate_payload_path_stays_within_root(root, &path, artifact)?;
let metadata = fs::metadata(&path).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to stat artifact payload `{}` at {}: {err}",
artifact.id,
path.display()
))
})?;
if !metadata.is_file() {
return Err(DagMlError::RuntimeValidation(format!(
"artifact payload `{}` at {} is not a regular file",
artifact.id,
path.display()
)));
}
let size_bytes = metadata.len();
if let Some(expected_size) = artifact.size_bytes {
if expected_size != size_bytes {
return Err(DagMlError::RuntimeValidation(format!(
"artifact payload `{}` size mismatch: expected {}, got {}",
artifact.id, expected_size, size_bytes
)));
}
}
let content_fingerprint =
sha256_file_hex(&path, &format!("artifact payload `{}`", artifact.id))?;
let expected_fingerprint = artifact
.content_fingerprint
.as_deref()
.expect("portable artifact validation requires content_fingerprint");
if !content_fingerprint.eq_ignore_ascii_case(expected_fingerprint) {
return Err(DagMlError::RuntimeValidation(format!(
"artifact payload `{}` content fingerprint mismatch",
artifact.id
)));
}
Ok(ArtifactPayloadMetadata {
uri,
content_fingerprint,
size_bytes,
})
}
pub(crate) fn validate_payload_path_stays_within_root(
root: &Path,
path: &Path,
artifact: &ArtifactRef,
) -> Result<()> {
let root = fs::canonicalize(root).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to canonicalize artifact payload root `{}`: {err}",
root.display()
))
})?;
let path = fs::canonicalize(path).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to canonicalize artifact payload `{}` at {}: {err}",
artifact.id,
path.display()
))
})?;
if !path.starts_with(&root) {
return Err(DagMlError::RuntimeValidation(format!(
"artifact payload `{}` resolves outside store root `{}`",
artifact.id,
root.display()
)));
}
Ok(())
}
pub(crate) fn sha256_file_hex(path: &Path, label: &str) -> Result<String> {
let mut file = fs::File::open(path).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to open {label} at {}: {err}",
path.display()
))
})?;
let mut hasher = Sha256::new();
let mut buffer = [0u8; 64 * 1024];
loop {
let read = file.read(&mut buffer).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to read {label} at {}: {err}",
path.display()
))
})?;
if read == 0 {
break;
}
hasher.update(&buffer[..read]);
}
Ok(bytes_to_hex(&hasher.finalize()))
}
#[cfg(test)]
pub(crate) fn sha256_bytes_hex(bytes: &[u8]) -> String {
bytes_to_hex(&Sha256::digest(bytes))
}
pub(crate) fn bytes_to_hex(bytes: &[u8]) -> String {
let mut out = String::with_capacity(bytes.len() * 2);
for byte in bytes {
use std::fmt::Write as _;
write!(&mut out, "{byte:02x}").expect("writing to String cannot fail");
}
out
}
pub(crate) fn validate_relative_artifact_uri(artifact_id: &ArtifactId, uri: &str) -> Result<()> {
if uri.is_empty() {
return Err(DagMlError::RuntimeValidation(format!(
"artifact `{artifact_id}` has empty uri"
)));
}
if uri.chars().any(char::is_control) {
return Err(DagMlError::RuntimeValidation(format!(
"artifact `{artifact_id}` uri has control characters"
)));
}
if uri.starts_with('/') || uri.starts_with('\\') {
return Err(DagMlError::RuntimeValidation(format!(
"artifact `{artifact_id}` uri `{uri}` must be a relative path"
)));
}
let mut prefix = uri.chars();
if let (Some(drive), Some(':')) = (prefix.next(), prefix.next()) {
if drive.is_ascii_alphabetic() {
return Err(DagMlError::RuntimeValidation(format!(
"artifact `{artifact_id}` uri `{uri}` must be a relative path"
)));
}
}
let first_segment = uri.split(['/', '\\']).next().unwrap_or(uri);
if first_segment.contains(':') {
return Err(DagMlError::RuntimeValidation(format!(
"artifact `{artifact_id}` uri `{uri}` must not include a scheme or colon in its first path segment"
)));
}
for segment in uri.split(['/', '\\']) {
if segment == ".." {
return Err(DagMlError::RuntimeValidation(format!(
"artifact `{artifact_id}` uri `{uri}` must not contain `..` components"
)));
}
}
Ok(())
}
pub(crate) fn validate_runtime_fingerprint(label: &str, value: &str) -> Result<()> {
if value.len() != 64 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
return Err(DagMlError::RuntimeValidation(format!(
"{label} fingerprint must be a 64-character hex digest"
)));
}
Ok(())
}
pub(crate) fn read_runtime_json<T: serde::de::DeserializeOwned>(
path: &Path,
label: &str,
) -> Result<T> {
let data = fs::read(path).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to read {label} at {}: {err}",
path.display()
))
})?;
serde_json::from_slice(&data).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to parse {label} at {}: {err}",
path.display()
))
})
}
pub(crate) fn write_runtime_json<T: Serialize>(path: &Path, value: &T, label: &str) -> Result<()> {
let mut data = serde_json::to_vec_pretty(value).map_err(|err| {
DagMlError::RuntimeValidation(format!("failed to serialize {label}: {err}"))
})?;
data.push(b'\n');
fs::write(path, data).map_err(|err| {
DagMlError::RuntimeValidation(format!(
"failed to write {label} at {}: {err}",
path.display()
))
})
}
#[derive(Clone, Debug, Default)]
pub struct InMemoryPredictionCacheStore {
payloads: BTreeMap<String, crate::bundle::BundlePredictionCachePayload>,
materialization_records: RefCell<Vec<PredictionCacheMaterializationRecord>>,
}
impl InMemoryPredictionCacheStore {
pub fn from_payloads(
bundle: &ExecutionBundle,
payloads: BundlePredictionCachePayloadSet,
) -> Result<Self> {
payloads.validate_against_bundle(bundle)?;
Ok(Self {
payloads: payloads
.caches
.into_iter()
.map(|payload| (payload.requirement_key.clone(), payload))
.collect(),
materialization_records: RefCell::new(Vec::new()),
})
}
pub fn payload_count(&self) -> usize {
self.payloads.len()
}
pub fn materialization_records(&self) -> Vec<PredictionCacheMaterializationRecord> {
self.materialization_records.borrow().clone()
}
}
impl RuntimePredictionCacheStore for InMemoryPredictionCacheStore {
fn load_blocks(&self, requirement_key: &str) -> Result<Vec<PredictionBlock>> {
let payload = self.payloads.get(requirement_key).ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"prediction cache store is missing requirement `{requirement_key}`"
))
})?;
payload.validate()?;
if payload.prediction_level != PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"prediction cache store requirement `{requirement_key}` contains {:?} predictions, not sample blocks",
payload.prediction_level
)));
}
Ok(payload.blocks.clone())
}
fn load_aggregated_blocks(
&self,
requirement_key: &str,
) -> Result<Vec<AggregatedPredictionBlock>> {
let payload = self.payloads.get(requirement_key).ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"prediction cache store is missing requirement `{requirement_key}`"
))
})?;
payload.validate()?;
if payload.prediction_level == PredictionLevel::Sample {
return Err(DagMlError::RuntimeValidation(format!(
"prediction cache store requirement `{requirement_key}` contains sample predictions, not aggregated blocks"
)));
}
Ok(payload.aggregated_blocks.clone())
}
fn materialize(&self, request: &PredictionCacheMaterializationRequest) -> Result<HandleRef> {
request.requirement.validate()?;
request.cache.validate()?;
if request.requirement.key() != request.cache.requirement_key {
return Err(DagMlError::RuntimeValidation(format!(
"prediction cache materialization request for `{}` uses cache `{}` with mismatched requirement `{}`",
request.requirement.key(),
request.cache.cache_id,
request.cache.requirement_key
)));
}
let payload = self
.payloads
.get(&request.cache.requirement_key)
.ok_or_else(|| {
DagMlError::RuntimeValidation(format!(
"prediction cache store is missing requirement `{}`",
request.cache.requirement_key
))
})?;
validate_prediction_cache_payload_matches_record(payload, &request.cache)?;
let fingerprint = stable_json_fingerprint(&(
&request.run_id,
&request.bundle_id,
request.phase,
&request.variant_id,
&request.cache.requirement_key,
&request.cache.cache_id,
request.cache.prediction_level,
&request.cache.content_fingerprint,
))?;
let handle = HandleRef {
handle: u64::from_str_radix(&fingerprint[..16], 16)
.expect("sha256 hex prefix should fit into u64"),
kind: HandleKind::Prediction,
owner_controller: request.producer_controller_id.clone(),
};
self.materialization_records
.borrow_mut()
.push(PredictionCacheMaterializationRecord {
run_id: request.run_id.clone(),
bundle_id: request.bundle_id.clone(),
phase: request.phase,
variant_id: request.variant_id.clone(),
requirement_key: request.cache.requirement_key.clone(),
cache_id: request.cache.cache_id.clone(),
handle: handle.clone(),
});
Ok(handle)
}
}