Displays accelerator node pool information as part of status output (#46)

* Displays accelerator pool info with status command

* removes unused method

Author: JyotinderSingh

keras_remote/cli/commands/status.py

@@ -2,13 +2,17 @@
 
 import click
 from pulumi.automation import CommandError
-from rich.table import Table
 
 from keras_remote.cli.config import InfraConfig
 from keras_remote.cli.constants import DEFAULT_ZONE
 from keras_remote.cli.infra.program import create_program
 from keras_remote.cli.infra.stack_manager import get_stack
-from keras_remote.cli.output import banner, console, warning
+from keras_remote.cli.output import (
+  banner,
+  console,
+  infrastructure_state,
+  warning,
+)
 from keras_remote.cli.prerequisites_check import check_all
 from keras_remote.cli.prompts import resolve_project
 
@@ -57,13 +61,4 @@ def status(project, zone):
     warning("No infrastructure found. Run 'keras-remote up' first.")
     return
 
-  table = Table(title="Infrastructure State")
-  table.add_column("Resource", style="bold")
-  table.add_column("Value", style="green")
-
-  for key, output in outputs.items():
-    table.add_row(key, str(output.value))
-
-  console.print()
-  console.print(table)
-  console.print()
+  infrastructure_state(outputs)

keras_remote/cli/infra/program.py

@@ -131,6 +131,35 @@ def pulumi_program():
       f"{ar_location}-docker.pkg.dev/{project_id}/keras-remote",
     )
 
+    # 7. Accelerator node pool exports
+    if isinstance(accelerator, GpuConfig):
+      pulumi.export(
+        "accelerator",
+        {
+          "type": "GPU",
+          "name": accelerator.name,
+          "count": accelerator.count,
+          "machine_type": accelerator.machine_type,
+          "node_pool": "gpu-pool",
+          "node_count": 1,
+        },
+      )
+    elif isinstance(accelerator, TpuConfig):
+      pulumi.export(
+        "accelerator",
+        {
+          "type": "TPU",
+          "name": accelerator.name,
+          "chips": accelerator.chips,
+          "topology": accelerator.topology,
+          "machine_type": accelerator.machine_type,
+          "node_pool": f"tpu-{accelerator.name}-pool",
+          "node_count": accelerator.num_nodes,
+        },
+      )
+    else:
+      pulumi.export("accelerator", None)
+
   return pulumi_program
 
 

keras_remote/cli/infra/program_test.py

@@ -1,10 +1,10 @@
-"""Tests for keras_remote.cli.infra.program — TPU node pool creation."""
+"""Tests for keras_remote.cli.infra.program — node pool and exports."""
 
 from unittest import mock
 
 from absl.testing import absltest, parameterized
 
-from keras_remote.core.accelerators import TpuConfig
+from keras_remote.core.accelerators import GpuConfig, TpuConfig
 
 # Patch the pulumi_gcp module before importing program, so the module-level
 # import inside program.py picks up the mock.
@@ -86,5 +86,65 @@ def test_pool_name_includes_tpu_name(self, gcp_mock):
     self.assertEqual(positional_args[0], "tpu-v5p-pool")
 
 
+def _make_config(accelerator=None):
+  """Create a mock InfraConfig for testing."""
+  config = mock.MagicMock()
+  config.project = "test-project"
+  config.zone = "us-central1-a"
+  config.cluster_name = "test-cluster"
+  config.accelerator = accelerator
+  return config
+
+
+def _run_program_and_get_exports(config):
+  """Run the Pulumi program and return a dict of exported key -> value."""
+  with (
+    mock.patch.object(program, "pulumi") as pulumi_mock,
+    mock.patch.object(program, "gcp"),
+  ):
+    program_fn = program.create_program(config)
+    program_fn()
+    return {
+      call.args[0]: call.args[1] for call in pulumi_mock.export.call_args_list
+    }
+
+
+class TestAcceleratorExports(absltest.TestCase):
+  """Verify accelerator metadata is exported correctly."""
+
+  def test_gpu_exports(self):
+    gpu = GpuConfig("l4", 1, "nvidia-l4", "g2-standard-4")
+    exports = _run_program_and_get_exports(_make_config(gpu))
+
+    self.assertIn("accelerator", exports)
+    accel = exports["accelerator"]
+    self.assertEqual(accel["type"], "GPU")
+    self.assertEqual(accel["name"], "l4")
+    self.assertEqual(accel["count"], 1)
+    self.assertEqual(accel["machine_type"], "g2-standard-4")
+    self.assertEqual(accel["node_pool"], "gpu-pool")
+    self.assertEqual(accel["node_count"], 1)
+
+  def test_tpu_exports(self):
+    tpu = TpuConfig("v5p", 8, "2x2x2", "tpu-v5p-slice", "ct5p-hightpu-4t", 2)
+    exports = _run_program_and_get_exports(_make_config(tpu))
+
+    self.assertIn("accelerator", exports)
+    accel = exports["accelerator"]
+    self.assertEqual(accel["type"], "TPU")
+    self.assertEqual(accel["name"], "v5p")
+    self.assertEqual(accel["chips"], 8)
+    self.assertEqual(accel["topology"], "2x2x2")
+    self.assertEqual(accel["machine_type"], "ct5p-hightpu-4t")
+    self.assertEqual(accel["node_pool"], "tpu-v5p-pool")
+    self.assertEqual(accel["node_count"], 2)
+
+  def test_cpu_only_exports_none(self):
+    exports = _run_program_and_get_exports(_make_config(None))
+
+    self.assertIn("accelerator", exports)
+    self.assertIsNone(exports["accelerator"])
+
+
 if __name__ == "__main__":
   absltest.main()

