Refactor time measurement in TimeTook class

microsoft · Mar 17, 2024 · ed048fc · ed048fc
1 parent 2d19261
commit ed048fc
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/rag_experiment_accelerator/utils/timetook.py b/rag_experiment_accelerator/utils/timetook.py
@@ -1,4 +1,4 @@
-from datetime import datetime
+import time
 from .logging import get_logger
 
 
@@ -18,12 +18,11 @@ def __init__(self, description, logger):
         self.end = None
 
     def __enter__(self):
-        self.start = datetime.now()
+        self.start = time.perf_counter()
         self.logger.info(f"Starting {self.description}")
 
     def __exit__(self, type, value, traceback):
-        self.end = datetime.now()
+        self.end = time.perf_counter()
         self.logger.info(
-            f"Time took for {self.description}: "
-            f"{(self.end - self.start).total_seconds()} seconds"
+            f"Time took for {self.description}: " f"{self.end - self.start} seconds"
         )