rusty_stack/installers/components/
textgen.rs

1//! text-generation-webui installer — ports `scripts/install_textgen.sh`.
2//!
3//! Clones the text-generation-webui repository, installs Python dependencies
4//! with ROCm-only filtering (excludes nvidia/CUDA packages), and creates
5//! launcher shim.
6//!
7//! # Validation Assertions
8//!
9//! - **VAL-INSTALL-024**: text-generation-webui correct git clone
10//! - **VAL-INSTALL-025**: App installers correct target directory
11
12use crate::installers::common::command_exists;
13
14// ===========================================================================
15// Types
16// ===========================================================================
17
18/// A constructed shell command.
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct ShellCommand {
21    /// The program to run.
22    pub program: String,
23    /// Arguments to pass.
24    pub args: Vec<String>,
25    /// Environment variables to set.
26    pub env: Vec<(String, String)>,
27}
28
29impl ShellCommand {
30    /// Format as a shell command string.
31    pub fn to_command_string(&self) -> String {
32        let env_prefix = self
33            .env
34            .iter()
35            .map(|(k, v)| format!("{k}={v}"))
36            .collect::<Vec<_>>()
37            .join(" ");
38        let cmd = if self.args.is_empty() {
39            self.program.clone()
40        } else {
41            format!("{} {}", self.program, self.args.join(" "))
42        };
43        if env_prefix.is_empty() {
44            cmd
45        } else {
46            format!("{env_prefix} {cmd}")
47        }
48    }
49}
50
51/// Configuration for the text-generation-webui installer.
52#[derive(Debug, Clone)]
53pub struct TextgenConfig {
54    /// Target installation directory.
55    pub install_dir: String,
56    /// Python binary to use (resolved from MLSTACK_PYTHON_BIN or platform detection).
57    pub python_bin: String,
58    /// Whether to use `uv pip` instead of `python -m pip`.
59    pub use_uv: bool,
60    /// Whether to add `--break-system-packages` for global pip installs.
61    pub break_system_packages: bool,
62    /// Whether to run in dry-run mode.
63    pub dry_run: bool,
64}
65
66impl Default for TextgenConfig {
67    fn default() -> Self {
68        Self {
69            install_dir: default_install_dir(),
70            python_bin: resolve_python_bin(),
71            use_uv: command_exists("uv"),
72            break_system_packages: true,
73            dry_run: false,
74        }
75    }
76}
77
78/// Resolve the Python binary from ML Stack context.
79///
80/// Priority:
81/// 1. `MLSTACK_PYTHON_BIN` env var
82/// 2. `UV_PYTHON` env var
83/// 3. Platform environment detection (venvs, conda, system)
84/// 4. Fallback to `python3`
85fn resolve_python_bin() -> String {
86    // Priority 1 & 2: ML Stack configured Python
87    for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
88        if let Ok(val) = std::env::var(key) {
89            let val = val.trim().to_string();
90            if !val.is_empty() {
91                return val;
92            }
93        }
94    }
95
96    // Priority 3: Platform environment detection
97    let interpreters = crate::platform::environment::python_interpreters();
98    if let Some(python) = interpreters.first() {
99        return python.to_string_lossy().to_string();
100    }
101
102    // Priority 4: Fallback
103    "python3".to_string()
104}
105
106/// text-generation-webui Git repository URL.
107pub const REPO_URL: &str = "https://github.com/oobabooga/text-generation-webui.git";
108
109/// Default branch for text-generation-webui.
110pub const DEFAULT_BRANCH: &str = "main";
111
112/// Default web port.
113pub const WEB_PORT: u16 = 7860;
114
115/// Directories to preserve during git update.
116pub const PRESERVE_DIRS: &[&str] = &[
117    "models",
118    "loras",
119    "embeddings",
120    "presets",
121    "characters",
122    "training",
123];
124
125/// Regex pattern for packages to exclude from requirements.txt.
126/// Original script excludes: nvidia-*, cuda*, tensorrt*, triton[versioned],
127/// xformers, flash-attn, torch/torchvision/torchaudio.
128pub const EXCLUDED_PATTERNS: &[&str] = &[
129    "nvidia-",
130    "cuda",
131    "tensorrt",
132    "triton",
133    "xformers",
134    "flash-attn",
135    "torch",
136    "torchvision",
137    "torchaudio",
138];
139
140/// Default text-generation-webui installation directory.
141pub fn default_install_dir() -> String {
142    format!(
143        "{}/text-generation-webui",
144        std::env::var("HOME").unwrap_or_default()
145    )
146}
147
148/// The text-generation-webui installer.
149pub struct TextgenInstaller {
150    config: TextgenConfig,
151}
152
153impl TextgenInstaller {
154    /// Create a new text-generation-webui installer with the given config.
155    pub fn new(config: TextgenConfig) -> Self {
156        Self { config }
157    }
158
159    /// Create with default config.
160    pub fn with_defaults() -> Self {
161        Self::new(TextgenConfig::default())
162    }
163
164    // -----------------------------------------------------------------------
165    // Command construction (VAL-INSTALL-024, VAL-INSTALL-025)
166    // -----------------------------------------------------------------------
167
168    /// Construct the git clone command for text-generation-webui.
169    ///
170    /// The original script runs:
171    /// `git clone "https://github.com/oobabooga/text-generation-webui.git" "$TEXTGEN_DIR"`
172    pub fn build_git_clone_command(&self) -> ShellCommand {
173        ShellCommand {
174            program: "git".to_string(),
175            args: vec![
176                "clone".to_string(),
177                REPO_URL.to_string(),
178                self.config.install_dir.clone(),
179            ],
180            env: vec![],
181        }
182    }
183
184    /// Construct the git update commands for an existing clone.
185    ///
186    /// The original script runs:
187    /// 1. `git -C "$DIR" remote set-head origin -a`
188    /// 2. `git -C "$DIR" checkout "$DEFAULT_BRANCH"`
189    /// 3. `git -C "$DIR" fetch --all`
190    /// 4. `git -C "$DIR" stash push -u` (if user data exists)
191    /// 5. `git -C "$DIR" reset --hard "origin/$DEFAULT_BRANCH"`
192    /// 6. `git -C "$DIR" stash pop` (if user data was stashed)
193    pub fn build_git_update_commands(&self, has_user_data: bool) -> Vec<ShellCommand> {
194        let dir = &self.config.install_dir;
195        let mut cmds = vec![
196            ShellCommand {
197                program: "git".to_string(),
198                args: vec![
199                    "-C".to_string(),
200                    dir.clone(),
201                    "remote".to_string(),
202                    "set-head".to_string(),
203                    "origin".to_string(),
204                    "-a".to_string(),
205                ],
206                env: vec![],
207            },
208            ShellCommand {
209                program: "git".to_string(),
210                args: vec![
211                    "-C".to_string(),
212                    dir.clone(),
213                    "checkout".to_string(),
214                    DEFAULT_BRANCH.to_string(),
215                ],
216                env: vec![],
217            },
218            ShellCommand {
219                program: "git".to_string(),
220                args: vec![
221                    "-C".to_string(),
222                    dir.clone(),
223                    "fetch".to_string(),
224                    "--all".to_string(),
225                ],
226                env: vec![],
227            },
228        ];
229
230        if has_user_data {
231            cmds.push(ShellCommand {
232                program: "git".to_string(),
233                args: vec![
234                    "-C".to_string(),
235                    dir.clone(),
236                    "stash".to_string(),
237                    "push".to_string(),
238                    "-u".to_string(),
239                    "-m".to_string(),
240                    "rusty-stack-preserve-user-data".to_string(),
241                    "--".to_string(),
242                ]
243                .into_iter()
244                .chain(PRESERVE_DIRS.iter().map(|s| s.to_string()))
245                .collect(),
246                env: vec![],
247            });
248        }
249
250        cmds.push(ShellCommand {
251            program: "git".to_string(),
252            args: vec![
253                "-C".to_string(),
254                dir.clone(),
255                "reset".to_string(),
256                "--hard".to_string(),
257                format!("origin/{DEFAULT_BRANCH}"),
258            ],
259            env: vec![],
260        });
261
262        if has_user_data {
263            cmds.push(ShellCommand {
264                program: "git".to_string(),
265                args: vec![
266                    "-C".to_string(),
267                    dir.clone(),
268                    "stash".to_string(),
269                    "pop".to_string(),
270                ],
271                env: vec![],
272            });
273        }
274
275        cmds
276    }
277
278    /// Construct the pip install command for filtered requirements.
279    ///
280    /// The original script filters requirements.txt excluding nvidia/CUDA packages:
281    /// `grep -v -E '^(nvidia-|cuda|tensorrt|triton([=<>!\[]|$)|xformers|flash-attn|torch([=<>! ]|$)|torchvision|torchaudio)'`
282    ///
283    /// Uses `uv pip install` when `use_uv` is true, otherwise `python -m pip install`.
284    /// Adds `--break-system-packages` for global installs when `break_system_packages` is true.
285    pub fn build_pip_install_command(&self, requirements_path: &str) -> ShellCommand {
286        let (program, mut args) = self.build_pip_prefix();
287        args.push("-r".to_string());
288        args.push(requirements_path.to_string());
289
290        ShellCommand {
291            program,
292            args,
293            env: vec![],
294        }
295    }
296
297    /// Construct the pip install command for AMD-specific requirements.
298    ///
299    /// Uses `uv pip install` when `use_uv` is true, otherwise `python -m pip install`.
300    /// Adds `--break-system-packages` for global installs when `break_system_packages` is true.
301    pub fn build_amd_requirements_command(&self, amd_req_path: &str) -> ShellCommand {
302        let (program, mut args) = self.build_pip_prefix();
303        args.push("-r".to_string());
304        args.push(amd_req_path.to_string());
305
306        ShellCommand {
307            program,
308            args,
309            env: vec![],
310        }
311    }
312
313    /// Build the pip install prefix (program + base args).
314    ///
315    /// Returns `(program, args)` ready for additional arguments to be appended.
316    fn build_pip_prefix(&self) -> (String, Vec<String>) {
317        if self.config.use_uv {
318            let mut args = vec!["pip".to_string(), "install".to_string()];
319            // uv pip needs --system for global installs (no venv)
320            args.push("--system".to_string());
321            // uv pip: respect UV_PYTHON env var or pass --python
322            if std::env::var("UV_PYTHON").is_err() {
323                args.push("--python".to_string());
324                args.push(self.config.python_bin.clone());
325            }
326            ("uv".to_string(), args)
327        } else {
328            let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
329            if self.config.break_system_packages {
330                args.push("--break-system-packages".to_string());
331            }
332            (self.config.python_bin.clone(), args)
333        }
334    }
335
336    /// Get the grep filter pattern for excluding packages from requirements.
337    /// Matches the regex from the original script.
338    pub fn requirements_filter_pattern() -> String {
339        r"^(nvidia-|cuda|tensorrt|triton([=<>!\[]|$)|xformers|flash-attn|torch([=<>! ]|$)|torchvision|torchaudio)".to_string()
340    }
341
342    /// Construct the launcher shim content.
343    ///
344    /// Creates a shell script at `~/.local/bin/textgen` that runs:
345    /// `cd "$TEXTGEN_DIR" && HIP_VISIBLE_DEVICES=$GPUS ... python3 server.py --chat`
346    pub fn build_launcher_content(&self, gpu_devices: &str) -> String {
347        format!(
348            "#!/bin/bash\n# text-generation-webui launcher for Stan's ML Stack\n# Auto-detected GPU devices: {}\ncd \"{}\" && \\\n    HIP_VISIBLE_DEVICES={} \\\n    CUDA_VISIBLE_DEVICES={} \\\n    {} server.py --chat \"$@\"\n",
349            gpu_devices,
350            self.config.install_dir,
351            gpu_devices,
352            gpu_devices,
353            self.config.python_bin,
354        )
355    }
356
357    /// Construct the systemd service unit content.
358    pub fn build_systemd_service(&self, gpu_devices: &str) -> String {
359        format!(
360            "[Unit]\nDescription=text-generation-webui - LLM chat interface with ROCm support\nAfter=network.target\n\n\
361             [Service]\nType=simple\nWorkingDirectory=\"{}\"\n\
362             Environment=\"HIP_VISIBLE_DEVICES={}\"\n\
363             Environment=\"CUDA_VISIBLE_DEVICES={}\"\n\
364             ExecStart={} \"{}\" --chat\nRestart=on-failure\n\n\
365             [Install]\nWantedBy=default.target\n",
366            self.config.install_dir,
367            gpu_devices,
368            gpu_devices,
369            self.config.python_bin,
370            self.config.install_dir,
371        )
372    }
373
374    /// Detect GPU devices — the canonical, iGPU-filtered discrete list.
375    ///
376    /// Routes through `crate::installer::detect_gpu_list` (Stage 1) so the
377    /// integrated GPU is never included and no dGPU is missed. Baked into the
378    /// launcher shim + systemd unit as `HIP_VISIBLE_DEVICES`/`CUDA_VISIBLE_DEVICES`.
379    pub fn detect_gpu_devices(&self) -> String {
380        crate::installer::detect_gpu_list()
381    }
382
383    /// Get the install directory.
384    pub fn install_dir(&self) -> &str {
385        &self.config.install_dir
386    }
387
388    /// Get the repository URL.
389    pub fn repo_url(&self) -> &str {
390        REPO_URL
391    }
392
393    /// Get the default branch.
394    pub fn default_branch(&self) -> &str {
395        DEFAULT_BRANCH
396    }
397}
398
399// ===========================================================================
400// Tests
401// ===========================================================================
402
403#[cfg(test)]
404mod tests {
405    use super::*;
406
407    fn make_installer(dir: &str) -> TextgenInstaller {
408        TextgenInstaller::new(TextgenConfig {
409            install_dir: dir.to_string(),
410            python_bin: "python3".to_string(),
411            use_uv: false,
412            break_system_packages: true,
413            dry_run: false,
414        })
415    }
416
417    fn make_installer_with_python(dir: &str, python_bin: &str) -> TextgenInstaller {
418        TextgenInstaller::new(TextgenConfig {
419            install_dir: dir.to_string(),
420            python_bin: python_bin.to_string(),
421            use_uv: false,
422            break_system_packages: true,
423            dry_run: false,
424        })
425    }
426
427    fn make_installer_uv(dir: &str) -> TextgenInstaller {
428        TextgenInstaller::new(TextgenConfig {
429            install_dir: dir.to_string(),
430            python_bin: "python3".to_string(),
431            use_uv: true,
432            break_system_packages: true,
433            dry_run: false,
434        })
435    }
436
437    // -----------------------------------------------------------------------
438    // VAL-INSTALL-024: text-generation-webui correct git clone
439    // -----------------------------------------------------------------------
440
441    #[test]
442    fn test_git_clone_url_matches_original_script() {
443        assert_eq!(
444            REPO_URL, "https://github.com/oobabooga/text-generation-webui.git",
445            "URL must match install_textgen.sh REPO_URL"
446        );
447    }
448
449    #[test]
450    fn test_git_clone_command() {
451        let installer = make_installer("/home/user/text-generation-webui");
452        let cmd = installer.build_git_clone_command();
453        assert_eq!(cmd.program, "git");
454        assert!(cmd.args.contains(&"clone".to_string()));
455        assert!(cmd.args.contains(&REPO_URL.to_string()));
456        assert!(cmd
457            .args
458            .contains(&"/home/user/text-generation-webui".to_string()));
459    }
460
461    #[test]
462    fn test_git_clone_command_string() {
463        let installer = make_installer("/home/user/text-generation-webui");
464        let cmd = installer.build_git_clone_command();
465        let s = cmd.to_command_string();
466        assert!(s.contains("git clone"));
467        assert!(s.contains(REPO_URL));
468    }
469
470    #[test]
471    fn test_default_branch_is_main() {
472        assert_eq!(DEFAULT_BRANCH, "main");
473    }
474
475    // -----------------------------------------------------------------------
476    // VAL-INSTALL-025: App installers correct target directory
477    // -----------------------------------------------------------------------
478
479    #[test]
480    fn test_default_install_dir_is_home_textgen() {
481        let config = TextgenConfig::default();
482        assert!(
483            config.install_dir.ends_with("/text-generation-webui"),
484            "Default dir must end with /text-generation-webui, got: {}",
485            config.install_dir,
486        );
487    }
488
489    #[test]
490    fn test_custom_install_dir() {
491        let installer = make_installer("/opt/apps/text-generation-webui");
492        assert_eq!(installer.install_dir(), "/opt/apps/text-generation-webui");
493    }
494
495    #[test]
496    fn test_pip_install_command() {
497        let installer = make_installer("/home/user/text-generation-webui");
498        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
499        assert_eq!(cmd.program, "python3");
500        assert!(cmd.args.contains(&"-m".to_string()));
501        assert!(cmd.args.contains(&"pip".to_string()));
502        assert!(cmd.args.contains(&"install".to_string()));
503        assert!(cmd.args.contains(&"/tmp/filtered_reqs.txt".to_string()));
504    }
505
506    #[test]
507    fn test_pip_install_command_uses_configured_python() {
508        let installer = make_installer_with_python(
509            "/home/user/text-generation-webui",
510            "/home/user/rocm_venv/bin/python",
511        );
512        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
513        assert_eq!(
514            cmd.program, "/home/user/rocm_venv/bin/python",
515            "Pip install must use configured Python interpreter, not bare 'python3'"
516        );
517    }
518
519    #[test]
520    fn test_pip_install_command_includes_break_system_packages() {
521        let installer = make_installer("/home/user/text-generation-webui");
522        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
523        assert!(
524            cmd.args.contains(&"--break-system-packages".to_string()),
525            "Global pip install must include --break-system-packages for PEP 668 compatibility"
526        );
527    }
528
529    #[test]
530    fn test_pip_install_command_no_break_system_packages_when_disabled() {
531        let installer = TextgenInstaller::new(TextgenConfig {
532            install_dir: "/home/user/text-generation-webui".to_string(),
533            python_bin: "python3".to_string(),
534            use_uv: false,
535            break_system_packages: false,
536            dry_run: false,
537        });
538        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
539        assert!(
540            !cmd.args.contains(&"--break-system-packages".to_string()),
541            "Venv pip install should NOT include --break-system-packages"
542        );
543    }
544
545    #[test]
546    fn test_pip_install_command_uses_uv_when_available() {
547        let installer = make_installer_uv("/home/user/text-generation-webui");
548        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
549        // When use_uv is true, the program is "uv" and args start with "pip"
550        assert_eq!(
551            cmd.program, "uv",
552            "When use_uv is true, program must be 'uv'"
553        );
554        assert!(
555            cmd.args.first().map(|s| s == "pip").unwrap_or(false),
556            "When use_uv is true, args must start with 'pip'"
557        );
558        assert!(
559            !cmd.args.contains(&"--break-system-packages".to_string()),
560            "uv pip does not need --break-system-packages"
561        );
562    }
563
564    #[test]
565    fn test_amd_requirements_command_uses_configured_python() {
566        let installer = make_installer_with_python(
567            "/home/user/text-generation-webui",
568            "/home/user/rocm_venv/bin/python",
569        );
570        let cmd = installer.build_amd_requirements_command(
571            "/home/user/text-generation-webui/requirements_amd.txt",
572        );
573        assert_eq!(
574            cmd.program, "/home/user/rocm_venv/bin/python",
575            "AMD requirements command must use configured Python interpreter"
576        );
577    }
578
579    #[test]
580    fn test_amd_requirements_command_includes_break_system_packages() {
581        let installer = make_installer("/home/user/text-generation-webui");
582        let cmd = installer.build_amd_requirements_command(
583            "/home/user/text-generation-webui/requirements_amd.txt",
584        );
585        assert!(
586            cmd.args.contains(&"--break-system-packages".to_string()),
587            "Global pip install must include --break-system-packages for AMD requirements"
588        );
589    }
590
591    #[test]
592    fn test_default_config_detects_python_from_env() {
593        // The default config should resolve python from MLSTACK_PYTHON_BIN or platform detection
594        let config = TextgenConfig::default();
595        assert!(
596            !config.python_bin.is_empty(),
597            "Default python_bin must not be empty"
598        );
599        // It should NOT be bare "python3" when MLSTACK_PYTHON_BIN is set
600        // (unless that's what MLSTACK_PYTHON_BIN says)
601    }
602
603    #[test]
604    fn test_default_config_detects_uv() {
605        let config = TextgenConfig::default();
606        // Should match whether uv is available
607        let uv_available = command_exists("uv");
608        assert_eq!(
609            config.use_uv, uv_available,
610            "Default use_uv should match uv availability"
611        );
612    }
613
614    #[test]
615    fn test_requirements_filter_excludes_nvidia_cuda() {
616        let pattern = TextgenInstaller::requirements_filter_pattern();
617        assert!(pattern.contains("nvidia-"));
618        assert!(pattern.contains("cuda"));
619        assert!(pattern.contains("tensorrt"));
620        assert!(pattern.contains("triton"));
621        assert!(pattern.contains("xformers"));
622        assert!(pattern.contains("flash-attn"));
623        assert!(pattern.contains("torch"));
624    }
625
626    #[test]
627    fn test_preserve_dirs_match_original_script() {
628        // Original: models, loras, embeddings, presets, characters, training
629        assert_eq!(
630            PRESERVE_DIRS,
631            [
632                "models",
633                "loras",
634                "embeddings",
635                "presets",
636                "characters",
637                "training"
638            ]
639        );
640    }
641
642    #[test]
643    fn test_git_update_without_user_data() {
644        let installer = make_installer("/home/user/text-generation-webui");
645        let cmds = installer.build_git_update_commands(false);
646        assert_eq!(cmds.len(), 4);
647        let last = cmds.last().unwrap();
648        assert!(last.args.contains(&"reset".to_string()));
649    }
650
651    #[test]
652    fn test_git_update_with_user_data() {
653        let installer = make_installer("/home/user/text-generation-webui");
654        let cmds = installer.build_git_update_commands(true);
655        assert_eq!(cmds.len(), 6);
656        let cmd_strings: Vec<String> = cmds.iter().map(|c| c.to_command_string()).collect();
657        let has_stash = cmd_strings.iter().any(|s| s.contains("stash"));
658        assert!(has_stash);
659    }
660
661    #[test]
662    fn test_launcher_content() {
663        let installer = make_installer("/home/user/text-generation-webui");
664        let content = installer.build_launcher_content("0,1");
665        assert!(content.starts_with("#!/bin/bash"));
666        assert!(content.contains("/home/user/text-generation-webui"));
667        assert!(content.contains("HIP_VISIBLE_DEVICES=0,1"));
668        assert!(content.contains("python3 server.py --chat"));
669    }
670
671    #[test]
672    fn test_launcher_content_uses_configured_python() {
673        let installer = make_installer_with_python(
674            "/home/user/text-generation-webui",
675            "/home/user/rocm_venv/bin/python",
676        );
677        let content = installer.build_launcher_content("0,1");
678        assert!(
679            content.contains("/home/user/rocm_venv/bin/python server.py --chat"),
680            "Launcher must use configured Python interpreter"
681        );
682    }
683
684    #[test]
685    fn test_systemd_service_content() {
686        let installer = make_installer("/home/user/text-generation-webui");
687        let svc = installer.build_systemd_service("0,1");
688        assert!(svc.contains("[Unit]"));
689        assert!(svc.contains("text-generation-webui"));
690        assert!(svc.contains("HIP_VISIBLE_DEVICES=0,1"));
691        assert!(svc.contains("--chat"));
692    }
693
694    #[test]
695    fn test_amd_requirements_command() {
696        let installer = make_installer("/home/user/text-generation-webui");
697        let cmd = installer.build_amd_requirements_command(
698            "/home/user/text-generation-webui/requirements_amd.txt",
699        );
700        assert_eq!(cmd.program, "python3");
701        assert!(cmd.args.contains(&"-r".to_string()));
702        assert!(cmd
703            .args
704            .contains(&"/home/user/text-generation-webui/requirements_amd.txt".to_string()));
705    }
706}
rusty_stack/installers/components/textgen.rs

rusty_stack/installers/components/
textgen.rs
