cleanup

Author: JyotinderSingh

AGENTS.md

@@ -25,7 +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
+  → ensure_credentials()            # Verify/auto-configure gcloud, ADC, kubeconfig
   → _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
@@ -41,7 +41,7 @@ keras_remote/
 | ---------------------------- | -------------------------------------------------------------------------------- |
 | `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)      |
+| `credentials.py`             | Credential verification & auto-setup (gcloud, ADC, kubeconfig)                   |
 | `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                                     |

keras_remote/backend/execution_test.py

@@ -193,34 +193,6 @@ 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/commands/up.py

@@ -8,7 +8,6 @@
 from keras_remote.cli.config import InfraConfig
 from keras_remote.cli.constants import DEFAULT_CLUSTER_NAME, DEFAULT_ZONE
 from keras_remote.cli.infra.post_deploy import (
-  configure_docker_auth,
   configure_kubectl,
   install_gpu_drivers,
   install_lws,
@@ -24,7 +23,6 @@
 )
 from keras_remote.cli.prerequisites_check import check_all
 from keras_remote.cli.prompts import prompt_accelerator, resolve_project
-from keras_remote.constants import zone_to_ar_location
 from keras_remote.core import accelerators
 from keras_remote.core.accelerators import GpuConfig
 
@@ -113,11 +111,9 @@ def up(project, zone, accelerator, cluster_name, yes):
     )
 
   # Post-deploy steps
