Adds CLI for Infrastructure Management (#14)

* Adds CLI to orchestrate infra setup/teardown

* Fix exceptions and err msgs

* fix down logic + improve maintainability

* Fixes issue when APIs are not enabled by skipping cleanup steps instead of hanging.

* address reviews + Moves GCS handling to pulumi.

* Uses py sdk instead of gcloud for container builds.

* update CLI to use new accelerator registry types

Replace dict-based accelerator configs with typed GpuConfig/TpuConfig
from the core registry. Introduces InfraConfig dataclass and removes
the now-redundant accelerator_configs.py lookup tables.

* disable allow_create for  command

* remove dead code

* rebasing fixes

* Adds check for kubectl and docker in the prerequisites check

* use a released nvidia driver instead of master

* Adds an error helper

* rename prerequisites to prerequisites_check

* log error when checking for image existence

* improve conditional logic

* Adds label based filtering for resource cleanup

Author: JyotinderSingh

.gitignore

@@ -206,4 +206,7 @@ marimo/_static/
 marimo/_lsp/
 __marimo__/
 
-.claude/
+.claude/
+
+# Pulumi
+.pulumi/

README.md

@@ -45,38 +45,74 @@ final_loss = train_model()
 
 ## Installation
 
-### From Source
+### Library Only
+
+Install the core package to use the `@keras_remote.run()` decorator in your code:
 
 ```bash
 git clone https://github.com/keras-team/keras-remote.git
 cd keras-remote
 pip install -e .
 ```
 
+This is sufficient if your infrastructure (GKE cluster, Artifact Registry, etc.) is already provisioned.
+
+### Library + CLI
+
+Install with the `cli` extra to also get the `keras-remote` command for managing infrastructure:
+
+```bash
+git clone https://github.com/keras-team/keras-remote.git
+cd keras-remote
+pip install -e ".[cli]"
+```
+
+This adds the `keras-remote up`, `keras-remote down`, `keras-remote status`, and `keras-remote config` commands for provisioning and tearing down cloud resources.
+
 ### Requirements
 
 - Python 3.11+
 - Google Cloud SDK (`gcloud`)
   - Run `gcloud auth login` and `gcloud auth application-default login`
+- [Pulumi CLI](https://www.pulumi.com/docs/install/) (required for `[cli]` install only)
 - A Google Cloud project with billing enabled
 
 ## Quick Start
 
 ### 1. Configure Google Cloud
 
-Run the automated setup script:
+Run the CLI setup command:
 
 ```bash
-./setup.sh
+keras-remote up
 ```
 
-The script will:
+This will interactively:
 
 - Prompt for your GCP project ID
+- Let you choose an accelerator type (CPU, GPU, or TPU)
 - Enable required APIs (Cloud Build, Artifact Registry, Cloud Storage, GKE)
 - Create the Artifact Registry repository
-- Configure Docker authentication
-- Verify the setup
+- Provision a GKE cluster with optional accelerator node pools
+- Configure Docker authentication and kubectl access
+
+You can also run non-interactively:
+
+```bash
+keras-remote up --project=my-project --accelerator=t4 --yes
+```
+
+To view current infrastructure state:
+
+```bash
+keras-remote status
+```
+
+To view configuration:
+
+```bash
+keras-remote config
+```
 
 ### 2. Set Environment Variables
 
@@ -124,7 +160,7 @@ def train():
 - Support for GPU accelerators (T4, L4, A100, V100, H100)
 - Lower overhead for iterative development
 
-**Setup:** Run `./setup.sh` and select GKE.
+**Setup:** Run `keras-remote up` and select a GPU accelerator.
 
 ### TPU VM
 
@@ -348,15 +384,18 @@ gcloud artifacts repositories describe keras-remote \
 Remove all Keras Remote resources to avoid charges:
 
 ```bash
-./cleanup.sh
+keras-remote down
 ```
 
 This removes:
 
-- Cloud Storage buckets
-- Artifact Registry repositories
-- GKE clusters (if created by setup)
-- TPU VMs
+- GKE cluster and accelerator node pools (via Pulumi)
+- Artifact Registry repository and container images
+- Cloud Storage buckets (jobs and builds)
+- TPU VMs and orphaned Compute Engine VMs
+
+Use `--yes` to skip the confirmation prompt, or `--pulumi-only` to only
+destroy Pulumi-managed resources.
 
 ## Contributing
 

cleanup.sh

@@ -1 +1,8 @@
+import os
+
+# Suppress noisy gRPC fork/logging messages before any gRPC imports
+os.environ.setdefault("GRPC_VERBOSITY", "NONE")
+os.environ.setdefault("GLOG_minloglevel", "3")
+os.environ.setdefault("GRPC_ENABLE_FORK_SUPPORT", "0")
+
 from keras_remote.core.core import run

keras_remote/__init__.py

@@ -13,6 +13,7 @@
 
 import cloudpickle
 
+from keras_remote.constants import get_default_zone, zone_to_region
 from keras_remote.infra import container_builder
 from keras_remote.backend import gke_client
 from keras_remote.infra import infra
@@ -54,11 +55,7 @@ class JobContext:
 
     def __post_init__(self):
         self.bucket_name = f"{self.project}-keras-remote-jobs"
-        self.region = (
-            self.zone.rsplit("-", 1)[0]
-            if self.zone and "-" in self.zone
-            else "us-central1"
-        )
+        self.region = zone_to_region(self.zone)
         self.display_name = f"keras-remote-{self.func.__name__}-{self.job_id}"
 
     @classmethod
@@ -75,7 +72,7 @@ def from_params(
     ) -> "JobContext":
         """Factory method with default resolution for zone/project."""
         if not zone:
-            zone = os.environ.get("KERAS_REMOTE_ZONE", "us-central1-a")
+            zone = get_default_zone()
         if not project:
             project = os.environ.get("KERAS_REMOTE_PROJECT")
             if not project:
@@ -212,7 +209,6 @@ def _upload_artifacts(ctx: JobContext) -> None:
         job_id=ctx.job_id,
         payload_path=ctx.payload_path,
         context_path=ctx.context_path,
-        location=ctx.region,
         project=ctx.project,
     )
 

keras_remote/backend/execution.py

@@ -6,7 +6,6 @@
 from kubernetes import client, config
 from kubernetes.client.rest import ApiException
 
-from keras_remote.core.accelerators import GpuConfig
 from keras_remote.core.accelerators import TpuConfig
 from keras_remote.core import accelerators
 from keras_remote.infra import infra
@@ -172,7 +171,7 @@ def cleanup_job(job_name, namespace="default"):
 
 def _parse_accelerator(accelerator):
     """Convert accelerator string to GKE pod spec fields."""
-    parsed = parse_accelerator(accelerator)
+    parsed = accelerators.parse_accelerator(accelerator)
 
     if parsed is None:
         return {
@@ -334,7 +333,7 @@ def _print_pod_logs(core_v1, job_name, namespace):
         pods = core_v1.list_namespaced_pod(
             namespace, label_selector=f"job-name={job_name}"
         )
-        
+
         for pod in pods.items:
             with suppress(ApiException):
                 logs = core_v1.read_namespaced_pod_log(