feat: improve accelerator validation and container caching

- Group accelerators by category (CPU, GPU, TPU) for container image sharing, reducing redundant builds.
- Implement preflight check to validate node pool existence before building containers.
- Improve error messages for node selector mismatches in pod scheduling.
- Update simple_demo.py to return training loss and use more idiomatic result retrieval.
- Add preflight check note to README with link to Quick Start.

Author: jeffcarp

README.md

@@ -228,6 +228,9 @@ See [examples/Dockerfile.prebuilt](examples/Dockerfile.prebuilt) for a template.
 
 ## Supported Accelerators
 
+Note: each accelerator and topology requires [setting up its own NodePool](#quick-start)
+as a prerequisite.
+
 ### TPUs
 
 | Type           | Configurations                              |

keras_remote/backend/execution.py

@@ -101,6 +101,10 @@ def __init__(self, cluster: Optional[str] = None, namespace: str = "default"):
     self.cluster = cluster
     self.namespace = namespace
 
+  def validate_preflight(self, ctx: JobContext) -> None:
+    """Perform preflight checks before building container or uploading artifacts."""
+    pass
+
   def submit_job(self, ctx: JobContext) -> Any:
     """Submit a job to the backend. Returns backend-specific job handle."""
     raise NotImplementedError
@@ -117,6 +121,16 @@ def cleanup_job(self, job: Any, ctx: JobContext) -> None:
 class GKEBackend(BaseK8sBackend):
   """Backend adapter for standard GKE Jobs."""
 
+  def validate_preflight(self, ctx: JobContext) -> None:
+    """Check if the required node pool exists for the accelerator."""
+    gke_client.validate_preflight(
+      accelerator=ctx.accelerator,
+      project=ctx.project,
+      cluster=self.cluster,
+      zone=ctx.zone,
+      namespace=self.namespace,
+    )
+
   def submit_job(self, ctx: JobContext) -> Any:
     """Submit job to GKE cluster."""
     return gke_client.submit_k8s_job(
@@ -142,6 +156,17 @@ def cleanup_job(self, job: Any, ctx: JobContext) -> None:
 class PathwaysBackend(BaseK8sBackend):
   """Backend adapter for ML Pathways using LeaderWorkerSet."""
 
+  def validate_preflight(self, ctx: JobContext) -> None:
+    """Preflight checks for Pathways (currently same as GKE)."""
+    # Pathways also runs on GKE nodes with specific labels
+    gke_client.validate_preflight(
+      accelerator=ctx.accelerator,
+      project=ctx.project,
+      cluster=self.cluster,
+      zone=ctx.zone,
+      namespace=self.namespace,
+    )
+
   def submit_job(self, ctx: JobContext) -> Any:
     """Submit LWS job to GKE cluster."""
     return pathways_client.submit_pathways_job(
@@ -289,6 +314,9 @@ def execute_remote(ctx: JobContext, backend: BaseK8sBackend) -> Any:
     cluster=backend.cluster,
   )
 
+  # Preflight check
+  backend.validate_preflight(ctx)
+
   with tempfile.TemporaryDirectory() as tmpdir:
     # Phase 1: Package artifacts
     _prepare_artifacts(ctx, tmpdir)

keras_remote/backend/gke_client.py

@@ -179,6 +179,46 @@ def cleanup_job(job_name, namespace="default"):
       logging.warning("Failed to delete job %s: %s", job_name, e.reason)
 
 
+def validate_preflight(accelerator, project, cluster, zone, namespace="default"):
+  """Check if the required node pool exists for the accelerator.
+
+  Args:
+      accelerator: Accelerator string (e.g., 'l4', 'v3-8')
+      project: GCP project ID
+      cluster: GKE cluster name
+      zone: GCP zone
+      namespace: Kubernetes namespace
+
+  Raises:
+      RuntimeError: If no nodes match the required accelerator selector.
+  """
+  _load_kube_config()
+  accel_config = _parse_accelerator(accelerator)
+  node_selector = accel_config.get("node_selector")
+
+  if not node_selector:
+    return  # CPU or no selector required
+
+  core_v1 = client.CoreV1Api()
+  try:
+    # Construct label selector string: "key1=val1,key2=val2"
+    label_selector = ",".join([f"{k}={v}" for k, v in node_selector.items()])
+    nodes = core_v1.list_node(label_selector=label_selector)
+
+    if not nodes.items:
+      selector_str = ", ".join([f"{k}: {v}" for k, v in node_selector.items()])
+      raise RuntimeError(
+        f"Preflight check failed: No nodes match the accelerator selector: {selector_str}. "
+        "Check that your GKE cluster has a node pool with the correct accelerator type. "
+        "See all supported accelerator symbols here: \n"
+        "https://github.com/keras-team/remote#supported-accelerators"
+      )
+  except ApiException as e:
+    # If we can't list nodes due to permissions, log a warning but proceed
+    # to avoid blocking users with restricted kubeconfig.
+    logging.warning("Preflight check: Failed to query nodes: %s", e.reason)
+
+
 def _parse_accelerator(accelerator):
   """Convert accelerator string to GKE pod spec fields."""
   parsed = accelerators.parse_accelerator(accelerator)
@@ -374,7 +414,15 @@ def _check_pod_scheduling(core_v1, job_name, namespace):
               "didn't match Pod's node affinity/selector" in msg
               or "node selector" in msg.lower()
             ):
+              selector = pod.spec.node_selector
+              selector_str = (
+                ", ".join([f"{k}: {v}" for k, v in selector.items()])
+                if selector
+                else "None"
+              )
               raise RuntimeError(
-                "No nodes match the GPU selector. Check that your node pool "
-                "has the correct GPU type label."
+                f"No nodes match the accelerator selector: {selector_str}. "
+                "Check that your node pool has the correct accelerator type label. "
+                "See all supported accelerator symbols here: \n"
+                "https://github.com/keras-team/remote#supported-accelerators"
               )

keras_remote/backend/gke_client_test.py

@@ -356,9 +356,10 @@ def test_kubeconfig_fallback(self):
 
 
 class TestCheckPodScheduling(parameterized.TestCase):
-  def _make_pending_pod(self, message):
+  def _make_pending_pod(self, message, node_selector=None):
     pod = MagicMock()
     pod.status.phase = "Pending"
+    pod.spec.node_selector = node_selector
     condition = MagicMock()
     condition.type = "PodScheduled"
     condition.status = "False"
@@ -371,16 +372,20 @@ def _make_pending_pod(self, message):
       testcase_name="insufficient_gpu",
       condition_message="Insufficient nvidia.com/gpu",
       error_match="No GPU nodes available",
+      node_selector=None,
     ),
     dict(
       testcase_name="node_selector_mismatch",
       condition_message="didn't match Pod's node affinity/selector",
-      error_match="No nodes match",
+      error_match="No nodes match the accelerator selector: cloud.google.com/gke-accelerator: nvidia-l4",
+      node_selector={"cloud.google.com/gke-accelerator": "nvidia-l4"},
     ),
   )
