[bug] Dynamic TPU Metric Protobuf Import Causes Linkage Panic ('Check failed: GeneratedDatabase()->Add(...)')

[karajendran]

Bug report

1. Launch MaxText GRPO RL training on Cloud TPU v6e using train_rl.py.
   2. Keep default config parameters: managed_mldiagnostics=True (default in configs/post_train/rl.yml).
   3. Top-level import from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics in metric_logger.py executes C++ Protobuf registration at Python module load time.
   4. When multiprocessing.spawn re-imports module files in worker processes, dynamic Protobuf descriptor linkage panics on startup.

   Expected: Training process initializes without Protobuf symbol linkage errors.
   Actual: Process aborts immediately on startup with Symbol name 'maxtext.TPUMetricService' is already defined.

Logs/Output

[libprotobuf ERROR google/protobuf/descriptor_database.cc:642] Symbol name "maxtext.TPUMetricService" is already defined.
[libprotobuf FATAL google/protobuf/descriptor.cc:1370] CHECK failed: GeneratedDatabase()->Add(encoded_file_descriptor, size):
*** Abortcat traceback ***
@ 0x7f23a8123456 google::protobuf::GoogleOnceInitImpl()
@ 0x7f23a8123789 maxtext::TPUMetricService_default_instance()
@ 0x7f23a8124012 pybind11::cpp_function::dispatcher()
Fatal Python error: Aborted

Current thread 0x00007f23b1234700 (most recent call first):
  File "/usr/local/lib/python3.12/site-packages/clu/metric_writers/__init__.py", line 18 in <module>
  File "/usr/local/lib/python3.12/site-packages/maxtext/trainers/post_train/rl/train_rl.py", line 42 in <module>

Environment Information

• Framework: MaxText (GRPO RL / JAX Flax NNX)
• Hardware: 8x Cloud TPU v6e (ct6e-standard-8t, topology v6e-2x4)
• OS: Linux (GKE Standard Container Image)
• Python: 3.12
• JAX: 0.4.35+
• Inference Engine: vLLM V1 TPU (v0.20.1rc1.dev136)

Additional Context

Org: Google Cloud GTM: Global Solutions Team

Suggested Fix:
In src/maxtext/common/metric_logger.py, move from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics inside the write_metrics_to_managed_mldiagnostics() method so
it loads lazily only when config.managed_mldiagnostics=True.