Coverage for qreward / utils / scheduler / pools.py: 100.00%

1"""Running task pool for concurrency monitoring.""" 

2 

3import threading (empty)

4import time (empty)

5import warnings (empty)

6from collections import OrderedDict (empty)

7 

8# Default number of time windows to track historical peak concurrency 

9DEFAULT_WINDOW_MAX_SIZE = 12 (empty)

10 

11# Default interval (seconds) for each time window 

12DEFAULT_WINDOW_INTERVAL = 60 (empty)

13 

14# Default concurrency threshold below which tasks are always allowed 

15DEFAULT_THRESHOLD = 3 (empty)

16 

17 

18class RunningTaskPool: (empty)

19 """Manager for monitoring and controlling task concurrency.""" 

20 

21 # Global lock for protecting global pool 

22 global_lock = threading.Lock() (empty)

23 

24 # Global pool dictionary, storing different pool instances by key 

25 global_task_pool: dict = {} (empty)

26 

27 @classmethod (empty)

28 def get_pool( (empty)

29 cls, 

30 key: str, 

31 window_max_size: int = DEFAULT_WINDOW_MAX_SIZE, 

32 window_interval: int = DEFAULT_WINDOW_INTERVAL, 

33 threshold: int = DEFAULT_THRESHOLD, 

34 ) -> "RunningTaskPool": 

35 """Get or create a task pool instance for the given key. 

36 

37 Uses singleton pattern to ensure one pool per key. 

38 

39 Args: 

40 key: Unique identifier for the pool 

41 window_max_size: Max number of time windows, default 12 

42 window_interval: Interval of each window in seconds, default 60 

43 threshold: Concurrency threshold, default 3 

44 

45 Returns: 

46 RunningTaskPool instance 

47 """ 

48 with cls.global_lock: (empty)

49 if key not in cls.global_task_pool: (empty)

50 cls.global_task_pool[key] = cls( (empty)

51 window_max_size=window_max_size, 

52 window_interval=window_interval, 

53 threshold=threshold, 

54 ) 

55 return cls.global_task_pool[key] (empty)

56 

57 def __init__( (empty)

58 self, 

59 window_max_size: int = DEFAULT_WINDOW_MAX_SIZE, 

60 window_interval: int = DEFAULT_WINDOW_INTERVAL, 

61 threshold: int = DEFAULT_THRESHOLD, 

62 ): 

63 """Initialize task pool. 

64 

65 Args: 

66 window_max_size: Max number of time windows 

67 window_interval: Interval of each window in seconds 

68 threshold: Concurrency threshold 

69 """ 

70 self._value = 0 (empty)

71 self._max_size_map = OrderedDict() (empty)

72 self._window_max_size = window_max_size (empty)

73 self._window_interval = window_interval (empty)

74 self._threshold = threshold (empty)

75 self._lock = threading.Lock() (empty)

76 

77 def add(self, value: int = 1) -> int: (empty)

78 """Update current running task count and record historical peak. 

79 

80 Args: 

81 value: Value to change 

82 

83 Returns: 

84 Current running task count 

85 """ 

86 with self._lock: (empty)

87 self._value += value (empty)

88 key = int(time.time()) // self._window_interval (empty)

89 if key in self._max_size_map: (empty)

90 if self._max_size_map[key] < self._value: (empty)

91 self._max_size_map[key] = self._value (empty)

92 else: 

93 while len(self._max_size_map) >= self._window_max_size: (empty)

94 self._max_size_map.popitem(last=False) (empty)

95 self._max_size_map[key] = self._value (empty)

96 return self._value (empty)

97 

98 def can_submit(self, multiply: float = 1) -> bool: (empty)

99 """Check if current task status allows execution. 

100 

101 Logic: 

102 1. If current task count is below threshold, allow execution 

103 2. If historical max exceeds current value * multiply, system is 

104 overloaded 

105 

106 Args: 

107 multiply: Multiplier for load calculation 

108 

109 Returns: 

110 True if can continue, False if should throttle 

111 """ 

112 with self._lock: (empty)

113 if self._value <= self._threshold: (empty)

114 return True (empty)

115 max_value = 0 (empty)

116 for v in self._max_size_map.values(): (empty)

117 if v > max_value: (empty)

118 max_value = v (empty)

119 if max_value > self._value * multiply: (empty)

120 return True (empty)

121 return False (empty)

122 

123 def less_than(self, multiply: float = 1) -> bool: (empty)

124 warnings.warn( (empty)

125 "less_than is deprecated, use can_submit instead", 

126 DeprecationWarning, 

127 stacklevel=2, 

128 ) 

129 return self.can_submit(multiply) (empty)