Adds auto-login support to @run()

Author: JyotinderSingh

AGENTS.md

@@ -16,6 +16,7 @@ keras_remote/
 ├── cli/            # CLI for infrastructure provisioning (Pulumi-based)
 │   ├── commands/   # up, down, status, config
 │   └── infra/      # Pulumi programs and stack management
+├── credentials.py  # Credential verification & auto-setup (shared by core & CLI)
 └── constants.py    # Zone/region utilities
 ```
 
@@ -24,6 +25,7 @@ keras_remote/
 ```python
 @keras_remote.run() called
   → JobContext.from_params()        # Resolve config from args/env vars
+  → ensure_credentials()            # Verify/auto-configure gcloud, ADC, kubeconfig, Docker auth
   → _prepare_artifacts()            # Serialize function (cloudpickle), zip working dir
   → _build_container()              # Build or retrieve cached Docker image
   → _upload_artifacts()             # Upload payload.pkl, context.zip to GCS
@@ -35,23 +37,24 @@ keras_remote/
 
 ## Key Modules
 
-| Module                       | Responsibility                                                                |
-| ---------------------------- | ----------------------------------------------------------------------------- |
-| `core/core.py`               | `@run()` decorator, backend routing, env var capture                          |
-| `core/accelerators.py`       | Accelerator registry (`GPUS`, `TPUS`), parser (`parse_accelerator`)           |
-| `backend/execution.py`       | `JobContext` dataclass, `BackendClient` protocol, `execute_remote()` pipeline |
-| `backend/gke_client.py`      | K8s Job creation, status polling, pod log retrieval                           |
-| `backend/pathways_client.py` | LeaderWorkerSet creation for multi-host TPUs                                  |
-| `infra/container_builder.py` | Content-hashed Docker image building via Cloud Build                          |
-| `utils/packager.py`          | `save_payload()` (cloudpickle), `zip_working_dir()`                           |
-| `utils/storage.py`           | GCS upload/download/cleanup for job artifacts                                 |
-| `runner/remote_runner.py`    | Runs inside container: deserialize, execute, upload result                    |
-| `cli/main.py`                | CLI entry point (`keras-remote` command)                                      |
+| Module                       | Responsibility                                                                   |
+| ---------------------------- | -------------------------------------------------------------------------------- |
+| `core/core.py`               | `@run()` decorator, backend routing, env var capture                             |
+| `core/accelerators.py`       | Accelerator registry (`GPUS`, `TPUS`), parser (`parse_accelerator`)              |
+| `credentials.py`             | Credential verification & auto-setup (gcloud, ADC, kubeconfig, Docker auth)      |
+| `backend/execution.py`       | `JobContext` dataclass, `BaseK8sBackend` base class, `execute_remote()` pipeline |
+| `backend/gke_client.py`      | K8s Job creation, status polling, pod log retrieval                              |
+| `backend/pathways_client.py` | LeaderWorkerSet creation for multi-host TPUs                                     |
+| `infra/container_builder.py` | Content-hashed Docker image building via Cloud Build                             |
+| `utils/packager.py`          | `save_payload()` (cloudpickle), `zip_working_dir()`                              |
+| `utils/storage.py`           | GCS upload/download/cleanup for job artifacts                                    |
+| `runner/remote_runner.py`    | Runs inside container: deserialize, execute, upload result                       |
+| `cli/main.py`                | CLI entry point (`keras-remote` command)                                         |
 
 ## Key Abstractions
 
 - **`JobContext`** (`backend/execution.py`): Mutable dataclass carrying all job state through the pipeline — inputs, generated IDs, artifact paths, image URI.
-- **`BackendClient`** protocol (`backend/execution.py`): Interface with `submit_job`, `wait_for_job`, `cleanup_job`. Implemented by `GKEBackend` and `PathwaysBackend`.
+- **`BaseK8sBackend`** (`backend/execution.py`): Base class with `submit_job`, `wait_for_job`, `cleanup_job`. Subclassed by `GKEBackend` and `PathwaysBackend`.
 - **`GpuConfig` / `TpuConfig`** (`core/accelerators.py`): Frozen dataclasses for accelerator metadata. Single source of truth used by runtime, container builder, and CLI.
 - **`InfraConfig`** (`cli/config.py`): CLI provisioning configuration (project, zone, cluster, accelerator).
 
@@ -62,7 +65,6 @@ keras_remote/
 - **Formatter/linter**: `ruff` (2-space indent, 80-char line length target, E501 ignored)
 - **Rules**: B, E, F, N, PYI, T20, TID, SIM, W, I, NPY
 - **Dataclasses**: Frozen for immutable configs, mutable for state objects
