mod artifact;
mod attachment;
pub mod config;
mod dynamic;
mod lifecycle;
mod logging;
mod parameter;
mod run;
mod sealed;
#[cfg(test)]
mod tests;
pub use attachment::{
detect_file_media_type, AttachmentNotFound, AttachmentTable, Compression,
DEFAULT_FILE_MEDIA_TYPE,
};
pub use dynamic::{ExperimentDyn, RunDyn, SamplingDyn, SealedRunDyn, SolveDyn};
pub use lifecycle::{ExperimentLifecycle, RunLifecycle};
pub use logging::AttachmentLogger;
pub use parameter::{ParameterValue, RunParameterCell};
pub use run::{FailedSampleRecord, FailedSolveRecord, FinishedSampleRecord, FinishedSolveRecord};
pub(crate) use sealed::experiment_manifest_record_from_artifact;
pub use sealed::{Sampling, SealedRun, Solve};
use crate::artifact::local_registry::{LocalRegistry, StoredDescriptor, TempLocalRegistry};
use crate::artifact::{media_types, ImageRef, LocalArtifact};
use anyhow::{ensure, Context, Result};
use oci_spec::image::Descriptor;
use parameter::ParameterSet;
use rmpv::Value as MessagePackValue;
use std::sync::{Mutex, MutexGuard};
use std::{
collections::{BTreeMap, HashMap},
io::Cursor,
time::{Duration, Instant},
};
const EXPERIMENT_STATUS_FINISHED: &str = "finished";
const EXPERIMENT_STATUS_DRAFT: &str = "draft";
const EXPERIMENT_STATUS_FAILED: &str = "failed";
const EXPERIMENT_STATUS_INTERRUPTED: &str = "interrupted";
const RUN_PARAMETERS_MEDIA_TYPE: &str = "application/org.ommx.v1.experiment.run-parameters+msgpack";
const EXPERIMENT_ARTIFACT_MEDIA_TYPE: &str = media_types::V1_EXPERIMENT_MEDIA_TYPE;
pub(crate) const EXPERIMENT_CONFIG_MEDIA_TYPE: &str =
"application/org.ommx.v1.experiment.config+json";
const RUN_STATUS_FINISHED: &str = "finished";
const RUN_STATUS_FAILED: &str = "failed";
const RUN_STATUS_INTERRUPTED: &str = "interrupted";
const SOLVE_STATUS_FINISHED: &str = "finished";
const SOLVE_STATUS_FAILED: &str = "failed";
const SOLVE_STATUS_INTERRUPTED: &str = "interrupted";
const SAMPLING_STATUS_FINISHED: &str = "finished";
const SAMPLING_STATUS_FAILED: &str = "failed";
const SAMPLING_STATUS_INTERRUPTED: &str = "interrupted";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExperimentStatus {
Finished,
Draft,
Failed,
Interrupted,
}
impl ExperimentStatus {
pub fn as_str(&self) -> &'static str {
match self {
Self::Finished => EXPERIMENT_STATUS_FINISHED,
Self::Draft => EXPERIMENT_STATUS_DRAFT,
Self::Failed => EXPERIMENT_STATUS_FAILED,
Self::Interrupted => EXPERIMENT_STATUS_INTERRUPTED,
}
}
pub(crate) fn from_config(status: &str) -> Result<Self> {
match status {
EXPERIMENT_STATUS_FINISHED => Ok(Self::Finished),
EXPERIMENT_STATUS_DRAFT => Ok(Self::Draft),
EXPERIMENT_STATUS_FAILED => Ok(Self::Failed),
EXPERIMENT_STATUS_INTERRUPTED => Ok(Self::Interrupted),
_ => {
crate::bail!(
"Experiment status is {status}, expected {EXPERIMENT_STATUS_FINISHED}, \
{EXPERIMENT_STATUS_DRAFT}, {EXPERIMENT_STATUS_FAILED}, or \
{EXPERIMENT_STATUS_INTERRUPTED}"
)
}
}
}
}
impl std::fmt::Display for ExperimentStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum AutosavePolicy {
#[default]
EveryRunClose,
EveryNRuns(u32),
MinInterval(Duration),
Disabled,
}
impl AutosavePolicy {
fn validate(self) -> Result<()> {
ensure!(
!matches!(self, Self::EveryNRuns(0)),
"AutosavePolicy::EveryNRuns requires a non-zero Run count"
);
Ok(())
}
}
#[derive(Debug, Clone)]
struct AutosaveController {
policy: AutosavePolicy,
last_autosaved_run_count: usize,
last_attempt_at: Option<Instant>,
}
impl AutosaveController {
fn new(current_run_count: usize) -> Self {
Self {
policy: AutosavePolicy::default(),
last_autosaved_run_count: current_run_count,
last_attempt_at: None,
}
}
fn set_policy(&mut self, policy: AutosavePolicy, current_run_count: usize) -> Result<()> {
policy.validate()?;
self.policy = policy;
self.last_autosaved_run_count = current_run_count;
self.last_attempt_at = None;
Ok(())
}
fn begin_autosave_attempt(&mut self, now: Instant, current_run_count: usize) -> bool {
let due = match self.policy {
AutosavePolicy::EveryRunClose => true,
AutosavePolicy::EveryNRuns(run_count) => {
current_run_count.saturating_sub(self.last_autosaved_run_count)
>= run_count as usize
}
AutosavePolicy::MinInterval(interval) => self
.last_attempt_at
.is_none_or(|last| now.saturating_duration_since(last) >= interval),
AutosavePolicy::Disabled => false,
};
if due && matches!(self.policy, AutosavePolicy::MinInterval(_)) {
self.last_attempt_at = Some(now);
}
due
}
fn record_forced_attempt(&mut self, now: Instant) {
if matches!(self.policy, AutosavePolicy::MinInterval(_)) {
self.last_attempt_at = Some(now);
}
}
fn mark_autosaved(&mut self, current_run_count: usize) {
self.last_autosaved_run_count = current_run_count;
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RunStatus {
Finished,
Failed,
Interrupted,
}
impl RunStatus {
pub fn as_str(&self) -> &'static str {
match self {
Self::Finished => RUN_STATUS_FINISHED,
Self::Failed => RUN_STATUS_FAILED,
Self::Interrupted => RUN_STATUS_INTERRUPTED,
}
}
}
impl std::fmt::Display for RunStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SolveStatus {
Finished,
Failed,
Interrupted,
}
impl SolveStatus {
pub fn as_str(&self) -> &'static str {
match self {
Self::Finished => SOLVE_STATUS_FINISHED,
Self::Failed => SOLVE_STATUS_FAILED,
Self::Interrupted => SOLVE_STATUS_INTERRUPTED,
}
}
fn from_config(status: &str) -> Result<Self> {
match status {
SOLVE_STATUS_FINISHED => Ok(Self::Finished),
SOLVE_STATUS_FAILED => Ok(Self::Failed),
SOLVE_STATUS_INTERRUPTED => Ok(Self::Interrupted),
_ => {
crate::bail!(
"Solve status is {status}, expected {SOLVE_STATUS_FINISHED}, \
{SOLVE_STATUS_FAILED}, or {SOLVE_STATUS_INTERRUPTED}"
)
}
}
}
}
impl std::fmt::Display for SolveStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SamplingStatus {
Finished,
Failed,
Interrupted,
}
impl SamplingStatus {
pub fn as_str(&self) -> &'static str {
match self {
Self::Finished => SAMPLING_STATUS_FINISHED,
Self::Failed => SAMPLING_STATUS_FAILED,
Self::Interrupted => SAMPLING_STATUS_INTERRUPTED,
}
}
fn from_config(status: &str) -> Result<Self> {
match status {
SAMPLING_STATUS_FINISHED => Ok(Self::Finished),
SAMPLING_STATUS_FAILED => Ok(Self::Failed),
SAMPLING_STATUS_INTERRUPTED => Ok(Self::Interrupted),
_ => {
crate::bail!(
"Sampling status is {status}, expected {SAMPLING_STATUS_FINISHED}, \
{SAMPLING_STATUS_FAILED}, or {SAMPLING_STATUS_INTERRUPTED}"
)
}
}
}
}
impl std::fmt::Display for SamplingStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug)]
pub struct Experiment<'reg> {
registry: &'reg LocalRegistry,
state: Mutex<UnsealedExperimentState<'reg>>,
unresolved_run_behavior: UnresolvedRunBehavior,
}
#[derive(Debug, Clone)]
pub struct SealedExperiment<'reg> {
lifecycle: ExperimentLifecycle,
artifact: LocalArtifact<'reg>,
attachments: AttachmentTable<StoredDescriptor<'reg>>,
runs: BTreeMap<u64, sealed::SealedRun<'reg>>,
run_parameters: parameter::RunParameterTable,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Trace {
bytes: Vec<u8>,
}
impl Trace {
pub fn from_bytes(bytes: impl Into<Vec<u8>>) -> Self {
Self {
bytes: bytes.into(),
}
}
pub fn as_bytes(&self) -> &[u8] {
&self.bytes
}
pub fn into_bytes(self) -> Vec<u8> {
self.bytes
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Name {
Named(ImageRef),
Anonymous,
}
impl Name {
fn resolve(self, registry: &LocalRegistry) -> Result<ImageRef> {
match self {
Self::Named(image_name) => Ok(image_name),
Self::Anonymous => registry.synthesize_anonymous_experiment_image_name(),
}
}
}
impl From<ImageRef> for Name {
fn from(image_name: ImageRef) -> Self {
Self::Named(image_name)
}
}
#[derive(Debug)]
pub struct Run<'exp, 'reg> {
experiment: &'exp Experiment<'reg>,
run_id: u64,
attachments: AttachmentTable<StoredDescriptor<'reg>>,
trace: Option<StoredDescriptor<'reg>>,
solves: Vec<SolveEntry<'reg>>,
next_solve_id: u64,
samplings: Vec<SamplingEntry<'reg>>,
next_sampling_id: u64,
parameters: ParameterSet,
interrupt_on_drop: bool,
closed: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum UnresolvedRunBehavior {
Abandon,
Interrupt,
}
struct ScopedExperimentGuard<'reg> {
experiment: Option<Experiment<'reg>>,
}
impl<'reg> ScopedExperimentGuard<'reg> {
fn new(experiment: Experiment<'reg>) -> Self {
Self {
experiment: Some(experiment),
}
}
fn experiment(&self) -> &Experiment<'reg> {
self.experiment
.as_ref()
.expect("scoped Experiment is present while its callback runs")
}
fn take(&mut self) -> Experiment<'reg> {
self.experiment
.take()
.expect("scoped Experiment has exactly one terminal transition")
}
}
impl Drop for ScopedExperimentGuard<'_> {
fn drop(&mut self) {
let Some(experiment) = self.experiment.take() else {
return;
};
if let Err(error) = experiment.commit_checkpoint(ExperimentLifecycle::Interrupted {
reason: Some("Experiment scope unwound".to_string()),
}) {
tracing::warn!(
error = %error,
"Failed to publish interrupted Experiment checkpoint during scoped unwind"
);
}
}
}
#[derive(Debug)]
struct RunEntry<'reg> {
run_id: u64,
lifecycle: RunLifecycle,
attachments: AttachmentTable<StoredDescriptor<'reg>>,
trace: Option<StoredDescriptor<'reg>>,
solves: Vec<SolveEntry<'reg>>,
samplings: Vec<SamplingEntry<'reg>>,
parameters: ParameterSet,
}
#[derive(Debug, Clone)]
struct SolveEntry<'reg> {
solve_id: u64,
status: SolveStatus,
input: StoredDescriptor<'reg>,
output: Option<StoredDescriptor<'reg>>,
adapter: String,
adapter_options: String,
diagnostics: Option<StoredDescriptor<'reg>>,
}
#[derive(Debug, Clone)]
struct SamplingEntry<'reg> {
sampling_id: u64,
status: SamplingStatus,
input: StoredDescriptor<'reg>,
output: Option<StoredDescriptor<'reg>>,
adapter: String,
adapter_options: String,
diagnostics: Option<StoredDescriptor<'reg>>,
}
#[derive(Debug, Clone)]
pub struct AdapterDiagnosticPayload {
value: MessagePackValue,
}
impl AdapterDiagnosticPayload {
pub fn new(bytes: Vec<u8>) -> Result<Self> {
let mut cursor = Cursor::new(&bytes);
let value = rmpv::decode::read_value(&mut cursor)
.context("Adapter diagnostic payload must be valid MessagePack")?;
ensure!(
cursor.position() == bytes.len() as u64,
"Adapter diagnostic payload must contain exactly one MessagePack value",
);
Self::from_value(value)
}
pub fn from_value(value: MessagePackValue) -> Result<Self> {
ensure!(
matches!(value, MessagePackValue::Array(_)),
"Adapter diagnostic payload must decode to a MessagePack array",
);
Ok(Self { value })
}
pub fn value(&self) -> &MessagePackValue {
&self.value
}
pub(crate) fn to_msgpack_bytes(&self) -> Result<Vec<u8>> {
let mut bytes = Vec::new();
rmpv::encode::write_value(&mut bytes, &self.value)
.context("Failed to encode Adapter diagnostic payload as MessagePack")?;
Ok(bytes)
}
}
fn read_adapter_diagnostic_payload(
record_kind: &str,
record_id: u64,
descriptor: &StoredDescriptor<'_>,
) -> Result<(Vec<u8>, AdapterDiagnosticPayload)> {
descriptor.ensure_media_type(&media_types::diagnostic_msgpack())?;
let bytes = descriptor.registry().get_blob(descriptor)?;
let payload = AdapterDiagnosticPayload::new(bytes.clone())
.with_context(|| format!("Invalid {record_kind} {record_id} diagnostic payload"))?;
Ok((bytes, payload))
}
#[derive(Debug)]
struct UnsealedExperimentState<'reg> {
image_name: ImageRef,
subject: Option<oci_spec::image::Descriptor>,
annotations: HashMap<String, String>,
attachments: AttachmentTable<StoredDescriptor<'reg>>,
runs: BTreeMap<u64, RunEntry<'reg>>,
next_run_id: u64,
autosave: AutosaveController,
}
impl<'reg> UnsealedExperimentState<'reg> {
fn autosave_after_run_close(
&mut self,
registry: &'reg LocalRegistry,
) -> Result<Option<LocalArtifact<'reg>>> {
let run_count = self.runs.len();
if !self
.autosave
.begin_autosave_attempt(Instant::now(), run_count)
{
return Ok(None);
}
let artifact = self.autosave_checkpoint(registry)?;
self.autosave.mark_autosaved(run_count);
Ok(Some(artifact))
}
}
impl Experiment<'static> {
pub fn new(name: impl Into<Name>) -> Result<Self> {
let registry = LocalRegistry::shared_default()?;
Self::with_registry(registry, name)
}
pub fn scoped(
name: impl Into<Name>,
f: impl FnOnce(&Experiment<'static>) -> Result<()>,
) -> Result<SealedExperiment<'static>> {
let registry = LocalRegistry::shared_default()?;
Self::scoped_with_registry(registry, name, f)
}
}
impl<'reg> Experiment<'reg> {
pub fn scoped_with_registry(
registry: &'reg LocalRegistry,
name: impl Into<Name>,
f: impl FnOnce(&Experiment<'reg>) -> Result<()>,
) -> Result<SealedExperiment<'reg>> {
let experiment =
Self::with_registry_and_run_behavior(registry, name, UnresolvedRunBehavior::Interrupt)?;
let mut guard = ScopedExperimentGuard::new(experiment);
match f(guard.experiment()) {
Ok(()) => guard.take().commit(),
Err(error) => {
if let Err(checkpoint_error) =
guard.take().commit_checkpoint(ExperimentLifecycle::Failed {
reason: Some(error.to_string()),
})
{
tracing::warn!(
error = %checkpoint_error,
"Failed to publish failed Experiment checkpoint after callback error"
);
}
Err(error)
}
}
}
pub fn with_temp_local_registry<T>(
name: impl Into<Name>,
f: impl FnOnce(Experiment<'_>) -> anyhow::Result<T>,
) -> Result<T> {
let temp = TempLocalRegistry::new()?;
let experiment = Experiment::with_registry(temp.registry(), name)?;
f(experiment)
}
pub fn with_registry(registry: &'reg LocalRegistry, name: impl Into<Name>) -> Result<Self> {
Self::with_registry_and_run_behavior(registry, name, UnresolvedRunBehavior::Abandon)
}
fn with_registry_and_run_behavior(
registry: &'reg LocalRegistry,
name: impl Into<Name>,
unresolved_run_behavior: UnresolvedRunBehavior,
) -> Result<Self> {
let image_name = name.into().resolve(registry)?;
Ok(Experiment {
registry,
state: Mutex::new(UnsealedExperimentState {
image_name,
subject: None,
annotations: HashMap::new(),
attachments: AttachmentTable::new(),
runs: BTreeMap::new(),
next_run_id: 0,
autosave: AutosaveController::new(0),
}),
unresolved_run_behavior,
})
}
pub fn image_name(&self) -> ImageRef {
self.lock_state().image_name.clone()
}
pub fn set_annotation(&self, key: impl Into<String>, value: impl Into<String>) -> Result<()> {
let key = key.into();
ensure!(
!crate::is_reserved_annotation_key(&key),
"Annotation key `{key}` is reserved for OMMX metadata"
);
self.lock_state().annotations.insert(key, value.into());
Ok(())
}
pub fn set_autosave_policy(&self, policy: AutosavePolicy) -> Result<()> {
let mut state = self.lock_state();
let run_count = state.runs.len();
state.autosave.set_policy(policy, run_count)
}
pub fn run(&self) -> Result<Run<'_, 'reg>> {
let mut state = self.lock_state();
let run_id = allocate_next_run_id(&mut state.next_run_id)?;
Ok(Run {
experiment: self,
run_id,
attachments: AttachmentTable::new(),
trace: None,
solves: Vec::new(),
next_solve_id: 0,
samplings: Vec::new(),
next_sampling_id: 0,
parameters: ParameterSet::new(),
interrupt_on_drop: self.unresolved_run_behavior == UnresolvedRunBehavior::Interrupt,
closed: false,
})
}
pub fn scoped_run<T>(&self, f: impl FnOnce(&mut Run<'_, 'reg>) -> Result<T>) -> Result<T> {
let mut run = self.run()?.interrupt_on_drop();
match f(&mut run) {
Ok(value) => {
run.finish()?;
Ok(value)
}
Err(error) => {
if let Err(finish_error) = run.finish_failed_with_reason(error.to_string()) {
tracing::warn!(
error = %finish_error,
"Failed to finish failed Run after callback error"
);
}
Err(error)
}
}
}
fn push_closed_run(&self, run: RunEntry<'reg>) -> Result<()> {
let mut state = self.lock_state();
if state.runs.contains_key(&run.run_id) {
crate::bail!("Run {} has already been registered", run.run_id);
}
state.runs.insert(run.run_id, run);
if let Err(error) = state.autosave_after_run_close(self.registry) {
tracing::warn!(
error = %error,
"Failed to publish Experiment autosave checkpoint after Run close"
);
}
Ok(())
}
fn lock_state(&self) -> MutexGuard<'_, UnsealedExperimentState<'reg>> {
match self.state.lock() {
Ok(state) => state,
Err(poisoned) => {
tracing::warn!("Experiment state mutex was poisoned; continuing with inner state");
poisoned.into_inner()
}
}
}
pub fn commit(self) -> Result<SealedExperiment<'reg>> {
let (registry, state) = self.into_parts();
let artifact = state.commit(registry)?;
SealedExperiment::from_artifact(artifact)
}
fn commit_checkpoint(self, lifecycle: ExperimentLifecycle) -> Result<LocalArtifact<'reg>> {
let (registry, state) = self.into_parts();
state.commit_checkpoint(registry, lifecycle)
}
fn into_parts(self) -> (&'reg LocalRegistry, UnsealedExperimentState<'reg>) {
let Experiment {
registry,
state,
unresolved_run_behavior: _,
} = self;
let state = match state.into_inner() {
Ok(state) => state,
Err(poisoned) => {
tracing::warn!("Experiment state mutex was poisoned; consuming inner state");
poisoned.into_inner()
}
};
(registry, state)
}
}
impl<'reg> logging::AttachmentLoggerStorage for &Experiment<'reg> {
type Descriptor = StoredDescriptor<'reg>;
fn with_local_registry<R>(&self, f: impl FnOnce(&LocalRegistry) -> Result<R>) -> Result<R> {
f(self.registry)
}
fn with_attachment_table<R>(
&mut self,
f: impl FnOnce(&mut AttachmentTable<Self::Descriptor>) -> Result<R>,
) -> Result<R> {
let mut state = self.lock_state();
f(&mut state.attachments)
}
fn descriptor_for_attachment_table(&self, descriptor: Descriptor) -> Result<Self::Descriptor> {
self.registry.stored_descriptor(descriptor)
}
}
impl<'reg> SealedExperiment<'reg> {
pub fn artifact(&self) -> LocalArtifact<'reg> {
self.artifact.clone()
}
pub fn into_artifact(self) -> LocalArtifact<'reg> {
self.artifact
}
pub fn fork(&self, name: impl Into<Name>) -> Result<Experiment<'reg>> {
let registry = self.artifact.registry();
let image_name = name.into().resolve(registry)?;
let subject = Some(self.artifact.stored_manifest_descriptor()?.into());
let mut runs = BTreeMap::new();
let mut parameters_by_run = self.run_parameters.parameter_sets()?;
for run in self.runs.values() {
let parameters = parameters_by_run
.remove(&run.run_id())
.unwrap_or_else(ParameterSet::new);
let solves = run
.solves()
.iter()
.map(|solve| SolveEntry {
solve_id: solve.solve_id(),
status: solve.status().clone(),
input: solve.input_descriptor().clone(),
output: solve.output_descriptor().cloned(),
adapter: solve.adapter().to_string(),
adapter_options: solve.adapter_options().to_string(),
diagnostics: solve.diagnostic_descriptor().cloned(),
})
.collect();
let samplings = run
.samplings()
.iter()
.map(|sampling| SamplingEntry {
sampling_id: sampling.sampling_id(),
status: sampling.status().clone(),
input: sampling.input_descriptor().clone(),
output: sampling.output_descriptor().cloned(),
adapter: sampling.adapter().to_string(),
adapter_options: sampling.adapter_options().to_string(),
diagnostics: sampling.diagnostic_descriptor().cloned(),
})
.collect();
runs.insert(
run.run_id(),
RunEntry {
run_id: run.run_id(),
lifecycle: run.lifecycle().clone(),
attachments: run.attachment_table().clone(),
trace: run.trace_descriptor().cloned(),
solves,
samplings,
parameters,
},
);
}
Ok(Experiment {
registry,
state: Mutex::new(UnsealedExperimentState {
image_name,
subject,
annotations: HashMap::new(),
attachments: self.attachments.clone(),
next_run_id: next_run_id(runs.keys().copied())?,
autosave: AutosaveController::new(runs.len()),
runs,
}),
unresolved_run_behavior: UnresolvedRunBehavior::Abandon,
})
}
}
fn next_run_id(run_ids: impl Iterator<Item = u64>) -> Result<u64> {
match run_ids.max() {
Some(max) => max
.checked_add(1)
.ok_or_else(|| anyhow::anyhow!("Run ID space is exhausted")),
None => Ok(0),
}
}
fn allocate_next_run_id(next_run_id: &mut u64) -> Result<u64> {
let run_id = *next_run_id;
*next_run_id = next_run_id
.checked_add(1)
.ok_or_else(|| anyhow::anyhow!("Run ID space is exhausted"))?;
Ok(run_id)
}