rusty_stack/installers/components/
pytorch_profiler.rs

1//! PyTorch Profiler installer — ports `scripts/install_pytorch_profiler.sh`.
2//!
3//! Constructs correct pip install command for PyTorch Profiler (torch-tb-profiler)
4//! and TensorBoard. Depends on PyTorch being installed first.
5//!
6//! # Validation Assertion
7//!
8//! - **VAL-INSTALL-020**: PyTorch Profiler correct pip command
9
10// ===========================================================================
11// Types
12// ===========================================================================
13
14/// Installation method.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum InstallMethod {
17    /// Install globally.
18    Global,
19    /// Install in a virtual environment.
20    Venv,
21    /// Try global, fallback to venv.
22    Auto,
23}
24
25impl std::fmt::Display for InstallMethod {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            InstallMethod::Global => write!(f, "global"),
29            InstallMethod::Venv => write!(f, "venv"),
30            InstallMethod::Auto => write!(f, "auto"),
31        }
32    }
33}
34
35/// Configuration for the PyTorch Profiler installer.
36#[derive(Debug, Clone)]
37pub struct PytorchProfilerConfig {
38    /// Python binary to use.
39    pub python_bin: String,
40    /// Installation method.
41    pub method: InstallMethod,
42    /// Whether to run in dry-run mode.
43    pub dry_run: bool,
44    /// Whether strict ROCm mode is enabled.
45    pub strict_rocm: bool,
46}
47
48impl Default for PytorchProfilerConfig {
49    fn default() -> Self {
50        Self {
51            python_bin: "python3".to_string(),
52            method: InstallMethod::Auto,
53            dry_run: false,
54            strict_rocm: true,
55        }
56    }
57}
58
59/// A constructed shell command.
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub struct ShellCommand {
62    /// The program to run.
63    pub program: String,
64    /// Arguments to pass.
65    pub args: Vec<String>,
66    /// Environment variables to set.
67    pub env: Vec<(String, String)>,
68}
69
70impl ShellCommand {
71    /// Format as a shell command string.
72    pub fn to_command_string(&self) -> String {
73        let env_prefix = self
74            .env
75            .iter()
76            .map(|(k, v)| format!("{k}={v}"))
77            .collect::<Vec<_>>()
78            .join(" ");
79        let cmd = if self.args.is_empty() {
80            self.program.clone()
81        } else {
82            format!("{} {}", self.program, self.args.join(" "))
83        };
84        if env_prefix.is_empty() {
85            cmd
86        } else {
87            format!("{env_prefix} {cmd}")
88        }
89    }
90}
91
92/// The PyTorch Profiler installer.
93pub struct PytorchProfilerInstaller {
94    config: PytorchProfilerConfig,
95}
96
97impl PytorchProfilerInstaller {
98    /// Create a new PyTorch Profiler installer with the given config.
99    pub fn new(config: PytorchProfilerConfig) -> Self {
100        Self { config }
101    }
102
103    /// Create with default config.
104    pub fn with_defaults() -> Self {
105        Self::new(PytorchProfilerConfig::default())
106    }
107
108    // -----------------------------------------------------------------------
109    // Dependencies
110    // -----------------------------------------------------------------------
111
112    /// PyTorch Profiler depends on PyTorch.
113    pub fn dependencies(&self) -> &[&str] {
114        &["pytorch"]
115    }
116
117    /// Validate that all dependencies are satisfied.
118    pub fn validate_dependencies(&self, installed_components: &[&str]) -> anyhow::Result<()> {
119        for dep in self.dependencies() {
120            if !installed_components.contains(dep) {
121                anyhow::bail!("PyTorch Profiler requires '{}' to be installed first", dep);
122            }
123        }
124        Ok(())
125    }
126
127    // -----------------------------------------------------------------------
128    // Command construction (VAL-INSTALL-020)
129    // -----------------------------------------------------------------------
130
131    /// Construct the pip install command for PyTorch Profiler.
132    ///
133    /// The original script installs:
134    /// `pip install torch-tb-profiler tensorboard`
135    pub fn build_install_command(&self) -> ShellCommand {
136        let use_break = self.config.method == InstallMethod::Global
137            || self.config.method == InstallMethod::Auto;
138
139        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
140        if use_break {
141            args.push("--break-system-packages".to_string());
142        }
143        args.extend(["torch-tb-profiler".to_string(), "tensorboard".to_string()]);
144
145        ShellCommand {
146            program: self.config.python_bin.clone(),
147            args,
148            env: vec![],
149        }
150    }
151}
152
153// ===========================================================================
154// Tests
155// ===========================================================================
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    // --- VAL-INSTALL-020: PyTorch Profiler correct pip command ---
162
163    #[test]
164    fn test_install_command_global() {
165        let installer = PytorchProfilerInstaller::new(PytorchProfilerConfig {
166            method: InstallMethod::Global,
167            ..Default::default()
168        });
169        let cmd = installer.build_install_command();
170        assert_eq!(cmd.program, "python3");
171        assert!(cmd.args.contains(&"-m".to_string()));
172        assert!(cmd.args.contains(&"pip".to_string()));
173        assert!(cmd.args.contains(&"install".to_string()));
174        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
175        assert!(cmd.args.contains(&"torch-tb-profiler".to_string()));
176        assert!(cmd.args.contains(&"tensorboard".to_string()));
177    }
178
179    #[test]
180    fn test_install_command_venv() {
181        let installer = PytorchProfilerInstaller::new(PytorchProfilerConfig {
182            method: InstallMethod::Venv,
183            ..Default::default()
184        });
185        let cmd = installer.build_install_command();
186        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
187        assert!(cmd.args.contains(&"torch-tb-profiler".to_string()));
188        assert!(cmd.args.contains(&"tensorboard".to_string()));
189    }
190
191    #[test]
192    fn test_dependencies() {
193        let installer = PytorchProfilerInstaller::with_defaults();
194        assert!(installer.dependencies().contains(&"pytorch"));
195    }
196
197    #[test]
198    fn test_validate_dependencies_success() {
199        let installer = PytorchProfilerInstaller::with_defaults();
200        assert!(installer.validate_dependencies(&["pytorch"]).is_ok());
201    }
202
203    #[test]
204    fn test_validate_dependencies_missing() {
205        let installer = PytorchProfilerInstaller::with_defaults();
206        let result = installer.validate_dependencies(&[]);
207        assert!(result.is_err());
208        assert!(result.unwrap_err().to_string().contains("pytorch"));
209    }
210}
rusty_stack/installers/components/pytorch_profiler.rs

rusty_stack/installers/components/
pytorch_profiler.rs
