TPU Monitoring Library

pip install libtpu 

pip install -U "jax[tpu]" 

pip install torch~=2.6.0 'torch_xla[tpu]~=2.6.0' \  -f https://storage.googleapis.com/libtpu-releases/index.html \  -f https://storage.googleapis.com/libtpu-wheels/index.html  # Optional: if you're using custom kernels, install pallas dependencies pip install 'torch_xla[pallas]' \  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html 

from libtpu.sdk import tpumonitoring 

 from libtpu.sdk import tpumonitoring tpumonitoring.help() " libtpu.sdk.monitoring.help():  List all supported functionality.  libtpu.sdk.monitoring.list_support_metrics()  List support metric names in the list of str format.  libtpu.sdk.monitoring.get_metric(metric_name:str)  Get metric data with metric name. It represents the snapshot mode.  The metric data is a object with `description()` and `data()` methods,  where the `description()` returns a string describe the format of data  and data unit, `data()` returns the metric data in the list in str format. " 

from libtpu.sdk import tpumonitoring tpumonitoring.list_supported_metrics() ["duty_cycle_pct", "tensorcore_util", "hbm_util", ...] 

from libtpu.sdk import tpumonitoring metric = tpumonitoring.get_metric("duty_cycle_pct") metric.description() "The metric provides a list of duty cycle percentages, one for each accelerator (from accelerator_0 to accelerator_x). The duty cycle represents the percentage of time an accelerator was actively processing during the last sample period, indicating TPU utilization." metric.data() ["0.00", "0.00", "0.00", "0.00"] # accelerator_0-3 

import jax import jax.numpy as jnp from libtpu.sdk import tpumonitoring import time  # --- Your JAX model and training setup would go here ---  # --- Example placeholder model and data (replace with your actual setup)--- def simple_model(x):  return jnp.sum(x) def loss_fn(params, x, y):  preds = simple_model(x)  return jnp.mean((preds - y)**2) def train_step(params, x, y, optimizer):  grads = jax.grad(loss_fn)(params, x, y)  return optimizer.update(grads, params) key = jax.random.PRNGKey(0) params = jnp.array([1.0, 2.0]) # Example params optimizer = ... # Your optimizer (for example, optax.adam) data_x = jnp.ones((10, 10)) data_y = jnp.zeros((10,)) num_epochs = 10 log_interval_steps = 2 # Log duty cycle every 2 steps for epoch in range(num_epochs):  for step in range(5): # Example steps per epoch  params = train_step(params, data_x, data_y, optimizer)  if (step + 1) % log_interval_steps == 0:  # --- Integrate TPU Monitoring Library here to get duty_cycle ---  duty_cycle_metric = tpumonitoring.get_metric("duty_cycle_pct")  duty_cycle_data = duty_cycle_metric.data  print(f"Epoch {epoch+1}, Step {step+1}: TPU Duty Cycle Data:")  print(f" Description: {duty_cycle_metric.description}")  print(f" Data: {duty_cycle_data}")  # --- End TPU Monitoring Library Integration ---  # --- Rest of your training loop logic ---  time.sleep(0.1) # Simulate some computation print("Training complete.") 

# The following code sample shows how to check HBM utilization before JAX inference: import jax import jax.numpy as jnp from libtpu.sdk import tpumonitoring  # --- Your JAX model and inference setup would go here ---  # --- Example placeholder model (replace with your actual model loading/setup)--- def simple_model(x):  return jnp.sum(x) key = jax.random.PRNGKey(0) params = ... # Load your trained parameters  # Integrate the TPU Monitoring Library to get HBM utilization before inference hbm_util_metric = tpumonitoring.get_metric("hbm_util") hbm_util_data = hbm_util_metric.data print("HBM Utilization Before Inference:") print(f" Description: {hbm_util_metric.description}") print(f" Data: {hbm_util_data}")  # End TPU Monitoring Library Integration  # Your Inference Logic input_data = jnp.ones((1, 10)) # Example input predictions = simple_model(input_data) print("Inference Predictions:", predictions) print("Inference complete.") 

