xref: /aosp_15_r20/external/pytorch/benchmarks/distributed/rpc/parameter_server/metrics/CPUMetric.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1import time
2
3from .MetricBase import MetricBase
4
5
6class CPUMetric(MetricBase):
7    def __init__(self, name: str):
8        self.name = name
9        self.start = None
10        self.end = None
11
12    def record_start(self):
13        self.start = time.time()
14
15    def record_end(self):
16        self.end = time.time()
17
18    def elapsed_time(self):
19        if self.start is None:
20            raise RuntimeError("start is None")
21        if self.end is None:
22            raise RuntimeError("end is None")
23        return self.end - self.start
24