Add animated loading display with tips to CLI infrastructure panels (#76)

Author: JyotinderSingh

keras_remote/backend/log_streaming.py

@@ -5,16 +5,14 @@
 job execution.
 """
 
-import sys
 import threading
-from collections import deque
 
 import urllib3
 from absl import logging
 from kubernetes.client.rest import ApiException
 from rich.console import Console
-from rich.live import Live
-from rich.panel import Panel
+
+from keras_remote.cli.output import LiveOutputPanel
 
 _MAX_DISPLAY_LINES = 25
 
@@ -27,14 +25,13 @@ def _stream_pod_logs(core_v1, pod_name, namespace):
 
   In interactive terminals, logs are displayed in a Rich Live panel.
   In non-interactive contexts (piped output, CI), logs are streamed
-  as raw lines with Rich Rule delimiters.
+  as plain lines with Rule delimiters.
 
   Args:
       core_v1: Kubernetes CoreV1Api client.
       pod_name: Name of the pod to stream logs from.
       namespace: Kubernetes namespace.
   """
-  console = Console()
   resp = None
   try:
     resp = core_v1.read_namespaced_pod_log(
@@ -43,10 +40,22 @@ def _stream_pod_logs(core_v1, pod_name, namespace):
       follow=True,
       _preload_content=False,
     )
-    if console.is_terminal:
-      _render_live_panel(resp, pod_name, console)
-    else:
-      _render_plain(resp, pod_name, console)
+    title = f"Remote logs \u2022 {pod_name}"
+    with LiveOutputPanel(
+      title,
+      max_lines=_MAX_DISPLAY_LINES,
+      target_console=Console(),
+      show_subtitle=False,
+    ) as panel:
+      buffer = ""
+      for chunk in resp.stream(decode_content=True):
+        buffer += chunk.decode("utf-8", errors="replace")
+        while "\n" in buffer:
+          line, buffer = buffer.split("\n", 1)
+          panel.on_output(line)
+      # Flush remaining partial line
+      if buffer.strip():
+        panel.on_output(buffer)
   except ApiException:
     pass  # Pod deleted or not found
   except urllib3.exceptions.ProtocolError:
@@ -60,45 +69,6 @@ def _stream_pod_logs(core_v1, pod_name, namespace):
       resp.release_conn()
 
 
-def _render_live_panel(resp, pod_name, console):
-  """Render streaming logs inside a Rich Live panel."""
-  lines = deque(maxlen=_MAX_DISPLAY_LINES)
-  title = f"Remote logs \u2022 {pod_name}"
-  buffer = ""
-
-  with Live(
-    _make_log_panel(lines, title),
-    console=console,
-    refresh_per_second=4,
-  ) as live:
-    for chunk in resp.stream(decode_content=True):
-      buffer += chunk.decode("utf-8", errors="replace")
-      while "\n" in buffer:
-        line, buffer = buffer.split("\n", 1)
-        lines.append(line)
-      live.update(_make_log_panel(lines, title))
-
-    # Flush remaining partial line
-    if buffer.strip():
-      lines.append(buffer)
-      live.update(_make_log_panel(lines, title))
-
-
-def _render_plain(resp, pod_name, console):
-  """Render streaming logs as raw lines with Rule delimiters."""
-  console.rule(f"Remote logs ({pod_name})", style="blue")
-  for chunk in resp.stream(decode_content=True):
-    sys.stdout.write(chunk.decode("utf-8", errors="replace"))
-    sys.stdout.flush()
-  console.rule("End remote logs", style="blue")
-
-
-def _make_log_panel(lines, title):
-  """Build a Panel renderable from accumulated log lines."""
-  content = "\n".join(lines) if lines else "Waiting for output..."
-  return Panel(content, title=title, border_style="blue")
-
-
 class LogStreamer:
   """Context manager that owns the log-streaming thread lifecycle.
 

keras_remote/backend/log_streaming_test.py

@@ -1,6 +1,5 @@
 """Tests for keras_remote.backend.log_streaming — live pod log streaming."""
 
-import io
 from unittest import mock
 from unittest.mock import MagicMock
 
@@ -9,8 +8,6 @@
 
 from keras_remote.backend.log_streaming import (
   LogStreamer,
-  _render_live_panel,
-  _render_plain,
   _stream_pod_logs,
 )
 
@@ -40,33 +37,23 @@ def test_calls_log_api_correctly(self):
       _preload_content=False,
     )
 
-  def test_routes_by_terminal(self):
-    for is_terminal in (True, False):
-      mock_core = MagicMock()
-      mock_core.read_namespaced_pod_log.return_value = self._make_mock_resp(
-        [b"hello\n"]
-      )
-
-      with (
-        mock.patch(
-          "keras_remote.backend.log_streaming.Console"
-        ) as mock_console_cls,
-        mock.patch(
-          "keras_remote.backend.log_streaming._render_live_panel"
-        ) as mock_live,
-        mock.patch(
-          "keras_remote.backend.log_streaming._render_plain"
-        ) as mock_plain,
-      ):
-        mock_console_cls.return_value.is_terminal = is_terminal
-        _stream_pod_logs(mock_core, "pod-1", "default")
-
-      if is_terminal:
-        mock_live.assert_called_once()
-        mock_plain.assert_not_called()
-      else:
-        mock_plain.assert_called_once()
-        mock_live.assert_not_called()
+  def test_handles_partial_lines(self):
+    mock_core = MagicMock()
+    # "hello\nwor" then "ld\n" — "world" is split across chunks
+    mock_core.read_namespaced_pod_log.return_value = self._make_mock_resp(
+      [b"hello\nwor", b"ld\n"]
+    )
+
+    with mock.patch(
+      "keras_remote.backend.log_streaming.LiveOutputPanel"
+    ) as mock_panel_cls:
+      mock_panel = MagicMock()
+      mock_panel_cls.return_value.__enter__ = MagicMock(return_value=mock_panel)
+      mock_panel_cls.return_value.__exit__ = MagicMock(return_value=False)
+      _stream_pod_logs(mock_core, "pod-1", "default")
+
+    lines = [call[0][0] for call in mock_panel.on_output.call_args_list]
+    self.assertEqual(lines, ["hello", "world"])
 
   def test_releases_conn_on_api_exception(self):
     mock_core = MagicMock()
@@ -111,66 +98,6 @@ def test_logs_warning_on_unexpected_error(self):
     mock_resp.release_conn.assert_called_once()
 
 
-class TestRenderPlain(absltest.TestCase):
-  """Tests for the non-terminal plain rendering path."""
-
-  def test_streams_chunks_to_stdout(self):
-    mock_resp = MagicMock()
-    mock_resp.stream.return_value = [b"line 1\n", b"line 2\n"]
-    console = MagicMock()
-
-    with mock.patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
-      _render_plain(mock_resp, "pod-1", console)
-
-    self.assertIn("line 1", mock_stdout.getvalue())
-    self.assertIn("line 2", mock_stdout.getvalue())
-
-  def test_prints_rule_delimiters(self):
-    mock_resp = MagicMock()
-    mock_resp.stream.return_value = []
-    console = MagicMock()
-
-    _render_plain(mock_resp, "pod-1", console)
-
-    self.assertEqual(console.rule.call_count, 2)
-    # Opening rule contains pod name
-    self.assertIn("pod-1", console.rule.call_args_list[0][0][0])
-
-  def test_handles_utf8_decode_errors(self):
-    mock_resp = MagicMock()
-    mock_resp.stream.return_value = [b"valid\n", b"\xff\xfe invalid\n"]
-    console = MagicMock()
-
-    with mock.patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
-      _render_plain(mock_resp, "pod-1", console)
-
-    output = mock_stdout.getvalue()
-    self.assertIn("valid", output)
-    self.assertIn("invalid", output)
-
-
-class TestRenderLivePanel(absltest.TestCase):
-  """Tests for the terminal Live panel rendering path."""
-
-  def test_handles_partial_lines(self):
-    mock_resp = MagicMock()
-    # "hello\nwor" then "ld\n" — "world" is split across chunks
-    mock_resp.stream.return_value = [b"hello\nwor", b"ld\n"]
-    console = MagicMock()
-
-    with mock.patch("keras_remote.backend.log_streaming.Live") as mock_live_cls:
-      mock_live = MagicMock()
-      mock_live_cls.return_value.__enter__ = MagicMock(return_value=mock_live)
-      mock_live_cls.return_value.__exit__ = MagicMock(return_value=False)
-      _render_live_panel(mock_resp, "pod-1", console)
-
-    # Check the final panel contains both complete lines
-    last_panel = mock_live.update.call_args_list[-1][0][0]
-    panel_content = last_panel.renderable
-    self.assertIn("hello", panel_content)
-    self.assertIn("world", panel_content)
-
-
 class TestLogStreamer(absltest.TestCase):
   def test_start_launches_daemon_thread(self):
     mock_core = MagicMock()

keras_remote/cli/commands/up.py

@@ -14,10 +14,10 @@
 from keras_remote.cli.infra.state import apply_update, load_state
 from keras_remote.cli.options import common_options
 from keras_remote.cli.output import (
+  LiveOutputPanel,
   banner,
   config_summary,
   console,
-  success,
   warning,
 )
 from keras_remote.cli.prerequisites_check import check_all
@@ -89,17 +89,13 @@ def up(project, zone, accelerator, cluster_name, yes):
 
   console.print()
 
-  # Run Pulumi
-  console.print("[bold]Provisioning infrastructure...[/bold]\n")
   pulumi_ok = apply_update(config)
   pulumi_failed = not pulumi_ok
 
   if pulumi_failed:
     warning("Attempting post-deploy configuration anyway...")
 
   # Post-deploy steps
-  console.print("\n[bold]Running post-deploy configuration...[/bold]\n")
-
   steps = [
     (
       "kubectl configuration",
@@ -115,14 +111,21 @@ def up(project, zone, accelerator, cluster_name, yes):
     steps.append(("GPU driver installation", install_gpu_drivers))
 
   failures = []
-  for name, fn in steps:
-    console.print(f"{name}...")
-    try:
-      fn()
-      success(f"{name} complete.")
-    except subprocess.CalledProcessError as e:
-      failures.append(name)
-      warning(f"{name} failed: {e}")
+  with LiveOutputPanel("Post-deploy configuration", transient=True) as panel:
+    for name, fn in steps:
+      panel.on_output(f"{name}...")
+      try:
+        fn()
+        panel.on_output(f"{name} complete.")
+      except subprocess.CalledProcessError as e:
+        failures.append(name)
+        panel.on_output(f"{name} failed: {e}")
+        if e.stderr:
+          stderr_text = e.stderr.decode("utf-8", errors="replace").strip()
+          if stderr_text:
+            for line in stderr_text.splitlines():
+              panel.on_output(f"  {line}")
+        panel.mark_error()
 
   # Final summary
   console.print()

keras_remote/cli/infra/post_deploy.py

@@ -33,6 +33,7 @@ def configure_kubectl(cluster_name, zone, project):
       f"--project={project}",
     ],
     check=True,
+    capture_output=True,
     env=env,
   )
 
@@ -46,6 +47,7 @@ def install_gpu_drivers():
   subprocess.run(
     ["kubectl", "apply", "-f", NVIDIA_DRIVER_DAEMONSET_URL],
     check=True,
+    capture_output=True,
   )
 
 
@@ -57,4 +59,5 @@ def install_lws():
   subprocess.run(
     ["kubectl", "apply", "--server-side", "-f", LWS_INSTALL_URL],
     check=True,
+    capture_output=True,
   )