reformat code

Author: divyashreepathihalli

keras_remote/backend/gke_client.py

@@ -1,9 +1,9 @@
 """GKE job submission for keras_remote."""
 
+import json
+import subprocess
 import time
 from contextlib import suppress
-import subprocess
-import json
 
 from absl import logging
 from kubernetes import client, config
@@ -400,9 +400,11 @@ def _validate_node_pool_exists(selector: dict) -> bool:
   """Use gcloud to verify that a GKE NodePool matches the pod node selector."""
   try:
     # Requires gcloud CLI and valid credentials.
-    out = subprocess.check_output([
-        "gcloud", "container", "node-pools", "list",
-        "--format", "json"], text=True, stderr=subprocess.DEVNULL)
+    out = subprocess.check_output(
+      ["gcloud", "container", "node-pools", "list", "--format", "json"],
+      text=True,
+      stderr=subprocess.DEVNULL,
+    )
     pools = json.loads(out)
     for pool in pools:
       config = pool.get("config", {})
@@ -415,6 +417,7 @@ def _validate_node_pool_exists(selector: dict) -> bool:
     # If gcloud is missing or unauthenticated, degrade gracefully and assume pool exists
     return True
 
+
 def _check_pod_scheduling(core_v1, job_name, namespace):
   """Check for pod scheduling issues and raise helpful errors."""
   with suppress(ApiException):
@@ -429,9 +432,15 @@ def _check_pod_scheduling(core_v1, job_name, namespace):
             if "Insufficient nvidia.com/gpu" in msg:
               selector = pod.spec.node_selector or {}
               if not _validate_node_pool_exists(selector):
-                  selector_str = ", ".join([f"{k}: {v}" for k, v in selector.items()]) if selector else "None"
-                  raise RuntimeError(f"No GKE node pool exists with selector '{selector_str}'. "
-                                     "Please use 'keras-remote pool add' to configure this accelerator.")
+                selector_str = (
+                  ", ".join([f"{k}: {v}" for k, v in selector.items()])
+                  if selector
+                  else "None"
+                )
+                raise RuntimeError(
+                  f"No GKE node pool exists with selector '{selector_str}'. "
+                  "Please use 'keras-remote pool add' to configure this accelerator."
+                )
               logging.info(
                 f"Pod {pod.metadata.name} is Pending: Insufficient nvidia.com/gpu. "
                 "Waiting for GKE Cluster Autoscaler to provision a new node... (scale-to-zero)\n"
@@ -442,12 +451,18 @@ def _check_pod_scheduling(core_v1, job_name, namespace):
               or "node selector" in msg.lower()
             ):
               selector = pod.spec.node_selector or {}
-              selector_str = ", ".join([f"{k}: {v}" for k, v in selector.items()]) if selector else "None"
-              
+              selector_str = (
+                ", ".join([f"{k}: {v}" for k, v in selector.items()])
+                if selector
+                else "None"
+              )
+
               if not _validate_node_pool_exists(selector):
-                  raise RuntimeError(f"No GKE node pool exists with selector '{selector_str}'. "
-                                     "Please use 'keras-remote pool add' to configure this accelerator.")
-              
+                raise RuntimeError(
+                  f"No GKE node pool exists with selector '{selector_str}'. "
+                  "Please use 'keras-remote pool add' to configure this accelerator."
+                )
+
               logging.info(
                 f"Pod {pod.metadata.name} is Pending: No currently running nodes "
                 f"match accelerator selector '{selector_str}'. "

keras_remote/cli/infra/program_test.py

@@ -350,12 +350,18 @@ def test_node_pool_scale_to_zero(
     is_multi_host = getattr(accelerator, "num_nodes", 1) > 1
 
     call_kwargs = gcp_mock.container.NodePool.call_args.kwargs
-    self.assertEqual(call_kwargs.get("initial_node_count"), expected_max_count if is_multi_host else 0)
+    self.assertEqual(
+      call_kwargs.get("initial_node_count"),
+      expected_max_count if is_multi_host else 0,
+    )
 
     autoscaling_kwargs = (
       gcp_mock.container.NodePoolAutoscalingArgs.call_args.kwargs
     )
-    self.assertEqual(autoscaling_kwargs.get("min_node_count"), expected_max_count if is_multi_host else 0)
+    self.assertEqual(
+      autoscaling_kwargs.get("min_node_count"),
+      expected_max_count if is_multi_host else 0,
+    )
     self.assertEqual(
       autoscaling_kwargs.get("max_node_count"), expected_max_count
     )