Skip to main content

use_ml_experiment/
lib.rs

1#![forbid(unsafe_code)]
2#![doc = include_str!("../README.md")]
3
4use core::{fmt, str::FromStr};
5use std::error::Error;
6
7pub mod prelude {
8    pub use crate::{
9        MlArtifactUri, MlExperimentError, MlExperimentId, MlExperimentName, MlExperimentStage,
10        MlParameterName, MlParameterValue, MlRunId, MlRunStatus, MlRunTag, MlTrackingBackendKind,
11    };
12}
13
14macro_rules! experiment_text_newtype {
15    ($name:ident) => {
16        #[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
17        pub struct $name(String);
18
19        impl $name {
20            pub fn new(value: impl AsRef<str>) -> Result<Self, MlExperimentError> {
21                non_empty_text(value).map(Self)
22            }
23
24            pub fn as_str(&self) -> &str {
25                &self.0
26            }
27        }
28
29        impl AsRef<str> for $name {
30            fn as_ref(&self) -> &str {
31                self.as_str()
32            }
33        }
34
35        impl fmt::Display for $name {
36            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
37                formatter.write_str(self.as_str())
38            }
39        }
40
41        impl FromStr for $name {
42            type Err = MlExperimentError;
43
44            fn from_str(value: &str) -> Result<Self, Self::Err> {
45                Self::new(value)
46            }
47        }
48
49        impl TryFrom<&str> for $name {
50            type Error = MlExperimentError;
51
52            fn try_from(value: &str) -> Result<Self, Self::Error> {
53                Self::new(value)
54            }
55        }
56    };
57}
58
59macro_rules! experiment_enum {
60    ($name:ident { $($variant:ident => $label:literal),+ $(,)? }) => {
61        #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
62        pub enum $name {
63            $($variant),+
64        }
65
66        impl $name {
67            pub const fn as_str(self) -> &'static str {
68                match self {
69                    $(Self::$variant => $label),+
70                }
71            }
72        }
73
74        impl fmt::Display for $name {
75            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
76                formatter.write_str(self.as_str())
77            }
78        }
79
80        impl FromStr for $name {
81            type Err = MlExperimentError;
82
83            fn from_str(value: &str) -> Result<Self, Self::Err> {
84                match normalized_label(value)?.as_str() {
85                    $($label => Ok(Self::$variant),)+
86                    _ => Err(MlExperimentError::UnknownLabel),
87                }
88            }
89        }
90    };
91}
92
93experiment_text_newtype!(MlExperimentName);
94experiment_text_newtype!(MlExperimentId);
95experiment_text_newtype!(MlRunId);
96experiment_text_newtype!(MlRunTag);
97experiment_text_newtype!(MlParameterName);
98experiment_text_newtype!(MlParameterValue);
99experiment_text_newtype!(MlArtifactUri);
100
101experiment_enum!(MlRunStatus {
102    Scheduled => "scheduled",
103    Running => "running",
104    Finished => "finished",
105    Failed => "failed",
106    Killed => "killed",
107    Cancelled => "cancelled",
108    Unknown => "unknown",
109});
110
111experiment_enum!(MlTrackingBackendKind {
112    Local => "local",
113    FileStore => "file-store",
114    Sql => "sql",
115    MlflowLike => "mlflow-like",
116    WeightsAndBiasesLike => "weights-and-biases-like",
117    NeptuneLike => "neptune-like",
118    Custom => "custom",
119});
120
121experiment_enum!(MlExperimentStage {
122    Exploration => "exploration",
123    Baseline => "baseline",
124    Tuning => "tuning",
125    Candidate => "candidate",
126    Validation => "validation",
127    ProductionCandidate => "production-candidate",
128    Archived => "archived",
129});
130
131#[derive(Clone, Copy, Debug, Eq, PartialEq)]
132pub enum MlExperimentError {
133    Empty,
134    UnknownLabel,
135}
136
137impl fmt::Display for MlExperimentError {
138    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
139        match self {
140            Self::Empty => formatter.write_str("ML experiment metadata text cannot be empty"),
141            Self::UnknownLabel => formatter.write_str("unknown ML experiment metadata label"),
142        }
143    }
144}
145
146impl Error for MlExperimentError {}
147
148fn non_empty_text(value: impl AsRef<str>) -> Result<String, MlExperimentError> {
149    let trimmed = value.as_ref().trim();
150    if trimmed.is_empty() {
151        Err(MlExperimentError::Empty)
152    } else {
153        Ok(trimmed.to_string())
154    }
155}
156
157fn normalized_label(value: &str) -> Result<String, MlExperimentError> {
158    let trimmed = value.trim();
159    if trimmed.is_empty() {
160        Err(MlExperimentError::Empty)
161    } else {
162        Ok(trimmed.to_ascii_lowercase().replace(['_', ' '], "-"))
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::{
169        MlExperimentError, MlExperimentName, MlExperimentStage, MlRunId, MlRunStatus,
170        MlTrackingBackendKind,
171    };
172
173    #[test]
174    fn validates_experiment_names_and_run_ids() -> Result<(), MlExperimentError> {
175        let experiment = MlExperimentName::new(" baseline ")?;
176        let run = MlRunId::new("run-001")?;
177
178        assert_eq!(experiment.as_str(), "baseline");
179        assert_eq!(run.as_str(), "run-001");
180        assert_eq!(MlExperimentName::new("  "), Err(MlExperimentError::Empty));
181        Ok(())
182    }
183
184    #[test]
185    fn displays_and_parses_experiment_enums() -> Result<(), MlExperimentError> {
186        assert_eq!("finished".parse::<MlRunStatus>()?, MlRunStatus::Finished);
187        assert_eq!(
188            "weights and biases like".parse::<MlTrackingBackendKind>()?,
189            MlTrackingBackendKind::WeightsAndBiasesLike
190        );
191        assert_eq!(
192            "production candidate".parse::<MlExperimentStage>()?,
193            MlExperimentStage::ProductionCandidate
194        );
195        Ok(())
196    }
197}