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    /// Construct the ROCm environment setup for the profiler.
153    pub fn build_rocm_env(&self) -> Vec<(String, String)> {
154        vec![
155            ("ROCM_PATH".to_string(), "/opt/rocm".to_string()),
156            (
157                "PATH".to_string(),
158                "/opt/rocm/bin:".to_string() + &std::env::var("PATH").unwrap_or_default(),
159            ),
160            (
161                "LD_LIBRARY_PATH".to_string(),
162                "/opt/rocm/lib:".to_string()
163                    + &std::env::var("LD_LIBRARY_PATH").unwrap_or_default(),
164            ),
165            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
166            ("PYTORCH_ROCM_ARCH".to_string(), "gfx1100".to_string()),
167        ]
168    }
169}
170
171// ===========================================================================
172// Tests
173// ===========================================================================
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    // --- VAL-INSTALL-020: PyTorch Profiler correct pip command ---
180
181    #[test]
182    fn test_install_command_global() {
183        let installer = PytorchProfilerInstaller::new(PytorchProfilerConfig {
184            method: InstallMethod::Global,
185            ..Default::default()
186        });
187        let cmd = installer.build_install_command();
188        assert_eq!(cmd.program, "python3");
189        assert!(cmd.args.contains(&"-m".to_string()));
190        assert!(cmd.args.contains(&"pip".to_string()));
191        assert!(cmd.args.contains(&"install".to_string()));
192        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
193        assert!(cmd.args.contains(&"torch-tb-profiler".to_string()));
194        assert!(cmd.args.contains(&"tensorboard".to_string()));
195    }
196
197    #[test]
198    fn test_install_command_venv() {
199        let installer = PytorchProfilerInstaller::new(PytorchProfilerConfig {
200            method: InstallMethod::Venv,
201            ..Default::default()
202        });
203        let cmd = installer.build_install_command();
204        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
205        assert!(cmd.args.contains(&"torch-tb-profiler".to_string()));
206        assert!(cmd.args.contains(&"tensorboard".to_string()));
207    }
208
209    #[test]
210    fn test_dependencies() {
211        let installer = PytorchProfilerInstaller::with_defaults();
212        assert!(installer.dependencies().contains(&"pytorch"));
213    }
214
215    #[test]
216    fn test_validate_dependencies_success() {
217        let installer = PytorchProfilerInstaller::with_defaults();
218        assert!(installer.validate_dependencies(&["pytorch"]).is_ok());
219    }
220
221    #[test]
222    fn test_validate_dependencies_missing() {
223        let installer = PytorchProfilerInstaller::with_defaults();
224        let result = installer.validate_dependencies(&[]);
225        assert!(result.is_err());
226        assert!(result.unwrap_err().to_string().contains("pytorch"));
227    }
228
229    #[test]
230    fn test_rocm_env() {
231        let installer = PytorchProfilerInstaller::with_defaults();
232        let env = installer.build_rocm_env();
233        assert!(env
234            .iter()
235            .any(|(k, v)| k == "ROCM_PATH" && v == "/opt/rocm"));
236        assert!(env
237            .iter()
238            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
239    }
240}
rusty_stack/installers/components/pytorch_profiler.rs

rusty_stack/installers/components/
pytorch_profiler.rs
