spot instance support

Author: JyotinderSingh

examples/example_gke.py

@@ -54,7 +54,7 @@ def simple_computation(x, y):
 
 
 # Example 2: Keras model training on CPU
-@keras_remote.run(accelerator="cpu")
+@keras_remote.run(accelerator="v6e-2x4", cluster="spot-tpu-nodes", spot=True)
 def train_simple_model_cpu():
   """Train a simple Keras model on remote CPU."""
 
@@ -111,10 +111,10 @@ def main():
   print("=" * 60)
 
   # Example 1: Simple computation (CPU)
-  print("\n--- Example 1: Simple Computation (CPU) ---")
-  print("Running simple_computation(10, 20) on GKE...")
-  result = simple_computation(10, 20)
-  print(f"Result: {result}")
+  # print("\n--- Example 1: Simple Computation (CPU) ---")
+  # print("Running simple_computation(10, 20) on GKE...")
+  # result = simple_computation(10, 20)
+  # print(f"Result: {result}")
 
   # Example 2: Model training on CPU
   print("\n--- Example 2: Keras Model Training (CPU) ---")

keras_remote/backend/execution.py

@@ -56,6 +56,9 @@ class JobContext:
   # Data volumes {mount_path: Data}
   volumes: Optional[dict] = None
 
+  # Configuration modifiers
+  spot: bool = False
+
   # Artifact paths (set during prepare phase)
   payload_path: Optional[str] = None
   context_path: Optional[str] = None
@@ -80,6 +83,7 @@ def from_params(
     env_vars: dict,
     cluster_name: Optional[str] = None,
     volumes: Optional[dict] = None,
+    spot: bool = False,
   ) -> "JobContext":
     """Factory method with default resolution for zone/project/cluster."""
     if not zone:
@@ -105,6 +109,7 @@ def from_params(
       project=project,
       cluster_name=cluster_name,
       volumes=volumes,
+      spot=spot,
     )
 
 
@@ -155,6 +160,7 @@ def submit_job(self, ctx: JobContext) -> Any:
       job_id=ctx.job_id,
       bucket_name=ctx.bucket_name,
       namespace=self.namespace,
+      spot=ctx.spot,
     )
 
   def wait_for_job(self, job: Any, ctx: JobContext) -> None:
@@ -191,6 +197,7 @@ def submit_job(self, ctx: JobContext) -> Any:
       job_id=ctx.job_id,
       bucket_name=ctx.bucket_name,
       namespace=self.namespace,
+      spot=ctx.spot,
     )
 
   def wait_for_job(self, job: Any, ctx: JobContext) -> None:

keras_remote/backend/gke_client.py