-  ar_location = zone_to_ar_location(zone)
   console.print("\n[bold]Running post-deploy configuration...[/bold]\n")
 
   steps = [
-    ("Docker authentication", lambda: configure_docker_auth(ar_location)),
     (
       "kubectl configuration",
       lambda: configure_kubectl(

keras_remote/cli/commands/up_test.py

@@ -27,9 +27,6 @@
     "keras_remote.cli.commands.up.create_program",
   ),
   "get_stack": mock.patch("keras_remote.cli.commands.up.get_stack"),
-  "configure_docker_auth": mock.patch(
-    "keras_remote.cli.commands.up.configure_docker_auth",
-  ),
   "configure_kubectl": mock.patch(
     "keras_remote.cli.commands.up.configure_kubectl",
   ),
@@ -68,7 +65,6 @@ def test_full_success(self):
     self.assertIn("Setup Complete", result.output)
     self.assertNotIn("Warnings", result.output)
     self.mocks["install_lws"].assert_called_once()
-    self.mocks["configure_docker_auth"].assert_called_once()
     self.mocks["configure_kubectl"].assert_called_once()
     self.mocks[
       "install_gpu_drivers"
@@ -83,27 +79,11 @@ def test_pulumi_failure_still_runs_post_deploy(self):
     result = self.runner.invoke(up, _CLI_ARGS)
 
     self.assertEqual(result.exit_code, 0, result.output)
-    self.mocks["configure_docker_auth"].assert_called_once()
     self.mocks["configure_kubectl"].assert_called_once()
     self.mocks["install_lws"].assert_called_once()
     self.assertIn("Setup Completed With Warnings", result.output)
     self.assertIn("Pulumi provisioning encountered errors", result.output)
 
-  def test_post_deploy_failure_does_not_block_others(self):
-    """One post-deploy step failing doesn't prevent the others from running."""
-    self.mocks[
-      "configure_docker_auth"
-    ].side_effect = subprocess.CalledProcessError(1, "gcloud")
-
-    result = self.runner.invoke(up, _CLI_ARGS)
-
-    self.assertEqual(result.exit_code, 0, result.output)
-    # Subsequent steps still called despite Docker auth failure.
-    self.mocks["configure_kubectl"].assert_called_once()
-    self.mocks["install_lws"].assert_called_once()
-    self.assertIn("Setup Completed With Warnings", result.output)
-    self.assertIn("Docker authentication", result.output)
-
   def test_multiple_post_deploy_failures(self):
     """Multiple post-deploy failures are all reported."""
     self.mocks["configure_kubectl"].side_effect = subprocess.CalledProcessError(

keras_remote/cli/constants.py

@@ -2,10 +2,12 @@
 
 import os
 
-from keras_remote.constants import DEFAULT_ZONE  # noqa: F401 — re-exported
+from keras_remote.constants import (
+  DEFAULT_CLUSTER_NAME,  # noqa: F401 — re-exported
+  DEFAULT_ZONE,  # noqa: F401 — re-exported
+)
 
 RESOURCE_NAME_PREFIX = "keras-remote"
-DEFAULT_CLUSTER_NAME = f"{RESOURCE_NAME_PREFIX}-cluster"
 STATE_DIR = os.environ.get(
   "KERAS_REMOTE_STATE_DIR",
   os.path.expanduser("~/.keras-remote/pulumi"),

keras_remote/cli/infra/post_deploy.py

@@ -1,7 +1,7 @@
 """Post-deploy steps that cannot be managed by Pulumi.
 
-These operations configure local machine state (Docker auth, kubectl)
-or apply Kubernetes manifests that depend on the cluster being ready.
+These operations configure local machine state (kubectl) or apply
+Kubernetes manifests that depend on the cluster being ready.
 """
 
 import os
@@ -13,24 +13,6 @@
 )
 
 
-def configure_docker_auth(ar_location):
-  """Configure Docker to authenticate with Artifact Registry.
-
-  Args:
-      ar_location: Multi-region location (e.g., "us", "europe", "asia").
-  """
-  subprocess.run(
-    [
-      "gcloud",
-      "auth",
-      "configure-docker",
-      f"{ar_location}-docker.pkg.dev",
-      "--quiet",
-    ],
-    check=True,
-  )
-
-
 def configure_kubectl(cluster_name, zone, project):
   """Configure kubectl to access the GKE cluster.
 

keras_remote/cli/prerequisites_check.py

@@ -2,8 +2,8 @@
 
 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.
+``click.ClickException``.  CLI-only tool checks (Pulumi, kubectl) remain
+here.
 """
 
 import shutil
@@ -37,14 +37,6 @@ def check_kubectl():
     )
 
 
-def check_docker():
-  """Verify Docker CLI is installed."""
-  if not shutil.which("docker"):
-    raise click.ClickException(
-      "Docker not found. Install from: https://docs.docker.com/get-docker/"
-    )
-
-
 def check_gke_auth_plugin():
   """Verify gke-gcloud-auth-plugin is installed; auto-install if missing."""
   try:
@@ -67,5 +59,4 @@ def check_all():
   check_pulumi()
   check_kubectl()
   check_gke_auth_plugin()
-  check_docker()
   check_gcloud_auth()

keras_remote/cli/prerequisites_check_test.py

@@ -6,7 +6,6 @@
 from absl.testing import absltest, parameterized
 
 from keras_remote.cli.prerequisites_check import (
-  check_docker,
   check_gcloud,
   check_gcloud_auth,
   check_gke_auth_plugin,
@@ -18,7 +17,7 @@
 
 
 class TestToolChecks(parameterized.TestCase):
-  """Tests for CLI-only tool checks (pulumi, kubectl, docker)."""
+  """Tests for CLI-only tool checks (pulumi, kubectl)."""
 
   @parameterized.named_parameters(
     dict(
@@ -31,11 +30,6 @@ class TestToolChecks(parameterized.TestCase):
       check_fn=check_kubectl,
       error_match="kubectl not found",
     ),
-    dict(
-      testcase_name="docker",
-      check_fn=check_docker,
-      error_match="Docker not found",
-    ),
   )
   def test_present(self, check_fn, error_match):
     with mock.patch("shutil.which", return_value="/usr/bin/tool"):
@@ -52,11 +46,6 @@ def test_present(self, check_fn, error_match):
       check_fn=check_kubectl,
       error_match="kubectl not found",
     ),
-    dict(
-      testcase_name="docker",
-      check_fn=check_docker,
-      error_match="Docker not found",
-    ),
   )
   def test_missing(self, check_fn, error_match):
     with (

keras_remote/constants.py

@@ -4,6 +4,7 @@
 
 ZONE_ENV_VAR = "KERAS_REMOTE_ZONE"
 DEFAULT_ZONE = "us-central1-a"
+DEFAULT_CLUSTER_NAME = "keras-remote-cluster"
 DEFAULT_REGION = DEFAULT_ZONE.rsplit("-", 1)[0]  # "us-central1"
 
 

keras_remote/core/core.py

@@ -7,6 +7,7 @@
   PathwaysBackend,
   execute_remote,
 )
+from keras_remote.constants import DEFAULT_CLUSTER_NAME
 from keras_remote.core import accelerators
 
 
@@ -117,7 +118,7 @@ def _execute_on_gke(
   """Execute function on GKE cluster with GPU/TPU nodes."""
   # Get GKE-specific defaults
   if not cluster:
-    cluster = os.environ.get("KERAS_REMOTE_CLUSTER")
+    cluster = os.environ.get("KERAS_REMOTE_CLUSTER", DEFAULT_CLUSTER_NAME)
   if not namespace:
     namespace = os.environ.get("KERAS_REMOTE_GKE_NAMESPACE", "default")
 
@@ -141,7 +142,7 @@ def _execute_on_pathways(
 ):
   """Execute function on GKE cluster via ML Pathways."""
   if not cluster:
-    cluster = os.environ.get("KERAS_REMOTE_CLUSTER")
+    cluster = os.environ.get("KERAS_REMOTE_CLUSTER", DEFAULT_CLUSTER_NAME)
   if not namespace:
     namespace = os.environ.get("KERAS_REMOTE_GKE_NAMESPACE", "default")