refactor slice_to_num_chips to adapt to Cloud config (#65)

Co-authored-by: Zhihao Shan <zhihaoshan@google.com>

Author: zhihaoshan-google

jetstream/core/config_lib.py

@@ -16,7 +16,6 @@
 
 import dataclasses
 import functools
-import math
 from typing import Any, Callable, List, Tuple, Type
 
 from jetstream.engine import engine_api
@@ -85,11 +84,12 @@ class InterleavedCPUTestServer(ServerConfig):
 
 
 def slice_to_num_chips(s: str) -> int:
-  """Converts a TPU spec like v5e=4x2 to the number of chips, 8."""
-  # Account for the case where it is written 'v5e:4x2'.
-  delim = "=" if "=" in s else ":"
-  i = math.prod([int(c) for c in s.split(delim)[1].split("x")])
-  return i
+  """Converts a TPU spec like v5e-8 or v5e=8 to the number of chips, 8."""
+  # Account for the case where it is written 'tpu=8' for compatibility.
+  delim = "-" if "-" in s else "="
+  # TODO: Support more accelerator type check.
+  accelerator_type, num_devices = s.split(delim)
+  return int(num_devices) if accelerator_type != "v4" else int(num_devices) // 2
 
 
 def _split_devices_by_slices(

jetstream/tests/core/test_config_lib.py

@@ -0,0 +1,35 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit test for config_lib.py."""
+
+from absl.testing import absltest, parameterized
+from jetstream.core import config_lib
+
+
+class TestConfigLib(parameterized.TestCase):
+
+  @parameterized.parameters(
+      ("tpu=8", 8),
+      ("v5e-8", 8),
+      ("v5e=4", 4),
+      ("v4-8", 4),
+  )
+  def test_slice_to_num_chips(self, accelerator_slice, expected_num_devices):
+    got = config_lib.slice_to_num_chips(accelerator_slice)
+    self.assertEqual(got, expected_num_devices)
+
+
+if __name__ == "__main__":
+  absltest.main()