ml_cellar/

ml_bin.rs

use std::env;
use std::path::PathBuf;

/// ML-bin is a specific release of a model stored under a rack (e.g., `vit-l/1.1/`).
/// ML-bin has artifacts like checkpoints, configs, logs.
/// It corresponds to a “vintage” in the wine analogy.
/// ML-bin is expressed as `{cellar_directory}/{rack_directory}/{project}/{version}/`.
pub struct MLBin {
    /// The absolute path to the root of the ML-cellar storage.
    pub cellar_directory: PathBuf,
    /// The name (relative directory path) of the rack.
    pub rack_name: PathBuf,
    /// The name (relative directory path) of the project.
    pub project_name: PathBuf,
    /// The version of the project.
    pub version: String,
}

impl MLBin {
    /// Creates a new MLBin instance.
    ///
    /// # Arguments
    ///
    /// - `cellar_directory` - The absolute path to the root of the ML-cellar storage
    /// - `rack_name` - The name (relative directory path) of the rack
    /// - `project_name` - The name (relative directory path) of the project
    /// - `version` - The version string of the model
    pub fn new(
        cellar_directory: PathBuf,
        rack_name: impl Into<PathBuf>,
        project_name: impl Into<PathBuf>,
        version: String,
    ) -> Self {
        Self {
            cellar_directory,
            rack_name: rack_name.into(),
            project_name: project_name.into(),
            version,
        }
    }

    /// Creates a new MLBin instance from an ML-bin path.
    ///
    /// This method parses the ML-bin path and extracts the rack name, project name,
    /// and version by analyzing the directory structure relative to the cellar and rack directories.
    ///
    /// # Arguments
    ///
    /// - `ml_bin_path` - The path to the ML-bin (can be relative or absolute)
    /// - `cellar_directory` - The absolute path to the root directory of the ML-cellar
    /// - `rack_directory` - The absolute path to the directory of the rack
    ///
    pub fn from_ml_bin_path(
        ml_bin_path: PathBuf,
        cellar_directory: PathBuf,
        rack_directory: PathBuf,
    ) -> Self {
        // Convert to absolute path
        // ml_bin_absolute_path: {cellar_directory(absolute)}/{rack_name}/{project_name}/{version}/
        let ml_bin_absolute_path = if ml_bin_path.is_absolute() {
            ml_bin_path
        } else {
            env::current_dir().unwrap().join(&ml_bin_path)
        };

        // rack_name: {rack_name}
        let rack_name = rack_directory
            .strip_prefix(&cellar_directory)
            .unwrap_or(&rack_directory)
            .to_path_buf();

        // ml_bin_path_from_rack: {project_name}/{version}
        let ml_bin_path_from_rack = ml_bin_absolute_path
            .strip_prefix(&rack_directory)
            .unwrap_or(&ml_bin_absolute_path)
            .to_path_buf();

        let project_name = match ml_bin_path_from_rack.parent() {
            Some(p) => p.to_path_buf(),
            None => PathBuf::new(),
        };

        let version = match ml_bin_path_from_rack.file_name().and_then(|n| n.to_str()) {
            Some(s) => s.to_string(),
            None => "".to_string(),
        };

        Self {
            cellar_directory,
            rack_name,
            project_name,
            version,
        }
    }

