trs-mlflow 0.7.0

//! Contains everything that is related model artifacts.
use crate::run::Run;
use builder_pattern::Builder;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

#[derive(Serialize, Deserialize, Debug, Clone, Builder)]
pub struct ListArtifacts {
    #[into]
    pub run_id: String,
    #[default(None)]
    pub path: Option<String>,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ListedArtifacts {
    pub root_uri: String,
    #[serde(default)]
    pub files: Vec<FileInfo>,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct FileInfo {
    pub path: String,
    pub is_dir: bool,
    #[serde(default)]
    pub file_size: u64,
}

#[derive(Serialize, Deserialize, Debug, Clone, Builder)]
pub struct Artifact {
    #[into]
    pub experiment_id: String,
    #[into]
    pub run_id: String,
    /// Path which is relative to mlflow root uri.
    #[into]
    pub path: PathBuf,
}

#[derive(Serialize, Deserialize, Debug, Clone, Builder)]
pub struct UploadArtifact {
    /// Represents a source of artifact on a disk.
    pub local: PathBuf,
    /// Determines where we want to store artifact.
    pub remote: Artifact,
}

/// Can be used to download multiple artifacts from a run.
#[derive(Serialize, Deserialize, Debug, Clone, Builder)]
pub struct DownloadRunArtifacts {
    #[into]
    pub experiment_id: String,
    #[into]
    pub run_id: String,
    /// All the artifacts of this run_id
    #[into]
    pub files: Vec<ArtifactFile>,
    /// Destination directory.
    #[into]
    pub destination: PathBuf,
}

#[derive(Serialize, Deserialize, Debug, Clone, Builder)]
pub struct ArtifactFile {
    /// Path which is relative to mlflow root uri.
    #[into]
    pub path: String,
    /// File's size in remote storage
    #[into]
    pub size: u64,
}

/// Can be used to download SINGLE artifact from a run.
#[derive(Serialize, Deserialize, Debug, Clone, Builder)]
pub struct DownloadRunArtifact {
    #[into]
    pub experiment_id: String,
    #[into]
    pub run_id: String,
    /// Path which is relative to mlflow root uri.
    #[into]
    pub file: String,
    /// Destination for [`Self::file].
    ///
    /// It could be a directory where we should store a file or a fullpath.
    #[into]
    pub destination: PathBuf,
    /// Expected file size (information obtained from remote storage)
    #[into]
    pub expected_size: u64,
}

/// Represents downloaded run artifacts.
#[derive(Serialize, Deserialize, Debug, Clone, Builder)]
pub struct RunArtifacts {
    #[into]
    pub experiment_id: String,
    #[into]
    pub run_id: String,
    #[into]
    pub paths: Vec<PathBuf>,
    /// Root for all [`Self::paths`].
    #[into]
    pub root: PathBuf,
}

impl DownloadRunArtifact {
    /// Returns a local path for [`Self::file`].
    pub fn path(&self) -> PathBuf {
        if !self.destination.ends_with(&self.file) {
            self.destination.join(&self.file)
        } else {
            self.destination.clone()
        }
    }
}

impl RunArtifacts {
    /// Returns a path to first file which name matches provided name.
    pub fn get_file(&self, name: impl AsRef<str>) -> Option<PathBuf> {
        for path in &self.paths {
            if path.ends_with(name.as_ref()) {
                return Some(path.clone());
            }
        }
        None
    }
}

impl DownloadRunArtifacts {
    /// Creates multiple single-download commands from [`Self`].
    pub fn as_single_downloads(self) -> Vec<DownloadRunArtifact> {
        let mut downloads = Vec::with_capacity(self.files.len());

        for file in self.files {
            let download = DownloadRunArtifact {
                experiment_id: self.experiment_id.clone(),
                run_id: self.run_id.clone(),
                file: file.path,
                destination: self.destination.clone(),
                expected_size: file.size,
            };

            downloads.push(download);
        }

        downloads
    }

    /// Creates download command for all artifacts in specified run.
    pub fn new_from_run(destination: impl AsRef<Path>, run: Run, list: ListedArtifacts) -> Self {
        Self {
            destination: destination.as_ref().to_owned(),
            run_id: run.info.run_id,
            experiment_id: run.info.experiment_id,
            files: list
                .files
                .into_iter()
                .map(|file| ArtifactFile {
                    path: file.path,
                    size: file.file_size,
                })
                .collect(),
        }
    }
}