[V1] Move usage stats to worker and start logging TPU hardware

vllm/usage/usage_lib.py

@@ -137,6 +137,9 @@ def __init__(self) -> None:
         self.gpu_type: Optional[str] = None
         self.gpu_memory_per_device: Optional[int] = None
         self.env_var_json: Optional[str] = None
+        self.tpu_count: Optional[int] = None
+        self.tpu_type: Optional[str] = None
+        self.tpu_memory_per_device: Optional[int] = None
 
         # vLLM Information
         self.model_architecture: Optional[str] = None
@@ -174,6 +177,19 @@ def _report_usage_once(self, model_architecture: str,
             self.gpu_memory_per_device = device_property.total_memory
         if current_platform.is_cuda():
             self.cuda_runtime = torch.version.cuda
+        if current_platform.is_tpu():
+            try:
+                import torch_xla.runtime as xr
+                from torch_xla.core import xla_model as xm
+                self.tpu_count = xr.world_size()
+                self.tpu_type = xm.xla_device_hw(xm.xla_device())
+                self.tpu_memory_per_device = xm.get_memory_info().bytes_limit
+            except ImportError:
+                logging.warning(
+                    "torch_xla not found, skipping TPU usage statistics.")
+                self.tpu_count = None
+                self.tpu_type = None
+                self.tpu_memory_per_device = None
         self.provider = _detect_cloud_provider()
         self.architecture = platform.machine()
         self.platform = platform.platform()