mlflow_client/
data.rs

1use std::time::SystemTime;
2
3use derive_ex::Ex;
4use ordered_float::OrderedFloat;
5use serde::{Deserialize, Serialize};
6use thiserror::Error;
7
8/// <https://mlflow.org/docs/latest/rest-api.html#fileinfo>
9#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
10pub struct FileInfo {
11    pub path: String,
12    pub is_dir: bool,
13    pub file_size: Option<i64>,
14}
15
16/// <https://mlflow.org/docs/latest/rest-api.html#run>
17#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
18pub struct Run {
19    pub info: RunInfo,
20    pub data: RunData,
21    pub inputs: RunInputs,
22}
23
24/// <https://mlflow.org/docs/latest/rest-api.html#runinfo>
25#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
26pub struct RunInfo {
27    pub run_id: String,
28    pub run_name: String,
29    pub experiment_id: String,
30    pub status: RunStatus,
31    pub start_time: Timestamp,
32    pub end_time: Option<Timestamp>,
33    pub artifact_uri: String,
34    pub lifecycle_stage: String,
35}
36
37/// <https://mlflow.org/docs/latest/rest-api.html#rundata>
38#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
39pub struct RunData {
40    #[serde(default)]
41    pub metrics: Vec<Metric>,
42    #[serde(default)]
43    pub params: Vec<Param>,
44    #[serde(default)]
45    pub tags: Vec<RunTag>,
46}
47
48/// <https://mlflow.org/docs/latest/rest-api.html#metric>
49#[derive(Serialize, Deserialize, Debug, Clone, Ex)]
50#[derive_ex(Eq, PartialEq, Ord, PartialOrd)]
51pub struct Metric {
52    pub key: String,
53    #[ord(key = OrderedFloat($))]
54    pub value: f64,
55    pub timestamp: Timestamp,
56    pub step: Option<i64>,
57}
58
59/// <https://mlflow.org/docs/latest/rest-api.html#param>
60#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
61pub struct Param {
62    pub key: String,
63    pub value: String,
64}
65
66/// <https://mlflow.org/docs/latest/rest-api.html#runinputs>
67#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
68pub struct RunInputs {
69    #[serde(default)]
70    pub dataset_inputs: Vec<DatasetInput>,
71}
72
73/// <https://mlflow.org/docs/latest/rest-api.html#datasetinput>
74#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
75pub struct DatasetInput {
76    pub tags: Vec<InputTag>,
77    pub dataset: Dataset,
78}
79
80/// <https://mlflow.org/docs/latest/rest-api.html#inputtag>
81#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
82pub struct InputTag {
83    pub key: String,
84    pub value: String,
85}
86
87/// <https://mlflow.org/docs/latest/rest-api.html#dataset>
88#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
89pub struct Dataset {
90    pub name: String,
91    pub digest: String,
92    pub source_type: String,
93    pub source: String,
94    pub schema: Option<String>,
95    pub profile: Option<String>,
96}
97
98/// <https://mlflow.org/docs/latest/rest-api.html#runstatus>
99#[derive(Debug, Serialize, Deserialize, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
100#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
101pub enum RunStatus {
102    Running,
103    Scheduled,
104    Finished,
105    Failed,
106    Killed,
107}
108
109/// <https://mlflow.org/docs/latest/rest-api.html#experiment>
110#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Ord, PartialOrd)]
111pub struct Experiment {
112    pub experiment_id: String,
113    pub name: String,
114    pub artifact_location: String,
115    pub lifecycle_stage: String,
116    pub last_update_time: Timestamp,
117    pub creation_time: Timestamp,
118    #[serde(default)]
119    pub tags: Vec<ExperimentTag>,
120}
121
122/// <https://mlflow.org/docs/latest/rest-api.html#viewtype>
123#[derive(Debug, Serialize, Deserialize, Clone, Copy, Default, Eq, PartialEq, Ord, PartialOrd)]
124#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
125pub enum ViewType {
126    #[default]
127    ActiveOnly,
128    DeletedOnly,
129    All,
130}
131
132/// <https://mlflow.org/docs/latest/rest-api.html#experimenttag>
133#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Ord, PartialOrd)]
134pub struct ExperimentTag {
135    pub key: String,
136    pub value: String,
137}
138
139/// <https://mlflow.org/docs/latest/rest-api.html#runtag>
140#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Ord, PartialOrd)]
141pub struct RunTag {
142    pub key: String,
143    pub value: String,
144}
145
146/// Unix timestamp in milliseconds.
147#[derive(Debug, Serialize, Deserialize, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
148#[serde(transparent)]
149pub struct Timestamp(pub i64);
150
151impl Timestamp {
152    pub fn now() -> Self {
153        SystemTime::now().try_into().unwrap()
154    }
155}
156
157impl From<i64> for Timestamp {
158    fn from(i: i64) -> Self {
159        Timestamp(i)
160    }
161}
162impl From<Timestamp> for i64 {
163    fn from(ts: Timestamp) -> i64 {
164        ts.0
165    }
166}
167
168impl TryFrom<SystemTime> for Timestamp {
169    type Error = TimestampError;
170    fn try_from(value: SystemTime) -> Result<Self, Self::Error> {
171        Ok(Timestamp(if value < SystemTime::UNIX_EPOCH {
172            let millis: i128 = SystemTime::UNIX_EPOCH
173                .duration_since(value)?
174                .as_millis()
175                .try_into()?;
176            (-millis).try_into()?
177        } else {
178            value
179                .duration_since(SystemTime::UNIX_EPOCH)?
180                .as_millis()
181                .try_into()?
182        }))
183    }
184}
185impl From<Timestamp> for SystemTime {
186    fn from(value: Timestamp) -> SystemTime {
187        if value.0 < 0 {
188            SystemTime::UNIX_EPOCH - std::time::Duration::from_millis(-(value.0 as i128) as u64)
189        } else {
190            SystemTime::UNIX_EPOCH + std::time::Duration::from_millis(value.0 as u64)
191        }
192    }
193}
194
195#[derive(Debug, Error)]
196pub enum TimestampError {
197    #[error("SystemTimeError: {0}")]
198    SystemTimeError(#[from] std::time::SystemTimeError),
199    #[error("TryFromIntError: {0}")]
200    TryFromIntError(#[from] std::num::TryFromIntError),
201}
202
203/// <https://mlflow.org/docs/latest/rest-api.html#request-structure>
204#[derive(Serialize, Debug, Clone, Copy, Default)]
205pub struct CreateExperimentOptions<'a> {
206    pub artifact_location: Option<&'a str>,
207    pub tags: &'a [ExperimentTag],
208}
209
210/// <https://mlflow.org/docs/latest/rest-api.html#mlflowsearchexperiments>
211#[derive(Serialize, Debug, Clone, Copy, Ex)]
212#[derive_ex(Default)]
213pub struct SearchExperimentsOptions<'a> {
214    pub filter: &'a str,
215    pub order_by: &'a [&'a str],
216    pub view_type: ViewType,
217}
218
219/// <https://mlflow.org/docs/latest/rest-api.html#mlflowcreaterun>
220#[derive(Serialize, Debug, Clone, Copy, Default)]
221pub struct CreateRunOptions<'a> {
222    pub start_time: Option<Timestamp>,
223    pub tags: &'a [RunTag],
224}
225
226/// <https://mlflow.org/docs/latest/rest-api.html#mlflowsearchruns>
227#[derive(Serialize, Debug, Clone, Copy, Ex)]
228#[derive_ex(Default)]
229pub struct SearchRunsOptions<'a> {
230    pub filter: &'a str,
231    pub run_view_type: ViewType,
232    pub order_by: &'a [&'a str],
233}
234
235/// <https://mlflow.org/docs/latest/rest-api.html#mlflowupdaterun>
236#[derive(Serialize, Debug, Clone, Copy, Default)]
237pub struct UpdateRunOptions<'a> {
238    pub status: Option<RunStatus>,
239    pub end_time: Option<Timestamp>,
240    pub run_name: Option<&'a str>,
241}
mlflow_client/data.rs

mlflow_client/
data.rs
