Update Pathways HBM calculation (#1097)

Signed-off-by: wenxindongwork <wenxindong@google.com>

Author: wenxindongwork

tests/test_utils.py

@@ -75,25 +75,34 @@ def test_hbm_usage_bytes_pathways_enabled(mock_devices, mock_live_arrays):
     mock_device2 = MagicMock()
     devices = [mock_device1, mock_device2]
 
-    # Create mock device buffers
-    mock_buffer1_dev1 = MagicMock()
-    mock_buffer1_dev1.device = mock_device1
-    mock_buffer1_dev1.nbytes = 2000  # 2000 bytes on device1
+    # Create mock addressable shards with data property
+    mock_data1_dev1 = MagicMock()
+    mock_data1_dev1.device = mock_device1
+    mock_data1_dev1.nbytes = 2000  # 2000 bytes on device1
 
-    mock_buffer1_dev2 = MagicMock()
-    mock_buffer1_dev2.device = mock_device2
-    mock_buffer1_dev2.nbytes = 2000  # 2000 bytes on device2
+    mock_data1_dev2 = MagicMock()
+    mock_data1_dev2.device = mock_device2
+    mock_data1_dev2.nbytes = 2000  # 2000 bytes on device2
 
-    mock_buffer2_dev1 = MagicMock()
-    mock_buffer2_dev1.device = mock_device1
-    mock_buffer2_dev1.nbytes = 1000  # 1000 bytes on device1
+    mock_data2_dev1 = MagicMock()
+    mock_data2_dev1.device = mock_device1
+    mock_data2_dev1.nbytes = 1000  # 1000 bytes on device1
 
-    # Create mock arrays with device buffers
+    mock_shard1_dev1 = MagicMock()
+    mock_shard1_dev1.data = mock_data1_dev1
+
+    mock_shard1_dev2 = MagicMock()
+    mock_shard1_dev2.data = mock_data1_dev2
+
+    mock_shard2_dev1 = MagicMock()
+    mock_shard2_dev1.data = mock_data2_dev1
+
+    # Create mock arrays with addressable_shards
     mock_array1 = MagicMock()
-    mock_array1.device_buffers = [mock_buffer1_dev1, mock_buffer1_dev2]
+    mock_array1.addressable_shards = [mock_shard1_dev1, mock_shard1_dev2]
 
     mock_array2 = MagicMock()
-    mock_array2.device_buffers = [mock_buffer2_dev1]
+    mock_array2.addressable_shards = [mock_shard2_dev1]
 
     mock_live_arrays.return_value = [mock_array1, mock_array2]
 

tpu_inference/utils.py

@@ -132,8 +132,8 @@ def pathways_hbm_usage_gb(devices: Any) -> List[Tuple[float, float]]:
     hbm_used = defaultdict(int)
     hbm_limit = get_device_hbm_limit()
     for array in live_arrays:
-        for buffer in array.device_buffers:
-            hbm_used[buffer.device] += buffer.nbytes
+        for buffer in array.addressable_shards:
+            hbm_used[buffer.data.device] += buffer.data.nbytes
     return [(hbm_used[device], hbm_limit) for device in devices]