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    /// Legacy script name (kept for documentation/test reference only).
211    #[allow(dead_code)]
212    pub fn legacy_script_name(&self) -> &'static str {
213        match self {
214            RepairStep::Rocm => "install_rocm.sh",
215            RepairStep::AmdgpuDrivers => "install_amdgpu_drivers.sh",
216            RepairStep::PyTorch => "install_pytorch_rocm.sh",
217            RepairStep::MlStackCore => "install_ml_stack.sh",
218            RepairStep::Aiter => "install_aiter.sh",
219            RepairStep::MigraphxPython => "install_migraphx_python.sh",
220            RepairStep::DeepSpeed => "install_deepspeed.sh",
221            RepairStep::EnvironmentVariables => "setup_permanent_rocm_env.sh",
222        }
223    }
224}
225
226/// Configuration for the repair tool.
227#[derive(Debug, Clone)]
228pub struct RepairConfig {
229    /// Python binary to use.
230    pub python_bin: String,
231    /// Whether to run in dry-run mode.
232    pub dry_run: bool,
233    /// Whether to run in non-interactive mode.
234    pub non_interactive: bool,
235}
236
237impl Default for RepairConfig {
238    fn default() -> Self {
239        Self {
240            python_bin: "python3".to_string(),
241            dry_run: false,
242            non_interactive: false,
243        }
244    }
245}
246
247/// The ML Stack repair tool.
248pub struct RepairInstaller {
249    #[allow(dead_code)]
250    config: RepairConfig,
251}
252
253impl RepairInstaller {
254    /// Create a new repair tool with the given config.
255    pub fn new(config: RepairConfig) -> Self {
256        Self { config }
257    }
258
259    /// Create with default config.
260    pub fn with_defaults() -> Self {
261        Self::new(RepairConfig::default())
262    }
263
264    // -----------------------------------------------------------------------
265    // Step enumeration (VAL-INSTALL-026)
266    // -----------------------------------------------------------------------
267
268    /// Get the repair step sequence in the correct order.
269    pub fn repair_sequence() -> Vec<RepairStep> {
270        RepairStep::all_in_order()
271    }
272
273    /// Get the native component IDs for all repair steps, in order.
274    pub fn native_component_ids() -> Vec<&'static str> {
275        Self::repair_sequence()
276            .iter()
277            .map(|s| s.native_component_id())
278            .collect()
279    }
280
281    // -----------------------------------------------------------------------
282    // Failure propagation (VAL-INSTALL-027)
283    // -----------------------------------------------------------------------
284
285    /// Run the repair sequence with mock step results.
286    /// Continues through all steps even if some fail, collecting results.
287    /// Returns a summary with all step outcomes.
288    pub fn run_repair_sequence<F>(&self, mut step_runner: F) -> RepairResult
289    where
290        F: FnMut(RepairStep) -> Result<(), String>,
291    {
292        let mut result = RepairResult::new();
293
294        for step in Self::repair_sequence() {
295            match step_runner(step) {
296                Ok(()) => result.add_step(step.name(), true, None),
297                Err(msg) => result.add_step(step.name(), false, Some(msg)),
298            }
299        }
300
301        result
302    }
303}
304
305// ===========================================================================
306// Tests
307// ===========================================================================
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312
313    #[test]
314    fn test_repair_sequence_order_matches_original_script() {
315        let sequence = RepairInstaller::repair_sequence();
316        assert_eq!(sequence.len(), 8);
317        assert_eq!(sequence[0], RepairStep::Rocm);
318        assert_eq!(sequence[1], RepairStep::AmdgpuDrivers);
319        assert_eq!(sequence[2], RepairStep::PyTorch);
320        assert_eq!(sequence[3], RepairStep::MlStackCore);
321        assert_eq!(sequence[4], RepairStep::Aiter);
322        assert_eq!(sequence[5], RepairStep::MigraphxPython);
323        assert_eq!(sequence[6], RepairStep::DeepSpeed);
324        assert_eq!(sequence[7], RepairStep::EnvironmentVariables);
325    }
326
327    #[test]
328    fn test_repair_step_names() {
329        assert_eq!(RepairStep::Rocm.name(), "ROCm Installation");
330        assert_eq!(RepairStep::AmdgpuDrivers.name(), "AMDGPU Drivers");
331        assert_eq!(RepairStep::PyTorch.name(), "PyTorch Installation");
332        assert_eq!(RepairStep::MlStackCore.name(), "ML Stack Core");
333        assert_eq!(RepairStep::Aiter.name(), "AITER Installation");
334        assert_eq!(RepairStep::MigraphxPython.name(), "MIGraphX Python");
335        assert_eq!(RepairStep::DeepSpeed.name(), "DeepSpeed Installation");
336        assert_eq!(
337            RepairStep::EnvironmentVariables.name(),
338            "Environment Variables"
339        );
340    }
341
342    #[test]
343    fn test_repair_step_native_component_ids() {
344        assert_eq!(RepairStep::Rocm.native_component_id(), "rocm");
345        assert_eq!(
346            RepairStep::AmdgpuDrivers.native_component_id(),
347            "amdgpu-drivers"
348        );
349        assert_eq!(RepairStep::PyTorch.native_component_id(), "pytorch");
350        assert_eq!(
351            RepairStep::MlStackCore.native_component_id(),
352            "ml-stack-core"
353        );
354        assert_eq!(RepairStep::Aiter.native_component_id(), "aiter");
355        assert_eq!(
356            RepairStep::MigraphxPython.native_component_id(),
357            "migraphx-python"
358        );
359        assert_eq!(RepairStep::DeepSpeed.native_component_id(), "deepspeed");
360        assert_eq!(
361            RepairStep::EnvironmentVariables.native_component_id(),
362            "permanent-env"
363        );
364    }
365
366    #[test]
367    fn test_native_component_ids_match_repair_sequence() {
368        let ids = RepairInstaller::native_component_ids();
369        let steps = RepairInstaller::repair_sequence();
370        assert_eq!(ids.len(), steps.len());
371        for (step, id) in steps.iter().zip(ids.iter()) {
372            assert_eq!(step.native_component_id(), *id);
373        }
374    }
375
376    #[test]
377    fn test_repair_uses_native_dispatch_not_scripts() {
378        // Verify every repair step maps to a native component ID
379        // (no bash script references needed for repair)
380        for step in RepairInstaller::repair_sequence() {
381            let id = step.native_component_id();
382            assert!(
383                crate::installers::components::is_native_component(id),
384                "Repair step '{}' maps to component '{}' which must be native",
385                step.name(),
386                id
387            );
388        }
389    }
390
391    #[test]
392    fn test_repair_propagates_failures() {
393        let repair = RepairInstaller::with_defaults();
394        let result = repair.run_repair_sequence(|step| match step {
395            RepairStep::PyTorch => Err("PyTorch installation failed".to_string()),
396            RepairStep::DeepSpeed => Err("DeepSpeed build error".to_string()),
397            _ => Ok(()),
398        });
399
400        assert_eq!(result.failed_count(), 2);
401        assert_eq!(result.successful_count(), 6);
402        assert!(!result.all_succeeded());
403
404        let failed = result.failed_steps();
405        assert!(failed.contains(&"PyTorch Installation"));
406        assert!(failed.contains(&"DeepSpeed Installation"));
407    }
408
409    #[test]
410    fn test_repair_all_succeed() {
411        let repair = RepairInstaller::with_defaults();
412        let result = repair.run_repair_sequence(|_| Ok(()));
413
414        assert!(result.all_succeeded());
415        assert_eq!(result.successful_count(), 8);
416        assert_eq!(result.failed_count(), 0);
417    }
418
419    #[test]
420    fn test_repair_summary_format() {
421        let repair = RepairInstaller::with_defaults();
422        let result = repair.run_repair_sequence(|step| match step {
423            RepairStep::PyTorch => Err("timeout".to_string()),
424            _ => Ok(()),
425        });
426
427        let summary = result.format_summary();
428        assert!(summary.contains("ML Stack Repair Summary"));
429        assert!(summary.contains("7 succeeded, 1 failed"));
430        assert!(summary.contains("✗ PyTorch Installation (timeout)"));
431        assert!(summary.contains("✓ ROCm Installation"));
432    }
433
434    #[test]
435    fn test_repair_continues_after_failure() {
436        let repair = RepairInstaller::with_defaults();
437        let mut visited = Vec::new();
438
439        let _result = repair.run_repair_sequence(|step| {
440            visited.push(step);
441            if step == RepairStep::Rocm {
442                Err("ROCm failed".to_string())
443            } else {
444                Ok(())
445            }
446        });
447
448        assert_eq!(visited.len(), 8);
449        assert_eq!(visited[0], RepairStep::Rocm);
450        assert_eq!(visited[7], RepairStep::EnvironmentVariables);
451    }
452
453    #[test]
454    fn test_repair_result_default() {
455        let result = RepairResult::default();
456        assert_eq!(result.steps.len(), 0);
457        assert!(result.all_succeeded());
458    }
459}
rusty_stack/installers/components/repair.rs

rusty_stack/installers/components/
repair.rs
