pg.time now supports aggregating sub-timers with the same name.

pyglove/core/utils/timing.py

@@ -56,6 +56,16 @@ def to_json(self, **kwargs) -> Dict[str, Any]:
           **kwargs,
       )
 
+    def merge(self, other: 'TimeIt.Status') -> 'TimeIt.Status':
+      """Merges the status of two `pg.timeit`."""
+      assert other.name == self.name, (self.name, other.name)
+      return TimeIt.Status(
+          name=self.name,
+          elapse=self.elapse + other.elapse,
+          has_ended=self.has_ended and other.has_ended,
+          error=self.error or other.error,
+      )
+
   @dataclasses.dataclass
   class StatusSummary(json_conversion.JSONConvertible):
     """Aggregated summary for repeated calls for `pg.timeit`."""
@@ -125,7 +135,7 @@ def __init__(self, name: str = ''):
     self._name: str = name
     self._start_time: Optional[float] = None
     self._end_time: Optional[float] = None
-    self._child_contexts: Dict[str, TimeIt] = {}
+    self._child_contexts: List[TimeIt] = []
     self._error: Optional[error_utils.ErrorInfo] = None
     self._parent: Optional[TimeIt] = None
 
@@ -137,13 +147,11 @@ def name(self) -> str:
   @property
   def children(self) -> List['TimeIt']:
     """Returns child contexts."""
-    return list(self._child_contexts.values())
+    return self._child_contexts
 
   def add(self, context: 'TimeIt'):
     """Adds a child context."""
-    if context.name in self._child_contexts:
-      raise ValueError(f'`timeit` with name {context.name!r} already exists.')
-    self._child_contexts[context.name] = context
+    self._child_contexts.append(context)
 
   def start(self):
     """Starts timing."""
@@ -206,11 +214,14 @@ def status(self) -> Dict[str, Status]:
             has_ended=self.has_ended, error=self._error,
         )
     }
-    for child in self._child_contexts.values():
+    for child in self._child_contexts:
       child_result = child.status()
       for k, v in child_result.items():
         key = f'{self.name}.{k}' if self.name else k
-        result[key] = v
+        if key in result:
+          result[key] = result[key].merge(v)
+        else:
+          result[key] = v
     return result
 
   def __enter__(self):

pyglove/core/utils/timing_test.py

@@ -62,27 +62,30 @@ def test_timeit(self):
         self.assertFalse(r['node.child'].has_ended)
         self.assertTrue(r['node.child.grandchild'].has_ended)
 
-        with self.assertRaisesRegex(ValueError, '.* already exists'):
-          with timing.timeit('grandchild'):
-            pass
+        with timing.timeit('grandchild') as t3:
+          time.sleep(0.5)
+        self.assertEqual(t1.children, [t2, t3])
 
     elapse2 = t.elapse
     self.assertTrue(t.has_ended)
     self.assertGreater(elapse2, elapse1)
     time.sleep(0.5)
     self.assertEqual(elapse2, t.elapse)
 
-    statuss = t.status()
+    status = t.status()
     self.assertEqual(
-        list(statuss.keys()),
+        list(status.keys()),
         ['node', 'node.child', 'node.child.grandchild']
     )
     self.assertEqual(
-        sorted([v.elapse for v in statuss.values()], reverse=True),
-        [v.elapse for v in statuss.values()],
+        status['node.child.grandchild'].elapse, t2.elapse + t3.elapse
+    )
+    self.assertEqual(
+        sorted([v.elapse for v in status.values()], reverse=True),
+        [v.elapse for v in status.values()],
     )
-    self.assertTrue(all(v.has_ended for v in statuss.values()))
-    self.assertFalse(any(v.has_error for v in statuss.values()))
+    self.assertTrue(all(v.has_ended for v in status.values()))
+    self.assertFalse(any(v.has_error for v in status.values()))
     status = t.status()
     json_dict = json_conversion.to_json(status)
     status2 = json_conversion.from_json(json_dict)