rusty_stack/installers/components/
bitsandbytes_multi.rs

1//! bitsandbytes installer — ports `scripts/install_bitsandbytes_multi.sh`.
2//!
3//! Constructs correct pip install with version constraints for ROCm.
4//! Builds from ROCm fork when PyPI version lacks HIP binaries.
5//!
6//! # Validation Assertion
7//!
8//! - **VAL-INSTALL-017**: bitsandbytes correct pip command
9
10// ===========================================================================
11// Types
12// ===========================================================================
13
14/// Installation method.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum InstallMethod {
17    /// Install globally.
18    Global,
19    /// Install in a virtual environment.
20    Venv,
21    /// Try global, fallback to venv.
22    Auto,
23}
24
25impl std::fmt::Display for InstallMethod {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            InstallMethod::Global => write!(f, "global"),
29            InstallMethod::Venv => write!(f, "venv"),
30            InstallMethod::Auto => write!(f, "auto"),
31        }
32    }
33}
34
35/// Configuration for the bitsandbytes installer.
36#[derive(Debug, Clone)]
37pub struct BitsAndBytesConfig {
38    /// ROCm version string (e.g., "7.2.0").
39    pub rocm_version: String,
40    /// GPU architecture (e.g., "gfx1100").
41    pub gpu_arch: String,
42    /// Python binary to use.
43    pub python_bin: String,
44    /// Installation method.
45    pub method: InstallMethod,
46    /// Whether to run in dry-run mode.
47    pub dry_run: bool,
48}
49
50impl Default for BitsAndBytesConfig {
51    fn default() -> Self {
52        Self {
53            rocm_version: "7.2.0".to_string(),
54            gpu_arch: "gfx1100".to_string(),
55            python_bin: "python3".to_string(),
56            method: InstallMethod::Auto,
57            dry_run: false,
58        }
59    }
60}
61
62/// A constructed shell command.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct ShellCommand {
65    /// The program to run.
66    pub program: String,
67    /// Arguments to pass.
68    pub args: Vec<String>,
69    /// Environment variables to set.
70    pub env: Vec<(String, String)>,
71}
72
73impl ShellCommand {
74    /// Format as a shell command string.
75    pub fn to_command_string(&self) -> String {
76        let env_prefix = self
77            .env
78            .iter()
79            .map(|(k, v)| format!("{k}={v}"))
80            .collect::<Vec<_>>()
81            .join(" ");
82        let cmd = if self.args.is_empty() {
83            self.program.clone()
84        } else {
85            format!("{} {}", self.program, self.args.join(" "))
86        };
87        if env_prefix.is_empty() {
88            cmd
89        } else {
90            format!("{env_prefix} {cmd}")
91        }
92    }
93}
94
95/// The bitsandbytes installer.
96pub struct BitsAndBytesInstaller {
97    config: BitsAndBytesConfig,
98}
99
100impl BitsAndBytesInstaller {
101    /// Create a new bitsandbytes installer with the given config.
102    pub fn new(config: BitsAndBytesConfig) -> Self {
103        Self { config }
104    }
105
106    /// Create with default config.
107    pub fn with_defaults() -> Self {
108        Self::new(BitsAndBytesConfig::default())
109    }
110
111    // -----------------------------------------------------------------------
112    // Git clone URLs
113    // -----------------------------------------------------------------------
114
115    /// Get the primary ROCm bitsandbytes fork URL.
116    pub fn rocm_fork_url(&self) -> &'static str {
117        "https://github.com/ROCm/bitsandbytes.git"
118    }
119
120    /// Get the fallback TimDettmers bitsandbytes fork URL.
121    pub fn fallback_fork_url(&self) -> &'static str {
122        "https://github.com/TimDettmers/bitsandbytes.git"
123    }
124
125    /// Construct the git clone command for the ROCm fork.
126    pub fn build_git_clone_command(&self, target_dir: &str) -> ShellCommand {
127        ShellCommand {
128            program: "git".to_string(),
129            args: vec![
130                "clone".to_string(),
131                "--recursive".to_string(),
132                self.rocm_fork_url().to_string(),
133                target_dir.to_string(),
134            ],
135            env: vec![],
136        }
137    }
138
139    // -----------------------------------------------------------------------
140    // Command construction (VAL-INSTALL-017)
141    // -----------------------------------------------------------------------
142
143    /// Construct the initial pip install command for bitsandbytes from PyPI.
144    ///
145    /// The original script tries:
146    /// `pip install bitsandbytes --break-system-packages --no-cache-dir --no-deps`
147    pub fn build_pypi_install_command(&self) -> ShellCommand {
148        let use_break = self.config.method == InstallMethod::Global
149            || self.config.method == InstallMethod::Auto;
150
151        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
152        if use_break {
153            args.push("--break-system-packages".to_string());
154        }
155        args.extend([
156            "bitsandbytes".to_string(),
157            "--no-cache-dir".to_string(),
158            "--no-deps".to_string(),
159        ]);
160
161        ShellCommand {
162            program: self.config.python_bin.clone(),
163            args,
164            env: vec![],
165        }
166    }
167
168    /// Construct the build environment variables for ROCm source build.
169    pub fn build_rocm_env(&self) -> Vec<(String, String)> {
170        let rocm_path = "/opt/rocm".to_string();
171        let version_short = self.rocm_version_short();
172
173        vec![
174            ("BNB_USE_ROCM".to_string(), "1".to_string()),
175            ("BNB_ROCM_ARCH".to_string(), self.config.gpu_arch.clone()),
176            ("ROCM_PATH".to_string(), rocm_path.clone()),
177            ("ROCM_HOME".to_string(), rocm_path.clone()),
178            ("HIP_PATH".to_string(), rocm_path),
179            ("BNB_ROCM_VERSION".to_string(), version_short),
180        ]
181    }
182
183    /// Extract version shortcode (e.g., "7.2.0" -> "72").
184    fn rocm_version_short(&self) -> String {
185        let parts: Vec<&str> = self.config.rocm_version.split('.').collect();
186        if parts.len() >= 2 {
187            format!("{}{}", parts[0], parts[1])
188        } else {
189            "72".to_string()
190        }
191    }
192
193    /// Construct the build dependencies pip install command.
194    pub fn build_build_deps_command(&self) -> ShellCommand {
195        let use_break = self.config.method == InstallMethod::Global
196            || self.config.method == InstallMethod::Auto;
197
198        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
199        if use_break {
200            args.push("--break-system-packages".to_string());
201        }
202        args.extend([
203            "scikit-build-core".to_string(),
204            "cmake".to_string(),
205            "ninja".to_string(),
206        ]);
207
208        ShellCommand {
209            program: self.config.python_bin.clone(),
210            args,
211            env: vec![],
212        }
213    }
214
215    /// Construct the source build pip install command.
216    ///
217    /// The original script uses:
218    /// `pip install . --no-build-isolation --break-system-packages
219    ///  -Ccmake.define.COMPUTE_BACKEND=hip --no-deps
220    ///  -Ccmake.define.ROCM_PATH=/opt/rocm`
221    pub fn build_source_install_command(&self) -> ShellCommand {
222        let use_break = self.config.method == InstallMethod::Global
223            || self.config.method == InstallMethod::Auto;
224
225        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
226        if use_break {
227            args.push("--break-system-packages".to_string());
228        }
229        args.extend([
230            ".".to_string(),
231            "--no-build-isolation".to_string(),
232            "--no-deps".to_string(),
233            "-Ccmake.define.COMPUTE_BACKEND=hip".to_string(),
234            "-Ccmake.define.ROCM_PATH=/opt/rocm".to_string(),
235        ]);
236
237        ShellCommand {
238            program: self.config.python_bin.clone(),
239            args,
240            env: self.build_rocm_env(),
241        }
242    }
243
244    /// Construct the uninstall command for bitsandbytes.
245    pub fn build_uninstall_command(&self) -> ShellCommand {
246        ShellCommand {
247            program: self.config.python_bin.clone(),
248            args: vec![
249                "-m".to_string(),
250                "pip".to_string(),
251                "uninstall".to_string(),
252                "-y".to_string(),
253                "bitsandbytes".to_string(),
254            ],
255            env: vec![],
256        }
257    }
258}
259
260// ===========================================================================
261// Tests
262// ===========================================================================
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    // --- VAL-INSTALL-017: bitsandbytes correct pip command ---
269
270    #[test]
271    fn test_pypi_install_command() {
272        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
273            method: InstallMethod::Global,
274            ..Default::default()
275        });
276        let cmd = installer.build_pypi_install_command();
277        assert_eq!(cmd.program, "python3");
278        assert!(cmd.args.contains(&"-m".to_string()));
279        assert!(cmd.args.contains(&"pip".to_string()));
280        assert!(cmd.args.contains(&"install".to_string()));
281        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
282        assert!(cmd.args.contains(&"bitsandbytes".to_string()));
283        assert!(cmd.args.contains(&"--no-cache-dir".to_string()));
284        assert!(cmd.args.contains(&"--no-deps".to_string()));
285    }
286
287    #[test]
288    fn test_pypi_install_command_venv() {
289        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
290            method: InstallMethod::Venv,
291            ..Default::default()
292        });
293        let cmd = installer.build_pypi_install_command();
294        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
295        assert!(cmd.args.contains(&"bitsandbytes".to_string()));
296    }
297
298    #[test]
299    fn test_rocm_fork_url() {
300        let installer = BitsAndBytesInstaller::with_defaults();
301        assert_eq!(
302            installer.rocm_fork_url(),
303            "https://github.com/ROCm/bitsandbytes.git"
304        );
305    }
306
307    #[test]
308    fn test_fallback_fork_url() {
309        let installer = BitsAndBytesInstaller::with_defaults();
310        assert_eq!(
311            installer.fallback_fork_url(),
312            "https://github.com/TimDettmers/bitsandbytes.git"
313        );
314    }
315
316    #[test]
317    fn test_git_clone_command() {
318        let installer = BitsAndBytesInstaller::with_defaults();
319        let cmd = installer.build_git_clone_command("/tmp/bnb");
320        assert_eq!(cmd.program, "git");
321        assert!(cmd.args.contains(&"clone".to_string()));
322        assert!(cmd.args.contains(&"--recursive".to_string()));
323        assert!(cmd.args.iter().any(|a| a.contains("ROCm/bitsandbytes")));
324    }
325
326    #[test]
327    fn test_build_rocm_env() {
328        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
329            rocm_version: "7.2.0".to_string(),
330            gpu_arch: "gfx1100".to_string(),
331            ..Default::default()
332        });
333        let env = installer.build_rocm_env();
334        assert!(env.iter().any(|(k, v)| k == "BNB_USE_ROCM" && v == "1"));
335        assert!(env
336            .iter()
337            .any(|(k, v)| k == "BNB_ROCM_ARCH" && v == "gfx1100"));
338        assert!(env
339            .iter()
340            .any(|(k, v)| k == "ROCM_PATH" && v == "/opt/rocm"));
341        assert!(env
342            .iter()
343            .any(|(k, v)| k == "BNB_ROCM_VERSION" && v == "72"));
344    }
345
346    #[test]
347    fn test_rocm_version_short() {
348        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
349            rocm_version: "7.2.0".to_string(),
350            ..Default::default()
351        });
352        assert_eq!(installer.rocm_version_short(), "72");
353    }
354
355    #[test]
356    fn test_rocm_version_short_legacy() {
357        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
358            rocm_version: "7.0.0".to_string(),
359            ..Default::default()
360        });
361        assert_eq!(installer.rocm_version_short(), "70");
362    }
363
364    #[test]
365    fn test_build_deps_command() {
366        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
367            method: InstallMethod::Global,
368            ..Default::default()
369        });
370        let cmd = installer.build_build_deps_command();
371        assert!(cmd.args.contains(&"scikit-build-core".to_string()));
372        assert!(cmd.args.contains(&"cmake".to_string()));
373        assert!(cmd.args.contains(&"ninja".to_string()));
374    }
375
376    #[test]
377    fn test_source_install_command() {
378        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
379            rocm_version: "7.2.0".to_string(),
380            gpu_arch: "gfx1100".to_string(),
381            method: InstallMethod::Global,
382            ..Default::default()
383        });
384        let cmd = installer.build_source_install_command();
385        assert_eq!(cmd.program, "python3");
386        assert!(cmd.args.contains(&"--no-build-isolation".to_string()));
387        assert!(cmd.args.contains(&"--no-deps".to_string()));
388        assert!(cmd.args.contains(&".".to_string()));
389        assert!(cmd.args.iter().any(|a| a.contains("COMPUTE_BACKEND=hip")));
390        assert!(cmd.args.iter().any(|a| a.contains("ROCM_PATH=/opt/rocm")));
391        // Verify ROCm env vars
392        assert!(cmd.env.iter().any(|(k, v)| k == "BNB_USE_ROCM" && v == "1"));
393    }
394
395    #[test]
396    fn test_uninstall_command() {
397        let installer = BitsAndBytesInstaller::with_defaults();
398        let cmd = installer.build_uninstall_command();
399        assert!(cmd.args.contains(&"uninstall".to_string()));
400        assert!(cmd.args.contains(&"-y".to_string()));
401        assert!(cmd.args.contains(&"bitsandbytes".to_string()));
402    }
403}
rusty_stack/installers/components/bitsandbytes_multi.rs

rusty_stack/installers/components/
bitsandbytes_multi.rs
