lean_ctx/core/patterns/
mlflow.rs

1//! MLflow CLI output compression.
2//!
3//! `mlflow run` drives conda/pip env builds that flood the output with
4//! `Collecting ...`, `Downloading ...` and progress bars before the actual
5//! run. We strip the python-logging timestamp prefix, drop env-build noise,
6//! deduplicate and keep the run lifecycle (`Run (ID '..') succeeded`),
7//! registered-model/version lines, metrics and errors.
8
9use crate::core::compressor::strip_ansi;
10use std::collections::HashSet;
11
12pub fn compress(_cmd: &str, output: &str) -> Option<String> {
13    let trimmed = output.trim();
14    if trimmed.is_empty() {
15        return Some("mlflow: ok".to_string());
16    }
17
18    let mut kept: Vec<String> = Vec::new();
19    let mut seen: HashSet<String> = HashSet::new();
20
21    for raw in trimmed.lines() {
22        let stripped = strip_ansi(raw);
23        let body = strip_log_prefix(stripped.trim());
24        if body.is_empty() || is_noise(body) {
25            continue;
26        }
27        if seen.insert(normalize(body)) {
28            kept.push(body.to_string());
29        }
30    }
31
32    if kept.is_empty() {
33        return Some("mlflow: ok".to_string());
34    }
35    Some(kept.join("\n"))
36}
37
38/// Drop a leading `YYYY/MM/DD HH:MM:SS LEVEL component:` python-logging prefix.
39fn strip_log_prefix(line: &str) -> &str {
40    let mut it = line.splitn(4, ' ');
41    let (Some(date), Some(time), Some(level), rest) = (it.next(), it.next(), it.next(), it.next())
42    else {
43        return line;
44    };
45    if is_date(date) && is_time(time) && is_level(level) {
46        rest.unwrap_or("").trim_start()
47    } else {
48        line
49    }
50}
51
52fn is_date(s: &str) -> bool {
53    let p: Vec<&str> = s.split('/').collect();
54    p.len() == 3
55        && p.iter()
56            .all(|x| !x.is_empty() && x.chars().all(|c| c.is_ascii_digit()))
57}
58
59fn is_time(s: &str) -> bool {
60    let p: Vec<&str> = s.split(':').collect();
61    p.len() == 3
62        && p.iter()
63            .all(|x| !x.is_empty() && x.chars().all(|c| c.is_ascii_digit()))
64}
65
66fn is_level(s: &str) -> bool {
67    matches!(
68        s,
69        "INFO" | "WARNING" | "WARN" | "ERROR" | "DEBUG" | "CRITICAL"
70    )
71}
72
73fn is_noise(line: &str) -> bool {
74    const PREFIXES: [&str; 9] = [
75        "Collecting ",
76        "Requirement already satisfied",
77        "Downloading ",
78        "Installing ",
79        "Building wheel",
80        "Preparing metadata",
81        "Using cached ",
82        "Channels:",
83        "Platform:",
84    ];
85    if PREFIXES.iter().any(|p| line.starts_with(p)) {
86        return true;
87    }
88    line.contains("Solving environment")
89        || line.contains("MB/s")
90        || line.contains("kB/s")
91        || line.contains("━")
92        || line.contains("it/s]")
93}
94
95/// Collapse digits so per-run IDs/metrics with the same shape dedupe sensibly
96/// while distinct messages survive.
97fn normalize(s: &str) -> String {
98    s.chars().filter(|c| !c.is_ascii_whitespace()).collect()
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    const RUN: &str = "2024/01/01 12:00:00 INFO mlflow.projects.utils: === Created directory /tmp/x ===\n2024/01/01 12:00:01 INFO mlflow.projects.backend.local: === Running command 'python train.py' ===\nCollecting numpy==1.26.0\nDownloading numpy-1.26.0.whl (18.2 MB)\n   ━━━━━━━━━━ 18.2/18.2 MB 25.1 MB/s\nRequirement already satisfied: scipy in /usr/lib\n2024/01/01 12:00:30 INFO mlflow.projects: === Run (ID 'abc123def456') succeeded ===\n";
106
107    #[test]
108    fn strips_prefix_and_env_noise_keeps_lifecycle() {
109        let r = compress("mlflow run .", RUN).unwrap();
110        assert!(r.contains("Run (ID 'abc123def456') succeeded"), "{r}");
111        assert!(r.contains("Running command"), "{r}");
112        assert!(!r.contains("2024/01/01"), "drops log timestamp: {r}");
113        assert!(!r.contains("Collecting"), "drops pip noise: {r}");
114        assert!(!r.contains("MB/s"), "drops download progress: {r}");
115    }
116
117    #[test]
118    fn shorter_than_input() {
119        let r = compress("mlflow run .", RUN).unwrap();
120        assert!(r.len() < RUN.len());
121    }
122
123    #[test]
124    fn empty_is_ok() {
125        assert_eq!(compress("mlflow run .", "").unwrap(), "mlflow: ok");
126    }
127}
lean_ctx/core/patterns/mlflow.rs

lean_ctx/core/patterns/
mlflow.rs
