Fix TPU7x chip counting to account for chiplet architecture

For TPU7x devices, each physical chip contains 2 chiplets that are
exposed to the host as separate devices. The previous
implementation counted these devices directly, resulting in 2x the
actual chip count being reported.

Example of the issue:
- tpu7x-8 has 4 physical chips with 8 chiplets total
- Previous behavior: reported num_chips=8 (incorrect)
- Fixed behavior: reports num_chips=4 (correct)

Changes:
- Add get_num_chiplets_per_chip() helper function that returns 2 for
  tpu7x devices and 1 for all other TPU types, following the same
  pattern as get_num_cores_per_chip()
- Modify get_num_chips() to divide device count by chiplets_per_chip
  using integer division
- Update logging to conditionally show chiplets_per_chip for tpu7x
  devices only (non-tpu7x devices don't have chiplets, so we avoid
  showing confusing information)
- Add comprehensive tests for tpu7x chip counting scenarios
- Update existing get_num_chips() tests to mock get_tpu_type() since
  the function now calls it

Test coverage:
- test_get_num_chiplets_per_chip: Tests tpu7x variants, other TPU
  types, and edge cases (None, empty string)
- test_get_num_chips_tpu7x_from_accel: Tests tpu7x-8 with /dev/accel*
- test_get_num_chips_tpu7x_4_from_accel: Tests tpu7x-4 with /dev/accel*
- test_get_num_chips_tpu7x_from_vfio: Tests tpu7x-8 with /dev/vfio
- test_get_num_chips_non_tpu7x_unchanged: Verifies backward compatibility

Backward compatibility:
All non-tpu7x devices return chiplets_per_chip=1, making the division
a no-op. Behavior is identical to before the fix.

Signed-off-by: burbajr <[email protected]>

Author: burbajr

tests/test_tpu_info.py

@@ -5,8 +5,9 @@
 import requests
 
 from tpu_inference.tpu_info import (get_node_name, get_node_worker_id,
-                                    get_num_chips, get_num_cores_per_chip,
-                                    get_tpu_metadata, get_tpu_type)
+                                    get_num_chiplets_per_chip, get_num_chips,
+                                    get_num_cores_per_chip, get_tpu_metadata,
+                                    get_tpu_type)
 
 
 # Mock requests.get for get_tpu_metadata tests
@@ -98,23 +99,105 @@ def test_get_num_cores_per_chip(mock_get_tpu_type, tpu_type, expected):
     assert get_num_cores_per_chip() == expected
 
 
+# Test get_num_chiplets_per_chip
+@pytest.mark.parametrize(
+    "tpu_type, expected",
+    [
+        ("tpu7x-8", 2),
+        ("tpu7x-4", 2),
+        ("TPU7x", 2),  # Case insensitive
+        ("v5litepod-4", 1),
+        ("v6e-8", 1),
+        ("v4-8", 1),
+        (None, 1),  # Default when tpu_type is None
+        ("", 1),  # Default when tpu_type is empty
+    ])
+@patch("tpu_inference.tpu_info.get_tpu_type")
+def test_get_num_chiplets_per_chip(mock_get_tpu_type, tpu_type, expected):
+    """Test get_num_chiplets_per_chip with different TPU types."""
+    mock_get_tpu_type.return_value = tpu_type
+    assert get_num_chiplets_per_chip() == expected
+
+
 # Test get_num_chips
+@patch("tpu_inference.tpu_info.get_tpu_type", return_value="v6e-8")
 @patch("tpu_inference.tpu_info.glob.glob",
        return_value=["/dev/accel0", "/dev/accel1"])
-def test_get_num_chips_from_accel(mock_glob):
+def test_get_num_chips_from_accel(mock_glob, mock_get_tpu_type):
     """Test get_num_chips when /dev/accel* files exist."""
     assert get_num_chips() == 2
 
 
+@patch("tpu_inference.tpu_info.get_tpu_type", return_value="v6e-8")
 @patch("tpu_inference.tpu_info.glob.glob", return_value=[])
 @patch("tpu_inference.tpu_info.os.listdir", return_value=["0", "1", "2"])
-def test_get_num_chips_from_vfio(mock_listdir, mock_glob):
+def test_get_num_chips_from_vfio(mock_listdir, mock_glob, mock_get_tpu_type):
     """Test get_num_chips when /dev/accel* files don't exist but /dev/vfio entries do."""
     assert get_num_chips() == 3
 
 
+@patch("tpu_inference.tpu_info.get_tpu_type", return_value="v6e-8")
 @patch("tpu_inference.tpu_info.glob.glob", return_value=[])
 @patch("tpu_inference.tpu_info.os.listdir", side_effect=FileNotFoundError)
-def test_get_num_chips_not_found(mock_listdir, mock_glob, caplog):
+def test_get_num_chips_not_found(mock_listdir, mock_glob, mock_get_tpu_type,
+                                 caplog):
     """Test get_num_chips when neither files nor directory are found."""
     assert get_num_chips() == 0
