rusty_stack/installers/components/
repair.rs

1//! Repair ML Stack utility — ports `scripts/repair_ml_stack.sh`.
2//!
3//! Calls correct sub-installers in sequence via native Rust modules,
4//! propagates failures, and reports a summary of repair results.
5//!
6//! # Validation Assertions
7//!
8//! - **VAL-INSTALL-026**: Repair invokes correct sub-installers in order
9//! - **VAL-INSTALL-027**: Repair propagates individual failures
10//! - **VAL-INSTALL-051**: Repair dispatches to native Rust modules, not bash scripts
11
12// ===========================================================================
13// Types
14// ===========================================================================
15
16/// A constructed shell command.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct ShellCommand {
19    /// The program to run.
20    pub program: String,
21    /// Arguments to pass.
22    pub args: Vec<String>,
23    /// Environment variables to set.
24    pub env: Vec<(String, String)>,
25}
26
27impl ShellCommand {
28    /// Format as a shell command string.
29    pub fn to_command_string(&self) -> String {
30        let env_prefix = self
31            .env
32            .iter()
33            .map(|(k, v)| format!("{k}={v}"))
34            .collect::<Vec<_>>()
35            .join(" ");
36        let cmd = if self.args.is_empty() {
37            self.program.clone()
38        } else {
39            format!("{} {}", self.program, self.args.join(" "))
40        };
41        if env_prefix.is_empty() {
42            cmd
43        } else {
44            format!("{env_prefix} {cmd}")
45        }
46    }
47}
48
49/// Result of a single repair step.
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct RepairStepResult {
52    /// Name of the repair step.
53    pub name: String,
54    /// Whether the step succeeded.
55    pub success: bool,
56    /// Optional error message.
57    pub error_message: Option<String>,
58}
59
60/// Result of the full repair sequence.
61#[derive(Debug, Clone)]
62pub struct RepairResult {
63    /// Individual step results in order.
64    pub steps: Vec<RepairStepResult>,
65}
66
67impl RepairResult {
68    /// Create a new empty repair result.
69    pub fn new() -> Self {
70        Self { steps: Vec::new() }
71    }
72
73    /// Add a step result.
74    pub fn add_step(&mut self, name: &str, success: bool, error_message: Option<String>) {
75        self.steps.push(RepairStepResult {
76            name: name.to_string(),
77            success,
78            error_message,
79        });
80    }
81
82    /// Get the number of successful steps.
83    pub fn successful_count(&self) -> usize {
84        self.steps.iter().filter(|s| s.success).count()
85    }
86
87    /// Get the number of failed steps.
88    pub fn failed_count(&self) -> usize {
89        self.steps.iter().filter(|s| !s.success).count()
90    }
91
92    /// Whether all steps succeeded.
93    pub fn all_succeeded(&self) -> bool {
94        self.steps.iter().all(|s| s.success)
95    }
96
97    /// Get failed step names.
98    pub fn failed_steps(&self) -> Vec<&str> {
99        self.steps
100            .iter()
101            .filter(|s| !s.success)
102            .map(|s| s.name.as_str())
103            .collect()
104    }
105
106    /// Format a summary report.
107    pub fn format_summary(&self) -> String {
108        let mut lines = Vec::new();
109        lines.push("ML Stack Repair Summary".to_string());
110        lines.push("=======================".to_string());
111        for step in &self.steps {
112            let status = if step.success { "✓" } else { "✗" };
113            let msg = match &step.error_message {
114                Some(m) => format!(" ({})", m),
115                None => String::new(),
116            };
117            lines.push(format!("{} {}{}", status, step.name, msg));
118        }
119        lines.push(String::new());
120        lines.push(format!(
121            "Total: {} succeeded, {} failed",
122            self.successful_count(),
123            self.failed_count()
124        ));
125        lines.join("\n")
126    }
127}
128
129impl Default for RepairResult {
130    fn default() -> Self {
131        Self::new()
132    }
133}
134
135/// Repair step identifier — matches the sequence in repair_ml_stack.sh.
136///
137/// The original script's main() repair sequence is:
138/// 1. ROCm installation (fix_rocm_installation)
139/// 2. AMDGPU drivers (fix_amdgpu_drivers)
140/// 3. PyTorch (fix_pytorch_installation)
141/// 4. ML Stack Core (fix_ml_stack_core)
142/// 5. AITER (fix_aiter_installation)
143/// 6. MIGraphX Python (fix_migraphx_python_installation)
144/// 7. DeepSpeed (fix_deepspeed_installation)
145/// 8. Environment variables (fix_environment_variables)
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
147pub enum RepairStep {
148    /// Fix ROCm installation.
149    Rocm,
150    /// Fix AMDGPU drivers.
151    AmdgpuDrivers,
152    /// Fix PyTorch installation.
153    PyTorch,
154    /// Fix ML Stack Core components.
155    MlStackCore,
156    /// Fix AITER installation.
157    Aiter,
158    /// Fix MIGraphX Python bindings.
159    MigraphxPython,
160    /// Fix DeepSpeed installation.
161    DeepSpeed,
162    /// Fix environment variables.
163    EnvironmentVariables,
164}
165
166impl RepairStep {
167    /// Get all repair steps in the correct order.
168    pub fn all_in_order() -> Vec<RepairStep> {
169        vec![
170            RepairStep::Rocm,
171            RepairStep::AmdgpuDrivers,
172            RepairStep::PyTorch,
173            RepairStep::MlStackCore,
174            RepairStep::Aiter,
175            RepairStep::MigraphxPython,
176            RepairStep::DeepSpeed,
177            RepairStep::EnvironmentVariables,
178        ]
179    }
180
181    /// Get the display name of the repair step.
182    pub fn name(&self) -> &'static str {
183        match self {
184            RepairStep::Rocm => "ROCm Installation",
185            RepairStep::AmdgpuDrivers => "AMDGPU Drivers",
186            RepairStep::PyTorch => "PyTorch Installation",
187            RepairStep::MlStackCore => "ML Stack Core",
188            RepairStep::Aiter => "AITER Installation",
189            RepairStep::MigraphxPython => "MIGraphX Python",
190            RepairStep::DeepSpeed => "DeepSpeed Installation",
191            RepairStep::EnvironmentVariables => "Environment Variables",
192        }
193    }
194
195    /// Get the native component ID that this repair step dispatches to.
196    /// Each repair step re-runs the corresponding native installer module.
197    pub fn native_component_id(&self) -> &'static str {
198        match self {
199            RepairStep::Rocm => "rocm",
200            RepairStep::AmdgpuDrivers => "amdgpu-drivers",
201            RepairStep::PyTorch => "pytorch",
202            RepairStep::MlStackCore => "ml-stack-core",
203            RepairStep::Aiter => "aiter",
204            RepairStep::MigraphxPython => "migraphx-python",
205            RepairStep::DeepSpeed => "deepspeed",
206            RepairStep::EnvironmentVariables => "permanent-env",
207        }
208    }
209}
210
211/// Configuration for the repair tool.
212#[derive(Debug, Clone)]
213pub struct RepairConfig {
214    /// Python binary to use.
215    pub python_bin: String,
216    /// Whether to run in dry-run mode.
217    pub dry_run: bool,
218    /// Whether to run in non-interactive mode.
219    pub non_interactive: bool,
220}
221
222impl Default for RepairConfig {
223    fn default() -> Self {
224        Self {
225            python_bin: "python3".to_string(),
226            dry_run: false,
227            non_interactive: false,
228        }
229    }
230}
231
232/// The ML Stack repair tool.
233pub struct RepairInstaller {
234    config: RepairConfig,
235}
236
237impl RepairInstaller {
238    /// Create a new repair tool with the given config.
239    pub fn new(config: RepairConfig) -> Self {
240        Self { config }
241    }
242
243    /// Create with default config.
244    pub fn with_defaults() -> Self {
245        Self::new(RepairConfig::default())
246    }
247
248    // -----------------------------------------------------------------------
249    // Step enumeration (VAL-INSTALL-026)
250    // -----------------------------------------------------------------------
251
252    /// Get the repair step sequence in the correct order.
253    pub fn repair_sequence() -> Vec<RepairStep> {
254        RepairStep::all_in_order()
255    }
256
257    /// Get the native component IDs for all repair steps, in order.
258    pub fn native_component_ids() -> Vec<&'static str> {
259        Self::repair_sequence()
260            .iter()
261            .map(|s| s.native_component_id())
262            .collect()
263    }
264
265    // -----------------------------------------------------------------------
266    // Failure propagation (VAL-INSTALL-027)
267    // -----------------------------------------------------------------------
268
269    /// Run the repair sequence with mock step results.
270    /// Continues through all steps even if some fail, collecting results.
271    /// Returns a summary with all step outcomes.
272    pub fn run_repair_sequence<F>(&self, mut step_runner: F) -> RepairResult
273    where
274        F: FnMut(RepairStep) -> Result<(), String>,
275    {
276        let mut result = RepairResult::new();
277
278        for step in Self::repair_sequence() {
279            // Honor dry-run: record each step as skipped without invoking the runner.
280            if self.config.dry_run {
281                result.add_step(step.name(), true, Some("dry-run: skipped".to_string()));
282                continue;
283            }
284            match step_runner(step) {
285                Ok(()) => result.add_step(step.name(), true, None),
286                Err(msg) => result.add_step(step.name(), false, Some(msg)),
287            }
288        }
289
290        result
291    }
292}
293
294// ===========================================================================
295// Tests
296// ===========================================================================
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    #[test]
303    fn test_repair_sequence_order_matches_original_script() {
304        let sequence = RepairInstaller::repair_sequence();
305        assert_eq!(sequence.len(), 8);
306        assert_eq!(sequence[0], RepairStep::Rocm);
307        assert_eq!(sequence[1], RepairStep::AmdgpuDrivers);
308        assert_eq!(sequence[2], RepairStep::PyTorch);
309        assert_eq!(sequence[3], RepairStep::MlStackCore);
310        assert_eq!(sequence[4], RepairStep::Aiter);
311        assert_eq!(sequence[5], RepairStep::MigraphxPython);
312        assert_eq!(sequence[6], RepairStep::DeepSpeed);
313        assert_eq!(sequence[7], RepairStep::EnvironmentVariables);
314    }
315
316    #[test]
317    fn test_repair_step_names() {
318        assert_eq!(RepairStep::Rocm.name(), "ROCm Installation");
319        assert_eq!(RepairStep::AmdgpuDrivers.name(), "AMDGPU Drivers");
320        assert_eq!(RepairStep::PyTorch.name(), "PyTorch Installation");
321        assert_eq!(RepairStep::MlStackCore.name(), "ML Stack Core");
322        assert_eq!(RepairStep::Aiter.name(), "AITER Installation");
323        assert_eq!(RepairStep::MigraphxPython.name(), "MIGraphX Python");
324        assert_eq!(RepairStep::DeepSpeed.name(), "DeepSpeed Installation");
325        assert_eq!(
326            RepairStep::EnvironmentVariables.name(),
327            "Environment Variables"
328        );
329    }
330
331    #[test]
332    fn test_repair_step_native_component_ids() {
333        assert_eq!(RepairStep::Rocm.native_component_id(), "rocm");
334        assert_eq!(
335            RepairStep::AmdgpuDrivers.native_component_id(),
336            "amdgpu-drivers"
337        );
338        assert_eq!(RepairStep::PyTorch.native_component_id(), "pytorch");
339        assert_eq!(
340            RepairStep::MlStackCore.native_component_id(),
341            "ml-stack-core"
342        );
343        assert_eq!(RepairStep::Aiter.native_component_id(), "aiter");
344        assert_eq!(
345            RepairStep::MigraphxPython.native_component_id(),
346            "migraphx-python"
347        );
348        assert_eq!(RepairStep::DeepSpeed.native_component_id(), "deepspeed");
349        assert_eq!(
350            RepairStep::EnvironmentVariables.native_component_id(),
351            "permanent-env"
352        );
353    }
354
355    #[test]
356    fn test_native_component_ids_match_repair_sequence() {
357        let ids = RepairInstaller::native_component_ids();
358        let steps = RepairInstaller::repair_sequence();
359        assert_eq!(ids.len(), steps.len());
360        for (step, id) in steps.iter().zip(ids.iter()) {
361            assert_eq!(step.native_component_id(), *id);
362        }
363    }
364
365    #[test]
366    fn test_repair_uses_native_dispatch_not_scripts() {
367        // Verify every repair step maps to a native component ID
368        // (no bash script references needed for repair)
369        for step in RepairInstaller::repair_sequence() {
370            let id = step.native_component_id();
371            assert!(
372                crate::installers::components::is_native_component(id),
373                "Repair step '{}' maps to component '{}' which must be native",
374                step.name(),
375                id
376            );
377        }
378    }
379
380    #[test]
381    fn test_repair_propagates_failures() {
382        let repair = RepairInstaller::with_defaults();
383        let result = repair.run_repair_sequence(|step| match step {
384            RepairStep::PyTorch => Err("PyTorch installation failed".to_string()),
385            RepairStep::DeepSpeed => Err("DeepSpeed build error".to_string()),
386            _ => Ok(()),
387        });
388
389        assert_eq!(result.failed_count(), 2);
390        assert_eq!(result.successful_count(), 6);
391        assert!(!result.all_succeeded());
392
393        let failed = result.failed_steps();
394        assert!(failed.contains(&"PyTorch Installation"));
395        assert!(failed.contains(&"DeepSpeed Installation"));
396    }
397
398    #[test]
399    fn test_repair_all_succeed() {
400        let repair = RepairInstaller::with_defaults();
401        let result = repair.run_repair_sequence(|_| Ok(()));
402
403        assert!(result.all_succeeded());
404        assert_eq!(result.successful_count(), 8);
405        assert_eq!(result.failed_count(), 0);
406    }
407
408    #[test]
409    fn test_repair_summary_format() {
410        let repair = RepairInstaller::with_defaults();
411        let result = repair.run_repair_sequence(|step| match step {
412            RepairStep::PyTorch => Err("timeout".to_string()),
413            _ => Ok(()),
414        });
415
416        let summary = result.format_summary();
417        assert!(summary.contains("ML Stack Repair Summary"));
418        assert!(summary.contains("7 succeeded, 1 failed"));
419        assert!(summary.contains("✗ PyTorch Installation (timeout)"));
420        assert!(summary.contains("✓ ROCm Installation"));
421    }
422
423    #[test]
424    fn test_repair_continues_after_failure() {
425        let repair = RepairInstaller::with_defaults();
426        let mut visited = Vec::new();
427
428        let _result = repair.run_repair_sequence(|step| {
429            visited.push(step);
430            if step == RepairStep::Rocm {
431                Err("ROCm failed".to_string())
432            } else {
433                Ok(())
434            }
435        });
436
437        assert_eq!(visited.len(), 8);
438        assert_eq!(visited[0], RepairStep::Rocm);
439        assert_eq!(visited[7], RepairStep::EnvironmentVariables);
440    }
441
442    #[test]
443    fn test_repair_result_default() {
444        let result = RepairResult::default();
445        assert_eq!(result.steps.len(), 0);
446        assert!(result.all_succeeded());
447    }
448}
rusty_stack/installers/components/repair.rs

rusty_stack/installers/components/
repair.rs
