output.py

"""Rich console output helpers for the keras-remote CLI."""

import random
import time

from rich.console import Console
from rich.live import Live
from rich.panel import Panel
from rich.table import Table
from rich.text import Text

from keras_remote.core.accelerators import GpuConfig, TpuConfig

console = Console()

_SPINNER_FRAMES = ("⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏")

_SUBTITLE_MESSAGES = (
  # Fun phrases and helpful tips, interleaved.
  "Painting the pods",
  "Tip: Pass Data('./dataset/') as a function arg to upload data",
  "Winding all the butterflies",
  "Tip: Use volumes={'/mnt': Data('./data/')} to mount data on the pod",
  "Warming the compute engine",
  "Tip: Data is content-hashed — identical data is uploaded only once",
  "Reticulating splines",
  "Tip: Data() accepts GCS URIs too, e.g. Data('gs://bucket/path/')",
  "Charging the flux capacitor",
  "Tip: Data objects nested in lists/dicts are recursively discovered",
  "Aligning the cloud crystals",
  "Tip: Container images are content-hashed — unchanged deps skip rebuilds",
  "Feeding the hamsters",
  "Tip: Add a requirements.txt or pyproject.toml to auto-install deps on the remote pod",
  "Consulting the oracle",
  "Tip: Use --cluster to manage multiple clusters in the same project",
  "Calibrating the widgets",
  "Tip: Run 'keras-remote pool add --accelerator v5p-8' to add a TPU pool",
  "Herding the containers",
  "Tip: Run 'keras-remote pool list' to see all accelerators on your cluster",
  "Polishing the tensors",
  "Tip: Pass --yes to 'pool add/remove' to skip the confirmation prompt",
  "Summoning the cluster spirits",
  "Tip: Use cluster= in @run() to pick a cluster (or env KERAS_REMOTE_CLUSTER)",
  "Untangling the neural pathways",
  "Tip: Set zone= in @run() to pick a GCP zone (or env KERAS_REMOTE_ZONE)",
  "Brewing the cloud juice",
  "Tip: Use capture_env_vars=['PREFIX_*'] in @run() to forward env vars to the worker",
  "Wrangling the cloud gremlins",
  "Tip: Multi-host TPUs (e.g. v6e-4x4) auto-select the Pathways backend",
  "Compiling the butterfly wings",
  "Tip: Your working directory is auto-zipped and sent to the pod",
  "Tuning the hyperparameters of the universe",
  "Tip: Remote exceptions are re-raised locally with original traceback",
  "Spinning up the hamster wheels",
  "Tip: Run 'keras-remote config show' to check your current settings",
  "Negotiating with the load balancer",
  "Tip: Use container_image= in @run() to bring your own Docker image",
  "Teaching the pods to dance",
  "Tip: Use namespace= in @run() to pick a K8s namespace",
  "Downloading more RAM",
  "Tip: Set KERAS_REMOTE_PROJECT or --project to pick a specific GCP project",
)


class LiveOutputPanel:
  """Context manager that displays streaming output in a Rich Live panel.

  Shows the last `max_lines` in a bordered box. Supports error state
  (yellow border) and optional transient mode (clears on success).

  In non-interactive terminals, falls back to plain console output.
  """

  def __init__(
    self,
    title,
    *,
    max_lines=7,
    target_console=None,
    transient=False,
    show_subtitle=True,
  ):
    self._title = title
    self._max_lines = max_lines
    self._lines = []
    self._has_error = False
    self._transient = transient
    self._show_subtitle = show_subtitle
    self._console = target_console or console
    self._live = None
    self._start_time = None
    self._phrase_order = None

  def __enter__(self):
    self._start_time = time.monotonic()
    self._phrase_order = list(range(len(_SUBTITLE_MESSAGES)))
    random.shuffle(self._phrase_order)
    if self._console.is_terminal:
      self._live = Live(
        self,
        console=self._console,
        refresh_per_second=5,
      )
      self._live.__enter__()
    else:
      self._console.rule(self._title, style="blue")
    return self

  def __exit__(self, exc_type, exc_val, exc_tb):
    if exc_type is not None:
      self._has_error = True
    if self._live:
      if self._transient and not self._has_error:
        self._live.update(Text(""))
      self._live.__exit__(exc_type, exc_val, exc_tb)
    else:
      style = "yellow" if self._has_error else "blue"
      self._console.rule(style=style)
    return False

  def __rich__(self):
    return self._make_panel()

  def on_output(self, line):
    """Append a line and refresh the display."""
    stripped = line.rstrip("\n")
    if self._live:
      self._lines.append(stripped)
    else:
      self._console.print(stripped)

  def mark_error(self):
    """Turn the panel border yellow to indicate an error."""
    self._has_error = True
    if self._live:
      self._live.refresh()

  def _make_subtitle(self):
    if self._start_time is None or self._phrase_order is None:
      return None
    elapsed = time.monotonic() - self._start_time
    spinner_idx = int(elapsed * 4) % len(_SPINNER_FRAMES)
    spinner = _SPINNER_FRAMES[spinner_idx]
    msg_idx = int(elapsed / 4) % len(_SUBTITLE_MESSAGES)
    message = _SUBTITLE_MESSAGES[self._phrase_order[msg_idx]]
    suffix = "" if message.startswith("Tip:") else "..."
    return f"[italic]{spinner} {message}{suffix}[/italic]"

  def _make_panel(self):
    if self._lines:
      visible = (
        self._lines if self._has_error else self._lines[-self._max_lines :]
      )
      content = "\n".join(visible)
    else:
      content = "Waiting..."
    style = "yellow" if self._has_error else "blue"
    return Panel(
      content,
      title=self._title,
      subtitle=self._make_subtitle()
      if self._show_subtitle and not self._has_error
      else None,
      border_style=style,
    )


