lean_ctx_patterns/
mlflow.rs

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

lean_ctx_patterns/
mlflow.rs
