Multi-stack support

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/infra/program.py

@@ -71,7 +71,7 @@ def pulumi_program():
     # 2. Artifact Registry docker repository
     repo = gcp.artifactregistry.Repository(
       "keras-remote-repo",
-      repository_id="keras-remote",
+      repository_id=f"kr-{cluster_name}",
       location=ar_location,
       format="DOCKER",
       description="keras-remote container images",
@@ -84,7 +84,7 @@ def pulumi_program():
 
     gcp.storage.Bucket(
       "keras-remote-jobs-bucket",
-      name=f"{project_id}-keras-remote-jobs",
+      name=f"{project_id}-kr-{cluster_name}-jobs",
       location=region,
       project=project_id,
       force_destroy=True,
@@ -93,7 +93,7 @@ def pulumi_program():
 
     gcp.storage.Bucket(
       "keras-remote-builds-bucket",
-      name=f"{project_id}-keras-remote-builds",
+      name=f"{project_id}-kr-{cluster_name}-builds",
       location=ar_location,
       project=project_id,
       force_destroy=True,
@@ -170,7 +170,7 @@ def pulumi_program():
     pulumi.export(
       "ar_registry",
       repo.name.apply(
-        lambda _: f"{ar_location}-docker.pkg.dev/{project_id}/keras-remote"
+        lambda _: f"{ar_location}-docker.pkg.dev/{project_id}/kr-{cluster_name}"
       ),
     )
 

keras_remote/cli/infra/stack_manager.py

@@ -33,8 +33,9 @@ def get_stack(program_fn, config):
     click.echo("Pulumi CLI not found. Installing...")
     pulumi_cmd = auto.PulumiCommand.install(root=PULUMI_ROOT)
 
-  # Use project ID as stack name so each GCP project gets its own stack
-  stack_name = config.project
+  # Each (project, cluster) pair gets its own stack, so multiple clusters
+  # within the same GCP project are fully independent.
+  stack_name = f"{config.project}-{config.cluster_name}"
 
   project_settings = auto.ProjectSettings(
     name=RESOURCE_NAME_PREFIX,

keras_remote/constants.py

@@ -13,6 +13,11 @@ def get_default_zone():
   return os.environ.get(ZONE_ENV_VAR, DEFAULT_ZONE)
 
 
+def get_default_cluster_name():
+  """Return cluster name from KERAS_REMOTE_CLUSTER env var, or DEFAULT_CLUSTER_NAME."""
+  return os.environ.get("KERAS_REMOTE_CLUSTER", DEFAULT_CLUSTER_NAME)
+
+
 def zone_to_region(zone):
   """Convert a GCP zone to its region (e.g. 'us-central1-a' -> 'us-central1')."""
   return zone.rsplit("-", 1)[0] if zone and "-" in zone else DEFAULT_REGION

keras_remote/core/core.py

@@ -154,6 +154,7 @@ def _execute_on_gke(
     zone,
     project,
     env_vars,
+    cluster_name=cluster,
     volumes=volumes,
   )
   return execute_remote(ctx, GKEBackend(cluster=cluster, namespace=namespace))
@@ -187,6 +188,7 @@ def _execute_on_pathways(
     zone,
     project,
     env_vars,
+    cluster_name=cluster,
     volumes=volumes,
   )
   return execute_remote(

keras_remote/infra/container_builder.py

@@ -14,7 +14,11 @@
 from google.cloud import artifactregistry_v1, storage
 from google.cloud.devtools import cloudbuild_v1
 
-from keras_remote.constants import get_default_zone, zone_to_ar_location
+from keras_remote.constants import (
+  get_default_cluster_name,
+  get_default_zone,
+  zone_to_ar_location,
+)
 from keras_remote.core import accelerators
 
 REMOTE_RUNNER_FILE_NAME = "remote_runner.py"
@@ -26,7 +30,12 @@
 
 
 def get_or_build_container(
-  base_image, requirements_path, accelerator_type, project, zone=None
+  base_image,
+  requirements_path,
+  accelerator_type,
+  project,
+  zone=None,
+  cluster_name=None,
 ):
   """Get existing container or build if requirements changed.
 
@@ -38,11 +47,13 @@ def get_or_build_container(
       accelerator_type: TPU/GPU type (e.g., 'v3-8')
       project: GCP project ID
       zone: GCP zone for region derivation (defaults to KERAS_REMOTE_ZONE)
+      cluster_name: GKE cluster name (defaults to KERAS_REMOTE_CLUSTER)
 
   Returns:
       Container image URI in Artifact Registry
   """
   ar_location = zone_to_ar_location(zone or get_default_zone())
+  cluster_name = cluster_name or get_default_cluster_name()
   category = accelerators.get_category(accelerator_type)
 
   # Generate deterministic hash from requirements + base image + category
@@ -53,8 +64,9 @@ def get_or_build_container(
   # Use category for image name (e.g., 'tpu-hash', 'gpu-hash')
   image_tag = f"{category}-{requirements_hash[:12]}"
 
-  # Use Artifact Registry
-  registry = f"{ar_location}-docker.pkg.dev/{project}/keras-remote"
+  # Use Artifact Registry (cluster-scoped repo)
+  repo_id = f"kr-{cluster_name}"
+  registry = f"{ar_location}-docker.pkg.dev/{project}/{repo_id}"
   image_uri = f"{registry}/base:{image_tag}"
 
   # Check if image exists
@@ -63,7 +75,7 @@ def get_or_build_container(
     ar_url = (
       "https://console.cloud.google.com/artifacts"
       f"/docker/{project}/{ar_location}"
-      f"/keras-remote/base?project={project}"
+      f"/{repo_id}/base?project={project}"
     )
     logging.info("View image: %s", ar_url)
     return image_uri
@@ -77,6 +89,7 @@ def get_or_build_container(
     project,
     image_uri,
     ar_location,
+    cluster_name,
   )
 
 
@@ -155,6 +168,7 @@ def _build_and_push(
   project,
   image_uri,
   ar_location="us",
+  cluster_name=None,
 ):
   """Build and push Docker image using Cloud Build.
 
@@ -200,8 +214,9 @@ def _build_and_push(
           os.path.join(tmpdir, "requirements.txt"), arcname="requirements.txt"
         )
 
-    # Upload source to GCS
-    bucket_name = f"{project}-keras-remote-builds"
+    # Upload source to GCS (cluster-scoped bucket)
+    cluster_name = cluster_name or get_default_cluster_name()
+    bucket_name = f"{project}-kr-{cluster_name}-builds"
     source_gcs = _upload_build_source(tarball_path, bucket_name, project)
 
     # Submit build to Cloud Build

keras_remote/infra/container_builder_test.py

@@ -236,10 +236,11 @@ def test_returns_cached_when_image_exists(self):
         accelerator_type="l4",
         project="test-proj",
         zone="us-central1-a",
+        cluster_name="my-cluster",
       )
 
     mock_build.assert_not_called()
-    self.assertIn("us-docker.pkg.dev/test-proj/keras-remote/base:", result)
+    self.assertIn("us-docker.pkg.dev/test-proj/kr-my-cluster/base:", result)
 
   def test_builds_when_image_missing(self):
     with (
@@ -249,7 +250,7 @@ def test_builds_when_image_missing(self):
       ),
       mock.patch(
         "keras_remote.infra.container_builder._build_and_push",
-        return_value="us-docker.pkg.dev/proj/keras-remote/base:gpu-bbbbbbbbbbbb",
+        return_value="us-docker.pkg.dev/proj/kr-my-cluster/base:gpu-bbbbbbbbbbbb",
       ) as mock_build,
     ):
       result = get_or_build_container(
@@ -258,11 +259,13 @@ def test_builds_when_image_missing(self):
         accelerator_type="l4",
         project="proj",
         zone="us-central1-a",
+        cluster_name="my-cluster",
       )
 
     mock_build.assert_called_once()
     self.assertEqual(
-      result, "us-docker.pkg.dev/proj/keras-remote/base:gpu-bbbbbbbbbbbb"
+      result,
+      "us-docker.pkg.dev/proj/kr-my-cluster/base:gpu-bbbbbbbbbbbb",
     )
 
   def _get_image_uri(self, accelerator_type, project, zone):
@@ -276,13 +279,14 @@ def _get_image_uri(self, accelerator_type, project, zone):
         accelerator_type=accelerator_type,
         project=project,
         zone=zone,
+        cluster_name="my-cluster",
       )
 
   def test_image_uri_format_tpu_europe(self):
     result = self._get_image_uri("v3-4", "my-proj", "europe-west4-b")
 
     self.assertTrue(
-      result.startswith("europe-docker.pkg.dev/my-proj/keras-remote/base:")
+      result.startswith("europe-docker.pkg.dev/my-proj/kr-my-cluster/base:")
     )
     tag = result.split(":")[-1]
     self.assertRegex(tag, r"^tpu-[0-9a-f]{12}$")
@@ -291,7 +295,7 @@ def test_image_uri_format_gpu_us(self):
     result = self._get_image_uri("a100-80gb", "proj", "us-central1-a")
 
     self.assertTrue(
-      result.startswith("us-docker.pkg.dev/proj/keras-remote/base:")
+      result.startswith("us-docker.pkg.dev/proj/kr-my-cluster/base:")
     )
     tag = result.split(":")[-1]
     self.assertRegex(tag, r"^gpu-[0-9a-f]{12}$")