zernel/commands/
env.rs

1// Copyright (C) 2026 Dyber, Inc. — Proprietary
2
3//! zernel env — Environment management
4
5use anyhow::Result;
6use clap::Subcommand;
7use std::process::Command;
8
9#[derive(Subcommand)]
10pub enum EnvCommands {
11    /// Snapshot current environment (Python, CUDA, driver, packages)
12    Snapshot {
13        /// Output file
14        #[arg(long, default_value = "zernel-env.yml")]
15        output: String,
16    },
17    /// Diff two environment snapshots
18    Diff {
19        /// First snapshot
20        a: String,
21        /// Second snapshot
22        b: String,
23    },
24    /// Reproduce an environment from a snapshot
25    Reproduce {
26        /// Snapshot file to reproduce
27        file: String,
28    },
29    /// Export environment as Dockerfile
30    Export {
31        /// Export format (docker, conda, pip)
32        #[arg(long, default_value = "docker")]
33        format: String,
34        /// Output file
35        #[arg(long, default_value = "Dockerfile.zernel")]
36        output: String,
37    },
38    /// Show current environment
39    Show,
40}
41
42fn run_cmd(cmd: &str, args: &[&str]) -> String {
43    Command::new(cmd)
44        .args(args)
45        .output()
46        .ok()
47        .filter(|o| o.status.success())
48        .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
49        .unwrap_or_else(|| "N/A".into())
50}
51
52pub async fn run(cmd: EnvCommands) -> Result<()> {
53    match cmd {
54        EnvCommands::Show | EnvCommands::Snapshot { .. } => {
55            let os_info = run_cmd("uname", &["-srm"]);
56            let python = run_cmd("python3", &["--version"]);
57            let pip_list = run_cmd("pip", &["list", "--format=freeze"]);
58            let cuda = run_cmd("nvcc", &["--version"]);
59            let driver = run_cmd(
60                "nvidia-smi",
61                &["--query-gpu=driver_version", "--format=csv,noheader"],
62            );
63            let gpu_name = run_cmd("nvidia-smi", &["--query-gpu=name", "--format=csv,noheader"]);
64            let torch_ver = run_cmd("python3", &["-c", "import torch; print(torch.__version__)"]);
65            let torch_cuda = run_cmd(
66                "python3",
67                &["-c", "import torch; print(torch.version.cuda)"],
68            );
69
70            let snapshot = format!(
71                "# Zernel Environment Snapshot\n\
72                 # Generated: {}\n\n\
73                 os: {}\n\
74                 python: {}\n\
75                 nvidia_driver: {}\n\
76                 gpu: {}\n\
77                 cuda_toolkit: {}\n\
78                 torch: {}\n\
79                 torch_cuda: {}\n\n\
80                 # pip packages\n\
81                 packages:\n{}\n",
82                chrono::Utc::now().to_rfc3339(),
83                os_info,
84                python,
85                driver.lines().next().unwrap_or("N/A"),
86                gpu_name.lines().next().unwrap_or("N/A"),
87                cuda.lines().last().unwrap_or("N/A"),
88                torch_ver,
89                torch_cuda,
90                pip_list
91                    .lines()
92                    .map(|l| format!("  - {l}"))
93                    .collect::<Vec<_>>()
94                    .join("\n"),
95            );
96
97            if let EnvCommands::Snapshot { output } = cmd {
98                std::fs::write(&output, &snapshot)?;
99                println!("Environment snapshot saved to: {output}");
100            } else {
101                print!("{snapshot}");
102            }
103        }
104
105        EnvCommands::Diff { a, b } => {
106            let a_content = std::fs::read_to_string(&a)?;
107            let b_content = std::fs::read_to_string(&b)?;
108
109            println!("Environment Diff: {a} vs {b}");
110            println!("{}", "=".repeat(60));
111
112            let a_lines: std::collections::HashSet<&str> = a_content.lines().collect();
113            let b_lines: std::collections::HashSet<&str> = b_content.lines().collect();
114
115            let only_a: Vec<&&str> = a_lines.difference(&b_lines).collect();
116            let only_b: Vec<&&str> = b_lines.difference(&a_lines).collect();
117
118            if !only_a.is_empty() {
119                println!("\nOnly in {a}:");
120                for line in &only_a {
121                    if !line.starts_with('#') && !line.is_empty() {
122                        println!("  - {line}");
123                    }
124                }
125            }
126
127            if !only_b.is_empty() {
128                println!("\nOnly in {b}:");
129                for line in &only_b {
130                    if !line.starts_with('#') && !line.is_empty() {
131                        println!("  + {line}");
132                    }
133                }
134            }
135
136            if only_a.is_empty() && only_b.is_empty() {
137                println!("Environments are identical.");
138            }
139        }
140
141        EnvCommands::Reproduce { file } => {
142            let content = std::fs::read_to_string(&file)?;
143            println!("Reproducing environment from: {file}");
144
145            // Extract pip packages
146            let mut in_packages = false;
147            let mut packages = Vec::new();
148            for line in content.lines() {
149                if line.starts_with("packages:") {
150                    in_packages = true;
151                    continue;
152                }
153                if in_packages {
154                    if let Some(pkg) = line.strip_prefix("  - ") {
155                        packages.push(pkg.to_string());
156                    } else {
157                        in_packages = false;
158                    }
159                }
160            }
161
162            if packages.is_empty() {
163                println!("No packages found in snapshot.");
164                return Ok(());
165            }
166
167            println!("Installing {} packages...", packages.len());
168            let reqs = packages.join("\n");
169            let req_file = "/tmp/zernel-requirements.txt";
170            std::fs::write(req_file, &reqs)?;
171
172            let status = Command::new("pip")
173                .args(["install", "-r", req_file])
174                .status()?;
175
176            if status.success() {
177                println!("Environment reproduced successfully.");
178            } else {
179                println!("Some packages failed to install.");
180            }
181        }
182
183        EnvCommands::Export { format, output } => match format.as_str() {
184            "docker" => {
185                let pip_list = run_cmd("pip", &["list", "--format=freeze"]);
186                let dockerfile =
187                        "# Generated by zernel env export\n\
188                         FROM nvidia/cuda:12.6.0-runtime-ubuntu24.04\n\n\
189                         RUN apt-get update && apt-get install -y python3 python3-pip && rm -rf /var/lib/apt/lists/*\n\n\
190                         COPY requirements.txt /tmp/\n\
191                         RUN pip install --no-cache-dir -r /tmp/requirements.txt\n\n\
192                         WORKDIR /workspace\n";
193                std::fs::write(&output, dockerfile)?;
194                std::fs::write("requirements.txt", &pip_list)?;
195                println!("Exported to: {output} + requirements.txt");
196                println!("Build: docker build -f {output} -t my-env .");
197            }
198            "pip" => {
199                let pip_list = run_cmd("pip", &["list", "--format=freeze"]);
200                std::fs::write(&output, &pip_list)?;
201                println!("Exported to: {output}");
202                println!("Reproduce: pip install -r {output}");
203            }
204            other => {
205                println!("Unknown format: {other}. Available: docker, pip");
206            }
207        },
208    }
209    Ok(())
210}
zernel/commands/env.rs

zernel/commands/
env.rs