-  def test_scheduling_failure_raises(self, condition_message, error_match):
+  def test_scheduling_failure_raises(
+    self, condition_message, error_match, node_selector
+  ):
     mock_core = MagicMock()
-    pod = self._make_pending_pod(condition_message)
+    pod = self._make_pending_pod(condition_message, node_selector=node_selector)
     mock_core.list_namespaced_pod.return_value.items = [pod]
 
     with self.assertRaisesRegex(RuntimeError, error_match):

keras_remote/infra/container_builder.py

@@ -42,15 +42,15 @@ def get_or_build_container(
       Container image URI in Artifact Registry
   """
   ar_location = zone_to_ar_location(zone or get_default_zone())
+  category = accelerators.get_category(accelerator_type)
 
-  # Generate deterministic hash from requirements + base image
+  # Generate deterministic hash from requirements + base image + category
   requirements_hash = _hash_requirements(
-    requirements_path, accelerator_type, base_image
+    requirements_path, category, base_image
   )
 
-  # Sanitize accelerator type for image name
-  sanitized_accel = accelerator_type.replace(":", "-").replace("/", "-")
-  image_tag = f"{sanitized_accel}-{requirements_hash[:12]}"
+  # 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"
@@ -72,25 +72,25 @@ def get_or_build_container(
   return _build_and_push(
     base_image,
     requirements_path,
-    accelerator_type,
+    category,
     project,
     image_uri,
     ar_location,
   )
 
 
-def _hash_requirements(requirements_path, accelerator_type, base_image):
-  """Create deterministic hash from requirements + accelerator + remote_runner + base image.
+def _hash_requirements(requirements_path, category, base_image):
+  """Create deterministic hash from requirements + category + remote_runner + base image.
 
   Args:
       requirements_path: Path to requirements.txt (or None)
-      accelerator_type: TPU/GPU type
+      category: Accelerator category ('cpu', 'gpu', 'tpu')
       base_image: Base Docker image (e.g., 'python:3.12-slim')
 
   Returns:
       SHA256 hex digest
   """
-  content = f"base_image={base_image}\naccelerator={accelerator_type}\n"
+  content = f"base_image={base_image}\ncategory={category}\n"
 
   if requirements_path and os.path.exists(requirements_path):
     with open(requirements_path, "r") as f:
@@ -150,7 +150,7 @@ def _image_exists(image_uri, project):
 def _build_and_push(
   base_image,
   requirements_path,
-  accelerator_type,
+  category,
   project,
   image_uri,
   ar_location="us",
@@ -160,7 +160,7 @@ def _build_and_push(
   Args:
       base_image: Base Docker image
       requirements_path: Path to requirements.txt (or None)
-      accelerator_type: TPU/GPU type
+      category: Accelerator category ('cpu', 'gpu', 'tpu')
       project: GCP project ID
       image_uri: Target image URI
       ar_location: Artifact Registry multi-region (e.g., 'us')
@@ -173,7 +173,7 @@ def _build_and_push(
     dockerfile_content = _generate_dockerfile(
       base_image=base_image,
       requirements_path=requirements_path,
-      accelerator_type=accelerator_type,
+      category=category,
     )
 
     dockerfile_path = os.path.join(tmpdir, "Dockerfile")
@@ -255,19 +255,18 @@ def _build_and_push(
       raise RuntimeError(f"Build failed with status: {result.status}")
 
 
-def _generate_dockerfile(base_image, requirements_path, accelerator_type):
+def _generate_dockerfile(base_image, requirements_path, category):
   """Generate Dockerfile content based on configuration.
 
   Args:
       base_image: Base Docker image
       requirements_path: Path to requirements.txt (or None)
-      accelerator_type: TPU/GPU type
+      category: Accelerator category ('cpu', 'gpu', 'tpu')
 
   Returns:
       Dockerfile content as string
   """
-  # Determine JAX installation command based on accelerator
-  category = accelerators.get_category(accelerator_type)
+  # Determine JAX installation command based on accelerator category
   if category == "cpu":
     jax_install = "RUN python3 -m pip install jax"
   elif category == "tpu":