Fix protobuf descriptor collision by delaying cli import in tpu_info

[Shuwen-Fang]

Fix protobuf descriptor collision by delaying cli import in tpu_info

Root Cause

The tpu-info package imported the cli module in its top-level tpu_info/__init__.py. This triggered a chain of imports loading generated proto definitions (like tpu_metric_service.proto), which import google/protobuf/timestamp.proto.
When other libraries (like JAX or libtpu) also load these protos, a descriptor collision occurs:

E0000 00:00:1782937239.214124 1476265 descriptor_database.cc:633] File already exists in database: google/protobuf/timestamp.proto
F0000 00:00:1782937239.214232 1476265 descriptor.cc:2236] Check failed: GeneratedDatabase()->Add(encoded_file_descriptor, size)

This causes a fatal check failure, crashing the initialization subprocess with exit code -6 (SIGABRT).

Fix

Delay importing cli in tpu_info/__init__.py by commenting it out. It is only needed when tpu-info is run as a CLI tool (which uses tpu_info.cli:print_chip_info entry point and imports it explicitly). Anyone importing tpu_info as a library (e.g., tpu_inference to get device info) will no longer load the CLI and its proto dependencies, avoiding the collision.

Verification

Verified on a TPU VM:

• Reproduced the crash with official tpu-info==0.7.1 (where cli is imported at top-level).
• Verified the fix by installing from this branch; the training script successfully bypassed the protobuf initialization phase without any descriptor collision errors.
• Ran Google3 unit tests for tpu_info and confirmed they all pass.

Bug: b/528372351