[jax-inference-offloading] Print CI metric (#1819)

Example metrics:
```
0: JIO_METRIC_TRANSFER|mean=0.061956|std=0.001554|min=0.059907|max=0.064638|unit=s
0: JIO_METRIC_HANDSHAKE|mean=28.636484|std=0.000000|min=28.636484|max=28.636484|unit=s
0: JIO_METRIC_LOADMODEL|mean=5.969427|std=0.000000|min=5.969427|max=5.969427|unit=s
```

Author: yhtang

jax-inference-offloading/examples/trainer.py

@@ -153,6 +153,16 @@
 
 if jax.process_index() == 0:
   timer.summary(sort_by='name', precision=3)
+  for metric_name, metric_key in [
+    ('JIO_METRIC_TRANSFER', r'transport\.run\d+$'),
+    ('JIO_METRIC_HANDSHAKE', r'create_bridge\.handshake$'),
+    ('JIO_METRIC_LOADMODEL', r'load_model$'),
+  ]:
+    print(timer.ci_metric(
+      metric_name,
+      timer.node_stat(metric_key, ('mean', 'std', 'min', 'max')),
+      unit='s'
+    ))
 
 if jax.process_index() == 0:
   bridge.gateway.shutdown()

jax-inference-offloading/jax_inference_offloading/timer.py

@@ -14,9 +14,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import re
 import time
+import numpy as np
 from collections import defaultdict, deque
 from contextlib import contextmanager
+from typing import Callable
 
 
 class Timer:
@@ -39,7 +42,6 @@ class Timer:
     parse            : 0.3017 s
       json           : 0.3017 s
   train              : 2.8344 s
-
   """
 
   def __init__(self) -> None:
@@ -113,6 +115,38 @@ def _print_by_name(self, precision: int, col_sep: str) -> None:
       indent = "  " * depth
       print(f"{indent + segment:<{col_width_name}}{col_sep}{t:{col_width_time}.{precision}f}")
 
+  def node_stat(
+      self,
+      pattern: str,
+      stats: str | Callable[[list[float]], float] | list[str | Callable[[list[float]], float]] = 'mean'
+  ) -> float | tuple[float, ...] | None:
+    matcher = re.compile(pattern)
+    vals = [t for name, t in self._times.items() if matcher.fullmatch(name)]
+    if not vals:
+      return None
+
+    def get_op(name):
+      try:
+        return getattr(np, name)
+      except AttributeError:
+        raise ValueError(f"Unsupported statistic function name: {name}")    
+
+    if isinstance(stats, (list, tuple)):
+      return {s: get_op(s)(vals) for s in stats}
+    else:
+      return {stats: get_op(stats)(vals)}
+
+  def ci_metric(
+      self,
+      tag,
+      stats: dict[str, float],
+      unit: str = 's',
+      sep: str = '|',
+  ):
+    return sep.join(
+        [tag] + [f'{k}={v:.6f}' for k, v in stats.items()] + [f'unit={unit}']
+    )
+
   def reset(self) -> None:
     """Clear all recorded data."""
     self._times.clear()