-- **Protocols**: Used for backend abstraction (not ABCs)
 
 ### Environment Variables
 

examples/example_gke.py

@@ -35,7 +35,8 @@
 """
 
 import os
-import subprocess
+
+os.environ["KERAS_BACKEND"] = "jax"
 
 import keras
 import numpy as np
@@ -135,31 +136,6 @@ def main():
 if __name__ == "__main__":
   # Prerequisites:
   # 1. Set KERAS_REMOTE_PROJECT environment variable to your GCP project ID
-  # 2. Configure kubectl: gcloud container clusters get-credentials <cluster> --zone <zone>
-  # 3. Ensure your GKE cluster has GPU nodes with the required accelerator type
-  if not os.environ.get("KERAS_REMOTE_PROJECT"):
-    print("ERROR: KERAS_REMOTE_PROJECT environment variable not set")
-    print("Please set it to your GCP project ID:")
-    print("  export KERAS_REMOTE_PROJECT=your-project-id")
-    exit(1)
-
-  # Verify kubectl is configured
-  try:
-    result = subprocess.run(
-      ["kubectl", "cluster-info"], capture_output=True, text=True, timeout=10
-    )
-    if result.returncode != 0:
-      print("ERROR: kubectl is not configured or cluster is not accessible")
-      print("Please configure kubectl:")
-      print(
-        "  gcloud container clusters get-credentials <cluster-name> --zone <zone>"
-      )
-      exit(1)
-  except FileNotFoundError:
-    print("ERROR: kubectl not found. Please install kubectl.")
-    exit(1)
-  except subprocess.TimeoutExpired:
-    print("ERROR: kubectl timed out. Check your cluster connectivity.")
-    exit(1)
-
+  #    (if `project` param is not provided in the decorator)
+  # 2. Ensure your GKE cluster has GPU nodes with the required accelerator type
   main()

keras_remote/backend/execution.py

@@ -9,14 +9,15 @@
 import tempfile
 import uuid
 from dataclasses import dataclass, field
-from typing import Any, Callable, Optional, Protocol
+from typing import Any, Callable, Optional
 
 import cloudpickle
 from absl import logging
 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.credentials import ensure_credentials
 from keras_remote.infra import container_builder
 from keras_remote.utils import packager, storage
 
@@ -90,28 +91,24 @@ def from_params(
     )
 
 
-class BackendClient(Protocol):
-  """Protocol defining the interface for backend clients."""
+class BaseK8sBackend:
+  """Base class for Kubernetes-based backends."""
+
+  def __init__(self, cluster: Optional[str] = None, namespace: str = "default"):
+    self.cluster = cluster
+    self.namespace = namespace
 
   def submit_job(self, ctx: JobContext) -> Any:
     """Submit a job to the backend. Returns backend-specific job handle."""
-    ...
+    raise NotImplementedError
 
   def wait_for_job(self, job: Any, ctx: JobContext) -> None:
     """Wait for job completion. Raises RuntimeError if job fails."""
-    ...
+    raise NotImplementedError
 
   def cleanup_job(self, job: Any, ctx: JobContext) -> None:
     """Optional cleanup after job completion."""
-    ...
-
-
-class BaseK8sBackend:
-  """Base class for Kubernetes-based backends."""
-
-  def __init__(self, cluster: Optional[str] = None, namespace: str = "default"):
-    self.cluster = cluster
-    self.namespace = namespace
+    raise NotImplementedError
 
 
 class GKEBackend(BaseK8sBackend):
@@ -264,22 +261,28 @@ def _cleanup_and_return(ctx: JobContext, result_payload: dict) -> Any:
     raise result_payload["exception"]
 
 
-def execute_remote(ctx: JobContext, backend: BackendClient) -> Any:
+def execute_remote(ctx: JobContext, backend: BaseK8sBackend) -> Any:
   """Execute a function remotely using the specified backend.
 
   This is the unified executor that handles all common phases
   and delegates backend-specific operations to the backend client.
 
   Args:
       ctx: Job context with function and configuration
