trs_mlflow/
modelversion.rs

1//! Contains everything related to model versioning.
2use crate::experiment::KeyValue;
3use builder_pattern::Builder;
4use serde::{Deserialize, Serialize};
5use strum::Display;
6
7#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
8#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
9pub enum ModelVersionStatus {
10    PendingRegistration,
11    FailedRegistration,
12    Ready,
13}
14
15#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
16pub struct CreateModelVersion {
17    #[serde(rename = "name")]
18    #[into]
19    pub registered_model_name: String,
20    #[serde(rename = "source")]
21    #[into]
22    pub artifacts_url: String,
23    #[into]
24    pub run_id: String,
25    #[default(vec![])]
26    #[serde(default)]
27    pub tags: Vec<KeyValue>,
28    // #[into]
29    // pub run_link: String,
30    #[into]
31    pub description: String,
32}
33
34#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
35pub struct ModelVersion {
36    pub name: String,
37    pub version: String,
38    pub creation_timestamp: i64,
39    pub last_updated_timestamp: i64,
40    #[serde(default)]
41    pub user_id: String,
42    pub current_stage: String,
43    pub description: String,
44    #[serde(rename = "source")]
45    pub artifacts_url: String,
46    pub run_id: String,
47    #[serde(default)]
48    pub status_message: String,
49    #[serde(default)]
50    pub tags: Vec<KeyValue>,
51    pub run_link: String,
52    #[serde(default)]
53    pub aliases: Vec<String>,
54}
55
56#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Display)]
57pub enum ModelVersionStage {
58    None,
59    Production,
60    Staging,
61    Archived,
62}
63
64#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Builder)]
65pub struct TransitionModelVersionStage {
66    #[into]
67    pub name: String,
68    #[into]
69    pub version: String,
70    pub stage: ModelVersionStage,
71    #[default(false)]
72    pub archive_existing_versions: bool,
73}