keras_remote/cli/output.py

@@ -29,6 +29,68 @@ def error(msg):
   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))
+
+  if "accelerator" not in outputs:
+    table.add_row(
+      "Accelerator",
+      "[dim]Unknown (run 'keras-remote up' to refresh)[/dim]",
+    )
+  elif 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]))
+
+  console.print()
+  console.print(table)
+  console.print()
+
+
 def config_summary(config):
   """Display a configuration summary table."""
   table = Table(title="Configuration Summary")

keras_remote/cli/output_test.py

@@ -0,0 +1,101 @@
+"""Tests for keras_remote.cli.output — infrastructure_state display."""
+
+from unittest import mock
+
+from absl.testing import absltest
+from rich.console import Console
+
+from keras_remote.cli import output
+
+
+def _make_output(value):
+  """Create a mock Pulumi OutputValue."""
+  ov = mock.MagicMock()
+  ov.value = value
+  return ov
+
+
+def _render_text(outputs):
+  """Render infrastructure_state and capture output as text."""
+  from io import StringIO
+
+  sio = StringIO()
+  buf = Console(file=sio, force_terminal=False, width=120)
+  with mock.patch.object(output, "console", buf):
+    output.infrastructure_state(outputs)
+  return sio.getvalue()
+
+
+class TestInfrastructureState(absltest.TestCase):
+  """Verify infrastructure_state renders correctly."""
+
+  def _base_outputs(self):
+    return {
+      "project": _make_output("my-project"),
+      "zone": _make_output("us-central1-a"),
+      "cluster_name": _make_output("keras-remote-cluster"),
+      "cluster_endpoint": _make_output("34.123.45.67"),
+      "ar_registry": _make_output("us-docker.pkg.dev/my-project/keras-remote"),
+    }
+
+  def test_gpu_accelerator(self):
+    outputs = self._base_outputs()
+    outputs["accelerator"] = _make_output(
+      {
+        "type": "GPU",
+        "name": "l4",
+        "count": 1,
+        "machine_type": "g2-standard-4",
+        "node_pool": "gpu-pool",
+        "node_count": 1,
+      }
+    )
+    text = _render_text(outputs)
+
+    self.assertIn("my-project", text)
+    self.assertIn("GPU", text)
+    self.assertIn("l4", text)
+    self.assertIn("g2-standard-4", text)
+    self.assertIn("gpu-pool", text)
+
+  def test_tpu_accelerator(self):
+    outputs = self._base_outputs()
+    outputs["accelerator"] = _make_output(
+      {
+        "type": "TPU",
+        "name": "v5p",
+        "chips": 8,
+        "topology": "2x2x2",
+        "machine_type": "ct5p-hightpu-4t",
+        "node_pool": "tpu-v5p-pool",
+        "node_count": 2,
+      }
+    )
+    text = _render_text(outputs)
+
+    self.assertIn("TPU", text)
+    self.assertIn("v5p", text)
+    self.assertIn("2x2x2", text)
+    self.assertIn("ct5p-hightpu-4t", text)
+    self.assertIn("tpu-v5p-pool", text)
+
+  def test_cpu_only(self):
+    outputs = self._base_outputs()
+    outputs["accelerator"] = _make_output(None)
+    text = _render_text(outputs)
+
+    self.assertIn("CPU only", text)
+    self.assertNotIn("GPU", text)
+    self.assertNotIn("TPU", text)
+
+  def test_missing_accelerator_key_backward_compat(self):
+    outputs = self._base_outputs()
+    # No "accelerator" key — simulates old stack
+    text = _render_text(outputs)
+
+    self.assertIn("Unknown", text)
+    self.assertIn("keras-remote up", text)
+
+
+if __name__ == "__main__":
+  absltest.main()