trs-mlflow 0.7.0

//! Contains everything that is related to mlflow run.
use crate::experiment::KeyValue;
use builder_pattern::Builder;
use chrono::Utc;
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
pub struct CreateRun {
    #[into]
    pub experiment_id: String,
    #[into]
    pub run_name: String,
    #[default(Utc::now().timestamp_millis())]
    pub start_time: i64,
    #[default(vec![])]
    #[serde(default)]
    pub tags: Vec<KeyValue>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
pub struct UpdateRun {
    #[into]
    pub run_id: String,
    #[into]
    pub experiment_id: String,
    pub status: RunStatus,
    #[default(Utc::now().timestamp_millis())]
    pub end_time: i64,
    #[default(None)]
    pub run_name: Option<String>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Run {
    pub info: RunInfo,
    pub data: RunData,
    #[serde(default)]
    pub inputs: RunInputs,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RunInfo {
    pub run_id: String,
    pub run_name: String,
    pub experiment_id: String,
    pub status: RunStatus,
    pub start_time: u64,
    pub end_time: Option<u64>,
    pub artifact_uri: String,
    pub lifecycle_stage: String,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
pub struct RunData {
    #[serde(default)]
    #[default(vec![])]
    pub metrics: Vec<Metric>,
    #[serde(default)]
    #[default(vec![])]
    pub params: Vec<KeyValue>,
    #[serde(default)]
    #[default(vec![])]
    pub tags: Vec<KeyValue>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum RunStatus {
    Running,
    Scheduled,
    Finished,
    Failed,
    Killed,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
pub struct Metric {
    #[into]
    pub key: String,
    #[into]
    pub value: f64,
    #[default(Utc::now().timestamp_millis())]
    pub timestamp: i64,
    pub step: u64,
}

#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct RunInputs {
    #[serde(rename = "dataset_inputs", default)]
    pub inputs: Vec<DataSetInput>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
pub struct DataSetInput {
    #[default(vec![])]
    pub tags: Vec<KeyValue>,
    pub dataset: DataSet,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
pub struct DataSet {
    #[into]
    pub name: String,
    #[into]
    pub digest: String,
    #[into]
    pub source_type: String,
    #[into]
    pub source: String,
    #[into]
    pub schema: String,
    #[into]
    pub profile: String,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
pub struct SearchRuns {
    /// List of experiment IDs to search over.
    #[into]
    pub experiment_ids: Vec<String>,
    /// A filter expression over params, metrics, and tags, that allows
    /// returning a subset of runs. The syntax is a subset of SQL that supports
    /// ANDing together binary operations between a param, metric, or tag and a
    /// constant.
    ///
    /// Example: metrics.rmse < 1 and params.model_class = 'LogisticRegression'
    #[into]
    #[default("".to_owned())]
    pub filter: String,
    /// Whether to display only active, only deleted, or all runs. Defaults to
    /// only active runs.
    #[default(ViewType::default())]
    pub view: ViewType,
    /// Maximum number of runs desired. If unspecified, defaults to 1000. All
    /// servers are guaranteed to support a max_results threshold of at least
    /// 50,000 but may support more. Callers of this endpoint are encouraged to
    /// pass max_results explicitly and leverage page_token to iterate through
    /// experiments.
    #[default(None)]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_results: Option<u16>,
    #[into]
    #[default(vec![])]
    /// List of columns to be ordered by, including attributes, params, metrics,
    /// and tags with an optional “DESC” or “ASC” annotation, where “ASC” is the
    /// default. Example: [“params.input DESC”, “metrics.alpha ASC”,
    /// “metrics.rmse”] Tiebreaks are done by start_time DESC followed by run_id
    /// for runs with the same start time (and this is the default ordering
    /// criterion if order_by is not provided).
    pub order_by: Vec<String>,
    #[default(None)]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub page_token: Option<String>,
}

#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum ViewType {
    #[default]
    ActiveOnly,
    DeletedOnly,
    All,
}