-      backend: Backend client implementing BackendClient protocol
+      backend: Backend instance (GKEBackend or PathwaysBackend)
 
   Returns:
       The result of the remote function execution
 
   Raises:
       Exception: Re-raised from remote execution if it failed
   """
+  ensure_credentials(
+    project=ctx.project,
+    zone=ctx.zone,
+    cluster=backend.cluster,
+  )
+
   with tempfile.TemporaryDirectory() as tmpdir:
     # Phase 1: Package artifacts
     _prepare_artifacts(ctx, tmpdir)

keras_remote/backend/execution_test.py

@@ -151,6 +151,7 @@ def _make_ctx(self, container_image=None):
 
   def test_success_flow(self):
     with (
+      mock.patch("keras_remote.backend.execution.ensure_credentials"),
       mock.patch("keras_remote.backend.execution._build_container"),
       mock.patch("keras_remote.backend.execution._upload_artifacts"),
       mock.patch(
@@ -174,6 +175,7 @@ def test_success_flow(self):
 
   def test_cleanup_on_wait_failure(self):
     with (
+      mock.patch("keras_remote.backend.execution.ensure_credentials"),
       mock.patch("keras_remote.backend.execution._build_container"),
       mock.patch("keras_remote.backend.execution._upload_artifacts"),
       mock.patch(
@@ -191,6 +193,34 @@ def test_cleanup_on_wait_failure(self):
       # cleanup_job is called in finally block even when wait fails
       backend.cleanup_job.assert_called_once()
 
+  def test_ensure_credentials_called_with_correct_args(self):
+    with (
+      mock.patch(
+        "keras_remote.backend.execution.ensure_credentials"
+      ) as mock_creds,
+      mock.patch("keras_remote.backend.execution._build_container"),
+      mock.patch("keras_remote.backend.execution._upload_artifacts"),
+      mock.patch(
+        "keras_remote.backend.execution._download_result",
+        return_value={"success": True, "result": 0},
+      ),
+      mock.patch(
+        "keras_remote.backend.execution._cleanup_and_return",
+        return_value=0,
+      ),
+    ):
+      ctx = self._make_ctx()
+      backend = MagicMock()
+      backend.cluster = "test-cluster"
+
+      execute_remote(ctx, backend)
+
+      mock_creds.assert_called_once_with(
+        project="proj",
+        zone="us-central1-a",
+        cluster="test-cluster",
+      )
+
 
 if __name__ == "__main__":
   absltest.main()

keras_remote/cli/prerequisites_check.py

@@ -1,20 +1,24 @@
-"""Prerequisite checks for the keras-remote CLI."""
+"""Prerequisite checks for the keras-remote CLI.
+
+Delegates common credential checks (gcloud, auth plugin, ADC) to
+:mod:`keras_remote.credentials` and converts ``RuntimeError`` into
+``click.ClickException``.  CLI-only tool checks (Pulumi, kubectl, Docker)
+remain here.
+"""
 
 import shutil
-import subprocess
 
 import click
 
-from keras_remote.cli.output import warning
+from keras_remote import credentials
 
 
 def check_gcloud():
   """Verify gcloud CLI is installed."""
-  if not shutil.which("gcloud"):
-    raise click.ClickException(
-      "gcloud CLI not found. "
-      "Install from: https://cloud.google.com/sdk/docs/install"
-    )
+  try:
+    credentials.ensure_gcloud()
+  except RuntimeError as e:
+    raise click.ClickException(str(e))  # noqa: B904
 
 
 def check_pulumi():
@@ -42,37 +46,19 @@ def check_docker():
 
 
 def check_gke_auth_plugin():
-  """Verify gke-gcloud-auth-plugin is installed."""
-  if not shutil.which("gke-gcloud-auth-plugin"):
-    warning("gke-gcloud-auth-plugin not found.")
-    if click.confirm(
-      "Install it via `gcloud components install gke-gcloud-auth-plugin`?",
-      default=True,
-    ):
-      subprocess.run(
-        ["gcloud", "components", "install", "gke-gcloud-auth-plugin"],
-        check=True,
-      )
-    else:
-      raise click.ClickException(
-        "gke-gcloud-auth-plugin is required. "
-        "Install with: gcloud components install gke-gcloud-auth-plugin"
-      )
+  """Verify gke-gcloud-auth-plugin is installed; auto-install if missing."""
+  try:
+    credentials.ensure_gke_auth_plugin()
+  except RuntimeError as e:
+    raise click.ClickException(str(e))  # noqa: B904
 
 
 def check_gcloud_auth():
   """Check if gcloud Application Default Credentials are configured."""
-  result = subprocess.run(
-    ["gcloud", "auth", "application-default", "print-access-token"],
-    capture_output=True,
-  )
-  if result.returncode != 0:
-    warning("Application Default Credentials not found.")
-    click.echo("Running: gcloud auth application-default login")
-    subprocess.run(
-      ["gcloud", "auth", "application-default", "login"],
-      check=True,
-    )
+  try:
+    credentials.ensure_adc()
+  except RuntimeError as e:
+    raise click.ClickException(str(e))  # noqa: B904
 
 
 def check_all():