rusty_stack/installers/components/
vllm_multi.rs

1//! vLLM installer — ports `scripts/install_vllm_multi.sh`.
2//!
3//! Constructs correct git clone URL/branch + pip install commands for vLLM
4//! with ROCm support. vLLM depends on PyTorch (validated via preflight check).
5//!
6//! # Validation Assertions
7//!
8//! - **VAL-INSTALL-015**: vLLM correct git clone and pip install
9//! - **VAL-INSTALL-042**: vLLM declares dependency on PyTorch
10
11use crate::installers::common::RocmEnv;
12
13// ===========================================================================
14// Types
15// ===========================================================================
16
17/// Installation method for vLLM.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum InstallMethod {
20    /// Install globally.
21    Global,
22    /// Install in a virtual environment.
23    Venv,
24    /// Try global, fallback to venv.
25    Auto,
26}
27
28impl std::fmt::Display for InstallMethod {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        match self {
31            InstallMethod::Global => write!(f, "global"),
32            InstallMethod::Venv => write!(f, "venv"),
33            InstallMethod::Auto => write!(f, "auto"),
34        }
35    }
36}
37
38/// Configuration for the vLLM installer.
39#[derive(Debug, Clone)]
40pub struct VllmConfig {
41    /// ROCm version string (e.g., "7.2.0").
42    pub rocm_version: String,
43    /// GPU architecture (e.g., "gfx1100").
44    pub gpu_arch: String,
45    /// Python binary to use.
46    pub python_bin: String,
47    /// Installation method.
48    pub method: InstallMethod,
49    /// Whether to run in dry-run mode.
50    pub dry_run: bool,
51    /// Whether strict ROCm mode is enabled.
52    pub strict_rocm: bool,
53}
54
55impl Default for VllmConfig {
56    fn default() -> Self {
57        Self {
58            rocm_version: "7.2.0".to_string(),
59            gpu_arch: "gfx1100".to_string(),
60            python_bin: "python3".to_string(),
61            method: InstallMethod::Auto,
62            dry_run: false,
63            strict_rocm: true,
64        }
65    }
66}
67
68/// A constructed shell command.
69#[derive(Debug, Clone, PartialEq, Eq)]
70pub struct ShellCommand {
71    /// The program to run.
72    pub program: String,
73    /// Arguments to pass.
74    pub args: Vec<String>,
75    /// Environment variables to set.
76    pub env: Vec<(String, String)>,
77}
78
79impl ShellCommand {
80    /// Format as a shell command string.
81    pub fn to_command_string(&self) -> String {
82        let env_prefix = self
83            .env
84            .iter()
85            .map(|(k, v)| format!("{k}={v}"))
86            .collect::<Vec<_>>()
87            .join(" ");
88        let cmd = if self.args.is_empty() {
89            self.program.clone()
90        } else {
91            format!("{} {}", self.program, self.args.join(" "))
92        };
93        if env_prefix.is_empty() {
94            cmd
95        } else {
96            format!("{env_prefix} {cmd}")
97        }
98    }
99}
100
101/// The vLLM installer.
102pub struct VllmInstaller {
103    config: VllmConfig,
104}
105
106impl VllmInstaller {
107    /// Create a new vLLM installer with the given config.
108    pub fn new(config: VllmConfig) -> Self {
109        Self { config }
110    }
111
112    /// Create with default config.
113    pub fn with_defaults() -> Self {
114        Self::new(VllmConfig::default())
115    }
116
117    // -----------------------------------------------------------------------
118    // Dependencies (VAL-INSTALL-042)
119    // -----------------------------------------------------------------------
120
121    /// vLLM depends on PyTorch.
122    pub fn dependencies(&self) -> &[&str] {
123        &["pytorch"]
124    }
125
126    /// Validate that all dependencies are satisfied.
127    pub fn validate_dependencies(&self, installed_components: &[&str]) -> anyhow::Result<()> {
128        for dep in self.dependencies() {
129            if !installed_components.contains(dep) {
130                anyhow::bail!("vLLM requires '{}' to be installed first", dep);
131            }
132        }
133        Ok(())
134    }
135
136    // -----------------------------------------------------------------------
137    // ROCm detection helpers
138    // -----------------------------------------------------------------------
139
140    /// Detect ROCm major.minor version.
141    pub fn detect_rocm_mm(&self) -> String {
142        let v = &self.config.rocm_version;
143        // Extract major.minor from version string
144        let parts: Vec<&str> = v.split('.').collect();
145        if parts.len() >= 2 {
146            format!("{}.{}", parts[0], parts[1])
147        } else {
148            "7.2".to_string()
149        }
150    }
151
152    /// Get the ROCm PyTorch wheel index URL.
153    pub fn rocm_index_url(&self) -> String {
154        let mm = self.detect_rocm_mm();
155        format!("https://repo.radeon.com/rocm/manylinux/rocm-rel-{mm}/")
156    }
157
158    /// Get the vLLM ROCm wheels extra index URL.
159    pub fn vllm_wheels_url(&self) -> &'static str {
160        "https://wheels.vllm.ai/rocm/"
161    }
162
163    // -----------------------------------------------------------------------
164    // Command construction (VAL-INSTALL-015)
165    // -----------------------------------------------------------------------
166
167    /// Construct the ROCm build environment variables.
168    pub fn rocm_build_env(&self, rocm_env: &RocmEnv) -> Vec<(String, String)> {
169        let rocm_path = rocm_env
170            .path()
171            .map(|p| p.to_string_lossy().to_string())
172            .unwrap_or_else(|| "/opt/rocm".to_string());
173
174        vec![
175            ("ROCM_HOME".to_string(), rocm_path.clone()),
176            ("ROCM_PATH".to_string(), rocm_path.clone()),
177            ("HIP_PATH".to_string(), rocm_path.clone()),
178            ("HIP_ROOT_DIR".to_string(), rocm_path),
179            (
180                "PYTORCH_ROCM_ARCH".to_string(),
181                self.config.gpu_arch.clone(),
182            ),
183            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
184            ("VLLM_TARGET_DEVICE".to_string(), "rocm".to_string()),
185        ]
186    }
187
188    /// Construct the pip install command for vLLM from ROCm wheels.
189    ///
190    /// The original script uses:
191    /// `pip install --no-cache-dir --no-deps vllm --extra-index-url https://wheels.vllm.ai/rocm/`
192    pub fn build_vllm_install_command(&self) -> ShellCommand {
193        let use_break = self.config.method == InstallMethod::Global
194            || self.config.method == InstallMethod::Auto;
195
196        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
197        if use_break {
198            args.push("--break-system-packages".to_string());
199        }
200        args.extend([
201            "--no-cache-dir".to_string(),
202            "--no-deps".to_string(),
203            "vllm".to_string(),
204            "--extra-index-url".to_string(),
205            self.vllm_wheels_url().to_string(),
206        ]);
207
208        ShellCommand {
209            program: self.config.python_bin.clone(),
210            args,
211            env: vec![],
212        }
213    }
214
215    /// Construct the pip install command for vLLM dependencies.
216    ///
217    /// Installs the core dependencies needed by vLLM (excluding torch and xformers).
218    pub fn build_deps_install_command(&self) -> ShellCommand {
219        let use_break = self.config.method == InstallMethod::Global
220            || self.config.method == InstallMethod::Auto;
221
222        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
223        if use_break {
224            args.push("--break-system-packages".to_string());
225        }
226        args.extend([
227            "--no-cache-dir".to_string(),
228            "--extra-index-url".to_string(),
229            self.vllm_wheels_url().to_string(),
230        ]);
231
232        // Core vLLM dependencies
233        let deps = [
234            "accelerate",
235            "aiohttp",
236            "cloudpickle",
237            "fastapi",
238            "msgspec",
239            "prometheus-client",
240            "psutil",
241            "py-cpuinfo",
242            "pyzmq",
243            "requests",
244            "sentencepiece",
245            "tiktoken",
246            "uvicorn",
247            "einops",
248            "transformers",
249            "huggingface-hub",
250            "cachetools",
251            "cbor2",
252            "gguf",
253            "pybase64",
254            "ijson",
255            "python-json-logger",
256            "setproctitle",
257            "watchfiles",
258            "six",
259            "openai",
260            "blake3",
261            "lark",
262            "amdsmi",
263            "lm-format-enforcer",
264            "partial-json-parser",
265            "prometheus-fastapi-instrumentator",
266            "datasets",
267            "diskcache",
268            "timm",
269            "peft",
270            "numba",
271        ];
272        for dep in deps {
273            args.push(dep.to_string());
274        }
275
276        ShellCommand {
277            program: self.config.python_bin.clone(),
278            args,
279            env: vec![],
280        }
281    }
282
283    /// Construct the pip install command for vLLM versioned dependencies.
284    ///
285    /// These have specific version pins from the original script.
286    pub fn build_versioned_deps_command(&self) -> ShellCommand {
287        let use_break = self.config.method == InstallMethod::Global
288            || self.config.method == InstallMethod::Auto;
289
290        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
291        if use_break {
292            args.push("--break-system-packages".to_string());
293        }
294        args.extend([
295            "--no-cache-dir".to_string(),
296            "--extra-index-url".to_string(),
297            self.vllm_wheels_url().to_string(),
298        ]);
299
300        // Versioned deps from original script
301        let deps = [
302            "openai-harmony>=0.0.3",
303            "mistral-common[image]>=1.9.0",
304            "triton-kernels==1.0.0",
305            "outlines-core==0.2.11",
306            "xgrammar==0.1.29",
307            "llguidance>=1.3.0,<1.4.0",
308        ];
309        for dep in deps {
310            args.push(dep.to_string());
311        }
312
313        ShellCommand {
314            program: self.config.python_bin.clone(),
315            args,
316            env: vec![],
317        }
318    }
319
320    /// Construct the Triton cache environment setup command.
321    ///
322    /// Sets up writable cache directories for Triton kernel compilation.
323    pub fn triton_cache_env(&self) -> Vec<(String, String)> {
324        let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
325        let triton_home = format!("{home}/.cache/mlstack/triton");
326
327        vec![
328            ("MLSTACK_TRITON_HOME".to_string(), triton_home.clone()),
329            ("TRITON_HOME".to_string(), triton_home.clone()),
330            (
331                "TRITON_CACHE_DIR".to_string(),
332                format!("{triton_home}/cache"),
333            ),
334            ("TRITON_DUMP_DIR".to_string(), format!("{triton_home}/dump")),
335            (
336                "TRITON_OVERRIDE_DIR".to_string(),
337                format!("{triton_home}/override"),
338            ),
339        ]
340    }
341
342    /// Construct the force-reinstall command for vLLM wheel repair.
343    pub fn build_force_reinstall_command(&self) -> ShellCommand {
344        let args = vec![
345            "-m".to_string(),
346            "pip".to_string(),
347            "install".to_string(),
348            "--break-system-packages".to_string(),
349            "--no-cache-dir".to_string(),
350            "--force-reinstall".to_string(),
351            "--no-deps".to_string(),
352            "--extra-index-url".to_string(),
353            self.vllm_wheels_url().to_string(),
354            "vllm".to_string(),
355        ];
356
357        ShellCommand {
358            program: self.config.python_bin.clone(),
359            args,
360            env: vec![],
361        }
362    }
363
364    /// Construct the source rebuild command for vLLM.
365    ///
366    /// Used as a fallback when wheel install fails.
367    pub fn build_source_rebuild_command(&self) -> ShellCommand {
368        let args = vec![
369            "-m".to_string(),
370            "pip".to_string(),
371            "install".to_string(),
372            "--break-system-packages".to_string(),
373            "--no-cache-dir".to_string(),
374            "--force-reinstall".to_string(),
375            "--no-deps".to_string(),
376            "--no-build-isolation".to_string(),
377            "--no-binary".to_string(),
378            "vllm".to_string(),
379        ];
380
381        ShellCommand {
382            program: self.config.python_bin.clone(),
383            args,
384            env: vec![
385                ("VLLM_TARGET_DEVICE".to_string(), "rocm".to_string()),
386                ("VLLM_USE_ROCM".to_string(), "1".to_string()),
387                ("USE_ROCM".to_string(), "1".to_string()),
388            ],
389        }
390    }
391}
392
393// ===========================================================================
394// Tests
395// ===========================================================================
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400    use std::path::PathBuf;
401
402    // --- VAL-INSTALL-015: vLLM correct git clone and pip install ---
403
404    #[test]
405    fn test_vllm_install_command() {
406        let installer = VllmInstaller::new(VllmConfig {
407            method: InstallMethod::Global,
408            ..Default::default()
409        });
410        let cmd = installer.build_vllm_install_command();
411        assert_eq!(cmd.program, "python3");
412        assert!(cmd.args.contains(&"-m".to_string()));
413        assert!(cmd.args.contains(&"pip".to_string()));
414        assert!(cmd.args.contains(&"install".to_string()));
415        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
416        assert!(cmd.args.contains(&"--no-cache-dir".to_string()));
417        assert!(cmd.args.contains(&"--no-deps".to_string()));
418        assert!(cmd.args.contains(&"vllm".to_string()));
419        assert!(cmd.args.contains(&"--extra-index-url".to_string()));
420        assert!(cmd.args.iter().any(|a| a == "https://wheels.vllm.ai/rocm/"));
421    }
422
423    #[test]
424    fn test_vllm_install_command_venv() {
425        let installer = VllmInstaller::new(VllmConfig {
426            method: InstallMethod::Venv,
427            ..Default::default()
428        });
429        let cmd = installer.build_vllm_install_command();
430        // Venv should NOT have --break-system-packages
431        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
432        assert!(cmd.args.contains(&"vllm".to_string()));
433    }
434
435    #[test]
436    fn test_vllm_wheels_url() {
437        let installer = VllmInstaller::with_defaults();
438        assert_eq!(installer.vllm_wheels_url(), "https://wheels.vllm.ai/rocm/");
439    }
440
441    #[test]
442    fn test_rocm_index_url() {
443        let installer = VllmInstaller::new(VllmConfig {
444            rocm_version: "7.2.0".to_string(),
445            ..Default::default()
446        });
447        assert_eq!(
448            installer.rocm_index_url(),
449            "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/"
450        );
451    }
452
453    #[test]
454    fn test_rocm_index_url_legacy_version() {
455        let installer = VllmInstaller::new(VllmConfig {
456            rocm_version: "7.0.0".to_string(),
457            ..Default::default()
458        });
459        assert_eq!(
460            installer.rocm_index_url(),
461            "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0/"
462        );
463    }
464
465    #[test]
466    fn test_detect_rocm_mm() {
467        let installer = VllmInstaller::new(VllmConfig {
468            rocm_version: "7.2.0".to_string(),
469            ..Default::default()
470        });
471        assert_eq!(installer.detect_rocm_mm(), "7.2");
472    }
473
474    #[test]
475    fn test_rocm_build_env() {
476        let installer = VllmInstaller::new(VllmConfig {
477            gpu_arch: "gfx1100".to_string(),
478            ..Default::default()
479        });
480        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
481        let env = installer.rocm_build_env(&rocm_env);
482        assert!(env
483            .iter()
484            .any(|(k, v)| k == "ROCM_HOME" && v == "/opt/rocm"));
485        assert!(env
486            .iter()
487            .any(|(k, v)| k == "ROCM_PATH" && v == "/opt/rocm"));
488        assert!(env.iter().any(|(k, v)| k == "HIP_PATH" && v == "/opt/rocm"));
489        assert!(env
490            .iter()
491            .any(|(k, v)| k == "PYTORCH_ROCM_ARCH" && v == "gfx1100"));
492        assert!(env
493            .iter()
494            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
495        assert!(env
496            .iter()
497            .any(|(k, v)| k == "VLLM_TARGET_DEVICE" && v == "rocm"));
498    }
499
500    #[test]
501    fn test_deps_install_command() {
502        let installer = VllmInstaller::new(VllmConfig {
503            method: InstallMethod::Global,
504            ..Default::default()
505        });
506        let cmd = installer.build_deps_install_command();
507        assert!(cmd.args.contains(&"accelerate".to_string()));
508        assert!(cmd.args.contains(&"transformers".to_string()));
509        assert!(cmd.args.contains(&"einops".to_string()));
510        assert!(cmd.args.contains(&"fastapi".to_string()));
511        assert!(cmd.args.contains(&"amdsmi".to_string()));
512    }
513
514    #[test]
515    fn test_versioned_deps_command() {
516        let installer = VllmInstaller::with_defaults();
517        let cmd = installer.build_versioned_deps_command();
518        assert!(cmd.args.iter().any(|a| a.starts_with("triton-kernels==")));
519        assert!(cmd.args.iter().any(|a| a.starts_with("xgrammar==")));
520        assert!(cmd.args.iter().any(|a| a.starts_with("outlines-core==")));
521    }
522
523    #[test]
524    fn test_force_reinstall_command() {
525        let installer = VllmInstaller::with_defaults();
526        let cmd = installer.build_force_reinstall_command();
527        assert!(cmd.args.contains(&"--force-reinstall".to_string()));
528        assert!(cmd.args.contains(&"--no-deps".to_string()));
529        assert!(cmd.args.contains(&"vllm".to_string()));
530    }
531
532    #[test]
533    fn test_source_rebuild_command() {
534        let installer = VllmInstaller::with_defaults();
535        let cmd = installer.build_source_rebuild_command();
536        assert!(cmd.args.contains(&"--no-build-isolation".to_string()));
537        assert!(cmd.args.contains(&"--no-binary".to_string()));
538        assert!(cmd
539            .env
540            .iter()
541            .any(|(k, v)| k == "VLLM_TARGET_DEVICE" && v == "rocm"));
542        assert!(cmd
543            .env
544            .iter()
545            .any(|(k, v)| k == "VLLM_USE_ROCM" && v == "1"));
546    }
547
548    #[test]
549    fn test_triton_cache_env() {
550        let installer = VllmInstaller::with_defaults();
551        let env = installer.triton_cache_env();
552        assert!(env.iter().any(|(k, _)| k == "TRITON_HOME"));
553        assert!(env.iter().any(|(k, _)| k == "TRITON_CACHE_DIR"));
554        assert!(env.iter().any(|(k, _)| k == "TRITON_DUMP_DIR"));
555    }
556
557    // --- VAL-INSTALL-042: vLLM declares dependency on PyTorch ---
558
559    #[test]
560    fn test_dependencies() {
561        let installer = VllmInstaller::with_defaults();
562        assert!(installer.dependencies().contains(&"pytorch"));
563    }
564
565    #[test]
566    fn test_validate_dependencies_success() {
567        let installer = VllmInstaller::with_defaults();
568        assert!(installer
569            .validate_dependencies(&["pytorch", "rocm"])
570            .is_ok());
571    }
572
573    #[test]
574    fn test_validate_dependencies_missing_pytorch() {
575        let installer = VllmInstaller::with_defaults();
576        let result = installer.validate_dependencies(&["rocm"]);
577        assert!(result.is_err());
578        assert!(result.unwrap_err().to_string().contains("pytorch"));
579    }
580}
rusty_stack/installers/components/vllm_multi.rs

rusty_stack/installers/components/
vllm_multi.rs
