GGroenendaal commited on
Commit
615dee0
1 Parent(s): e9df5ab

add code for timing

Browse files
Files changed (2) hide show
  1. .env.example +2 -0
  2. src/utils/timing.py +57 -0
.env.example CHANGED
@@ -4,3 +4,5 @@ ELASTIC_HOST=https://localhost:9200
4
 
5
  LOG_LEVEL=INFO
6
  TRANSFORMERS_NO_ADVISORY_WARNINGS
 
 
 
4
 
5
  LOG_LEVEL=INFO
6
  TRANSFORMERS_NO_ADVISORY_WARNINGS
7
+
8
+ ENABLE_TIMING=TRUE
src/utils/timing.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import time
2
+ from typing import Dict
3
+ from dotenv import load_dotenv
4
+ import os
5
+ from src.utils.log import get_logger
6
+
7
+
8
+ logger = get_logger()
9
+
10
+
11
+ load_dotenv()
12
+
13
+
14
+ ENABLE_TIMING = os.getenv("ENABLE_TIMING", "false").lower() == "true"
15
+
16
+ if ENABLE_TIMING:
17
+ logger.info("Timing is enabled")
18
+
19
+
20
+ TimingType = Dict[str, float]
21
+
22
+ TIMES: TimingType = {}
23
+
24
+
25
+ def timeit(name: str):
26
+ def _timeit(func):
27
+ def wrapper(*args, **kwargs):
28
+ start = time.time()
29
+ result = func(*args, **kwargs)
30
+ end = time.time()
31
+
32
+ if name not in TIMES:
33
+ TIMES[name] = []
34
+ TIMES[name].append(end - start)
35
+
36
+ return result
37
+
38
+ if ENABLE_TIMING:
39
+ return wrapper
40
+ return func
41
+ return _timeit
42
+
43
+
44
+ def get_times() -> TimingType:
45
+ _warn_if_timing_disabled()
46
+ return TIMES
47
+
48
+
49
+ def reset_times() -> None:
50
+ _warn_if_timing_disabled()
51
+ TIMES.clear()
52
+
53
+
54
+ def _warn_if_timing_disabled() -> None:
55
+ if not ENABLE_TIMING:
56
+ logger.warning(
57
+ "Timing is disabled, please set ENABLE_TIMING to true in the .env file")