Centralize Pulumi state management, eliminate CLI duplication, and scope all resources to cluster (#74)

Author: JyotinderSingh

keras_remote/backend/execution.py

@@ -16,7 +16,11 @@
 from google.api_core import exceptions as google_exceptions
 
 from keras_remote.backend import gke_client, pathways_client
-from keras_remote.constants import get_default_zone, zone_to_region
+from keras_remote.constants import (
+  get_default_cluster_name,
+  get_default_zone,
+  zone_to_region,
+)
 from keras_remote.credentials import ensure_credentials
 from keras_remote.data import _make_data_ref
 from keras_remote.infra import container_builder
@@ -39,6 +43,7 @@ class JobContext:
   container_image: Optional[str]
   zone: str
   project: str
+  cluster_name: str
 
   # Generated identifiers
   job_id: str = field(default_factory=lambda: f"job-{uuid.uuid4().hex[:8]}")
@@ -58,7 +63,7 @@ class JobContext:
   image_uri: Optional[str] = None
 
   def __post_init__(self):
-    self.bucket_name = f"{self.project}-keras-remote-jobs"
+    self.bucket_name = f"{self.project}-kr-{self.cluster_name}-jobs"
     self.region = zone_to_region(self.zone)
     self.display_name = f"keras-remote-{self.func.__name__}-{self.job_id}"
 
@@ -73,9 +78,10 @@ def from_params(
     zone: Optional[str],
     project: Optional[str],
     env_vars: dict,
+    cluster_name: Optional[str] = None,
     volumes: Optional[dict] = None,
   ) -> "JobContext":
-    """Factory method with default resolution for zone/project."""
+    """Factory method with default resolution for zone/project/cluster."""
     if not zone:
       zone = get_default_zone()
     if not project:
@@ -85,6 +91,8 @@ def from_params(
           "project must be specified or set KERAS_REMOTE_PROJECT"
           " (or GOOGLE_CLOUD_PROJECT) environment variable"
         )
+    if not cluster_name:
+      cluster_name = get_default_cluster_name()
 
     return cls(
       func=func,
@@ -95,6 +103,7 @@ def from_params(
       container_image=container_image,
       zone=zone,
       project=project,
+      cluster_name=cluster_name,
       volumes=volumes,
     )
 
@@ -303,6 +312,7 @@ def _build_container(ctx: JobContext) -> None:
       accelerator_type=ctx.accelerator,
       project=ctx.project,
       zone=ctx.zone,
+      cluster_name=ctx.cluster_name,
     )
 
 

keras_remote/backend/execution_test.py

@@ -40,8 +40,9 @@ def test_post_init_derived_fields(self):
       container_image=None,
       zone="europe-west4-b",
       project="my-proj",
+      cluster_name="my-cluster",
     )
-    self.assertEqual(ctx.bucket_name, "my-proj-keras-remote-jobs")
+    self.assertEqual(ctx.bucket_name, "my-proj-kr-my-cluster-jobs")
     self.assertEqual(ctx.region, "europe-west4")
     self.assertTrue(ctx.display_name.startswith("keras-remote-my_train-"))
     self.assertRegex(ctx.job_id, r"^job-[0-9a-f]{8}$")
@@ -171,6 +172,7 @@ def _make_ctx(self, container_image=None):
       container_image=container_image,
       zone="us-central1-a",
       project="proj",
+      cluster_name="keras-remote-cluster",
     )
 
   def test_success_flow(self):

keras_remote/cli/commands/down.py

@@ -1,37 +1,18 @@
 """keras-remote down command — tear down infrastructure."""
 
 import click
-import pulumi.automation as auto
 
 from keras_remote.cli.config import InfraConfig
 from keras_remote.cli.constants import DEFAULT_CLUSTER_NAME, 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, success, warning
+from keras_remote.cli.infra.state import apply_destroy
+from keras_remote.cli.options import common_options
+from keras_remote.cli.output import banner, console, warning
 from keras_remote.cli.prerequisites_check import check_all
 from keras_remote.cli.prompts import resolve_project
 
 
 @click.command()
-@click.option(
-  "--project",
-  envvar="KERAS_REMOTE_PROJECT",
-  default=None,
-  help="GCP project ID [env: KERAS_REMOTE_PROJECT]",
-)
-@click.option(
-  "--zone",
-  envvar="KERAS_REMOTE_ZONE",
-  default=None,
-  help=(f"GCP zone [env: KERAS_REMOTE_ZONE, default: {DEFAULT_ZONE}]"),
-)
-@click.option(
-  "--cluster",
-  "cluster_name",
-  envvar="KERAS_REMOTE_CLUSTER",
-  default=None,
-  help="GKE cluster name [default: keras-remote-cluster]",
-)
+@common_options
 @click.option("--yes", "-y", is_flag=True, help="Skip confirmation prompt")
 def down(project, zone, cluster_name, yes):
   """Tear down keras-remote GCP infrastructure."""
