Merge pull request #15 from twosigma/tiny-progress-optimization

Tiny progress optimization
twosigma · Jan 1, 2023 · cb227c3 · cb227c3
2 parents d9fc79e + 3db2c9c
commit cb227c3
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/src/uberjob/progress/_simple_progress_observer.py b/src/uberjob/progress/_simple_progress_observer.py
@@ -84,7 +84,7 @@ class State:
     def __init__(self, start_time):
         self.section_scope_mapping = {}
         self.running_count = 0
-        self._running_section_scopes = set()
+        self._running_scope_states = set()
         self._prev_time = start_time
 
     def increment_total(self, section, scope, amount: int):
@@ -95,8 +95,7 @@ def increment_total(self, section, scope, amount: int):
     def increment_running(self, section, scope):
         self.update_weighted_elapsed()
         scope_state = self.section_scope_mapping[section][scope]
-        if not scope_state.running:
-            self._running_section_scopes.add((section, scope))
+        self._running_scope_states.add(scope_state)
         scope_state.running += 1
         self.running_count += 1
 
@@ -106,7 +105,7 @@ def increment_completed(self, section, scope):
         scope_state.running -= 1
         self.running_count -= 1
         if not scope_state.running:
-            self._running_section_scopes.remove((section, scope))
+            self._running_scope_states.remove(scope_state)
         scope_state.completed += 1
 
     def increment_failed(self, section, scope):
@@ -115,16 +114,15 @@ def increment_failed(self, section, scope):
         scope_state.running -= 1
         self.running_count -= 1
         if not scope_state.running:
-            self._running_section_scopes.remove((section, scope))
+            self._running_scope_states.remove(scope_state)
         scope_state.failed += 1
 
     def update_weighted_elapsed(self):
         t = time.time()
         if self.running_count:
             elapsed = t - self._prev_time
             multiplier = elapsed / self.running_count
-            for section, scope in self._running_section_scopes:
-                scope_state = self.section_scope_mapping[section][scope]
+            for scope_state in self._running_scope_states:
                 scope_state.weighted_elapsed += scope_state.running * multiplier
         self._prev_time = t