Skip to main content

rusty_stack/
state.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum Stage {
5    Welcome,
6    HardwareDetect,
7    Preflight,
8    ComponentSelect,
9    Configuration,
10    Confirm,
11    Installing,
12    Complete,
13    Benchmarks,
14    Recovery,
15}
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum Category {
19    Foundation,
20    Core,
21    UiUx,
22    Extension,
23    Environment,
24    Verification,
25    Performance,
26}
27
28#[derive(Debug, Clone)]
29pub struct Component {
30    pub id: String,
31    pub name: String,
32    pub description: String,
33    /// Script path relative to scripts_dir, or empty string for native Rust modules.
34    ///
35    /// For native Rust components, this field is empty and the installer
36    /// dispatches to the corresponding Rust module via
37    /// `installers::components::is_native_component()`.
38    ///
39    /// # Validation Assertions
40    ///
41    /// - **VAL-INSTALL-037**: Component.script no longer holds .sh filenames for ported components
42    /// - **VAL-INSTALL-038**: state.rs supports native module routing
43    /// - **VAL-INSTALL-040**: Verification/performance components use native routing
44    pub script: String,
45    pub category: Category,
46    pub required: bool,
47    pub selected: bool,
48    pub installed: bool,
49    pub progress: f32,
50    pub estimate: String,
51    /// Whether this component requires sudo to install
52    pub needs_sudo: bool,
53}
54
55impl Component {
56    /// Returns `true` if this component uses a native Rust installer
57    /// (no shell script involved).
58    ///
59    /// # Validation Assertions
60    ///
61    /// - **VAL-INSTALL-038**: state.rs supports native module routing
62    pub fn is_native(&self) -> bool {
63        crate::installers::components::is_native_component(&self.id)
64    }
65}
66
67#[derive(Debug, Clone, Default, Serialize, Deserialize)]
68pub struct GPUInfo {
69    pub model: String,
70    pub driver: String,
71    pub architecture: String,
72    pub rocm_version: String,
73    pub gpu_count: usize,
74    pub memory_gb: f32,
75    pub temperature_c: Option<f32>,
76    pub power_watts: Option<f32>,
77}
78
79#[derive(Debug, Clone, Default, Serialize, Deserialize)]
80pub struct SystemInfo {
81    pub os: String,
82    pub distribution: String,
83    pub kernel: String,
84    pub cpu_model: String,
85    pub memory_gb: f32,
86    pub storage_gb: f32,
87    pub storage_available_gb: f32,
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum PreflightStatus {
92    Passed,
93    Warning,
94    Failed,
95}
96
97impl PreflightStatus {
98    pub fn label(self) -> &'static str {
99        match self {
100            PreflightStatus::Passed => "passed",
101            PreflightStatus::Warning => "warning",
102            PreflightStatus::Failed => "failed",
103        }
104    }
105}
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub enum PreflightType {
109    Critical,
110    Warning,
111    Info,
112}
113
114#[derive(Debug, Clone)]
115pub struct PreflightCheck {
116    pub name: String,
117    pub status: PreflightStatus,
118    pub check_type: PreflightType,
119    pub message: String,
120    pub details: String,
121    pub score: i32,
122}
123
124#[derive(Debug, Clone, Default)]
125pub struct PreflightResult {
126    pub passed: bool,
127    pub summary: String,
128    pub checks: Vec<PreflightCheck>,
129    pub passed_count: usize,
130    pub failed_count: usize,
131    pub warning_count: usize,
132    pub total_score: i32,
133    pub can_continue: bool,
134}
135
136#[derive(Debug, Clone, Default)]
137pub struct HardwareState {
138    pub gpu: GPUInfo,
139    pub system: SystemInfo,
140    pub status: String,
141    pub progress: f32,
142}
143
144#[derive(Debug, Clone, Default)]
145pub struct InstallStatus {
146    pub progress: f32,
147    pub message: String,
148    pub completed: bool,
149}
150
151pub fn default_components() -> Vec<Component> {
152    use Category::*;
153
154    vec![
155        // Environment first
156        Component {
157            id: "permanent-env".into(),
158            name: "Permanent ROCm Env".into(),
159            description: "Unified permanent environment for Python 3.12".into(),
160            script: String::new(), // Native Rust installer
161            category: Environment,
162            required: false,
163            selected: true,
164            installed: false,
165            progress: 0.0,
166            estimate: "1-2 min".into(),
167            needs_sudo: true,
168        },
169        // Foundation components
170        Component {
171            id: "rocm".into(),
172            name: "ROCm Platform".into(),
173            description: "AMD ROCm GPU computing platform".into(),
174            script: String::new(), // Native Rust installer
175            category: Foundation,
176            required: true,
177            selected: true,
178            installed: false,
179            progress: 0.0,
180            estimate: "30-45 min".into(),
181            needs_sudo: true,
182        },
183        Component {
184            id: "pytorch".into(),
185            name: "PyTorch with ROCm".into(),
186            description: "PyTorch optimized for AMD GPUs".into(),
187            script: String::new(), // Native Rust installer
188            category: Foundation,
189            required: true,
190            selected: true,
191            installed: false,
192            progress: 0.0,
193            estimate: "10-15 min".into(),
194            needs_sudo: true,
195        },
196        Component {
197            id: "triton".into(),
198            name: "Triton".into(),
199            description: "Compiler for parallel programming".into(),
200            script: String::new(), // Native Rust installer
201            category: Foundation,
202            required: true,
203            selected: true,
204            installed: false,
205            progress: 0.0,
206            estimate: "5-10 min".into(),
207            needs_sudo: true,
208        },
209        Component {
210            id: "mpi4py".into(),
211            name: "MPI4Py".into(),
212            description: "MPI bindings for Python".into(),
213            script: String::new(), // Native Rust installer
214            category: Foundation,
215            required: true,
216            selected: true,
217            installed: false,
218            progress: 0.0,
219            estimate: "3-5 min".into(),
220            needs_sudo: true,
221        },
222        Component {
223            id: "deepspeed".into(),
224            name: "DeepSpeed".into(),
225            description: "Deep learning optimization library".into(),
226            script: String::new(), // Native Rust installer
227            category: Foundation,
228            required: true,
229            selected: true,
230            installed: false,
231            progress: 0.0,
232            estimate: "8-12 min".into(),
233            needs_sudo: true,
234        },
235        // Core components
236        Component {
237            id: "ml-stack-core".into(),
238            name: "ML Stack Core".into(),
239            description: "Core ML Stack components".into(),
240            script: String::new(), // Native Rust installer
241            category: Core,
242            required: false,
243            selected: false,
244            installed: false,
245            progress: 0.0,
246            estimate: "15-20 min".into(),
247            needs_sudo: true,
248        },
249        Component {
250            id: "flash-attn".into(),
251            name: "Flash Attention".into(),
252            description: "Efficient attention computation".into(),
253            script: String::new(), // Native Rust installer
254            category: Core,
255            required: false,
256            selected: true,
257            installed: false,
258            progress: 0.0,
259            estimate: "20-30 min".into(),
260            needs_sudo: true,
261        },
262        Component {
263            id: "repair-stack".into(),
264            name: "Repair ML Stack".into(),
265            description: "Repair ML Stack installation".into(),
266            script: String::new(), // Native Rust installer
267            category: Core,
268            required: false,
269            selected: false,
270            installed: false,
271            progress: 0.0,
272            estimate: "5-10 min".into(),
273            needs_sudo: true,
274        },
275        // Extension components
276        Component {
277            id: "megatron".into(),
278            name: "Megatron-LM".into(),
279            description: "Large-scale training framework".into(),
280            script: String::new(), // Native Rust installer
281            category: Extension,
282            required: false,
283            selected: false,
284            installed: false,
285            progress: 0.0,
286            estimate: "25-40 min".into(),
287            needs_sudo: true,
288        },
289        Component {
290            id: "vllm".into(),
291            name: "vLLM".into(),
292            description: "High-throughput inference engine".into(),
293            script: String::new(), // Native Rust installer
294            category: Extension,
295            required: false,
296            selected: false,
297            installed: false,
298            progress: 0.0,
299            estimate: "15-25 min".into(),
300            needs_sudo: true,
301        },
302        Component {
303            id: "aiter".into(),
304            name: "AITER".into(),
305            description: "AMD AITER optimization tooling".into(),
306            script: String::new(), // Native Rust installer
307            category: Extension,
308            required: false,
309            selected: false,
310            installed: false,
311            progress: 0.0,
312            estimate: "10-20 min".into(),
313            needs_sudo: true,
314        },
315        Component {
316            id: "vllm-studio".into(),
317            name: "vLLM Studio".into(),
318            description: "Model lifecycle manager for vLLM/SGLang".into(),
319            script: String::new(), // Native Rust installer
320            category: UiUx,
321            required: false,
322            selected: false,
323            installed: false,
324            progress: 0.0,
325            estimate: "10-20 min".into(),
326            needs_sudo: false,
327        },
328        Component {
329            id: "comfyui".into(),
330            name: "ComfyUI".into(),
331            description: "Node-based AI image generation UI with ROCm support".into(),
332            script: String::new(), // Native Rust installer
333            category: UiUx,
334            required: false,
335            selected: false,
336            installed: false,
337            progress: 0.0,
338            estimate: "5-10 min".into(),
339            needs_sudo: false,
340        },
341        Component {
342            id: "textgen".into(),
343            name: "text-generation-webui".into(),
344            description: "LLM chat/inference web UI with ROCm support (oobabooga)".into(),
345            script: String::new(), // Native Rust installer
346            category: UiUx,
347            required: false,
348            selected: false,
349            installed: false,
350            progress: 0.0,
351            estimate: "10-20 min".into(),
352            needs_sudo: false,
353        },
354        Component {
355            id: "onnx".into(),
356            name: "ONNX Runtime".into(),
357            description: "Cross-platform inference accelerator".into(),
358            script: String::new(), // Native Rust installer
359            category: Extension,
360            required: false,
361            selected: false,
362            installed: false,
363            progress: 0.0,
364            estimate: "2-5 min".into(),
365            needs_sudo: false,
366        },
367        Component {
368            id: "bitsandbytes".into(),
369            name: "BITSANDBYTES".into(),
370            description: "Efficient quantization for deep learning".into(),
371            script: String::new(), // Native Rust installer
372            category: Extension,
373            required: false,
374            selected: false,
375            installed: false,
376            progress: 0.0,
377            estimate: "3-5 min".into(),
378            needs_sudo: true,
379        },
380        Component {
381            id: "rocm-smi".into(),
382            name: "ROCm SMI".into(),
383            description: "System monitoring for AMD GPUs".into(),
384            script: String::new(), // Native Rust installer
385            category: Extension,
386            required: false,
387            selected: true,
388            installed: false,
389            progress: 0.0,
390            estimate: "2-3 min".into(),
391            needs_sudo: true,
392        },
393        Component {
394            id: "migraphx".into(),
395            name: "MIGraphX".into(),
396            description: "AMD graph inference engine".into(),
397            script: String::new(), // Native Rust installer
398            category: Extension,
399            required: false,
400            selected: false,
401            installed: false,
402            progress: 0.0,
403            estimate: "10-15 min".into(),
404            needs_sudo: true,
405        },
406        Component {
407            id: "pytorch-profiler".into(),
408            name: "PyTorch Profiler".into(),
409            description: "Performance analysis for PyTorch".into(),
410            script: String::new(), // Native Rust installer
411            category: Extension,
412            required: false,
413            selected: false,
414            installed: false,
415            progress: 0.0,
416            estimate: "3-5 min".into(),
417            needs_sudo: true,
418        },
419        Component {
420            id: "wandb".into(),
421            name: "Weights & Biases".into(),
422            description: "Experiment tracking and visualization".into(),
423            script: String::new(), // Native Rust installer
424            category: Extension,
425            required: false,
426            selected: false,
427            installed: false,
428            progress: 0.0,
429            estimate: "5-8 min".into(),
430            needs_sudo: true,
431        },
432        Component {
433            id: "llama-cpp".into(),
434            name: "llama.cpp (HIP/ROCm)".into(),
435            description: "llama.cpp with HIP/ROCm GPU acceleration".into(),
436            script: String::new(), // Native Rust installer
437            category: Extension,
438            required: false,
439            selected: false,
440            installed: false,
441            progress: 0.0,
442            estimate: "10-15 min".into(),
443            needs_sudo: true,
444        },
445        Component {
446            id: "fastvideo".into(),
447            name: "FastVideo".into(),
448            description: "Video generation framework with ROCm gfx11 optimizations".into(),
449            script: String::new(), // Native Rust installer
450            category: Extension,
451            required: false,
452            selected: false,
453            installed: false,
454            progress: 0.0,
455            estimate: "10-20 min".into(),
456            needs_sudo: false,
457        },
458        // Verification components (native Rust via verification module)
459        Component {
460            id: "verify-basic".into(),
461            name: "Verify Installation".into(),
462            description: "Basic verification".into(),
463            script: String::new(), // Native Rust verification
464            category: Verification,
465            required: false,
466            selected: false,
467            installed: false,
468            progress: 0.0,
469            estimate: "2-5 min".into(),
470            needs_sudo: true,
471        },
472        Component {
473            id: "verify-enhanced".into(),
474            name: "Enhanced Verify Installation".into(),
475            description: "Advanced verification".into(),
476            script: String::new(), // Native Rust verification
477            category: Verification,
478            required: false,
479            selected: false,
480            installed: false,
481            progress: 0.0,
482            estimate: "5-10 min".into(),
483            needs_sudo: true,
484        },
485        Component {
486            id: "verify-build".into(),
487            name: "Verify and Build".into(),
488            description: "Verify + build components".into(),
489            script: String::new(), // Native Rust verification
490            category: Verification,
491            required: false,
492            selected: false,
493            installed: false,
494            progress: 0.0,
495            estimate: "10-15 min".into(),
496            needs_sudo: true,
497        },
498        // Performance components (native Rust via benchmark_runners module)
499        Component {
500            id: "mlperf-inference".into(),
501            name: "MLPerf Inference".into(),
502            description: "MLPerf benchmark suite for inference performance".into(),
503            script: String::new(), // Native Rust benchmark
504            category: Performance,
505            required: false,
506            selected: false,
507            installed: false,
508            progress: 0.0,
509            estimate: "30-60 min".into(),
510            needs_sudo: false,
511        },
512        Component {
513            id: "rocm-benchmarks".into(),
514            name: "ROCm Benchmarks".into(),
515            description: "ROCm-specific performance benchmarks".into(),
516            script: String::new(), // Native Rust benchmark
517            category: Performance,
518            required: false,
519            selected: false,
520            installed: false,
521            progress: 0.0,
522            estimate: "20-40 min".into(),
523            needs_sudo: false,
524        },
525        Component {
526            id: "gpu-memory-bandwidth".into(),
527            name: "GPU Memory Bandwidth".into(),
528            description: "Memory bandwidth performance testing".into(),
529            script: String::new(), // Native Rust benchmark
530            category: Performance,
531            required: false,
532            selected: false,
533            installed: false,
534            progress: 0.0,
535            estimate: "10-15 min".into(),
536            needs_sudo: false,
537        },
538        Component {
539            id: "rocm-smi-bench".into(),
540            name: "ROCm SMI Benchmarks".into(),
541            description: "ROCm SMI performance monitoring".into(),
542            script: String::new(), // Native Rust benchmark
543            category: Performance,
544            required: false,
545            selected: false,
546            installed: false,
547            progress: 0.0,
548            estimate: "15-25 min".into(),
549            needs_sudo: false,
550        },
551        Component {
552            id: "vllm-performance".into(),
553            name: "vLLM Performance".into(),
554            description: "High-throughput vLLM inference benchmark".into(),
555            script: String::new(), // Native Rust benchmark
556            category: Performance,
557            required: false,
558            selected: false,
559            installed: false,
560            progress: 0.0,
561            estimate: "15-25 min".into(),
562            needs_sudo: false,
563        },
564        Component {
565            id: "deepspeed-performance".into(),
566            name: "DeepSpeed Performance".into(),
567            description: "DeepSpeed ZeRO optimization throughput benchmark".into(),
568            script: String::new(), // Native Rust benchmark
569            category: Performance,
570            required: false,
571            selected: false,
572            installed: false,
573            progress: 0.0,
574            estimate: "10-20 min".into(),
575            needs_sudo: false,
576        },
577        Component {
578            id: "megatron-performance".into(),
579            name: "Megatron-LM Performance".into(),
580            description: "Megatron-LM import and throughput benchmark".into(),
581            script: String::new(), // Native Rust benchmark
582            category: Performance,
583            required: false,
584            selected: false,
585            installed: false,
586            progress: 0.0,
587            estimate: "10-20 min".into(),
588            needs_sudo: false,
589        },
590        Component {
591            id: "all-benchmarks".into(),
592            name: "Full Suite Benchmark".into(),
593            description: "Run all post-installation performance tests".into(),
594            script: String::new(), // Native Rust benchmark
595            category: Performance,
596            required: false,
597            selected: false,
598            installed: false,
599            progress: 0.0,
600            estimate: "45-90 min".into(),
601            needs_sudo: false,
602        },
603    ]
604}