diff options
Diffstat (limited to 'app/backends/sglomni/wizard.py')
| -rw-r--r-- | app/backends/sglomni/wizard.py | 103 |
1 files changed, 57 insertions, 46 deletions
diff --git a/app/backends/sglomni/wizard.py b/app/backends/sglomni/wizard.py index 96679d6..f39e754 100644 --- a/app/backends/sglomni/wizard.py +++ b/app/backends/sglomni/wizard.py @@ -13,7 +13,7 @@ upstream recipes + HuggingFace weight pre-download). It is driven by ``audiobook.py``'s hub but can also be run directly: Usage: - python -m backends.sglomni [--models KEY[,KEY...]] [--all] + python -m backends.sglomni [KEY ...] [--models KEY[,KEY...]] [--all] [--skip-install] [--skip-python] The Configure screen (``models_screen``) manages models after the fact — @@ -32,9 +32,9 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from backends import common, envs, servers, setup from backends.sglomni import catalog as sg_catalog +from backends.sglomni import gpu as sg_gpu from backends.sglomni import models as sg_models -from backends.sglomni.constants import SERVER_NAME, SGLOMNI_PIP_PKG, \ - UV_PIP_PKG +from backends.sglomni.constants import SERVER_NAME, SGLOMNI_PIP_PKG from backends.sglomni.pythonenv import SGLOMNI_ENV, prepare_env from backends.sglomni.status import _is_installed from ui import taskview, tui @@ -43,10 +43,11 @@ _GO_BACK = object() def _nvidia_gpu_present() -> bool: - """True when an NVIDIA driver answers nvidia-smi (best effort).""" - proc = common.run_console_subprocess_quiet( - ["nvidia-smi", "-L"], timeout=10) - return proc is not None and proc.returncode == 0 + """True when an NVIDIA driver answers nvidia-smi (best effort). + + The same probe the launch decisions use (sglomni.gpu): nvidia-smi + names GPU 0, or there is no usable answer.""" + return sg_gpu.describe() is not None def _preflight() -> List[str]: @@ -240,6 +241,10 @@ def run_tui(args: Optional[argparse.Namespace] = None) -> int: settings = curses.wrapper(lambda scr: _wizard(scr, args)) except tui.WizardCancelled: return 1 + try: + curses.curs_set(1) # restore the text cursor hidden by the TUI + except curses.error: + pass if settings is None: return 1 return _execute(settings) @@ -258,21 +263,27 @@ def _collect_from_flags(args: argparse.Namespace, print(f"[WARNING] {warning}") if args.all: keys = [entry.key for entry in sg_catalog.ENTRIES] - elif args.models: - keys = [] - for part in args.models.split(","): + else: + # Positional keys and --models both feed the same list (the + # positional form is the --models shorthand's space-separated + # twin); duplicates collapse, unknown keys stop the run with the + # known set. + keys: List[str] = [] + for part in list(args.models_pos or []) + \ + (args.models or "").split(","): key = part.strip() if not key: continue if sg_catalog.entry_by_key(key) is None: known = ", ".join(e.key for e in sg_catalog.ENTRIES) parser.error(f"unknown model key {key!r} (known: {known})") - keys.append(key) - else: - keys = [] - print("[INFO] No --models given: installing the package only " - "(use --models KEY[,KEY...] or --all to add models, or the " - "TUI's Configure screen).") + if key not in keys: + keys.append(key) + if not keys: + print("[INFO] No models given: installing the package only " + "(pass model keys — positional or --models KEY[,KEY…] — " + "or --all to add models, or use the TUI's Configure " + "screen).") return { "keys": keys, "do_python": not args.skip_python, @@ -357,33 +368,20 @@ def models_screen(stdscr) -> int: def uninstall(*, emit=None, cancel=None) -> int: """Remove the SGLang-Omni backend entirely. - Phases: stop the managed server, pip-uninstall sglang-omni and every - catalog model's companion packages, delete every cached weight - snapshot, then remove the tool-owned venv (app/envs/sglomni — the - heavyweight CUDA stack is the install, so unlike the lighter backends - the whole environment goes) and the uv-managed interpreters under - app/envs/pythons. CANCEL is honored between phases only. Returns the - exit code (130 when cancelled before a remaining phase). + Phases: stop the managed server, delete every catalog model's cached + weight snapshot, then remove the tool-owned venv (app/envs/sglomni — + the heavyweight CUDA stack is the install, so unlike the lighter + backends the whole environment goes) and the uv-managed interpreters + under app/envs/pythons. No pip-uninstall phase: the venv removal IS + the cleanup, and pip-ing the package plus every companion out of an + environment that is about to be deleted is minutes of pure wait time. + CANCEL is honored between phases only. Returns the exit code (130 + when cancelled before a remaining phase). """ if servers.pid_for(SERVER_NAME) is not None: servers.stop(SERVER_NAME) if common.cancel_requested(cancel): return 130 - packages = [SGLOMNI_PIP_PKG] - for entry in sg_catalog.ENTRIES: - for spec, _no_deps in entry.extras: - name = spec.split("=")[0].split("<")[0].split(">")[0].strip() - if name and name not in packages: - packages.append(name) - if envs.env_exists(SGLOMNI_ENV): - rc = common.pip_uninstall(packages, emit=emit, env_dir=SGLOMNI_ENV) - if rc != 0: - print(f"[WARNING] pip uninstall failed (exit {rc}); the venv " - "is removed below anyway") - else: - rc = 0 - if common.cancel_requested(cancel): - return 130 sg_models.delete_model_weights() if common.cancel_requested(cancel): return 130 @@ -395,17 +393,23 @@ def uninstall(*, emit=None, cancel=None) -> int: print(f"[WARNING] Could not fully remove {directory}") else: print(f"[OK] {directory} removed.") - return rc + return 0 def update(*, emit=None, cancel=None) -> int: """Update the sglang-omni backend: pip install -U in its venv. A managed server that is running is stopped first (best-effort): it - imports the very package being upgraded. CANCEL is honored between - phases only. Model weights are untouched (they live in the shared - HuggingFace cache and survive package upgrades). When the venv does - not exist there is nothing to update. Returns the exit code. + imports the very package being upgraded. The upgrade is followed by a + companion refresh — every installed model's extras re-run (a pin + already satisfied is a pip no-op, so this is cheap when nothing + drifted) — so a newer sglang-omni's companion requirements are met + the way a fresh install would meet them; a failing extra warns and + leaves the update successful (the import probe re-heals it at the + next model install or server start). Model weights are untouched + (they live in the shared HuggingFace cache and survive package + upgrades). When the venv does not exist there is nothing to update. + CANCEL is honored between phases only. Returns the exit code. """ if servers.pid_for(SERVER_NAME) is not None: servers.stop(SERVER_NAME) @@ -420,9 +424,16 @@ def update(*, emit=None, cancel=None) -> int: if rc != 0: print(f"[WARNING] pip install -U failed (exit {rc}); update " f"{SGLOMNI_PIP_PKG} manually") - else: - print(f"[OK] {SGLOMNI_PIP_PKG} is up to date (or just upgraded).") - return rc + return rc + print(f"[OK] {SGLOMNI_PIP_PKG} is up to date (or just upgraded).") + for entry in sg_models.installed_entries(): + crc = sg_models.install_companions(entry, emit=emit, cancel=cancel, + force=True) + if crc != 0: + print(f"[WARNING] {entry.label}'s companion packages could " + "not all be refreshed; the next install or server start " + "retries what the import probe finds missing.") + return 0 def main() -> int: |