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            "--upgrade".to_string(),
157            "bitsandbytes".to_string(),
158            "--no-cache-dir".to_string(),
159            "--no-deps".to_string(),
160        ]);
161
162        ShellCommand {
163            program: self.config.python_bin.clone(),
164            args,
165            env: vec![],
166        }
167    }
168
169    /// Construct the build environment variables for ROCm source build.
170    pub fn build_rocm_env(&self) -> Vec<(String, String)> {
171        let rocm_path = "/opt/rocm".to_string();
172        let version_short = self.rocm_version_short();
173
174        vec![
175            ("BNB_USE_ROCM".to_string(), "1".to_string()),
176            ("BNB_ROCM_ARCH".to_string(), self.config.gpu_arch.clone()),
177            ("ROCM_PATH".to_string(), rocm_path.clone()),
178            ("ROCM_HOME".to_string(), rocm_path.clone()),
179            ("HIP_PATH".to_string(), rocm_path),
180            ("BNB_ROCM_VERSION".to_string(), version_short),
181        ]
182    }
183
184    /// Extract version shortcode (e.g., "7.2.0" -> "72").
185    fn rocm_version_short(&self) -> String {
186        let parts: Vec<&str> = self.config.rocm_version.split('.').collect();
187        if parts.len() >= 2 {
188            format!("{}{}", parts[0], parts[1])
189        } else {
190            "72".to_string()
191        }
192    }
193
194    /// Construct the build dependencies pip install command.
195    pub fn build_build_deps_command(&self) -> ShellCommand {
196        let use_break = self.config.method == InstallMethod::Global
197            || self.config.method == InstallMethod::Auto;
198
199        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
200        if use_break {
201            args.push("--break-system-packages".to_string());
202        }
203        args.extend([
204            "scikit-build-core".to_string(),
205            "cmake".to_string(),
206            "ninja".to_string(),
207        ]);
208
209        ShellCommand {
210            program: self.config.python_bin.clone(),
211            args,
212            env: vec![],
213        }
214    }
215
216    /// Construct the source build pip install command.
217    ///
218    /// The original script uses:
219    /// `pip install . --no-build-isolation --break-system-packages
220    ///  -Ccmake.define.COMPUTE_BACKEND=hip --no-deps
221    ///  -Ccmake.define.ROCM_PATH=/opt/rocm`
222    pub fn build_source_install_command(&self) -> ShellCommand {
223        let use_break = self.config.method == InstallMethod::Global
224            || self.config.method == InstallMethod::Auto;
225
226        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
227        if use_break {
228            args.push("--break-system-packages".to_string());
229        }
230        args.extend([
231            ".".to_string(),
232            "--no-build-isolation".to_string(),
233            "--no-deps".to_string(),
234            "-Ccmake.define.COMPUTE_BACKEND=hip".to_string(),
235            "-Ccmake.define.ROCM_PATH=/opt/rocm".to_string(),
236        ]);
237
238        ShellCommand {
239            program: self.config.python_bin.clone(),
240            args,
241            env: self.build_rocm_env(),
242        }
243    }
244
245    /// Construct the uninstall command for bitsandbytes.
246    pub fn build_uninstall_command(&self) -> ShellCommand {
247        ShellCommand {
248            program: self.config.python_bin.clone(),
249            args: vec![
250                "-m".to_string(),
251                "pip".to_string(),
252                "uninstall".to_string(),
253                "-y".to_string(),
254                "bitsandbytes".to_string(),
255            ],
256            env: vec![],
257        }
258    }
259}
260
261// ===========================================================================
262// Tests
263// ===========================================================================
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    // --- VAL-INSTALL-017: bitsandbytes correct pip command ---
270
271    #[test]
272    fn test_pypi_install_command() {
273        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
274            method: InstallMethod::Global,
275            ..Default::default()
276        });
277        let cmd = installer.build_pypi_install_command();
278        assert_eq!(cmd.program, "python3");
279        assert!(cmd.args.contains(&"-m".to_string()));
280        assert!(cmd.args.contains(&"pip".to_string()));
281        assert!(cmd.args.contains(&"install".to_string()));
282        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
283        assert!(cmd.args.contains(&"bitsandbytes".to_string()));
284        assert!(cmd.args.contains(&"--no-cache-dir".to_string()));
285        assert!(cmd.args.contains(&"--no-deps".to_string()));
286    }
287
288    #[test]
289    fn test_pypi_install_command_venv() {
290        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
291            method: InstallMethod::Venv,
292            ..Default::default()
293        });
294        let cmd = installer.build_pypi_install_command();
295        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
296        assert!(cmd.args.contains(&"bitsandbytes".to_string()));
297    }
298
299    #[test]
300    fn test_rocm_fork_url() {
301        let installer = BitsAndBytesInstaller::with_defaults();
302        assert_eq!(
303            installer.rocm_fork_url(),
304            "https://github.com/ROCm/bitsandbytes.git"
305        );
306    }
307
308    #[test]
309    fn test_fallback_fork_url() {
310        let installer = BitsAndBytesInstaller::with_defaults();
311        assert_eq!(
312            installer.fallback_fork_url(),
313            "https://github.com/TimDettmers/bitsandbytes.git"
314        );
315    }
316
317    #[test]
318    fn test_git_clone_command() {
319        let installer = BitsAndBytesInstaller::with_defaults();
320        let cmd = installer.build_git_clone_command("/tmp/bnb");
321        assert_eq!(cmd.program, "git");
322        assert!(cmd.args.contains(&"clone".to_string()));
323        assert!(cmd.args.contains(&"--recursive".to_string()));
324        assert!(cmd.args.iter().any(|a| a.contains("ROCm/bitsandbytes")));
325    }
326
327    #[test]
328    fn test_build_rocm_env() {
329        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
330            rocm_version: "7.2.0".to_string(),
331            gpu_arch: "gfx1100".to_string(),
332            ..Default::default()
333        });
334        let env = installer.build_rocm_env();
335        assert!(env.iter().any(|(k, v)| k == "BNB_USE_ROCM" && v == "1"));
336        assert!(env
337            .iter()
338            .any(|(k, v)| k == "BNB_ROCM_ARCH" && v == "gfx1100"));
339        assert!(env
340            .iter()
341            .any(|(k, v)| k == "ROCM_PATH" && v == "/opt/rocm"));
342        assert!(env
343            .iter()
344            .any(|(k, v)| k == "BNB_ROCM_VERSION" && v == "72"));
345    }
346
347    #[test]
348    fn test_rocm_version_short() {
349        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
350            rocm_version: "7.2.0".to_string(),
351            ..Default::default()
352        });
353        assert_eq!(installer.rocm_version_short(), "72");
354    }
355
356    #[test]
357    fn test_rocm_version_short_legacy() {
358        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
359            rocm_version: "7.0.0".to_string(),
360            ..Default::default()
361        });
362        assert_eq!(installer.rocm_version_short(), "70");
363    }
364
365    #[test]
366    fn test_build_deps_command() {
367        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
368            method: InstallMethod::Global,
369            ..Default::default()
370        });
371        let cmd = installer.build_build_deps_command();
372        assert!(cmd.args.contains(&"scikit-build-core".to_string()));
373        assert!(cmd.args.contains(&"cmake".to_string()));
374        assert!(cmd.args.contains(&"ninja".to_string()));
375    }
376
377    #[test]
378    fn test_source_install_command() {
379        let installer = BitsAndBytesInstaller::new(BitsAndBytesConfig {
380            rocm_version: "7.2.0".to_string(),
381            gpu_arch: "gfx1100".to_string(),
382            method: InstallMethod::Global,
383            ..Default::default()
384        });
385        let cmd = installer.build_source_install_command();
386        assert_eq!(cmd.program, "python3");
387        assert!(cmd.args.contains(&"--no-build-isolation".to_string()));
388        assert!(cmd.args.contains(&"--no-deps".to_string()));
389        assert!(cmd.args.contains(&".".to_string()));
390        assert!(cmd.args.iter().any(|a| a.contains("COMPUTE_BACKEND=hip")));
391        assert!(cmd.args.iter().any(|a| a.contains("ROCM_PATH=/opt/rocm")));
392        // Verify ROCm env vars
393        assert!(cmd.env.iter().any(|(k, v)| k == "BNB_USE_ROCM" && v == "1"));
394    }
395
396    #[test]
397    fn test_uninstall_command() {
398        let installer = BitsAndBytesInstaller::with_defaults();
399        let cmd = installer.build_uninstall_command();
400        assert!(cmd.args.contains(&"uninstall".to_string()));
401        assert!(cmd.args.contains(&"-y".to_string()));
402        assert!(cmd.args.contains(&"bitsandbytes".to_string()));
403    }
404}
rusty_stack/installers/components/bitsandbytes_multi.rs

rusty_stack/installers/components/
bitsandbytes_multi.rs