+
+
+# Test get_num_chips with tpu7x
+@patch("tpu_inference.tpu_info.get_tpu_type", return_value="tpu7x-8")
+@patch("tpu_inference.tpu_info.glob.glob",
+       return_value=[
+           "/dev/accel0", "/dev/accel1", "/dev/accel2", "/dev/accel3",
+           "/dev/accel4", "/dev/accel5", "/dev/accel6", "/dev/accel7"
+       ])
+def test_get_num_chips_tpu7x_from_accel(mock_glob, mock_get_tpu_type):
+    """Test get_num_chips for tpu7x-8 when /dev/accel* files exist.
+
+    tpu7x-8 has 4 physical chips with 2 chiplets each (8 total devices).
+    Should return 4 chips, not 8.
+    """
+    assert get_num_chips() == 4
+
+
+@patch("tpu_inference.tpu_info.get_tpu_type", return_value="tpu7x-4")
+@patch(
+    "tpu_inference.tpu_info.glob.glob",
+    return_value=["/dev/accel0", "/dev/accel1", "/dev/accel2", "/dev/accel3"])
+def test_get_num_chips_tpu7x_4_from_accel(mock_glob, mock_get_tpu_type):
+    """Test get_num_chips for tpu7x-4 when /dev/accel* files exist.
+
+    tpu7x-4 has 2 physical chips with 2 chiplets each (4 total devices).
+    Should return 2 chips, not 4.
+    """
+    assert get_num_chips() == 2
+
+
+@patch("tpu_inference.tpu_info.get_tpu_type", return_value="tpu7x-8")
+@patch("tpu_inference.tpu_info.glob.glob", return_value=[])
+@patch("tpu_inference.tpu_info.os.listdir",
+       return_value=["0", "1", "2", "3", "4", "5", "6", "7"])
+def test_get_num_chips_tpu7x_from_vfio(mock_listdir, mock_glob,
+                                       mock_get_tpu_type):
+    """Test get_num_chips for tpu7x when using /dev/vfio entries.
+
+    tpu7x-8 has 4 physical chips with 2 chiplets each (8 total devices).
+    Should return 4 chips when counting vfio entries.
+    """
+    assert get_num_chips() == 4
+
+
+@patch("tpu_inference.tpu_info.get_tpu_type", return_value="v6e-8")
+@patch("tpu_inference.tpu_info.glob.glob",
+       return_value=[
+           "/dev/accel0", "/dev/accel1", "/dev/accel2", "/dev/accel3",
+           "/dev/accel4", "/dev/accel5", "/dev/accel6", "/dev/accel7"
+       ])
+def test_get_num_chips_non_tpu7x_unchanged(mock_glob, mock_get_tpu_type):
+    """Test get_num_chips for non-tpu7x devices remains unchanged.
+
+    v6e-8 has 8 physical chips (1 chiplet each). Should return 8.
+    This verifies backward compatibility.
+    """
+    assert get_num_chips() == 8

tpu_inference/__init__.py

@@ -42,10 +42,15 @@
 else:
     # Either running on TPU or CPU
     try:
+        tpu_type = ti.get_tpu_type()
+        # For tpu7x, show chiplets_per_chip since each chip has 2 chiplets.
+        # Other TPU types don't have chiplets, so we don't show this info.
+        chiplets = f" | num_chiplets_per_chip={ti.get_num_chiplets_per_chip()}" if (
+            tpu_type and "tpu7x" in tpu_type.lower()) else ""
         logger.info(f"TPU info: node_name={ti.get_node_name()} | "
-                    f"tpu_type={ti.get_tpu_type()} | "
+                    f"tpu_type={tpu_type} | "
                     f"worker_id={ti.get_node_worker_id()} | "
-                    f"num_chips={ti.get_num_chips()} | "
+                    f"num_chips={ti.get_num_chips()}{chiplets} | "
                     f"num_cores_per_chip={ti.get_num_cores_per_chip()}")
     except Exception as e:
         logger.error(

tpu_inference/tpu_info.py

@@ -63,16 +63,41 @@ def get_num_cores_per_chip() -> int:
     return 2
 
 
+def get_num_chiplets_per_chip() -> int:
+    """Returns the number of chiplets per physical chip for the current TPU type.
+
+    For tpu7x devices, each physical chip contains 2 chiplets that are exposed
+    to the host as separate devices. For all other TPU types, each chip contains
+    1 chiplet.
+
+    Returns:
+        int: Number of chiplets per physical chip (1 or 2)
+    """
+    tpu_type = get_tpu_type()
+    if tpu_type and "tpu7x" in tpu_type.lower():
+        return 2
+    return 1
+
+
 def get_num_chips() -> int:
+    """Returns the number of physical TPU chips available on the current host.
+
+    For tpu7x devices, this correctly accounts for the fact that each physical
+    chip contains 2 chiplets that are exposed as separate /dev/accel* devices.
+
+    Returns:
+        int: Number of physical TPU chips
+    """
+    chiplets_per_chip = get_num_chiplets_per_chip()
     accel_files = glob.glob("/dev/accel*")
     if accel_files:
-        return len(accel_files)
+        return len(accel_files) // chiplets_per_chip
     try:
         vfio_entries = os.listdir("/dev/vfio")
         numeric_entries = [
             int(entry) for entry in vfio_entries if entry.isdigit()
         ]
-        return len(numeric_entries)
+        return len(numeric_entries) // chiplets_per_chip
     except FileNotFoundError as e:
         logger.error("Failed to detect number of TPUs: %s", e)
         return 0