Add an access bandwidth test script for 'ucm_connector'

Author: zhou-haitao

test/test_ucm_connector_save_load.py

@@ -56,11 +56,9 @@
     UCMConnector,
     UCMConnectorMetadata,
 )
+from ucm.logger import init_logger
 
-
-def log(msg: str):
-    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
-    print(f"[{timestamp}] {msg}", flush=True)
+logger = init_logger(__name__)
 
 
 def make_aligned_tensor(shape, dtype, device, alignment=4096):
@@ -92,7 +90,7 @@ def make_buffers(
     kv: int,
     is_mla: bool,
 ) -> Tuple[List[str], Dict[str, torch.Tensor]]:
-    log(f"Allocating buffers: blocks={block_number}, batch_size={batch_size}")
+    logger.info(f"Allocating buffers: blocks={block_number}, batch_size={batch_size}")
     hashes = [secrets.token_hex(16) for _ in range(block_number)]
     device = f"cuda:{device_id}"
     kv_caches: Dict[str, torch.Tensor] = {}
@@ -373,7 +371,7 @@ def broadcast(self, tensor, src):
     r_sizes, r_times, r_bws = [], [], []
 
     for round_idx in range(repeat):
-        log(f"Round {round_idx + 1}: start write test")
+        logger.info(f"Round {round_idx + 1}: start write test")
         start_hash_idx = round_idx * batch_size
         end_hash_idx = start_hash_idx + batch_size
         round_hashes = hashes[start_hash_idx:end_hash_idx]
@@ -401,7 +399,7 @@ def avg(values: List[float]) -> float:
     avg_r_time = avg(r_times)
     avg_r_bw = avg(r_bws)
 
-    log(
+    logger.info(
         "\n=== Summary ===\n"
         f"Write : size={avg_w_size:.4f} GB | time={avg_w_time:.4f} s | bw={avg_w_bw:.4f} GB/s\n"
         f"Read  : size={avg_r_size:.4f} GB | time={avg_r_time:.4f} s | bw={avg_r_bw:.4f} GB/s\n"