Sets min_nodes to 0 multi-host configs (#88)

* Sets min_nodes to 0 multi-host configs

* fix tests

Author: JyotinderSingh

keras_remote/cli/infra/program.py

@@ -246,7 +246,8 @@ def _create_tpu_node_pool(cluster, tpu: TpuConfig, zone, project_id, pool_name):
   # Single-host TPU slices (1 node) must not specify placement_policy;
   # multi-host slices require COMPACT placement with an explicit topology.
   is_multi_host = tpu.num_nodes > 1
-  min_nodes = tpu.num_nodes if is_multi_host else 0
+  # Autoscaling is enabled, so we need to set the min_node_count to 0.
+  min_nodes = 0
 
   placement = (
     gcp.container.NodePoolPlacementPolicyArgs(

keras_remote/cli/infra/program_test.py

@@ -80,12 +80,12 @@ def test_node_count_matches_config(self, gcp_mock):
 
     # Due to multi-host TPU workaround, initial_node_count is equal to num_nodes
     call_kwargs = gcp_mock.container.NodePool.call_args.kwargs
-    self.assertEqual(call_kwargs.get("initial_node_count"), 4)
+    self.assertEqual(call_kwargs.get("initial_node_count"), 0)
     autoscaling_kwargs = (
       gcp_mock.container.NodePoolAutoscalingArgs.call_args.kwargs
     )
     self.assertEqual(autoscaling_kwargs.get("max_node_count"), 4)
-    self.assertEqual(autoscaling_kwargs.get("min_node_count"), 4)
+    self.assertEqual(autoscaling_kwargs.get("min_node_count"), 0)
 
   @mock.patch.object(program, "gcp")
   def test_pool_name_passed_through(self, gcp_mock):
@@ -352,21 +352,13 @@ def test_node_pool_scale_to_zero(
         cluster, accelerator, "us-central2-b", "my-project", "test-pool"
       )
 
-    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"), 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"), 0)
     self.assertEqual(
       autoscaling_kwargs.get("max_node_count"), expected_max_count
     )