@@ -23,6 +23,7 @@ def submit_k8s_job(
   job_id,
   bucket_name,
   namespace="default",
+  spot=False,
 ):
   """Submit a Kubernetes Job to GKE cluster.
 
@@ -42,7 +43,7 @@ def submit_k8s_job(
   _load_kube_config()
 
   # Parse accelerator configuration
-  accel_config = _parse_accelerator(accelerator)
+  accel_config = _parse_accelerator(accelerator, spot=spot)
 
   # Create job specification
   job_name = f"keras-remote-{job_id}"
@@ -224,9 +225,9 @@ def validate_preflight(
     logging.warning("Preflight check: Failed to query nodes: %s", e.reason)
 
 
-def _parse_accelerator(accelerator):
+def _parse_accelerator(accelerator, spot=False):
   """Convert accelerator string to GKE pod spec fields."""
-  parsed = accelerators.parse_accelerator(accelerator)
+  parsed = accelerators.parse_accelerator(accelerator, spot=spot)
 
   if parsed is None:
     return {
@@ -241,7 +242,7 @@ def _parse_accelerator(accelerator):
     # For TPU Podslices (multi-node), resource requests must be per-node.
     # num_nodes is 1 for single-host TPUs (v3-8, v4-8, v5litepod-1/4/8).
     chips_per_node = parsed.chips // parsed.num_nodes
-    return {
+    config = {
       "node_selector": {
         "cloud.google.com/gke-tpu-accelerator": parsed.gke_accelerator,
         "cloud.google.com/gke-tpu-topology": parsed.topology,
@@ -254,8 +255,20 @@ def _parse_accelerator(accelerator):
       "jax_platform": "tpu",
     }
 
+    if parsed.spot:
+      config["node_selector"]["cloud.google.com/gke-spot"] = "true"
+      config["tolerations"].append(
+        {
+          "key": "cloud.google.com/gke-spot",
+          "operator": "Equal",
+          "value": "true",
+          "effect": "NoSchedule",
+        }
+      )
+    return config
+
   # GpuConfig
-  return {
+  config = {
     "node_selector": {"cloud.google.com/gke-accelerator": parsed.gke_label},
     "resource_limits": {"nvidia.com/gpu": str(parsed.count)},
     "resource_requests": {"nvidia.com/gpu": str(parsed.count)},
@@ -264,6 +277,17 @@ def _parse_accelerator(accelerator):
     ],
     "jax_platform": "gpu",
   }
+  if parsed.spot:
+    config["node_selector"]["cloud.google.com/gke-spot"] = "true"
+    config["tolerations"].append(
+      {
+        "key": "cloud.google.com/gke-spot",
+        "operator": "Equal",
+        "value": "true",
+        "effect": "NoSchedule",
+      }
+    )
+  return config
 
 
 def _load_kube_config():
@@ -441,6 +465,10 @@ def _check_node_pool_exists_cached(selector_items) -> bool:
       config_dict = pool.get("config", {})
       pool_labels = config_dict.get("labels", {}).copy()
 
+      # Spot VM mapping
+      if config_dict.get("spot"):
+        pool_labels["cloud.google.com/gke-spot"] = "true"
+
       # Map GKE injected node labels for accelerators mapping
       accel_config_list = config_dict.get("accelerators", [])
       if accel_config_list:
@@ -450,6 +478,13 @@ def _check_node_pool_exists_cached(selector_items) -> bool:
         else:
           pool_labels["cloud.google.com/gke-accelerator"] = accel_type
 
+      # TPU topology mapping from placement policy
+      placement_policy = pool.get("placementPolicy", {})
+      if placement_policy and placement_policy.get("tpuTopology"):
+        pool_labels["cloud.google.com/gke-tpu-topology"] = placement_policy[
+          "tpuTopology"
+        ]
+
       # TPU mapping fallback
       machine_type = config_dict.get("machineType", "")
 
@@ -460,7 +495,9 @@ def _check_node_pool_exists_cached(selector_items) -> bool:
           "goog-gke-accelerator-type"
         ]
 
-      if machine_type.startswith("ct"):
+      if machine_type.startswith("ct") and not pool_labels.get(
+        "cloud.google.com/gke-tpu-topology"
+      ):
         # We roughly map TPU topology presence for preflight
         pool_labels["cloud.google.com/gke-tpu-topology"] = selector.get(
           "cloud.google.com/gke-tpu-topology", ""

keras_remote/backend/gke_client_test.py

@@ -78,6 +78,38 @@ def test_tpu_v5litepod_4(self):
     )
     self.assertEqual(result["resource_limits"], {"google.com/tpu": "4"})
 
+  def test_spot_gpu(self):
+    result = _parse_accelerator("l4:spot")
+    self.assertEqual(
+      result["node_selector"]["cloud.google.com/gke-spot"], "true"
+    )
+    # Check for spot toleration
+    spot_tol = [
+      t
+      for t in result["tolerations"]
+      if t.get("key") == "cloud.google.com/gke-spot"
+    ]
+    self.assertLen(spot_tol, 1)
+    self.assertEqual(spot_tol[0]["value"], "true")
+
+  def test_spot_tpu(self):
+    result = _parse_accelerator("v6e-8:spot")
+    self.assertEqual(
+      result["node_selector"]["cloud.google.com/gke-spot"], "true"
+    )
+    # Check for spot toleration
+    spot_tol = [
+      t
+      for t in result["tolerations"]
+      if t.get("key") == "cloud.google.com/gke-spot"
+    ]
+    self.assertLen(spot_tol, 1)
+    self.assertEqual(spot_tol[0]["value"], "true")
+    # Should still have TPU toleration
+    self.assertTrue(
+      any(t.get("key") == "google.com/tpu" for t in result["tolerations"])
+    )
+
 
 class TestCreateJobSpec(absltest.TestCase):
   def _make_gpu_config(self):

keras_remote/backend/pathways_client.py

@@ -53,6 +53,7 @@ def submit_pathways_job(
   job_id,
   bucket_name,
   namespace="default",
+  spot=False,
 ):
   """Submit a LeaderWorkerSet to GKE cluster.
 
@@ -71,12 +72,10 @@ def submit_pathways_job(
   _load_kube_config()
   lws_version = _get_lws_version()
 
-  accel_config = _parse_accelerator(accelerator)
+  parsed_config = accelerators.parse_accelerator(accelerator, spot=spot)
+  accel_config = _parse_accelerator(accelerator, spot=spot)
   job_name = _get_job_name(job_id)
 
-  # Extract num nodes from the TPU configuration
-
-  parsed_config = accelerators.parse_accelerator(accelerator)
   if (
     isinstance(parsed_config, accelerators.TpuConfig)
     and parsed_config.num_nodes > 1
@@ -263,10 +262,12 @@ def _create_lws_spec(
     {"name": "TPU_WORKER_ID", "value": "$(LWS_WORKER_INDEX)"},
   ]
 
-  tolerations = [
-    {"key": t["key"], "operator": t["operator"], "effect": t["effect"]}
-    for t in accel_config["tolerations"]
-  ]
+  tolerations = []
+  for t in accel_config["tolerations"]:
+    entry = {"key": t["key"], "operator": t["operator"], "effect": t["effect"]}
+    if "value" in t:
+      entry["value"] = t["value"]
+    tolerations.append(entry)
 
   pod_template = {
     "metadata": {

keras_remote/backend/pathways_client_test.py

@@ -143,6 +143,33 @@ def test_env_vars(self):
     self.assertEqual(env["MEGASCALE_NUM_SLICES"], "4")
     self.assertEqual(env["TPU_WORKER_ID"], "$(LWS_WORKER_INDEX)")
 
+  def test_spot_spec(self):
+    """Test that spot selectors and tolerations are added when present."""
+    accel_config = self._make_tpu_accel_config()
+    accel_config["node_selector"]["cloud.google.com/gke-spot"] = "true"
+    accel_config["tolerations"].append(
+      {
+        "key": "cloud.google.com/gke-spot",
+        "operator": "Equal",
+        "value": "true",
+        "effect": "NoSchedule",
+      }
+    )
+
+    spec = self._make_spec(accel_config=accel_config)
+    pod_spec = spec["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]
+
+    self.assertEqual(
+      pod_spec["nodeSelector"]["cloud.google.com/gke-spot"], "true"
+    )
+    spot_tol = [
+      t
+      for t in pod_spec["tolerations"]
+      if t.get("key") == "cloud.google.com/gke-spot"
+    ]
+    self.assertLen(spot_tol, 1)
+    self.assertEqual(spot_tol[0]["value"], "true")
+
   def test_tpu_accel_config(self):
     """Test resources, tolerations, and node selector for TPU config."""
     spec = self._make_spec(accel_config=self._make_tpu_accel_config())

keras_remote/cli/commands/pool.py

@@ -29,13 +29,14 @@ def pool():
   "v5litepod, v5p, v6e, v3 (with optional count/topology)",
 )
 @click.option("--yes", "-y", is_flag=True, help="Skip confirmation prompt")
-def pool_add(project, zone, cluster_name, accelerator, yes):
+@click.option("--spot", is_flag=True, help="Use Spot VMs for node pool")
+def pool_add(project, zone, cluster_name, accelerator, yes, spot):
   """Add an accelerator node pool to the cluster."""
   banner("keras-remote Pool Add")
 
   # Parse the accelerator spec first to fail fast on bad input.
   try:
-    accel_config = accelerators.parse_accelerator(accelerator)
+    accel_config = accelerators.parse_accelerator(accelerator, spot=spot)
   except ValueError as e:
     raise click.BadParameter(str(e), param_hint="--accelerator") from e
 

keras_remote/cli/infra/program.py

@@ -237,6 +237,7 @@ def _create_gpu_node_pool(cluster, gpu: GpuConfig, zone, project_id, pool_name):
       ],
       labels={RESOURCE_NAME_PREFIX: "true"},
       max_run_duration=f"{NODE_MAX_RUN_DURATION_SECONDS}s",  # 24 hours
+      spot=gpu.spot,
     ),
   )
 
@@ -276,7 +277,10 @@ def _create_tpu_node_pool(cluster, tpu: TpuConfig, zone, project_id, pool_name):
       machine_type=tpu.machine_type,
       oauth_scopes=_BASE_OAUTH_SCOPES,
       labels={RESOURCE_NAME_PREFIX: "true"},
-      max_run_duration=f"{NODE_MAX_RUN_DURATION_SECONDS}s",  # 24 hours
+      max_run_duration=None
+      if tpu.spot
+      else f"{NODE_MAX_RUN_DURATION_SECONDS}s",  # 24 hours
+      spot=tpu.spot,
     ),
     placement_policy=placement,
   )