fix

Author: t00939662

ucm/integration/vllm/ucm_connector.py

@@ -1008,45 +1008,45 @@ def is_empty(self) -> bool:
 
     def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
         """Aggregate stats from another worker, preserving per-worker separation"""
-        if not other.is_empty():
-            assert isinstance(
-                other, UCMKVConnectorStats
-            ), "Expected UCMKVConnectorStats"
-            for worker_rank, worker_data in other.data.items():
-                if worker_rank not in self.data:
-                    self.data[worker_rank] = copy.deepcopy(worker_data)
-                else:
-                    for metric_name, values in worker_data.items():
-                        if metric_name not in self.data[worker_rank]:
-                            self.data[worker_rank][metric_name] = []
-                        self.data[worker_rank][metric_name].extend(values)
+        if other.is_empty():
+            return self
+        
+        assert isinstance(other, UCMKVConnectorStats), "Expected UCMKVConnectorStats"
+        
+        for worker_rank, worker_data in other.data.items():
+            if worker_rank not in self.data:
+                self.data[worker_rank] = copy.deepcopy(worker_data)
+                continue
+            
+            for metric_name, values in worker_data.items():
+                self.data[worker_rank].setdefault(metric_name, []).extend(values)
+        
         return self
 
     def reduce(self) -> dict[str, int | float]:
-        """Reduce the observations to representative values for CLI logging"""
+        """Reduce the observations to representative values for CLI logging."""
         result = {}
         is_count_metric = lambda k: any(x in k for x in ["num", "requests", "blocks"])
         is_hit_rate_metric = lambda k: "hit_rate" in k or "hit_rates" in k
-
+        
         for worker_rank, worker_data in sorted(self.data.items()):
             for metric_name, values in worker_data.items():
-                # For hit rate metrics, only show worker_0
                 if is_hit_rate_metric(metric_name) and worker_rank != "0":
                     continue
-
+                
                 suffix = "(total)" if is_count_metric(metric_name) else "(avg)"
                 if is_hit_rate_metric(metric_name) and worker_rank == "0":
                     worker_key = f"{metric_name} {suffix}"
                 else:
                     worker_key = f"worker_{worker_rank}_{metric_name} {suffix}"
-
+                
                 if not values:
                     result[worker_key] = 0 if is_count_metric(metric_name) else 0.0
                 elif is_count_metric(metric_name):
                     result[worker_key] = int(sum(values))
                 else:
                     result[worker_key] = round(sum(values) / len(values), 3)
-
+        
         return result