@@ -60,19 +41,7 @@ def down(project, zone, cluster_name, yes):
   console.print()
 
   config = InfraConfig(project=project, zone=zone, cluster_name=cluster_name)
-
-  # Pulumi destroy
-  try:
-    # Minimal config to load the stack — accelerator is not
-    # needed for destroy since the stack already has its state.
-    program = create_program(config)
-    stack = get_stack(program, config)
-    console.print("[bold]Destroying Pulumi-managed resources...[/bold]\n")
-    result = stack.destroy(on_output=print)
-    console.print()
-    success(f"Pulumi destroy complete. {result.summary.resource_changes}")
-  except auto.errors.CommandError as e:
-    warning(f"Pulumi destroy encountered an issue: {e}")
+  apply_destroy(config)
 
   # Summary
   console.print()

keras_remote/cli/commands/down_test.py

@@ -0,0 +1,110 @@
+"""Tests for keras_remote.cli.commands.down — destroy infrastructure."""
+
+from unittest import mock
+
+from absl.testing import absltest
+from click.testing import CliRunner
+
+from keras_remote.cli.commands.down import down
+
+# Shared CLI args that skip interactive prompts.
+_CLI_ARGS = [
+  "--project",
+  "test-project",
+  "--zone",
+  "us-central2-b",
+  "--yes",
+]
+
+# Patches applied to every test to bypass prerequisites and infrastructure.
+_BASE_PATCHES = {
+  "check_all": mock.patch("keras_remote.cli.commands.down.check_all"),
+  "resolve_project": mock.patch(
+    "keras_remote.cli.commands.down.resolve_project",
+    return_value="test-project",
+  ),
+  "apply_destroy": mock.patch(
+    "keras_remote.cli.commands.down.apply_destroy", return_value=True
+  ),
+}
+
+
+def _start_patches(test_case):
+  """Start all base patches and return a dict of mock objects."""
+  mocks = {}
+  for name, patcher in _BASE_PATCHES.items():
+    mocks[name] = test_case.enterContext(patcher)
+  return mocks
+
+
+class DownCommandTest(absltest.TestCase):
+  def setUp(self):
+    super().setUp()
+    self.runner = CliRunner()
+    self.mocks = _start_patches(self)
+
+  def test_successful_destroy(self):
+    """Successful destroy — exit code 0, 'Cleanup Complete' shown."""
+    result = self.runner.invoke(down, _CLI_ARGS)
+
+    self.assertEqual(result.exit_code, 0, result.output)
+    self.assertIn("Cleanup Complete", result.output)
+    self.mocks["apply_destroy"].assert_called_once()
+
+  def test_destroy_failure_still_shows_summary(self):
+    """apply_destroy returns False — summary still displayed."""
+    self.mocks["apply_destroy"].return_value = False
+
+    result = self.runner.invoke(down, _CLI_ARGS)
+
+    self.assertEqual(result.exit_code, 0, result.output)
+    self.assertIn("Cleanup Complete", result.output)
+    self.assertIn("Check manually", result.output)
+
+  def test_abort_on_no_confirmation(self):
+    """User declines confirmation — apply_destroy not called."""
+    args = ["--project", "test-project", "--zone", "us-central2-b"]
+
+    self.runner.invoke(down, args, input="n\n")
+
+    self.mocks["apply_destroy"].assert_not_called()
+
+  def test_yes_flag_skips_confirmation(self):
+    """--yes flag skips confirmation prompt."""
+    result = self.runner.invoke(down, _CLI_ARGS)
+
+    self.assertEqual(result.exit_code, 0, result.output)
+    self.mocks["apply_destroy"].assert_called_once()
+
+  def test_config_passed_correctly(self):
+    """InfraConfig args match CLI options."""
+    args = [
+      "--project",
+      "my-proj",
+      "--zone",
+      "europe-west1-b",
+      "--cluster",
+      "my-cluster",
+      "--yes",
+    ]
+
+    result = self.runner.invoke(down, args)
+
+    self.assertEqual(result.exit_code, 0, result.output)
+    config = self.mocks["apply_destroy"].call_args[0][0]
+    self.assertEqual(config.project, "my-proj")
+    self.assertEqual(config.zone, "europe-west1-b")
+    self.assertEqual(config.cluster_name, "my-cluster")
+
+  def test_resolve_project_allow_create_false(self):
+    """When --project not given, resolve_project(allow_create=False) is called."""
+    args = ["--zone", "us-central1-a", "--yes"]
+
+    result = self.runner.invoke(down, args)
+
+    self.assertEqual(result.exit_code, 0, result.output)
+    self.mocks["resolve_project"].assert_called_once_with(allow_create=False)
+
+
+if __name__ == "__main__":
+  absltest.main()