diff --git a/src/gpu_rent/cli.py b/src/gpu_rent/cli.py index 7f0fd5d..5de3da7 100644 --- a/src/gpu_rent/cli.py +++ b/src/gpu_rent/cli.py @@ -10,7 +10,6 @@ from datetime import datetime, timezone from typing import Optional import typer -from rich.console import Console from rich.table import Table from gpu_rent import __version__ @@ -23,22 +22,15 @@ from gpu_rent.ssh_ops import interactive_ssh, run_ssh from gpu_rent.state import load_state, preempt_window_end from gpu_rent.provision import ensure_swarmui_running, seed_civitai, seed_extensions from gpu_rent.sync_files import pull_tree, push_tree +from gpu_rent.term import console, err, log, ok, warn from gpu_rent.tunnel import run_tunnel -if sys.platform == "win32": - for _stream in (sys.stdout, sys.stderr): - try: - _stream.reconfigure(encoding="utf-8", errors="replace") - except (AttributeError, OSError): - pass - app = typer.Typer( no_args_is_help=True, pretty_exceptions_enable=False, add_completion=False, help="Прерываемый GPU Selectel + SwarmUI на localhost:17801. Сначала: gpu-rent setup / doctor. Ключи: docs/setup.md", ) -console = Console(highlight=False, legacy_windows=False) _DEBUG = False @@ -54,7 +46,7 @@ def _root( def _die(exc: BaseException) -> None: if _DEBUG: traceback.print_exc() - console.print(f"[red]{exc}[/red]") + err(str(exc)) raise typer.Exit(1) @@ -368,7 +360,7 @@ def setup( install_watchdog=watchdog if watchdog is not None else (False if yes else None), confirm=None if yes and watchdog is None else confirm, ask=None if yes and llm is not None else ask, - log=lambda m: console.print(m), + log=log, ) except Exception as exc: _die(exc) @@ -412,7 +404,6 @@ def up( from dataclasses import replace from gpu_rent.llm_runtime import ( - PRESET_HELP, append_vars_llm_runtime, decide_runtime, ensure_ollama_manifest_from_example, @@ -438,11 +429,24 @@ def up( asked_model_preset = False if not yes and runtime == "none" and not llm and not ollama and not llamacpp: - choice = typer.prompt( - "Поднять LLM рядом со SwarmUI? [none/ollama/llamacpp]", - default="none", + from gpu_rent.llm_runtime import ( + llamacpp_preset_menu, + llm_runtime_menu, + ollama_preset_menu, ) + from gpu_rent.prompts import prompt_menu + + def _ask(msg: str, default: str = "") -> str: + return typer.prompt(msg, default=default) + try: + choice = prompt_menu( + "LLM рядом со SwarmUI", + llm_runtime_menu(), + default="none", + ask=_ask, + show=log, + ) runtime = decide_runtime( flag=choice, ollama_flag=False, llamacpp_flag=False, from_config="none" ) @@ -452,63 +456,82 @@ def up( append_vars_llm_runtime(vars_path(), runtime) if runtime == "ollama": ensure_ollama_manifest_from_example() - console.print(PRESET_HELP) - preset = typer.prompt("Ollama preset", default="recommended") - if preset.strip().lower() not in {"keep", "example"}: - write_ollama_models_preset( - cfg.ollama_models_manifest, preset.strip().lower() + try: + preset = prompt_menu( + "Ollama preset", + ollama_preset_menu(include_keep=False), + default="recommended", + ask=_ask, + show=log, ) + except ValueError as exc: + raise GpuRentError(str(exc)) from exc + if preset not in {"keep", "example"}: + write_ollama_models_preset(cfg.ollama_models_manifest, preset) asked_model_preset = True elif runtime == "llamacpp": from gpu_rent.llm_runtime import ( - LLAMACPP_PRESET_HELP, ensure_llamacpp_manifest_from_example, write_llamacpp_models_preset, ) ensure_llamacpp_manifest_from_example() - console.print(LLAMACPP_PRESET_HELP) - preset = typer.prompt( - "llama.cpp GGUF preset [recommended/light/stock/empty]", - default="recommended", - ) - if preset.strip().lower() not in {"keep", "example"}: - write_llamacpp_models_preset( - cfg.llamacpp_models_manifest, preset.strip().lower() + try: + preset = prompt_menu( + "llama.cpp GGUF", + llamacpp_preset_menu(include_keep=False), + default="recommended", + ask=_ask, + show=log, ) + except ValueError as exc: + raise GpuRentError(str(exc)) from exc + if preset not in {"keep", "example"}: + write_llamacpp_models_preset(cfg.llamacpp_models_manifest, preset) asked_model_preset = True - # Runtime уже в vars (напр. llamacpp) — всё равно спросить модель, default=keep. - if not yes and not asked_model_preset and runtime == "llamacpp": + # Runtime уже в vars — спросить пресет, default=keep. + if not yes and not asked_model_preset and runtime in {"llamacpp", "ollama"}: from gpu_rent.llm_runtime import ( - LLAMACPP_PRESET_HELP, ensure_llamacpp_manifest_from_example, + llamacpp_preset_menu, + ollama_preset_menu, write_llamacpp_models_preset, ) + from gpu_rent.prompts import prompt_menu - ensure_llamacpp_manifest_from_example() - console.print(LLAMACPP_PRESET_HELP) - preset = typer.prompt( - "llama.cpp GGUF preset [recommended/light/stock/empty/keep]", - default="keep", - ) - key = preset.strip().lower() - if key not in {"keep", "example", ""}: - write_llamacpp_models_preset(cfg.llamacpp_models_manifest, key) - elif not yes and not asked_model_preset and runtime == "ollama": - ensure_ollama_manifest_from_example() - console.print(PRESET_HELP) - preset = typer.prompt( - "Ollama preset [recommended/light/stock/alt/empty/keep]", - default="keep", - ) - key = preset.strip().lower() - if key not in {"keep", "example", ""}: - write_ollama_models_preset(cfg.ollama_models_manifest, key) + def _ask2(msg: str, default: str = "") -> str: + return typer.prompt(msg, default=default) + + try: + if runtime == "llamacpp": + ensure_llamacpp_manifest_from_example() + key = prompt_menu( + "llama.cpp GGUF", + llamacpp_preset_menu(include_keep=True), + default="keep", + ask=_ask2, + show=log, + ) + if key not in {"keep", "example", ""}: + write_llamacpp_models_preset(cfg.llamacpp_models_manifest, key) + else: + ensure_ollama_manifest_from_example() + key = prompt_menu( + "Ollama preset", + ollama_preset_menu(include_keep=True), + default="keep", + ask=_ask2, + show=log, + ) + if key not in {"keep", "example", ""}: + write_ollama_models_preset(cfg.ollama_models_manifest, key) + except ValueError as exc: + raise GpuRentError(str(exc)) from exc cfg = replace(cfg, llm_runtime=runtime) if runtime != "none": - console.print(f"LLM runtime: {runtime}") + ok(f"LLM runtime: {runtime}") def confirm(msg: str) -> bool: return typer.confirm(msg) @@ -525,7 +548,7 @@ def up( update=False if no_update else None, confirm=confirm, ask=None if yes else ask, - log=lambda m: console.print(m), + log=log, ) if no_tunnel: from gpu_rent.access_card import print_access_card @@ -550,7 +573,7 @@ def up( cfg, state.floating_ip, open_browser=open_browser, - log=lambda m: console.print(m), + log=log, ) except GpuRentError as exc: _die(exc) @@ -563,7 +586,7 @@ def stop( """Удалить compute и FIP, диски оставить.""" try: cfg = load_config(require_auth=True) - cmd_stop(cfg, no_pull=no_pull, log=lambda m: console.print(m)) + cmd_stop(cfg, no_pull=no_pull, log=log) except GpuRentError as exc: _die(exc) @@ -579,7 +602,7 @@ def destroy( raise typer.Exit(1) try: cfg = load_config(require_auth=True) - cmd_stop(cfg, destroy_disks=True, no_pull=no_pull, log=lambda m: console.print(m)) + cmd_stop(cfg, destroy_disks=True, no_pull=no_pull, log=log) except GpuRentError as exc: _die(exc) @@ -636,7 +659,7 @@ def tunnel( cfg, state.floating_ip, open_browser=open_browser, - log=lambda m: console.print(m), + log=log, ) except GpuRentError as exc: _die(exc) @@ -654,9 +677,9 @@ def hold( cfg, host = _live() if clear: - clear_hold(cfg, host, log=lambda m: console.print(m)) + clear_hold(cfg, host, log=log) return - set_hold(cfg, host, minutes=minutes, until=until, log=lambda m: console.print(m)) + set_hold(cfg, host, minutes=minutes, until=until, log=log) except GpuRentError as exc: _die(exc) @@ -666,7 +689,7 @@ def seed_models() -> None: """Докачать Civitai-манифест на живой диск.""" try: cfg, host = _live() - seed_civitai(cfg, host, log=lambda m: console.print(m)) + seed_civitai(cfg, host, log=log) except GpuRentError as exc: _die(exc) @@ -677,9 +700,6 @@ def push_all() -> None: try: cfg, host = _live() - def log(msg: str) -> None: - console.print(msg) - push_tree(cfg, host, cfg.local_models_dir, "/mnt/swarm_data/Models", log, models=True) push_tree(cfg, host, cfg.local_wildcards_dir, "/mnt/swarm_data/Data/Wildcards", log, models=False) push_tree(cfg, host, cfg.local_workflows_dir, "/mnt/swarm_data/CustomWorkflows", log, models=False) @@ -697,7 +717,7 @@ def push_models() -> None: host, cfg.local_models_dir, "/mnt/swarm_data/Models", - lambda m: console.print(m), + log, models=True, ) except GpuRentError as exc: @@ -714,7 +734,7 @@ def pull_output_cmd() -> None: host, "/mnt/swarm_data/Output", cfg.local_output_dir, - lambda m: console.print(m), + log, ) except GpuRentError as exc: _die(exc) @@ -726,9 +746,6 @@ def seed_extensions_cmd() -> None: try: cfg, host = _live() - def log(msg: str) -> None: - console.print(msg) - seed_extensions(cfg, host, log) ensure_swarmui_running(cfg, host, log, restart=True) except GpuRentError as exc: @@ -768,9 +785,9 @@ def capture_root( host, dry_run=dry_run, kind_filter=kind, - log=lambda m: console.print(m), + log=log, ) - print_report(report, lambda m: console.print(m), dry_run=dry_run) + print_report(report, log, dry_run=dry_run) except GpuRentError as exc: _die(exc) @@ -796,11 +813,11 @@ def capture_models_cmd( None, dry_run=dry_run, kind_filter=kind, - log=lambda m: console.print(m), + log=log, ) print_report( report, - lambda m: console.print(m), + log, dry_run=dry_run, show_models=True, show_extensions=False, @@ -823,11 +840,11 @@ def capture_extensions_cmd( host, None, dry_run=dry_run, - log=lambda m: console.print(m), + log=log, ) print_report( report, - lambda m: console.print(m), + log, dry_run=dry_run, show_models=False, show_extensions=True, @@ -854,9 +871,9 @@ def capture_all_cmd( host, dry_run=dry_run, kind_filter=kind, - log=lambda m: console.print(m), + log=log, ) - print_report(report, lambda m: console.print(m), dry_run=dry_run) + print_report(report, log, dry_run=dry_run) except GpuRentError as exc: _die(exc) @@ -868,7 +885,7 @@ def resize_data(gb: int = typer.Option(..., "--gb", help="Новый разме from gpu_rent.resize import resize_data_volume cfg = load_config(require_auth=True) - resize_data_volume(cfg, gb, log=lambda m: console.print(m)) + resize_data_volume(cfg, gb, log=log) except GpuRentError as exc: _die(exc) @@ -905,9 +922,9 @@ def watchdog_install( try: from gpu_rent.local_watchdog import install_watchdog - install_watchdog(interval_minutes=interval, log=lambda m: console.print(m)) + install_watchdog(interval_minutes=interval, log=log) except Exception as exc: - console.print(f"[red]install fail:[/red] {exc}") + err(f"install fail: {exc}") raise typer.Exit(1) from exc @@ -917,9 +934,9 @@ def watchdog_uninstall() -> None: try: from gpu_rent.local_watchdog import uninstall_watchdog - uninstall_watchdog(log=lambda m: console.print(m)) + uninstall_watchdog(log=log) except Exception as exc: - console.print(f"[red]uninstall fail:[/red] {exc}") + err(f"uninstall fail: {exc}") raise typer.Exit(1) from exc @@ -946,7 +963,7 @@ def watchdog_tick( try: from gpu_rent.local_watchdog import run_tick - decision = run_tick(dry_run=dry_run, log=lambda m: console.print(m)) + decision = run_tick(dry_run=dry_run, log=log) if decision.kind == "stop" and not dry_run: raise typer.Exit(0) except GpuRentError as exc: diff --git a/src/gpu_rent/llm_runtime.py b/src/gpu_rent/llm_runtime.py index 2564010..ed4d769 100644 --- a/src/gpu_rent/llm_runtime.py +++ b/src/gpu_rent/llm_runtime.py @@ -26,12 +26,18 @@ OLLAMA_PRESETS: dict[str, list[str]] = { "empty": [], } -PRESET_HELP = ( - "recommended — Qwen2.5 7B abliterate (RU/EN, мало отказов, ~5GB)\n" - "light — qwen2.5:3b (быстрее, слабее)\n" - "stock — официальный qwen2.5:7b (больше цензуры)\n" - "alt — другой abliterate-пак 7B\n" - "empty — только runtime, без pull" +OLLAMA_PRESET_LABELS: dict[str, str] = { + "recommended": "Qwen2.5 7B abliterate (RU/EN, мало отказов, ~5GB)", + "light": "qwen2.5:3b (быстрее, слабее)", + "stock": "официальный qwen2.5:7b (больше цензуры)", + "alt": "другой abliterate-пак 7B", + "empty": "только runtime, без pull", + "keep": "не менять ollama-models.yaml", +} + +# Deprecated text blob — prefer menu helpers below. +PRESET_HELP = "\n".join( + f"{k} — {v}" for k, v in OLLAMA_PRESET_LABELS.items() if k != "keep" ) LLAMACPP_PRESETS: dict[str, list[str]] = { @@ -47,13 +53,48 @@ LLAMACPP_PRESETS: dict[str, list[str]] = { "empty": [], } -LLAMACPP_PRESET_HELP = ( - "recommended — Qwen2.5 7B abliterate GGUF Q4_K_M (~4.7GB, мало отказов)\n" - "light — Qwen2.5 3B Instruct Q4_K_M (~2GB)\n" - "stock — официальный Qwen2.5 7B Instruct Q4_K_M\n" - "empty — только llama-server, GGUF положи вручную / правь llamacpp-models.yaml" +LLAMACPP_PRESET_LABELS: dict[str, str] = { + "recommended": "Qwen2.5 7B abliterate GGUF Q4_K_M (~4.7GB, мало отказов)", + "light": "Qwen2.5 3B Instruct Q4_K_M (~2GB)", + "stock": "официальный Qwen2.5 7B Instruct Q4_K_M", + "empty": "только llama-server, GGUF вручную", + "keep": "не менять llamacpp-models.yaml", +} + +LLAMACPP_PRESET_HELP = "\n".join( + f"{k} — {v}" for k, v in LLAMACPP_PRESET_LABELS.items() if k != "keep" ) +LLM_RUNTIME_LABELS: dict[str, str] = { + "none": "только SwarmUI", + "ollama": "Ollama (+ pull моделей)", + "llamacpp": "llama.cpp server (+ GGUF)", +} + + +def llm_runtime_menu() -> list: + from gpu_rent.prompts import MenuItem + + return [MenuItem(k, f"{k} — {LLM_RUNTIME_LABELS[k]}") for k in ("none", "ollama", "llamacpp")] + + +def ollama_preset_menu(*, include_keep: bool = False) -> list: + from gpu_rent.prompts import MenuItem + + keys = list(OLLAMA_PRESETS.keys()) + if include_keep: + keys.append("keep") + return [MenuItem(k, OLLAMA_PRESET_LABELS.get(k, k)) for k in keys] + + +def llamacpp_preset_menu(*, include_keep: bool = False) -> list: + from gpu_rent.prompts import MenuItem + + keys = list(LLAMACPP_PRESETS.keys()) + if include_keep: + keys.append("keep") + return [MenuItem(k, LLAMACPP_PRESET_LABELS.get(k, k)) for k in keys] + @dataclass(frozen=True) class OllamaModelEntry: diff --git a/src/gpu_rent/pools.py b/src/gpu_rent/pools.py index 19e50a5..ed0a594 100644 --- a/src/gpu_rent/pools.py +++ b/src/gpu_rent/pools.py @@ -186,31 +186,35 @@ def scan_pools(cfg: Config, pools: tuple[str, ...] | None = None) -> list[PoolGp def format_pool_scan(offers: list[PoolGpuOffer], preference: tuple[str, ...]) -> list[str]: lines = [ - "скан пулов (мультизональные + кандидаты) × характеристики GPU:", - " (серые кнопки панели ≠ Nova; available_count бывает только у части flavors)", + "[bold cyan]скан пулов[/bold cyan] (мультизональные + кандидаты) × характеристики GPU:", + " [dim](серые кнопки панели ≠ Nova; available_count бывает только у части flavors)[/dim]", ] for offer in offers: tag = "multizone" if offer.multizone else "pool" - prefix = f" {offer.region} ({tag}):" + prefix = f" [bold]{offer.region}[/bold] ({tag}):" if offer.error: - lines.append(f"{prefix} ОШИБКА — {offer.error}") + lines.append(f"{prefix} [red]ОШИБКА[/red] — {offer.error}") continue types_s = ", ".join(offer.gpu_types) if offer.gpu_types else "—" lines.append(f"{prefix} типы: {types_s}") if offer.gpu_labels_found: best = offer.matched[0] if offer.matched else None - best_s = f" → лучший {best.name} ({best.label})" if best else "" + best_s = ( + f" → [green]лучший {best.name} ({best.label})[/green]" if best else "" + ) lines.append( f" preference: {', '.join(offer.gpu_labels_found)}{best_s}" ) else: - lines.append(f" preference: нет совпадений с {','.join(preference)}") + lines.append( + f" [yellow]preference: нет совпадений с {','.join(preference)}[/yellow]" + ) for offer in offers: if offer.matched and not offer.error: lines.append( - f"рекомендация: OS_REGION_NAME={offer.region} " + f"[green]рекомендация:[/green] OS_REGION_NAME={offer.region} " f"GPU_RENT_AZ={offer.region}a " - f"(сегмент a/b/c — в панели; диски и VM в одном AZ)" + f"[dim](сегмент a/b/c — в панели; диски и VM в одном AZ)[/dim]" ) break return lines diff --git a/src/gpu_rent/prompts.py b/src/gpu_rent/prompts.py new file mode 100644 index 0000000..2f1e984 --- /dev/null +++ b/src/gpu_rent/prompts.py @@ -0,0 +1,72 @@ +"""Numbered interactive menus (Enter = default index).""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass + +from rich.markup import escape + +Ask = Callable[[str, str], str] +Show = Callable[[str], None] + + +@dataclass(frozen=True) +class MenuItem: + key: str + label: str + + +def resolve_menu_choice( + raw: str, + items: Sequence[MenuItem], + *, + default_key: str, +) -> str: + """Accept 1-based index or exact key (case-insensitive). Empty → default.""" + text = (raw or "").strip().lower() + if not text: + return default_key + if text.isdigit(): + idx = int(text) + if 1 <= idx <= len(items): + return items[idx - 1].key + raise ValueError(f"выбор: жду 1…{len(items)}, получили {idx}") + for item in items: + if item.key.lower() == text: + return item.key + keys = ", ".join(f"{i}:{it.key}" for i, it in enumerate(items, 1)) + raise ValueError(f"выбор: жду номер или ключ ({keys}), получили {raw!r}") + + +def prompt_menu( + title: str, + items: Sequence[MenuItem], + *, + default: str, + ask: Ask, + show: Show | None = None, +) -> str: + """Print numbered list, ask for index (default = current recommendation).""" + if not items: + raise ValueError(f"{title}: пустое меню") + default_idx = 1 + for i, item in enumerate(items, 1): + if item.key == default: + default_idx = i + break + safe_title = escape(title) + lines = [f"[bold cyan]{safe_title}[/bold cyan]"] + for i, item in enumerate(items, 1): + label = escape(item.label) + if i == default_idx: + lines.append(f" [cyan]{i}.[/cyan] {label} [green]←[/green]") + else: + lines.append(f" [dim]{i}.[/dim] {label}") + emit = show + if emit is None: + from gpu_rent.term import log as emit # type: ignore[assignment] + for line in lines: + emit(line) + raw = ask(f"{title} [1-{len(items)}]", str(default_idx)) + return resolve_menu_choice(raw, items, default_key=items[default_idx - 1].key) diff --git a/src/gpu_rent/setup_wizard.py b/src/gpu_rent/setup_wizard.py index 69dbadd..d36aeb4 100644 --- a/src/gpu_rent/setup_wizard.py +++ b/src/gpu_rent/setup_wizard.py @@ -7,12 +7,13 @@ from collections.abc import Callable from pathlib import Path from gpu_rent.llm_runtime import ( - LLAMACPP_PRESET_HELP, - PRESET_HELP, append_vars_llm_runtime, ensure_llamacpp_manifest_from_example, ensure_ollama_manifest_from_example, + llamacpp_preset_menu, + llm_runtime_menu, normalize_runtime, + ollama_preset_menu, write_llamacpp_models_preset, write_ollama_models_preset, ) @@ -28,6 +29,7 @@ from gpu_rent.paths import ( vars_example_path, vars_path, ) +from gpu_rent.prompts import prompt_menu Log = Callable[[str], None] @@ -74,9 +76,12 @@ def run_setup( runtime = llm if runtime is None: if ask: - runtime = ask( - "LLM runtime [none/ollama/llamacpp]", - "none", + runtime = prompt_menu( + "LLM runtime", + llm_runtime_menu(), + default="none", + ask=ask, + show=log, ) else: runtime = "none" @@ -87,8 +92,13 @@ def run_setup( if runtime == "ollama": preset = ollama_preset if preset is None and ask: - log(PRESET_HELP) - preset = ask("Ollama preset [recommended/light/stock/alt/empty]", "recommended") + preset = prompt_menu( + "Ollama preset", + ollama_preset_menu(include_keep=False), + default="recommended", + ask=ask, + show=log, + ) if preset is None: preset = "recommended" if preset.strip().lower() in {"keep", "example", ""}: @@ -98,12 +108,14 @@ def run_setup( write_ollama_models_preset(ollama_models_manifest_path(), preset) log(f"ollama-models.yaml пресет={preset}") elif runtime == "llamacpp": - preset = ollama_preset # reuse --ollama-preset flag as generic LLM preset in setup + preset = ollama_preset if preset is None and ask: - log(LLAMACPP_PRESET_HELP) - preset = ask( - "llama.cpp GGUF preset [recommended/light/stock/empty]", - "recommended", + preset = prompt_menu( + "llama.cpp GGUF", + llamacpp_preset_menu(include_keep=False), + default="recommended", + ask=ask, + show=log, ) if preset is None: preset = "recommended" diff --git a/src/gpu_rent/term.py b/src/gpu_rent/term.py new file mode 100644 index 0000000..f8ccaaa --- /dev/null +++ b/src/gpu_rent/term.py @@ -0,0 +1,105 @@ +"""Colored terminal output (Rich). Plain log lines get severity colors.""" + +from __future__ import annotations + +import re +import sys + +from rich.console import Console +from rich.markup import escape + +if sys.platform == "win32": + for _stream in (sys.stdout, sys.stderr): + try: + _stream.reconfigure(encoding="utf-8", errors="replace") + except (AttributeError, OSError): + pass + +console = Console(highlight=False, legacy_windows=False) + +# Our own markup — leave as-is. +_MARKUP = re.compile( + r"\[(/?)(bold|dim|red|green|yellow|cyan|magenta|blue|white|" + r"bold [a-z]+|dim [a-z]+)\]" +) + +_ERR = re.compile( + r"(?i)\b(" + r"fail|failed|error|exception|traceback|отмен|ошиб|не\s+удал|" + r"блокирующ|refuse|denied|abort" + r")\b" +) +_WARN = re.compile( + r"(?i)(" + r"^!\s|" + r"\bwarn|\bwarning|предупрежд|внимание|" + r"тарифицир|idle-killer|preemptible:|" + r"₽:|диск data|серые кнопки|не\s+гасит|" + r"сейчас OS_REGION|лучше закрывается" + r")" +) +_OK = re.compile( + r"(?i)(" + r"^ok\b|готово|создал|создан|doctor ok|SSH ок|" + r"туннель 127|← выберем|рекомендация:|" + r"выбрано:|armed|bootstrap уже|слушает 127|" + r"pull Output:|push .+: \d|wrote " + r")" +) + + +def has_markup(text: str) -> bool: + return bool(_MARKUP.search(text)) + + +def paint(msg: str) -> str: + """Wrap a plain log line in Rich markup by severity heuristics.""" + if msg == "" or msg.isspace(): + return msg + if has_markup(msg): + return msg + + stripped = msg.lstrip() + body = escape(msg) + + if stripped.startswith("!"): + return f"[yellow]{body}[/yellow]" + if _ERR.search(stripped): + return f"[red]{body}[/red]" + if _WARN.search(stripped): + return f"[yellow]{body}[/yellow]" + if _OK.search(stripped): + return f"[green]{body}[/green]" + + # Numbered menus / lists + if re.match(r"^\s*\d+\.\s", msg): + return f"[white]{body}[/white]" + if stripped.endswith(":") and len(stripped) < 80 and not stripped.startswith("http"): + return f"[bold cyan]{body}[/bold cyan]" + + return body + + +def log(msg: str = "") -> None: + """CLI log callback — colored print.""" + console.print(paint(msg)) + + +def ok(msg: str) -> None: + console.print(f"[green]{escape(msg)}[/green]") + + +def warn(msg: str) -> None: + console.print(f"[yellow]{escape(msg)}[/yellow]") + + +def err(msg: str) -> None: + console.print(f"[red]{escape(msg)}[/red]") + + +def info(msg: str) -> None: + console.print(f"[cyan]{escape(msg)}[/cyan]") + + +def dim(msg: str) -> None: + console.print(f"[dim]{escape(msg)}[/dim]") diff --git a/src/gpu_rent/ux.py b/src/gpu_rent/ux.py index 22a3208..3b751a9 100644 --- a/src/gpu_rent/ux.py +++ b/src/gpu_rent/ux.py @@ -6,6 +6,8 @@ from collections.abc import Callable from dataclasses import dataclass from typing import Any +from rich.markup import escape + from gpu_rent.config import Config from gpu_rent.inventory import FlavorInfo, looks_like_gpu, rank_flavors, resolve_flavor @@ -21,14 +23,19 @@ def list_ranked_flavors(flavors: list[Any], cfg: Config) -> list[FlavorInfo]: def format_flavor_lines(ranked: list[FlavorInfo], picked: FlavorInfo | None = None) -> list[str]: if not ranked: - return ["flavors: пусто (проверь регион / FLAVOR_PREFERENCE)"] - lines = ["flavors (по FLAVOR_PREFERENCE):"] + return ["[yellow]flavors: пусто (проверь регион / FLAVOR_PREFERENCE)[/yellow]"] + lines = ["[bold cyan]flavors (по FLAVOR_PREFERENCE):[/bold cyan]"] for i, info in enumerate(ranked, 1): - mark = " ← выберем" if picked and info.id == picked.id else "" + mark = " [green]← выберем[/green]" if picked and info.id == picked.id else "" ram = f"{info.ram_mb // 1024}GB" if info.ram_mb else "?" vcpu = str(info.vcpus) if info.vcpus is not None else "?" - label = info.label or "—" - lines.append(f" {i}. [{label}] {info.name} vCPU={vcpu} RAM={ram} id={info.id}{mark}") + label = escape(info.label or "—") + name = escape(info.name) + fid = escape(info.id) + lines.append( + f" [cyan]{i}.[/cyan] [bold]{label}[/bold] {name} " + f"vCPU={vcpu} RAM={ram} id={fid}{mark}" + ) return lines @@ -114,21 +121,18 @@ def prompt_server_plan( if data_gb < 20: raise ValueError("Data disk GB: минимум 20") - spot_default = "Y" if spot else "n" - raw_spot = ( - ask( - "Preemptible GPU (дешевле, могут усыпить ~24ч)? [Y/n]", - spot_default, - ) - .strip() - .lower() + from gpu_rent.prompts import MenuItem, prompt_menu + + spot_key = prompt_menu( + "Тариф GPU", + [ + MenuItem("yes", "preemptible (дешевле, могут усыпить ~24ч)"), + MenuItem("no", "обычный on-demand"), + ], + default="yes" if spot else "no", + ask=ask, ) - if raw_spot in {"", "y", "yes", "1", "true", "on"}: - use_spot = True - elif raw_spot in {"n", "no", "0", "false", "off"}: - use_spot = False - else: - raise ValueError(f"preemptible: жду Y/n, получили {raw_spot!r}") + use_spot = spot_key == "yes" changed = ( chosen.id != picked.id diff --git a/tests/test_prompts.py b/tests/test_prompts.py new file mode 100644 index 0000000..b00c908 --- /dev/null +++ b/tests/test_prompts.py @@ -0,0 +1,32 @@ +from gpu_rent.prompts import MenuItem, prompt_menu, resolve_menu_choice + + +def test_resolve_by_number(): + items = [MenuItem("a", "A"), MenuItem("b", "B"), MenuItem("keep", "K")] + assert resolve_menu_choice("2", items, default_key="keep") == "b" + assert resolve_menu_choice("", items, default_key="keep") == "keep" + assert resolve_menu_choice("keep", items, default_key="a") == "keep" + + +def test_prompt_menu_prints_and_picks(): + items = [ + MenuItem("recommended", "rec"), + MenuItem("light", "lit"), + MenuItem("keep", "keep yaml"), + ] + shown: list[str] = [] + + def ask(msg: str, default: str = "") -> str: + assert default == "3" # keep is default + return "1" + + key = prompt_menu( + "GGUF", + items, + default="keep", + ask=ask, + show=shown.append, + ) + assert key == "recommended" + assert any("1. rec" in line for line in shown) + assert any("←" in line for line in shown) diff --git a/tests/test_server_plan.py b/tests/test_server_plan.py index e46a087..c931a60 100644 --- a/tests/test_server_plan.py +++ b/tests/test_server_plan.py @@ -31,7 +31,7 @@ def test_prompt_server_plan_with_ranked(monkeypatch, tmp_path): "gpu_rent.ux.list_ranked_flavors", lambda flavors, cfg: [f1, f2], ) - answers = iter(["2", "200", "n"]) + answers = iter(["2", "200", "2"]) def ask(msg: str, default: str = "") -> str: return next(answers)