rusty_stack/installers/components/
megatron.rs

1//! Megatron-LM installer — ports `scripts/install_megatron.sh`.
2//!
3//! Constructs correct git clone URL/branch and pip install commands.
4//! Megatron depends on PyTorch (ROCm-enabled) and MPI4Py.
5//!
6//! # Validation Assertions
7//!
8//! - **VAL-INSTALL-012**: Megatron correct git clone
9//! - **VAL-INSTALL-013**: Megatron correct pip install post-clone
10//! - **VAL-INSTALL-041**: Megatron declares dependency on PyTorch and MPI4Py
11//! - **VAL-INSTALL-043**: DeepSpeed declares dependency on PyTorch
12
13use crate::installers::common::RocmEnv;
14use std::fmt;
15use std::path::PathBuf;
16
17// ===========================================================================
18// Types
19// ===========================================================================
20
21/// Installation method.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum InstallMethod {
24    /// Install globally.
25    Global,
26    /// Install in a virtual environment.
27    Venv,
28    /// Try global, fallback to venv.
29    Auto,
30}
31
32impl fmt::Display for InstallMethod {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        match self {
35            InstallMethod::Global => write!(f, "global"),
36            InstallMethod::Venv => write!(f, "venv"),
37            InstallMethod::Auto => write!(f, "auto"),
38        }
39    }
40}
41
42/// Configuration for the Megatron-LM installer.
43#[derive(Debug, Clone)]
44pub struct MegatronConfig {
45    /// Installation method.
46    pub method: InstallMethod,
47    /// Python binary to use.
48    pub python_bin: String,
49    /// Whether to force reinstall.
50    pub force_reinstall: bool,
51    /// Whether to run in dry-run mode.
52    pub dry_run: bool,
53    /// Installation directory (defaults to `$HOME/Megatron-LM`).
54    pub install_dir: Option<PathBuf>,
55}
56
57impl Default for MegatronConfig {
58    fn default() -> Self {
59        Self {
60            method: InstallMethod::Auto,
61            python_bin: "python3".to_string(),
62            force_reinstall: false,
63            dry_run: false,
64            install_dir: None,
65        }
66    }
67}
68
69impl MegatronConfig {
70    /// Get the installation directory path.
71    ///
72    /// The original script uses `$HOME/Megatron-LM` as primary,
73    /// with `$HOME/.mlstack/src/Megatron-LM` as fallback.
74    pub fn install_dir(&self) -> PathBuf {
75        self.install_dir
76            .clone()
77            .unwrap_or_else(|| dirs_home().join("Megatron-LM"))
78    }
79
80    /// Get the fallback installation directory.
81    pub fn fallback_install_dir(&self) -> PathBuf {
82        dirs_home().join(".mlstack/src/Megatron-LM")
83    }
84}
85
86/// A constructed shell command.
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct ShellCommand {
89    /// The program to run.
90    pub program: String,
91    /// Arguments to pass.
92    pub args: Vec<String>,
93    /// Environment variables to set.
94    pub env: Vec<(String, String)>,
95    /// Working directory for the command.
96    pub working_dir: Option<PathBuf>,
97}
98
99impl ShellCommand {
100    /// Format as a shell command string.
101    pub fn to_command_string(&self) -> String {
102        let env_prefix = self
103            .env
104            .iter()
105            .map(|(k, v)| format!("{k}={v}"))
106            .collect::<Vec<_>>()
107            .join(" ");
108        let cmd = if self.args.is_empty() {
109            self.program.clone()
110        } else {
111            format!("{} {}", self.program, self.args.join(" "))
112        };
113        if env_prefix.is_empty() {
114            cmd
115        } else {
116            format!("{env_prefix} {cmd}")
117        }
118    }
119}
120
121/// The Megatron-LM installer.
122pub struct MegatronInstaller {
123    config: MegatronConfig,
124}
125
126impl MegatronInstaller {
127    /// Create a new Megatron installer with the given config.
128    pub fn new(config: MegatronConfig) -> Self {
129        Self { config }
130    }
131
132    /// Create with default config.
133    pub fn with_defaults() -> Self {
134        Self::new(MegatronConfig::default())
135    }
136
137    // -----------------------------------------------------------------------
138    // Dependencies (VAL-INSTALL-041)
139    // -----------------------------------------------------------------------
140
141    /// Get the list of required dependencies.
142    ///
143    /// Megatron-LM depends on:
144    /// - PyTorch (ROCm-enabled)
145    /// - MPI4Py
146    pub fn dependencies(&self) -> &[&str] {
147        &["pytorch", "mpi4py"]
148    }
149
150    // -----------------------------------------------------------------------
151    // ROCm environment
152    // -----------------------------------------------------------------------
153
154    /// Get the ROCm environment variables needed for Megatron.
155    ///
156    /// The original script sets:
157    /// - AMD_LOG_LEVEL=0
158    /// - HIP_VISIBLE_DEVICES
159    /// - ROCR_VISIBLE_DEVICES
160    /// - HSA_OVERRIDE_GFX_VERSION=11.0.0
161    /// - PYTORCH_ROCM_ARCH=gfx1100
162    /// - ROCM_PATH
163    /// - HSA_TOOLS_LIB
164    /// - PYTORCH_ALLOC_CONF=expandable_segments:True
165    pub fn rocm_env(&self, rocm_env: &RocmEnv) -> Vec<(String, String)> {
166        let rocm_path = rocm_env
167            .path()
168            .map(|p| p.to_string_lossy().to_string())
169            .unwrap_or_else(|| "/opt/rocm".to_string());
170
171        let mut env = vec![
172            ("AMD_LOG_LEVEL".to_string(), "0".to_string()),
173            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
174            ("PYTORCH_ROCM_ARCH".to_string(), "gfx1100".to_string()),
175            ("ROCM_PATH".to_string(), rocm_path.clone()),
176            (
177                "PYTORCH_ALLOC_CONF".to_string(),
178                "expandable_segments:True".to_string(),
179            ),
180        ];
181
182        // HSA_TOOLS_LIB - check for rocprofiler library
183        // Try the ROCm 7.x layout first (lib/rocprofiler-sdk/), then the old layout (lib/)
184        let profiler_lib = rocm_env.path().and_then(|p| {
185            let new_layout = p.join("lib/rocprofiler-sdk/librocprofiler-sdk-tool.so");
186            let old_layout = p.join("lib/librocprofiler-sdk-tool.so");
187            if new_layout.exists() {
188                Some(new_layout)
189            } else if old_layout.exists() {
190                Some(old_layout)
191            } else {
192                None
193            }
194        });
195        if let Some(lib) = profiler_lib {
196            env.push((
197                "HSA_TOOLS_LIB".to_string(),
198                lib.to_string_lossy().to_string(),
199            ));
200        } else {
201            env.push(("HSA_TOOLS_LIB".to_string(), "0".to_string()));
202        }
203
204        env
205    }
206
207    // -----------------------------------------------------------------------
208    // Git clone command (VAL-INSTALL-012)
209    // -----------------------------------------------------------------------
210
211    /// Construct the git clone command for Megatron-LM.
212    ///
213    /// The original script clones:
214    /// `git clone https://github.com/NVIDIA/Megatron-LM.git <dir>`
215    ///
216    /// It retries up to 3 times on failure.
217    pub fn build_git_clone_command(&self) -> ShellCommand {
218        ShellCommand {
219            program: "git".to_string(),
220            args: vec![
221                "clone".to_string(),
222                "https://github.com/NVIDIA/Megatron-LM.git".to_string(),
223                self.config.install_dir().to_string_lossy().to_string(),
224            ],
225            env: vec![],
226            working_dir: None,
227        }
228    }
229
230    /// Get the maximum number of clone retries.
231    pub fn max_clone_retries(&self) -> u32 {
232        3
233    }
234
235    // -----------------------------------------------------------------------
236    // Pip install command (VAL-INSTALL-013)
237    // -----------------------------------------------------------------------
238
239    /// Construct the pip install command for Megatron-LM.
240    ///
241    /// The original script installs with:
242    /// `pip install -e . --no-deps`
243    ///
244    /// The editable install is done from the cloned directory.
245    /// Dependencies are resolved separately (safe requirements filtering).
246    pub fn build_pip_install_command(&self) -> ShellCommand {
247        let is_global = self.config.method == InstallMethod::Global
248            || self.config.method == InstallMethod::Auto;
249
250        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
251
252        if is_global {
253            args.push("--break-system-packages".to_string());
254        }
255
256        args.push("-e".to_string());
257        args.push(".".to_string());
258        args.push("--no-deps".to_string());
259
260        ShellCommand {
261            program: self.config.python_bin.clone(),
262            args,
263            env: vec![],
264            working_dir: Some(self.config.install_dir()),
265        }
266    }
267
268    /// Construct the pip install command for a specific dependency.
269    ///
270    /// Used for installing individual safe dependencies.
271    pub fn build_dep_install_command(&self, package: &str) -> ShellCommand {
272        let is_global = self.config.method == InstallMethod::Global
273            || self.config.method == InstallMethod::Auto;
274
275        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
276
277        if is_global {
278            args.push("--break-system-packages".to_string());
279        }
280
281        args.push(package.to_string());
282
283        ShellCommand {
284            program: self.config.python_bin.clone(),
285            args,
286            env: vec![],
287            working_dir: None,
288        }
289    }
290
291    /// Construct a pip install command for a requirements file.
292    pub fn build_requirements_install_command(&self, requirements_path: &str) -> ShellCommand {
293        let is_global = self.config.method == InstallMethod::Global
294            || self.config.method == InstallMethod::Auto;
295
296        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
297
298        if is_global {
299            args.push("--break-system-packages".to_string());
300        }
301
302        args.push("-r".to_string());
303        args.push(requirements_path.to_string());
304
305        ShellCommand {
306            program: self.config.python_bin.clone(),
307            args,
308            env: vec![],
309            working_dir: None,
310        }
311    }
312
313    // -----------------------------------------------------------------------
314    // Safe dependency filtering
315    // -----------------------------------------------------------------------
316
317    /// Check if a Python package is safe to install (non-CUDA/non-NVIDIA).
318    ///
319    /// The original script filters out packages that would conflict with
320    /// the ROCm environment: nvidia-*, cuda*, torch, torchvision, etc.
321    pub fn is_safe_package(&self, package_name: &str) -> bool {
322        let lower = package_name.to_lowercase();
323        let lower = lower.split(';').next().unwrap_or(&lower);
324        let lower = lower.split('[').next().unwrap_or(lower);
325        let pkg_name = lower
326            .split(&['<', '>', '=', '~', '!', ' '][..])
327            .next()
328            .unwrap_or(lower)
329            .trim();
330
331        let blocked_prefixes = [
332            "nvidia",
333            "cuda",
334            "cudnn",
335            "cublas",
336            "cufft",
337            "curand",
338            "cusolver",
339            "cusparse",
340            "nccl",
341            "nvtx",
342            "nvjitlink",
343            "tensorrt",
344        ];
345
346        let blocked_exact = [
347            "torch",
348            "torchvision",
349            "torchaudio",
350            "triton",
351            "xformers",
352            "pytorch-cuda",
353            "torch-cuda",
354        ];
355
356        if blocked_exact.contains(&pkg_name) {
357            return false;
358        }
359
360        for prefix in &blocked_prefixes {
361            if pkg_name.starts_with(prefix) {
362                return false;
363            }
364        }
365
366        // Check for patterns like cupy-cuda*
367        if pkg_name.contains("cuda") {
368            return false;
369        }
370
371        true
372    }
373
374    // -----------------------------------------------------------------------
375    // Build error detection (VAL-INSTALL-014)
376    // -----------------------------------------------------------------------
377
378    /// Check build output for errors.
379    pub fn check_build_output(
380        &self,
381        stdout: &str,
382        stderr: &str,
383        exit_code: i32,
384    ) -> Result<(), String> {
385        if exit_code != 0 {
386            let error_patterns = [
387                "error:",
388                "Error:",
389                "FAILED:",
390                "fatal error:",
391                "undefined reference",
392                "No such file or directory",
393                "Permission denied",
394                "ImportError:",
395                "ModuleNotFoundError:",
396            ];
397
398            let combined = format!("{}\n{}", stdout, stderr);
399            let detected_errors: Vec<&str> = error_patterns
400                .iter()
401                .filter(|p| combined.contains(*p))
402                .copied()
403                .collect();
404
405            if detected_errors.is_empty() {
406                return Err(format!(
407                    "Megatron operation failed with exit code {}",
408                    exit_code
409                ));
410            }
411
412            return Err(format!(
413                "Megatron operation failed with exit code {}. Detected errors: {}",
414                exit_code,
415                detected_errors.join(", ")
416            ));
417        }
418
419        Ok(())
420    }
421}
422
423/// Get the home directory.
424fn dirs_home() -> PathBuf {
425    std::env::var("HOME")
426        .map(PathBuf::from)
427        .unwrap_or_else(|_| PathBuf::from("/root"))
428}
429
430// ===========================================================================
431// Tests
432// ===========================================================================
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437    use std::path::PathBuf;
438
439    // --- VAL-INSTALL-012: Megatron correct git clone ---
440
441    #[test]
442    fn test_git_clone_command() {
443        let installer = MegatronInstaller::with_defaults();
444        let cmd = installer.build_git_clone_command();
445
446        assert_eq!(cmd.program, "git");
447        assert!(cmd.args.contains(&"clone".to_string()));
448        assert!(cmd.args.iter().any(|a| a.contains("NVIDIA/Megatron-LM")));
449    }
450
451    #[test]
452    fn test_git_clone_target_dir() {
453        let installer = MegatronInstaller::new(MegatronConfig {
454            install_dir: Some(PathBuf::from("/home/user/Megatron-LM")),
455            ..Default::default()
456        });
457        let cmd = installer.build_git_clone_command();
458        assert!(cmd
459            .args
460            .iter()
461            .any(|a| a.contains("/home/user/Megatron-LM")));
462    }
463
464    #[test]
465    fn test_git_clone_retries() {
466        let installer = MegatronInstaller::with_defaults();
467        assert_eq!(installer.max_clone_retries(), 3);
468    }
469
470    // --- VAL-INSTALL-013: Megatron correct pip install post-clone ---
471
472    #[test]
473    fn test_pip_install_command() {
474        let installer = MegatronInstaller::new(MegatronConfig {
475            method: InstallMethod::Global,
476            ..Default::default()
477        });
478        let cmd = installer.build_pip_install_command();
479
480        assert_eq!(cmd.program, "python3");
481        assert!(cmd.args.contains(&"-m".to_string()));
482        assert!(cmd.args.contains(&"pip".to_string()));
483        assert!(cmd.args.contains(&"install".to_string()));
484        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
485        assert!(cmd.args.contains(&"-e".to_string()));
486        assert!(cmd.args.contains(&".".to_string()));
487        assert!(cmd.args.contains(&"--no-deps".to_string()));
488    }
489
490    #[test]
491    fn test_pip_install_venv_no_break_system() {
492        let installer = MegatronInstaller::new(MegatronConfig {
493            method: InstallMethod::Venv,
494            ..Default::default()
495        });
496        let cmd = installer.build_pip_install_command();
497        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
498    }
499
500    #[test]
501    fn test_pip_install_working_dir() {
502        let installer = MegatronInstaller::new(MegatronConfig {
503            install_dir: Some(PathBuf::from("/home/user/Megatron-LM")),
504            ..Default::default()
505        });
506        let cmd = installer.build_pip_install_command();
507        assert_eq!(
508            cmd.working_dir,
509            Some(PathBuf::from("/home/user/Megatron-LM"))
510        );
511    }
512
513    #[test]
514    fn test_dep_install_command() {
515        let installer = MegatronInstaller::new(MegatronConfig {
516            method: InstallMethod::Global,
517            ..Default::default()
518        });
519        let cmd = installer.build_dep_install_command("numpy");
520        assert!(cmd.args.contains(&"numpy".to_string()));
521        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
522    }
523
524    #[test]
525    fn test_requirements_install_command() {
526        let installer = MegatronInstaller::new(MegatronConfig {
527            method: InstallMethod::Global,
528            ..Default::default()
529        });
530        let cmd = installer.build_requirements_install_command("/tmp/req.txt");
531        assert!(cmd.args.contains(&"-r".to_string()));
532        assert!(cmd.args.contains(&"/tmp/req.txt".to_string()));
533        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
534    }
535
536    // --- VAL-INSTALL-041: Megatron declares dependency on PyTorch and MPI4Py ---
537
538    #[test]
539    fn test_dependencies() {
540        let installer = MegatronInstaller::with_defaults();
541        let deps = installer.dependencies();
542        assert!(deps.contains(&"pytorch"));
543        assert!(deps.contains(&"mpi4py"));
544    }
545
546    // --- ROCm environment ---
547
548    #[test]
549    fn test_rocm_env() {
550        let installer = MegatronInstaller::with_defaults();
551        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
552        let env = installer.rocm_env(&rocm_env);
553        assert!(env.iter().any(|(k, v)| k == "AMD_LOG_LEVEL" && v == "0"));
554        assert!(env
555            .iter()
556            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
557        assert!(env
558            .iter()
559            .any(|(k, v)| k == "PYTORCH_ROCM_ARCH" && v == "gfx1100"));
560        assert!(env
561            .iter()
562            .any(|(k, v)| k == "ROCM_PATH" && v == "/opt/rocm"));
563        assert!(env
564            .iter()
565            .any(|(k, v)| k == "PYTORCH_ALLOC_CONF" && v == "expandable_segments:True"));
566        assert!(env.iter().any(|(k, _)| k == "HSA_TOOLS_LIB"));
567    }
568
569    // --- Safe package filtering ---
570
571    #[test]
572    fn test_safe_package_allows_megatron_deps() {
573        let installer = MegatronInstaller::with_defaults();
574        // These should be safe
575        assert!(installer.is_safe_package("numpy"));
576        assert!(installer.is_safe_package("regex"));
577        assert!(installer.is_safe_package("sentencepiece"));
578        assert!(installer.is_safe_package("pyyaml"));
579        assert!(installer.is_safe_package("transformers"));
580        assert!(installer.is_safe_package("tokenizers"));
581        assert!(installer.is_safe_package("tqdm"));
582        assert!(installer.is_safe_package("einops"));
583        assert!(installer.is_safe_package("packaging"));
584        assert!(installer.is_safe_package("omegaconf"));
585    }
586
587    #[test]
588    fn test_safe_package_blocks_cuda_nvidia() {
589        let installer = MegatronInstaller::with_defaults();
590        // These should be blocked
591        assert!(!installer.is_safe_package("torch"));
592        assert!(!installer.is_safe_package("torchvision"));
593        assert!(!installer.is_safe_package("torchaudio"));
594        assert!(!installer.is_safe_package("triton"));
595        assert!(!installer.is_safe_package("nvidia-cublas-cu12"));
596        assert!(!installer.is_safe_package("nvidia-cuda-runtime-cu12"));
597        assert!(!installer.is_safe_package("pytorch-cuda"));
598        assert!(!installer.is_safe_package("torch-cuda"));
599        assert!(!installer.is_safe_package("xformers"));
600        assert!(!installer.is_safe_package("cupy-cuda12x"));
601    }
602
603    #[test]
604    fn test_safe_package_with_version_spec() {
605        let installer = MegatronInstaller::with_defaults();
606        assert!(installer.is_safe_package("numpy>=1.21"));
607        assert!(installer.is_safe_package("packaging>=20.0"));
608        assert!(!installer.is_safe_package("torch>=2.0"));
609    }
610
611    // --- Build error detection ---
612
613    #[test]
614    fn test_build_error_import() {
615        let installer = MegatronInstaller::with_defaults();
616        let result = installer.check_build_output("", "ImportError: No module named 'megatron'", 1);
617        assert!(result.is_err());
618        assert!(result.unwrap_err().contains("ImportError"));
619    }
620
621    #[test]
622    fn test_build_error_success() {
623        let installer = MegatronInstaller::with_defaults();
624        let result = installer.check_build_output("OK", "", 0);
625        assert!(result.is_ok());
626    }
627
628    // --- Config defaults ---
629
630    #[test]
631    fn test_config_defaults() {
632        let config = MegatronConfig::default();
633        assert_eq!(config.method, InstallMethod::Auto);
634        assert_eq!(config.python_bin, "python3");
635        assert!(!config.force_reinstall);
636        assert!(!config.dry_run);
637        assert!(config.install_dir.is_none());
638    }
639
640    #[test]
641    fn test_config_install_dir_default() {
642        let config = MegatronConfig::default();
643        let dir = config.install_dir();
644        assert!(dir.to_string_lossy().contains("Megatron-LM"));
645    }
646
647    #[test]
648    fn test_config_fallback_dir() {
649        let config = MegatronConfig::default();
650        let fallback = config.fallback_install_dir();
651        assert!(fallback
652            .to_string_lossy()
653            .contains(".mlstack/src/Megatron-LM"));
654    }
655
656    #[test]
657    fn test_command_string_format() {
658        let cmd = ShellCommand {
659            program: "python3".to_string(),
660            args: vec![
661                "-m".to_string(),
662                "pip".to_string(),
663                "install".to_string(),
664                "numpy".to_string(),
665            ],
666            env: vec![("ROCM_PATH".to_string(), "/opt/rocm".to_string())],
667            working_dir: None,
668        };
669        let s = cmd.to_command_string();
670        assert!(s.contains("ROCM_PATH=/opt/rocm"));
671        assert!(s.contains("python3 -m pip install numpy"));
672    }
673}
rusty_stack/installers/components/megatron.rs

rusty_stack/installers/components/
megatron.rs
