tpu_info_quickstart_print_chip_utilization_and_memory.py

1import tpu_info
2
3def main():
4    # Get all TPU chips available on the host
5    chips = tpu_info.get_chips()
6    
7    if not chips:
8        print("No TPU chips found.")
9        return
10
11    print(f"Found {len(chips)} TPU chip(s):")
12    for chip in chips:
13        # Print chip details such as chip ID and utilization
14        print(f"Chip ID: {chip.chip_id}")
15        print(f"  Utilization: {chip.utilization}%")
16        print(f"  Memory Usage: {chip.memory_usage_bytes / (1024**2):.2f} MB")
17
18if __name__ == "__main__":
19    main()