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}