code reformat

Author: divyashreepathihalli

examples/pathways_example.py

@@ -3,46 +3,49 @@
 os.environ["KERAS_BACKEND"] = "jax"
 
 import keras
-from keras import layers
 import numpy as np
+from keras import layers
+
 import keras_remote
 
+
 # A simple model that will be executed remotely
 @keras_remote.run(
-    accelerator="v5litepod-1", 
+  accelerator="v5litepod-1",
 )
 def train_simple_model():
-    print("Running Pathways job on JAX Backend!")
-    
-    # Create a simple dataset
-    x = np.random.rand(1000, 10)
-    y = np.random.randint(0, 2, size=(1000, 1))
-
-    # A simple sequential model
-    model = keras.Sequential([
-        keras.Input(shape=(10,)),
-        layers.Dense(32, activation='relu'),
-        layers.Dense(16, activation='relu'),
-        layers.Dense(1, activation='sigmoid')
-    ])
-
-    model.compile(
-        optimizer='adam',
-        loss='binary_crossentropy',
-        metrics=['accuracy']
-    )
-
-    print("Model Architecture:")
-    model.summary()
-    
-    # Train the model
-    print("\nStarting Training...")
-    history = model.fit(x, y, epochs=5, batch_size=32, validation_split=0.2)
-    
-    print("\nTraining completed successfully on Pathways!")
-    return history.history
+  print("Running Pathways job on JAX Backend!")
+
+  # Create a simple dataset
+  x = np.random.rand(1000, 10)
+  y = np.random.randint(0, 2, size=(1000, 1))
+
+  # A simple sequential model
+  model = keras.Sequential(
+    [
+      keras.Input(shape=(10,)),
+      layers.Dense(32, activation="relu"),
+      layers.Dense(16, activation="relu"),
+      layers.Dense(1, activation="sigmoid"),
+    ]
+  )
+
+  model.compile(
+    optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"]
+  )
+
+  print("Model Architecture:")
+  model.summary()
+
+  # Train the model
+  print("\nStarting Training...")
+  history = model.fit(x, y, epochs=5, batch_size=32, validation_split=0.2)
+
+  print("\nTraining completed successfully on Pathways!")
+  return history.history
+
 
 if __name__ == "__main__":
-    print("Submitting Pathways training job...")
-    result_history = train_simple_model()
-    print("Final validation accuracy:", result_history['val_accuracy'][-1])
+  print("Submitting Pathways training job...")
+  result_history = train_simple_model()
+  print("Final validation accuracy:", result_history["val_accuracy"][-1])

keras_remote/backend/execution.py

@@ -14,7 +14,7 @@
 import cloudpickle
 from absl import logging
 
-from keras_remote.backend import gke_client
+from keras_remote.backend import gke_client, pathways_client
 from keras_remote.constants import get_default_zone, zone_to_region
 from keras_remote.infra import container_builder
 from keras_remote.utils import packager, storage
@@ -106,59 +106,61 @@ def cleanup_job(self, job: Any, ctx: JobContext) -> None:
 
 
 class BaseK8sBackend:
-    """Base class for Kubernetes-based backends."""
+  """Base class for Kubernetes-based backends."""
+
+  def __init__(self, cluster: Optional[str] = None, namespace: str = "default"):
+    self.cluster = cluster
+    self.namespace = namespace
 
-    def __init__(self, cluster: Optional[str] = None, namespace: str = "default"):
-        self.cluster = cluster
-        self.namespace = namespace
 
 class GKEBackend(BaseK8sBackend):
-    """Backend adapter for standard GKE Jobs."""
-    def submit_job(self, ctx: JobContext) -> Any:
-        """Submit job to GKE cluster."""
-        return gke_client.submit_k8s_job(
-        display_name=ctx.display_name,
-        container_uri=ctx.image_uri,
-        accelerator=ctx.accelerator,
-        project=ctx.project,
-        job_id=ctx.job_id,
-        bucket_name=ctx.bucket_name,
-        namespace=self.namespace,
-        )
+  """Backend adapter for standard GKE Jobs."""
 
-    def wait_for_job(self, job: Any, ctx: JobContext) -> None:
-        """Wait for GKE job completion."""
-        gke_client.wait_for_job(job, namespace=self.namespace)
+  def submit_job(self, ctx: JobContext) -> Any:
+    """Submit job to GKE cluster."""
+    return gke_client.submit_k8s_job(
+      display_name=ctx.display_name,
+      container_uri=ctx.image_uri,
+      accelerator=ctx.accelerator,
+      project=ctx.project,
+      job_id=ctx.job_id,
+      bucket_name=ctx.bucket_name,
+      namespace=self.namespace,
+    )
+
+  def wait_for_job(self, job: Any, ctx: JobContext) -> None:
+    """Wait for GKE job completion."""
+    gke_client.wait_for_job(job, namespace=self.namespace)
 
-    def cleanup_job(self, job: Any, ctx: JobContext) -> None:
-        """Clean up K8s job resources."""
-        job_name = job.metadata.name
-        gke_client.cleanup_job(job_name, namespace=self.namespace)
+  def cleanup_job(self, job: Any, ctx: JobContext) -> None:
+    """Clean up K8s job resources."""
+    job_name = job.metadata.name
+    gke_client.cleanup_job(job_name, namespace=self.namespace)
 
 
 class PathwaysBackend(BaseK8sBackend):
-    """Backend adapter for ML Pathways using LeaderWorkerSet."""
-    def submit_job(self, ctx: JobContext) -> Any:
-        """Submit LWS job to GKE cluster."""
-        return pathways_client.submit_pathways_job(
-            display_name=ctx.display_name,
-            container_uri=ctx.image_uri,
-            accelerator=ctx.accelerator,
-            project=ctx.project,
-            job_id=ctx.job_id,
-            bucket_name=ctx.bucket_name,
-            namespace=self.namespace,
-        )
+  """Backend adapter for ML Pathways using LeaderWorkerSet."""
 
-    def wait_for_job(self, job: Any, ctx: JobContext) -> None:
-        """Wait for Pathways LWS completion."""
-        pathways_client.wait_for_job(job, ctx.job_id, namespace=self.namespace)
+  def submit_job(self, ctx: JobContext) -> Any:
+    """Submit LWS job to GKE cluster."""
+    return pathways_client.submit_pathways_job(
+      display_name=ctx.display_name,
+      container_uri=ctx.image_uri,
+      accelerator=ctx.accelerator,
+      project=ctx.project,
+      job_id=ctx.job_id,
+      bucket_name=ctx.bucket_name,
+      namespace=self.namespace,
+    )
 
-    def cleanup_job(self, job: Any, ctx: JobContext) -> None:
-        """Clean up LWS resources."""
-        job_name = f"keras-pathways-{ctx.job_id}"
-        pathways_client.cleanup_job(job_name, namespace=self.namespace)
+  def wait_for_job(self, job: Any, ctx: JobContext) -> None:
+    """Wait for Pathways LWS completion."""
+    pathways_client.wait_for_job(job, ctx.job_id, namespace=self.namespace)
 
+  def cleanup_job(self, job: Any, ctx: JobContext) -> None:
+    """Clean up LWS resources."""
+    job_name = f"keras-pathways-{ctx.job_id}"
+    pathways_client.cleanup_job(job_name, namespace=self.namespace)
 
 
 def _find_requirements(start_dir: str) -> Optional[str]: