mcai-benchmark 0.3.0

Library to benchmark MCAI workers.
Documentation
use crate::error::Result;
use itertools::iproduct;
use serde::{Deserialize, Serialize};
use serde_valid::Validate;
use serde_yaml;
use std::{collections::HashMap, fs};

#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(tag = "version")]
pub enum Configuration {
  #[serde(rename = "1")]
  Version1(Version1),
}

impl Configuration {
  pub fn get_worker_docker_image(&self) -> String {
    let Configuration::Version1(config) = self;
    config.worker_docker_image.clone()
  }

  pub fn get_output_folder(&self) -> Option<String> {
    let Configuration::Version1(config) = self;
    config.output_folder.clone()
  }

  pub fn get_envs_for_job(&self, job_name: &str) -> HashMap<String, String> {
    let Configuration::Version1(config) = self;
    config
      .benchmarks
      .get(job_name)
      .map(|config| config.envs.clone())
      .unwrap_or_default()
  }

  pub fn get_volumes_for_job(&self, job_name: &str) -> Vec<VolumeConfig> {
    let Configuration::Version1(config) = self;
    config
      .benchmarks
      .get(job_name)
      .map(|config| config.volumes.clone())
      .unwrap_or_default()
  }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Version1 {
  pub version: usize,
  pub worker_docker_image: String,
  pub output_folder: Option<String>,
  #[serde(flatten)]
  pub benchmarks: HashMap<String, BenchmarkConfig>,
}

impl Version1 {
  pub fn read_from_file(filepath: &str) -> Result<Version1> {
    let content = fs::read_to_string(filepath)?;
    let config = serde_yaml::from_str(&content)?;
    Ok(config)
  }
}

fn default_iterations() -> i64 {
  10
}

#[derive(Clone, Debug, Deserialize, Serialize, Validate)]
pub struct HardwareConfig {
  #[validate(min_items = 1)]
  pub cpu: Vec<f32>,
  #[validate(min_items = 1)]
  pub memory: Vec<i64>,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct VolumeConfig {
  pub host: String,
  pub container: String,
  #[serde(default)]
  pub readonly: bool,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct BenchmarkConfig {
  pub source_order: String,
  #[serde(default = "default_iterations")]
  pub iterations: i64,
  pub hardware: HardwareConfig,
  #[serde(default)]
  pub envs: HashMap<String, String>,
  #[serde(default)]
  pub volumes: Vec<VolumeConfig>,
}

impl BenchmarkConfig {
  pub fn get_hardware_configurations(&self) -> Result<Vec<(i64, f32)>> {
    let configurations: Vec<(i64, f32)> =
      iproduct!(self.hardware.memory.clone(), self.hardware.cpu.clone())
        .map(|(memory, cpu)| (memory, cpu))
        .collect();
    Ok(configurations)
  }
}