import jax import jax.numpy as jnp from libtpu.sdk import tpumonitoring # --- Your JAX model and training setup goes here --- # --- Example placeholder model and data (replace with your actual setup)--- def simple_model(x): return jnp.sum(x) # --- details here --- # ============================================================================== # Metric 1: TCP Delivery Rate # ============================================================================== # This metric reflects the delivery rate of the TCP connection (bytes delivered / elapsed time). # The output is a list of strings representing latency statistics: # [mean, p50, p90, p95, p99.9] # Example: ['100.00', '200.00', '300.00', '400.00', '500.00'] (Values in Mbps) # Get the metric object delivery_rate_metric = tpumonitoring.get_metric(metric_name_rate) # Print the description provided by the library print("Description:", delivery_rate_metric.description()) # Print the actual data payload print("Data:", delivery_rate_metric.data()) # ============================================================================== # Metric 2: TCP Minimum Round Trip Time (RTT) # ============================================================================== # This metric reflects the minimum RTT measured between sending a TCP packet # and receiving the acknowledgement. # The output is a list of strings representing latency statistics: # [mean, p50, p90, p95, p99.9] # Example: ['100.00', '200.00', '300.00', '400.00', '500.00'] (Values in us - microseconds) # Get the metric object min_rtt_metric = tpumonitoring.get_metric(metric_name_rtt) # Print the description provided by the library print("Description:", min_rtt_metric.description()) # Print the actual data payload print("Data:", min_rtt_metric.data()) 

summary = sdk.tpuz.get_core_state_summary() 

