Fix(tpu-info): Prevents crash on Python 3.12 incompatibility

PiperOrigin-RevId: 789376103

Author: Googler

tpu_info/pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
  "grpcio>=1.65.5",
  "protobuf",
  "rich",
+ "packaging",
 ]
 authors = [
  { name="Cloud TPU Team", email="cloud-tpu-eng@google.com" },

tpu_info/tpu_info/cli.py

@@ -73,14 +73,22 @@ def _get_runtime_info(rate: float) -> align.Align:
 def print_chip_info():
  """Print local TPU devices and libtpu runtime metrics."""
  cli_args = args.parse_arguments()
+ console_obj = console.Console()
+ is_incompatible = cli_helper.is_incompatible_python_version()
+ if is_incompatible:
+ console_obj.print(cli_helper.get_py_compat_warning_panel())
+
  if cli_args.version:
  print(f"- tpu-info version: {cli_helper.fetch_cli_version()}")
- print(f"- libtpu version: {cli_helper.fetch_libtpu_version()}")
- print(f"- accelerator type: {cli_helper.fetch_accelerator_type()}")
+ if is_incompatible:
+ print("- libtpu version: N/A (incompatible environment)")
+ print("- accelerator type: N/A (incompatible environment)")
+ else:
+ print(f"- libtpu version: {cli_helper.fetch_libtpu_version()}")
+ print(f"- accelerator type: {cli_helper.fetch_accelerator_type()}")
  return
 
  if cli_args.list_metrics:
- console_obj = console.Console()
  console_obj.print(
  panel.Panel(
  "\n".join(
@@ -93,14 +101,15 @@ def print_chip_info():
  )
  return
 
+ if is_incompatible:
+ return
+
  # TODO(wcromar): Merge all of this info into one table
  chip_type, count = device.get_local_chips()
  if not chip_type:
  print("No TPU chips found.")
  return
 
- console_obj = console.Console()
-
  if cli_args.process:
  table = cli_helper.fetch_process_table(chip_type, count)
  console_obj.print(table)

tpu_info/tpu_info/cli_helper.py

@@ -16,7 +16,9 @@
 
 import importlib.metadata
 import subprocess
+import sys
 from typing import Any, Dict, List, Optional, Tuple
+
 from tpu_info import args_helper
 from tpu_info import device
 from tpu_info import metrics
@@ -33,6 +35,38 @@ def _bytes_to_gib(size: int) -> float:
  return size / (1 << 30)
 
 
+def is_incompatible_python_version() -> bool:
+ """Checks if the current Python version is 3.12 or newer."""
+ if sys.version_info < (3, 12):
+ return False
+ try:
+ from packaging.version import parse as parse_version
+ import libtpu # pytype: disable=import-error
+
+ first_compatible_version = parse_version("3.20")
+ current_version = parse_version(libtpu.__version__)
+ if current_version < first_compatible_version:
+ return True
+ else:
+ return False
+ except (ImportError, AttributeError):
+ return True
+
+
+def get_py_compat_warning_panel() -> panel.Panel:
+ """Returns a Rich Panel with a Python compatibility warning."""
+ warning_text = (
+ "Some features are disabled due to an incompatibility between the libtpu"
+ " SDK and Python 3.12+.\n\nFor full functionality, please use a"
+ " different Python environment."
+ )
+ return panel.Panel(
+ f"[yellow]{warning_text}[/yellow]",
+ title="[bold yellow]Compatibility Warning[/bold yellow]",
+ border_style="yellow",
+ )
+
+
 def fetch_cli_version() -> str:
  """Returns the version of the current TPU CLI."""
  try: