Skip to main content

ferrum_cli/
layer_split_pipeline.rs

1use ferrum_types::{FerrumError, Result, RuntimeConfigEntry, RuntimeConfigSnapshot};
2use serde_json::Value;
3use std::collections::HashMap;
4
5pub const LAYER_SPLIT_PIPELINE_MODE_KEY: &str = "FERRUM_LAYER_SPLIT_PIPELINE_MODE";
6pub const LAYER_SPLIT_PIPELINE_MODE_BACKEND_OPTION: &str = "layer_split_pipeline_mode";
7
8#[derive(Clone, Copy, Debug, PartialEq, Eq, clap::ValueEnum)]
9pub enum LayerSplitPipelineModeArg {
10    Batch,
11    Overlapped,
12}
13
14impl LayerSplitPipelineModeArg {
15    pub fn as_str(self) -> &'static str {
16        match self {
17            Self::Batch => "batch",
18            Self::Overlapped => "overlapped",
19        }
20    }
21}
22
23pub fn push_cli_runtime_entry(
24    entries: &mut Vec<RuntimeConfigEntry>,
25    mode: Option<LayerSplitPipelineModeArg>,
26) {
27    if let Some(mode) = mode {
28        entries.push(RuntimeConfigEntry::new(
29            LAYER_SPLIT_PIPELINE_MODE_KEY,
30            mode.as_str(),
31            ferrum_types::RuntimeConfigSource::Cli,
32        ));
33    }
34}
35
36pub fn normalize_layer_split_pipeline_mode(value: &str) -> Result<&'static str> {
37    match value.trim().to_ascii_lowercase().as_str() {
38        "batch" => Ok("batch"),
39        "overlapped" => Ok("overlapped"),
40        other => Err(FerrumError::config(format!(
41            "{LAYER_SPLIT_PIPELINE_MODE_KEY} must be batch or overlapped, got {other:?}"
42        ))),
43    }
44}
45
46pub fn runtime_layer_split_pipeline_mode(
47    snapshot: &RuntimeConfigSnapshot,
48) -> Result<Option<&'static str>> {
49    snapshot
50        .entries
51        .iter()
52        .find(|entry| entry.key == LAYER_SPLIT_PIPELINE_MODE_KEY)
53        .map(|entry| normalize_layer_split_pipeline_mode(&entry.effective_value))
54        .transpose()
55}
56
57pub fn insert_backend_option_from_runtime(
58    snapshot: &RuntimeConfigSnapshot,
59    backend_options: &mut HashMap<String, Value>,
60) -> Result<()> {
61    if let Some(mode) = runtime_layer_split_pipeline_mode(snapshot)? {
62        backend_options.insert(
63            LAYER_SPLIT_PIPELINE_MODE_BACKEND_OPTION.to_string(),
64            Value::String(mode.to_string()),
65        );
66    }
67    Ok(())
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use ferrum_types::{RuntimeConfigSnapshot, RuntimeConfigSource};
74
75    #[test]
76    fn cli_mode_entry_records_product_pipeline_mode() {
77        let mut entries = Vec::new();
78        push_cli_runtime_entry(&mut entries, Some(LayerSplitPipelineModeArg::Batch));
79
80        assert_eq!(entries.len(), 1);
81        assert_eq!(entries[0].key, LAYER_SPLIT_PIPELINE_MODE_KEY);
82        assert_eq!(entries[0].effective_value, "batch");
83        assert_eq!(entries[0].source, RuntimeConfigSource::Cli);
84    }
85
86    #[test]
87    fn invalid_runtime_mode_is_rejected() {
88        let snapshot = RuntimeConfigSnapshot::from_entries([RuntimeConfigEntry::new(
89            LAYER_SPLIT_PIPELINE_MODE_KEY,
90            "serial",
91            RuntimeConfigSource::ConfigFile,
92        )]);
93
94        let err = runtime_layer_split_pipeline_mode(&snapshot)
95            .unwrap_err()
96            .to_string();
97        assert!(err.contains("must be batch or overlapped"));
98    }
99}