use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use serde::Deserialize;
use thiserror::Error;
use workload_spec::ImageRef;
use velveteen::TaskRuntime;
pub const ENV_TRANSFORM_IN_0: &str = "YAH_TRANSFORM_IN_0";
pub const ENV_TRANSFORM_OUT: &str = "YAH_TRANSFORM_OUT";
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct TransformRecipe {
pub name: String,
pub label: String,
pub placement: RecipePlacement,
pub image: ImageRef,
#[serde(default)]
pub steps: Vec<RecipeStep>,
}
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct RecipePlacement {
pub location: RecipeLocation,
pub runtime: TaskRuntime,
#[serde(default)]
pub platform: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RecipeLocation {
Local,
}
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct RecipeStep {
pub name: String,
pub argv: Vec<String>,
#[serde(default)]
pub timeout: u64,
}
#[derive(Error, Debug)]
pub enum RecipeError {
#[error("recipe {name:?} not found at {path}")]
NotFound { name: String, path: PathBuf },
#[error("reading {path}: {source}")]
Io {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("parsing recipe {path}: {source}")]
Parse {
path: PathBuf,
#[source]
source: toml::de::Error,
},
#[error(
"recipe {name:?} at {path} uses a bare-tag image; recipe images must \
be digest-pinned (e.g. `image = \"...:v1@sha256:<hex>\"`) for \
reproducibility (W164)"
)]
ImageNotPinned { name: String, path: PathBuf },
}
pub struct TransformRecipeLoader {
transforms_dir: PathBuf,
}
impl TransformRecipeLoader {
pub fn new(transforms_dir: impl AsRef<Path>) -> Self {
Self {
transforms_dir: transforms_dir.as_ref().to_path_buf(),
}
}
pub fn recipe_path(&self, name: &str) -> PathBuf {
self.transforms_dir.join(format!("{name}.toml"))
}
pub fn list_all(&self) -> Result<Vec<String>, RecipeError> {
let entries = match fs::read_dir(&self.transforms_dir) {
Ok(e) => e,
Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(source) => {
return Err(RecipeError::Io {
path: self.transforms_dir.clone(),
source,
})
}
};
let mut names: Vec<String> = entries
.filter_map(|e| e.ok())
.map(|e| e.path())
.filter(|p| p.extension().is_some_and(|ext| ext == "toml"))
.filter_map(|p| p.file_stem().map(|s| s.to_string_lossy().into_owned()))
.collect();
names.sort();
Ok(names)
}
pub fn load(&self, name: &str) -> Result<TransformRecipe, RecipeError> {
let path = self.recipe_path(name);
if !path.exists() {
return Err(RecipeError::NotFound { name: name.to_string(), path });
}
self.load_from_path(&path)
}
pub fn load_from_path(&self, path: &Path) -> Result<TransformRecipe, RecipeError> {
let content = fs::read_to_string(path).map_err(|source| RecipeError::Io {
path: path.to_path_buf(),
source,
})?;
let recipe: TransformRecipe = toml::from_str(&content).map_err(|source| RecipeError::Parse {
path: path.to_path_buf(),
source,
})?;
if recipe.image.digest.is_empty() {
return Err(RecipeError::ImageNotPinned {
name: recipe.name.clone(),
path: path.to_path_buf(),
});
}
Ok(recipe)
}
}
pub fn substitute_argv(template: &[String], params: &BTreeMap<String, String>) -> Vec<String> {
template.iter().map(|elem| substitute_one(elem, params)).collect()
}
fn substitute_one(elem: &str, params: &BTreeMap<String, String>) -> String {
let mut out = String::with_capacity(elem.len());
let mut rest = elem;
while let Some(start) = rest.find("{{") {
out.push_str(&rest[..start]);
let after = &rest[start + 2..];
let Some(end) = after.find("}}") else {
out.push_str("{{");
out.push_str(after);
return out;
};
let key = after[..end].trim();
if let Some(val) = params.get(key) {
out.push_str(val);
} else {
out.push_str("{{");
out.push_str(&after[..end]);
out.push_str("}}");
}
rest = &after[end + 2..];
}
out.push_str(rest);
out
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
const HASH_64: &str = "abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890";
fn whisper_quantize_toml() -> String {
format!(
r#"
name = "whisper-quantize"
label = "Quantize a whisper GGML model"
image = "ghcr.io/ggerganov/whisper.cpp:v1.7.4@sha256:{HASH_64}"
[placement]
location = "local"
runtime = "container"
[[steps]]
name = "quantize"
argv = ["./quantize", "{{{{YAH_TRANSFORM_IN_0}}}}", "{{{{YAH_TRANSFORM_OUT}}}}", "{{{{quant}}}}"]
timeout = 600
"#
)
}
#[test]
fn loads_sample_recipe_round_trip() {
let dir = tempdir().unwrap();
let transforms = dir.path().join("transforms");
fs::create_dir_all(&transforms).unwrap();
fs::write(
transforms.join("whisper-quantize.toml"),
whisper_quantize_toml(),
)
.unwrap();
let loader = TransformRecipeLoader::new(&transforms);
let recipe = loader.load("whisper-quantize").expect("load recipe");
assert_eq!(recipe.name, "whisper-quantize");
assert_eq!(recipe.placement.location, RecipeLocation::Local);
assert_eq!(recipe.placement.runtime, TaskRuntime::Container);
assert_eq!(recipe.image.registry, "ghcr.io");
assert_eq!(recipe.image.repository, "ggerganov/whisper.cpp");
assert_eq!(recipe.image.tag, "v1.7.4");
assert_eq!(recipe.image.digest, format!("sha256:{HASH_64}"));
assert_eq!(recipe.steps.len(), 1);
assert_eq!(recipe.steps[0].name, "quantize");
assert_eq!(recipe.steps[0].timeout, 600);
assert_eq!(
recipe.steps[0].argv,
vec![
"./quantize",
"{{YAH_TRANSFORM_IN_0}}",
"{{YAH_TRANSFORM_OUT}}",
"{{quant}}",
]
);
}
#[test]
fn list_all_returns_sorted_stems_and_ignores_non_toml() {
let dir = tempdir().unwrap();
let transforms = dir.path().join("transforms");
fs::create_dir_all(&transforms).unwrap();
fs::write(transforms.join("zeta.toml"), "").unwrap();
fs::write(transforms.join("alpha.toml"), "").unwrap();
fs::write(transforms.join("README.md"), "ignore me").unwrap();
let loader = TransformRecipeLoader::new(&transforms);
let names = loader.list_all().expect("list_all ok");
assert_eq!(names, vec!["alpha".to_string(), "zeta".to_string()]);
}
#[test]
fn list_all_missing_dir_is_empty_not_error() {
let dir = tempdir().unwrap();
let loader = TransformRecipeLoader::new(dir.path().join("does-not-exist"));
assert_eq!(loader.list_all().expect("ok"), Vec::<String>::new());
}
#[test]
fn rejects_recipe_with_bare_tag_string_image() {
let dir = tempdir().unwrap();
let transforms = dir.path().join("transforms");
fs::create_dir_all(&transforms).unwrap();
let bad = r#"
name = "bare-tag"
label = "no pin"
image = "node:20"
[placement]
location = "local"
runtime = "container"
[[steps]]
name = "noop"
argv = ["true"]
"#;
fs::write(transforms.join("bare-tag.toml"), bad).unwrap();
let loader = TransformRecipeLoader::new(&transforms);
let err = loader.load("bare-tag").expect_err("bare-tag must reject");
assert!(matches!(err, RecipeError::Parse { .. }), "got {err:?}");
}
#[test]
fn rejects_recipe_with_struct_image_missing_digest() {
let dir = tempdir().unwrap();
let transforms = dir.path().join("transforms");
fs::create_dir_all(&transforms).unwrap();
let bad = r#"
name = "struct-bare"
label = "struct-form bare tag"
[placement]
location = "local"
runtime = "container"
[image]
registry = "ghcr.io"
repository = "foo/bar"
tag = "v1"
[[steps]]
name = "noop"
argv = ["true"]
"#;
fs::write(transforms.join("struct-bare.toml"), bad).unwrap();
let loader = TransformRecipeLoader::new(&transforms);
let err = loader
.load("struct-bare")
.expect_err("struct-form bare tag must reject");
assert!(matches!(err, RecipeError::Parse { .. }), "got {err:?}");
}
#[test]
fn missing_recipe_reports_path() {
let dir = tempdir().unwrap();
let loader = TransformRecipeLoader::new(dir.path().join("transforms"));
let err = loader.load("nope").expect_err("missing recipe");
match err {
RecipeError::NotFound { name, .. } => assert_eq!(name, "nope"),
other => panic!("expected NotFound, got {other:?}"),
}
}
#[test]
fn substitute_argv_replaces_known_placeholders() {
let template = vec![
"./tool".to_string(),
"{{YAH_TRANSFORM_IN_0}}".to_string(),
"{{YAH_TRANSFORM_OUT}}".to_string(),
"--mode={{quant}}".to_string(),
];
let mut params = BTreeMap::new();
params.insert(ENV_TRANSFORM_IN_0.into(), "/cache/fetch/abc.bin".into());
params.insert(ENV_TRANSFORM_OUT.into(), "/tmp/out.bin".into());
params.insert("quant".into(), "q5_1".into());
let resolved = substitute_argv(&template, ¶ms);
assert_eq!(
resolved,
vec![
"./tool",
"/cache/fetch/abc.bin",
"/tmp/out.bin",
"--mode=q5_1",
]
);
}
#[test]
fn substitute_argv_preserves_unknown_placeholders() {
let template = vec!["{{unknown}}".to_string()];
let resolved = substitute_argv(&template, &BTreeMap::new());
assert_eq!(resolved, vec!["{{unknown}}".to_string()]);
}
#[test]
fn substitute_argv_does_not_split_values_with_spaces() {
let template = vec!["{{flag}}".to_string()];
let mut params = BTreeMap::new();
params.insert("flag".into(), "--a --b --c".into());
let resolved = substitute_argv(&template, ¶ms);
assert_eq!(resolved, vec!["--a --b --c"]);
assert_eq!(resolved.len(), 1, "must not split on spaces");
}
#[test]
fn substitute_argv_handles_unterminated_braces() {
let template = vec!["{{never_closed".to_string()];
let resolved = substitute_argv(&template, &BTreeMap::new());
assert_eq!(resolved, vec!["{{never_closed".to_string()]);
}
#[test]
fn substitute_argv_trims_whitespace_in_key() {
let template = vec!["{{ key }}".to_string()];
let mut params = BTreeMap::new();
params.insert("key".into(), "value".into());
let resolved = substitute_argv(&template, ¶ms);
assert_eq!(resolved, vec!["value".to_string()]);
}
}