    /// Returns the full absolute path to the ML-bin directory.
    /// Constructs the path by joining: `{cellar_directory}/{rack_name}/{project_name}/{version}`
    pub fn get_full_path(&self) -> PathBuf {
        self.cellar_directory
            .join(&self.rack_name)
            .join(&self.project_name)
            .join(&self.version)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    #[test]
    fn test_new() {
        let ml_bin = MLBin::new(
            PathBuf::from("/path/to/cellar"),
            "vit-l",
            "base",
            "1.1".to_string(),
        );

        assert_eq!(ml_bin.cellar_directory, PathBuf::from("/path/to/cellar"));
        assert_eq!(ml_bin.rack_name, PathBuf::from("vit-l"));
        assert_eq!(ml_bin.project_name, PathBuf::from("base"));
        assert_eq!(ml_bin.version, "1.1");
    }

    #[test]
    fn test_new_with_pathbuf() {
        let ml_bin = MLBin::new(
            PathBuf::from("/path/to/cellar"),
            PathBuf::from("llm-model"),
            PathBuf::from("projectA"),
            "2.0.3".to_string(),
        );

        assert_eq!(ml_bin.cellar_directory, PathBuf::from("/path/to/cellar"));
        assert_eq!(ml_bin.rack_name, PathBuf::from("llm-model"));
        assert_eq!(ml_bin.project_name, PathBuf::from("projectA"));
        assert_eq!(ml_bin.version, "2.0.3");
    }

    #[test]
    fn test_get_full_path() {
        let ml_bin = MLBin::new(
            PathBuf::from("/path/to/cellar"),
            "vit-l",
            "base",
            "1.1".to_string(),
        );

        let full_path = ml_bin.get_full_path();
        assert_eq!(full_path, PathBuf::from("/path/to/cellar/vit-l/base/1.1"));
    }

    #[test]
    fn test_get_full_path_with_nested_project() {
        let ml_bin = MLBin::new(
            PathBuf::from("/cellar"),
            "model",
            "project/subproject",
            "0.1.0".to_string(),
        );

        let full_path = ml_bin.get_full_path();
        assert_eq!(
            full_path,
            PathBuf::from("/cellar/model/project/subproject/0.1.0")
        );
    }

    #[test]
    fn test_from_ml_bin_path_absolute() {
        let cellar_dir = PathBuf::from("/path/to/cellar");
        let rack_dir = PathBuf::from("/path/to/cellar/vit-l");
        let ml_bin_path = PathBuf::from("/path/to/cellar/vit-l/base/1.1");

        let ml_bin = MLBin::from_ml_bin_path(ml_bin_path, cellar_dir.clone(), rack_dir);

        assert_eq!(ml_bin.cellar_directory, cellar_dir);
        assert_eq!(ml_bin.rack_name, PathBuf::from("vit-l"));
        assert_eq!(ml_bin.project_name, PathBuf::from("base"));
        assert_eq!(ml_bin.version, "1.1");
    }

    #[test]
    fn test_from_ml_bin_path_with_nested_project() {
        let cellar_dir = PathBuf::from("/cellar");
        let rack_dir = PathBuf::from("/cellar/model");
        let ml_bin_path = PathBuf::from("/cellar/model/project/subproject/2.0.0");

        let ml_bin = MLBin::from_ml_bin_path(ml_bin_path, cellar_dir.clone(), rack_dir);

        assert_eq!(ml_bin.cellar_directory, cellar_dir);
        assert_eq!(ml_bin.rack_name, PathBuf::from("model"));
        assert_eq!(ml_bin.project_name, PathBuf::from("project/subproject"));
        assert_eq!(ml_bin.version, "2.0.0");
    }

    /// Test that creating an MLBin and getting its path results in the same path
    #[test]
    fn test_from_ml_bin_path_roundtrip() {
        let cellar_dir = PathBuf::from("/test/cellar");
        let rack_dir = PathBuf::from("/test/cellar/my-rack");
        let original_path = PathBuf::from("/test/cellar/my-rack/my-project/1.2.3");

        let ml_bin =
            MLBin::from_ml_bin_path(original_path.clone(), cellar_dir.clone(), rack_dir.clone());
        let reconstructed_path = ml_bin.get_full_path();

        assert_eq!(original_path, reconstructed_path);
    }

    /// Test edge case where ml_bin_path_from_rack has no parent
    #[test]
    fn test_from_ml_bin_path_version_only() {
        let cellar_dir = PathBuf::from("/cellar");
        let rack_dir = PathBuf::from("/cellar/rack");
        let ml_bin_path = PathBuf::from("/cellar/rack/1.0");

        let ml_bin = MLBin::from_ml_bin_path(ml_bin_path, cellar_dir.clone(), rack_dir);

        assert_eq!(ml_bin.cellar_directory, cellar_dir);
        assert_eq!(ml_bin.rack_name, PathBuf::from("rack"));
        assert_eq!(ml_bin.project_name, PathBuf::from(""));
        assert_eq!(ml_bin.version, "1.0");
    }
}