def banner(text):
  """Display a styled banner."""
  console.print(Panel(f"  {text}", style="bold blue"))


def success(msg):
  """Display a success message."""
  console.print(f"[green]{msg}[/green]")


def warning(msg):
  """Display a warning message."""
  console.print(f"[yellow]{msg}[/yellow]")


def error(msg):
  """Display an error message."""
  console.print(f"[red]{msg}[/red]")


_INFRA_LABELS = {
  "project": "Project",
  "zone": "Zone",
  "cluster_name": "Cluster Name",
  "cluster_endpoint": "Cluster Endpoint",
  "ar_registry": "Artifact Registry",
}

_GPU_LABELS = {
  "name": "GPU Type",
  "count": "GPU Count",
  "machine_type": "Machine Type",
  "node_pool": "Node Pool",
  "node_count": "Node Count",
}

_TPU_LABELS = {
  "name": "TPU Type",
  "chips": "TPU Chips",
  "topology": "Topology",
  "machine_type": "Machine Type",
  "node_pool": "Node Pool",
  "node_count": "Node Count",
}


def infrastructure_state(outputs):
  """Display infrastructure state from Pulumi stack outputs.

  Args:
      outputs: dict of key -> pulumi.automation.OutputValue from stack.outputs().
  """
  table = Table(title="Infrastructure State")
  table.add_column("Resource", style="bold")
  table.add_column("Value", style="green")

  for key, label in _INFRA_LABELS.items():
    if key in outputs:
      table.add_row(label, str(outputs[key].value))

  # New format: "accelerators" (list of dicts)
  if "accelerators" in outputs:
    accel_list = outputs["accelerators"].value
    if not accel_list:
      table.add_row("Accelerators", "CPU only (no accelerator pools)")
    else:
      table.add_row("", "")
      table.add_row(f"Accelerator Pools ({len(accel_list)})", "")
      for i, accel in enumerate(accel_list, 1):
        _render_accelerator(table, accel, index=i)

  # Legacy format: "accelerator" (single dict or None)
  elif "accelerator" in outputs:
    if outputs["accelerator"].value is None:
      table.add_row("Accelerator", "CPU only")
    else:
      accel = outputs["accelerator"].value
      accel_type = accel.get("type", "Unknown")
      table.add_row("", "")
      table.add_row("Accelerator", accel_type)
      labels = _GPU_LABELS if accel_type == "GPU" else _TPU_LABELS
      for key, label in labels.items():
        if key in accel:
          table.add_row(f"  {label}", str(accel[key]))

  else:
    table.add_row(
      "Accelerator",
      "[dim]Unknown (run 'keras-remote up' to refresh)[/dim]",
    )

  console.print()
  console.print(table)
  console.print()


def _render_accelerator(table, accel, index=None):
  """Render a single accelerator pool entry in the status table."""
  accel_type = accel.get("type", "Unknown")
  pool_name = accel.get("node_pool", "")
  prefix = f"  Pool {index}" if index else "  Pool"
  table.add_row(f"{prefix}: {accel_type}", pool_name)
  labels = _GPU_LABELS if accel_type == "GPU" else _TPU_LABELS
  for key, label in labels.items():
    if key in accel and key != "node_pool":
      table.add_row(f"    {label}", str(accel[key]))


def config_summary(config):
  """Display a configuration summary table."""
  table = Table(title="Configuration Summary")
  table.add_column("Setting", style="bold")
  table.add_column("Value", style="green")

  table.add_row("Project", config.project)
  table.add_row("Zone", config.zone)
  table.add_row("Cluster Name", config.cluster_name)

  if not config.node_pools:
    table.add_row("Accelerators", "CPU only")
  else:
    accel_strs = []
    for np in config.node_pools:
      accel = np.accelerator
      if isinstance(accel, GpuConfig):
        accel_strs.append(f"GPU ({accel.gke_label})")
      elif isinstance(accel, TpuConfig):
        accel_strs.append(f"TPU ({accel.name}, {accel.topology})")
    table.add_row("Accelerators", ", ".join(accel_strs))

  console.print()
  console.print(table)