{ "host_name": "my-tpu-host-vm", "core_states": { "1": { "core_id": { "global_core_id": 1, "chip_id": 0, "core_on_chip": { "type": "TPU_CORE_TYPE_TENSOR_CORE", "index": 1 } }, "sequencer_info": [ { "sequencer_type": "TPU_SEQUENCER_TYPE_TENSOR_CORE_SEQUENCER", "sequencer_index": 0, "pc": 4490, "program_id": 3274167277388825310, "run_id": 3 } ], "program_fingerprint": "b'\\xbefB\\xc6\\x1eb\\xc1#\\xd0...'", "queued_program_info": [], "error_message": "" } // ... } } 

summary = sdk.tpuz.get_core_state_summary(include_hlo_info=True) 

"1": { "core_id": { "global_core_id": 1, "chip_id": 0, "core_on_chip": { "type": "TPU_CORE_TYPE_TENSOR_CORE", "index": 1 } }, "sequencer_info": [ { "sequencer_type": "TPU_SEQUENCER_TYPE_TENSOR_CORE_SEQUENCER", "sequencer_index": 0, "pc": 17776, "tag": 3, "tracemark": 2147483646, "program_id": 3230481660274331500, "run_id": 81, "hlo_location": "HLO: fusion.11; HLO computation: main.126_spmd", "hlo_detailed_info": "[{\"details\":\"HloModule fusion.11, entry_computation_layout={(bf16[20>..." } ], "program_fingerprint": "\ufffdU\ufffd4j\u7c6e\ufffd\ufffd{\u0017\ufffd\ufffdHHV\ufffdD\ufffde\uff>", "launch_id": 1394130914, "queued_program_info": [ { "run_id": 81, "launch_id": 1394130914, "program_fingerprint": "\ufffdU\ufffd4j\u7c6e\ufffd\ufffd{\u0017\ufffd\ufffdHHV\ufffdD\ufffde>" } ] } 

import jax import jax.numpy as jnp import time import threading from functools import partial from libtpu import sdk 

def monitor_tpu_status():  """Monitors TPU status in a background thread.""" while monitoring_active: try: summary = sdk.tpuz.get_core_state_summary(include_hlo_info=True) if summary and 'core_states' in summary: print(summary) else: print('WARNING: Call returned an empty or invalid summary.') except RuntimeError as e: print(f'FAIL: Error calling API: {e}') except Exception as e: print(f'FAIL: Unexpected error in monitor thread: {e}') for _ in range(MONITORING_INTERVAL_SECONDS * 2): if not monitoring_active: break time.sleep(0.5) print('✅ Monitoring thread stopped.') 

@partial(jax.jit, static_argnames=['num_heads']) def transformer_block(params, x, num_heads=32):  """A simplified but computationally intensive Transformer block.""" # Multi-head Self-Attention qkv = jnp.dot(x, params['qkv_kernel']) q, k, v = jnp.array_split(qkv, 3, axis=-1) # Reshape for multi-head attention q = q.reshape(q.shape[0], q.shape[1], num_heads, -1).transpose(0, 2, 1, 3) k = k.reshape(k.shape[0], k.shape[1], num_heads, -1).transpose(0, 2, 1, 3) v = v.reshape(v.shape[0], v.shape[1], num_heads, -1).transpose(0, 2, 1, 3) # Scaled dot-product attention attention_scores = jnp.einsum('nhqd,nhkd->nhqk', q, k) / jnp.sqrt(q.shape[-1]) attention_weights = jax.nn.softmax(attention_scores, axis=-1) attention_output = jnp.einsum('nhqk,nhvd->nhqd', attention_weights, v) attention_output = attention_output.transpose(0, 2, 1, 3).reshape(x.shape) attention_output = jnp.dot(attention_output, params['o_kernel']) # Residual connection and Layer Normalization 1 h1 = x + attention_output h1_norm = h1 - jnp.mean(h1, axis=-1, keepdims=True) h1_norm = h1_norm / jnp.sqrt( jnp.mean(jnp.square(h1_norm), axis=-1, keepdims=True) + 1e-5 ) # Feed-Forward Network ffn_hidden = jax.nn.gelu(jnp.dot(h1_norm, params['ffn1_kernel'])) ffn_output = jnp.dot(ffn_hidden, params['ffn2_kernel']) # Residual connection and Layer Normalization 2 h2 = h1_norm + ffn_output h2_norm = h2 - jnp.mean(h2, axis=-1, keepdims=True) h2_norm = h2_norm / jnp.sqrt( jnp.mean(jnp.square(h2_norm), axis=-1, keepdims=True) + 1e-5 ) return h2_norm 

def main(): num_devices = jax.device_count() print(f"Running on {num_devices} devices.") batch_size = 128 * num_devices seq_len = 512 embed_dim = 1024 ffn_dim = embed_dim * 4 key = jax.random.PRNGKey(0) params = { 'qkv_kernel': jax.random.normal( key, (embed_dim, embed_dim * 3), dtype=jnp.bfloat16 ), 'o_kernel': jax.random.normal( key, (embed_dim, embed_dim), dtype=jnp.bfloat16 ), 'ffn1_kernel': jax.random.normal( key, (embed_dim, ffn_dim), dtype=jnp.bfloat16 ), 'ffn2_kernel': jax.random.normal( key, (ffn_dim, embed_dim), dtype=jnp.bfloat16 ), } input_data = jax.random.normal( key, (batch_size, seq_len, embed_dim), dtype=jnp.bfloat16 ) input_data = jax.device_put(input_data) monitor_thread = threading.Thread(target=monitor_tpu_status) monitor_thread.start() print("Starting JAX computation loop...") start_time = time.time() iterations = 0 while time.time() - start_time < JOB_DURATION_SECONDS: result = transformer_block(params, input_data) result.block_until_ready() iterations += 1 print(f' -> Jax iteration {iterations} complete.', end='\r') print(f"\nCompleted {iterations} iterations in {JOB_DURATION_SECONDS} seconds.") global monitoring_active monitoring_active = False monitor_thread.join() if __name__ == '__main__': main() 

pip show libtpu 

import libtpu print(libtpu.__version__) 

pip install --upgrade libtpu