1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
|
# Copyright 2022 The Matrix.org Foundation C.I.C.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import time
from typing import Any, Callable, Dict, Optional
from scalene import scalene_profiler
class ProfilingDecider:
INSTANCES: Dict[str, "ProfilingDecider"] = {}
def __init__(self, name: str, cond: Callable[[], bool]) -> None:
ProfilingDecider.INSTANCES[name] = self
# Default to being armed if SCALENE is available as an env var.
self.armed = b"SCALENE" in os.environb
self._cond = cond
def decide(self) -> bool:
if not self.armed:
return False
if not self._cond():
return False
self.armed = False
return True
class CpuUtimeTracker:
def __init__(self) -> None:
self._update_times(time.time())
def _update_times(self, now_wall: float) -> None:
utime, _, _, _, elapsed = os.times()
self._last_utime = utime
self._last_elapsed = elapsed
self._last_wall = now_wall
self.min_elapse = 0.5
self.max_elapse = 120.0
def update_return_utime(self) -> Optional[float]:
"""
Returns CPU usage over this period, provided at least `min_elapse` have
elapsed.
"""
wall = time.time()
elapsed = wall - self._last_wall
if elapsed < self.min_elapse:
return None
last_utime = self._last_utime
last_elapsed = self._last_elapsed
self._update_times(wall)
if elapsed > self.max_elapse:
# the average will be a bit skewy if so much time has elapsed. Ignore.
return None
usage = (self._last_utime - last_utime) / (self._last_elapsed - last_elapsed)
return usage
class SelectiveProfiling:
def __init__(self, decider: ProfilingDecider, enable: bool = False):
self._decider = decider
self._enable = enable
def __enter__(self) -> None:
if not self._enable:
return
if not self._decider.decide():
self._enable = False
return
scalene_profiler.start()
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
if not self._enable:
scalene_profiler.stop()
|
