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 via rocm-smi.
375    pub fn detect_gpu_devices(&self) -> String {
376        let _ = command_exists("rocm-smi");
377        "0,1".to_string()
378    }
379
380    /// Get the install directory.
381    pub fn install_dir(&self) -> &str {
382        &self.config.install_dir
383    }
384
385    /// Get the repository URL.
386    pub fn repo_url(&self) -> &str {
387        REPO_URL
388    }
389
390    /// Get the default branch.
391    pub fn default_branch(&self) -> &str {
392        DEFAULT_BRANCH
393    }
394}
395
396// ===========================================================================
397// Tests
398// ===========================================================================
399
400#[cfg(test)]
401mod tests {
402    use super::*;
403
404    fn make_installer(dir: &str) -> TextgenInstaller {
405        TextgenInstaller::new(TextgenConfig {
406            install_dir: dir.to_string(),
407            python_bin: "python3".to_string(),
408            use_uv: false,
409            break_system_packages: true,
410            dry_run: false,
411        })
412    }
413
414    fn make_installer_with_python(dir: &str, python_bin: &str) -> TextgenInstaller {
415        TextgenInstaller::new(TextgenConfig {
416            install_dir: dir.to_string(),
417            python_bin: python_bin.to_string(),
418            use_uv: false,
419            break_system_packages: true,
420            dry_run: false,
421        })
422    }
423
424    fn make_installer_uv(dir: &str) -> TextgenInstaller {
425        TextgenInstaller::new(TextgenConfig {
426            install_dir: dir.to_string(),
427            python_bin: "python3".to_string(),
428            use_uv: true,
429            break_system_packages: true,
430            dry_run: false,
431        })
432    }
433
434    // -----------------------------------------------------------------------
435    // VAL-INSTALL-024: text-generation-webui correct git clone
436    // -----------------------------------------------------------------------
437
438    #[test]
439    fn test_git_clone_url_matches_original_script() {
440        assert_eq!(
441            REPO_URL, "https://github.com/oobabooga/text-generation-webui.git",
442            "URL must match install_textgen.sh REPO_URL"
443        );
444    }
445
446    #[test]
447    fn test_git_clone_command() {
448        let installer = make_installer("/home/user/text-generation-webui");
449        let cmd = installer.build_git_clone_command();
450        assert_eq!(cmd.program, "git");
451        assert!(cmd.args.contains(&"clone".to_string()));
452        assert!(cmd.args.contains(&REPO_URL.to_string()));
453        assert!(cmd
454            .args
455            .contains(&"/home/user/text-generation-webui".to_string()));
456    }
457
458    #[test]
459    fn test_git_clone_command_string() {
460        let installer = make_installer("/home/user/text-generation-webui");
461        let cmd = installer.build_git_clone_command();
462        let s = cmd.to_command_string();
463        assert!(s.contains("git clone"));
464        assert!(s.contains(REPO_URL));
465    }
466
467    #[test]
468    fn test_default_branch_is_main() {
469        assert_eq!(DEFAULT_BRANCH, "main");
470    }
471
472    // -----------------------------------------------------------------------
473    // VAL-INSTALL-025: App installers correct target directory
474    // -----------------------------------------------------------------------
475
476    #[test]
477    fn test_default_install_dir_is_home_textgen() {
478        let config = TextgenConfig::default();
479        assert!(
480            config.install_dir.ends_with("/text-generation-webui"),
481            "Default dir must end with /text-generation-webui, got: {}",
482            config.install_dir,
483        );
484    }
485
486    #[test]
487    fn test_custom_install_dir() {
488        let installer = make_installer("/opt/apps/text-generation-webui");
489        assert_eq!(installer.install_dir(), "/opt/apps/text-generation-webui");
490    }
491
492    #[test]
493    fn test_pip_install_command() {
494        let installer = make_installer("/home/user/text-generation-webui");
495        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
496        assert_eq!(cmd.program, "python3");
497        assert!(cmd.args.contains(&"-m".to_string()));
498        assert!(cmd.args.contains(&"pip".to_string()));
499        assert!(cmd.args.contains(&"install".to_string()));
500        assert!(cmd.args.contains(&"/tmp/filtered_reqs.txt".to_string()));
501    }
502
503    #[test]
504    fn test_pip_install_command_uses_configured_python() {
505        let installer = make_installer_with_python(
506            "/home/user/text-generation-webui",
507            "/home/user/rocm_venv/bin/python",
508        );
509        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
510        assert_eq!(
511            cmd.program, "/home/user/rocm_venv/bin/python",
512            "Pip install must use configured Python interpreter, not bare 'python3'"
513        );
514    }
515
516    #[test]
517    fn test_pip_install_command_includes_break_system_packages() {
518        let installer = make_installer("/home/user/text-generation-webui");
519        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
520        assert!(
521            cmd.args.contains(&"--break-system-packages".to_string()),
522            "Global pip install must include --break-system-packages for PEP 668 compatibility"
523        );
524    }
525
526    #[test]
527    fn test_pip_install_command_no_break_system_packages_when_disabled() {
528        let installer = TextgenInstaller::new(TextgenConfig {
529            install_dir: "/home/user/text-generation-webui".to_string(),
530            python_bin: "python3".to_string(),
531            use_uv: false,
532            break_system_packages: false,
533            dry_run: false,
534        });
535        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
536        assert!(
537            !cmd.args.contains(&"--break-system-packages".to_string()),
538            "Venv pip install should NOT include --break-system-packages"
539        );
540    }
541
542    #[test]
543    fn test_pip_install_command_uses_uv_when_available() {
544        let installer = make_installer_uv("/home/user/text-generation-webui");
545        let cmd = installer.build_pip_install_command("/tmp/filtered_reqs.txt");
546        // When use_uv is true, the program is "uv" and args start with "pip"
547        assert_eq!(
548            cmd.program, "uv",
549            "When use_uv is true, program must be 'uv'"
550        );
551        assert!(
552            cmd.args.first().map(|s| s == "pip").unwrap_or(false),
553            "When use_uv is true, args must start with 'pip'"
554        );
555        assert!(
556            !cmd.args.contains(&"--break-system-packages".to_string()),
557            "uv pip does not need --break-system-packages"
558        );
559    }
560
561    #[test]
562    fn test_amd_requirements_command_uses_configured_python() {
563        let installer = make_installer_with_python(
564            "/home/user/text-generation-webui",
565            "/home/user/rocm_venv/bin/python",
566        );
567        let cmd = installer.build_amd_requirements_command(
568            "/home/user/text-generation-webui/requirements_amd.txt",
569        );
570        assert_eq!(
571            cmd.program, "/home/user/rocm_venv/bin/python",
572            "AMD requirements command must use configured Python interpreter"
573        );
574    }
575
576    #[test]
577    fn test_amd_requirements_command_includes_break_system_packages() {
578        let installer = make_installer("/home/user/text-generation-webui");
579        let cmd = installer.build_amd_requirements_command(
580            "/home/user/text-generation-webui/requirements_amd.txt",
581        );
582        assert!(
583            cmd.args.contains(&"--break-system-packages".to_string()),
584            "Global pip install must include --break-system-packages for AMD requirements"
585        );
586    }
587
588    #[test]
589    fn test_default_config_detects_python_from_env() {
590        // The default config should resolve python from MLSTACK_PYTHON_BIN or platform detection
591        let config = TextgenConfig::default();
592        assert!(
593            !config.python_bin.is_empty(),
594            "Default python_bin must not be empty"
595        );
596        // It should NOT be bare "python3" when MLSTACK_PYTHON_BIN is set
597        // (unless that's what MLSTACK_PYTHON_BIN says)
598    }
599
600    #[test]
601    fn test_default_config_detects_uv() {
602        let config = TextgenConfig::default();
603        // Should match whether uv is available
604        let uv_available = command_exists("uv");
605        assert_eq!(
606            config.use_uv, uv_available,
607            "Default use_uv should match uv availability"
608        );
609    }
610
611    #[test]
612    fn test_requirements_filter_excludes_nvidia_cuda() {
613        let pattern = TextgenInstaller::requirements_filter_pattern();
614        assert!(pattern.contains("nvidia-"));
615        assert!(pattern.contains("cuda"));
616        assert!(pattern.contains("tensorrt"));
617        assert!(pattern.contains("triton"));
618        assert!(pattern.contains("xformers"));
619        assert!(pattern.contains("flash-attn"));
620        assert!(pattern.contains("torch"));
621    }
622
623    #[test]
624    fn test_preserve_dirs_match_original_script() {
625        // Original: models, loras, embeddings, presets, characters, training
626        assert_eq!(
627            PRESERVE_DIRS,
628            [
629                "models",
630                "loras",
631                "embeddings",
632                "presets",
633                "characters",
634                "training"
635            ]
636        );
637    }
638
639    #[test]
640    fn test_git_update_without_user_data() {
641        let installer = make_installer("/home/user/text-generation-webui");
642        let cmds = installer.build_git_update_commands(false);
643        assert_eq!(cmds.len(), 4);
644        let last = cmds.last().unwrap();
645        assert!(last.args.contains(&"reset".to_string()));
646    }
647
648    #[test]
649    fn test_git_update_with_user_data() {
650        let installer = make_installer("/home/user/text-generation-webui");
651        let cmds = installer.build_git_update_commands(true);
652        assert_eq!(cmds.len(), 6);
653        let cmd_strings: Vec<String> = cmds.iter().map(|c| c.to_command_string()).collect();
654        let has_stash = cmd_strings.iter().any(|s| s.contains("stash"));
655        assert!(has_stash);
656    }
657
658    #[test]
659    fn test_launcher_content() {
660        let installer = make_installer("/home/user/text-generation-webui");
661        let content = installer.build_launcher_content("0,1");
662        assert!(content.starts_with("#!/bin/bash"));
663        assert!(content.contains("/home/user/text-generation-webui"));
664        assert!(content.contains("HIP_VISIBLE_DEVICES=0,1"));
665        assert!(content.contains("python3 server.py --chat"));
666    }
667
668    #[test]
669    fn test_launcher_content_uses_configured_python() {
670        let installer = make_installer_with_python(
671            "/home/user/text-generation-webui",
672            "/home/user/rocm_venv/bin/python",
673        );
674        let content = installer.build_launcher_content("0,1");
675        assert!(
676            content.contains("/home/user/rocm_venv/bin/python server.py --chat"),
677            "Launcher must use configured Python interpreter"
678        );
679    }
680
681    #[test]
682    fn test_systemd_service_content() {
683        let installer = make_installer("/home/user/text-generation-webui");
684        let svc = installer.build_systemd_service("0,1");
685        assert!(svc.contains("[Unit]"));
686        assert!(svc.contains("text-generation-webui"));
687        assert!(svc.contains("HIP_VISIBLE_DEVICES=0,1"));
688        assert!(svc.contains("--chat"));
689    }
690
691    #[test]
692    fn test_amd_requirements_command() {
693        let installer = make_installer("/home/user/text-generation-webui");
694        let cmd = installer.build_amd_requirements_command(
695            "/home/user/text-generation-webui/requirements_amd.txt",
696        );
697        assert_eq!(cmd.program, "python3");
698        assert!(cmd.args.contains(&"-r".to_string()));
699        assert!(cmd
700            .args
701            .contains(&"/home/user/text-generation-webui/requirements_amd.txt".to_string()));
702    }
703}
rusty_stack/installers/components/textgen.rs

rusty_stack/installers/components/
textgen.rs
