Coverage for o2/models/constraints/size_rule_constraints.py: 86%

1from dataclasses import dataclass 

2from typing import TYPE_CHECKING, Callable, Optional, TypeGuard 

3 

4from dataclass_wizard import JSONWizard 

5from sympy import Symbol, lambdify 

6 

7from o2.models.constraints.batching_constraints import BatchingConstraints 

8from o2.models.timetable import RULE_TYPE, FiringRule 

9 

10if TYPE_CHECKING: 

11 from o2.models.timetable import TimetableType 

12 

13 

14@dataclass(frozen=True) 

15class SizeRuleConstraints(BatchingConstraints, JSONWizard): 

16 """Size rule constraints for batching.""" 

17 

18 duration_fn: str 

19 cost_fn: str 

20 min_size: Optional[int] 

21 max_size: Optional[int] 

22 

23 class _(JSONWizard.Meta): # noqa: N801 

24 key_transform_with_dump = "SNAKE" 

25 tag = RULE_TYPE.SIZE.value 

26 tag_key = "rule_type" 

27 

28 def verify_timetable(self, timetable: "TimetableType") -> bool: 

29 """Check if the timetable is valid against the constraints.""" 

30 rules: list[FiringRule[int]] = timetable.get_firing_rules_for_tasks( 

31 self.tasks, rule_type=RULE_TYPE.SIZE, batch_type=self.batch_type 

32 ) 

33 return all(self._verify_firing_rule(rule) for rule in rules) 

34 

35 def _verify_firing_rule(self, firing_rule: FiringRule[int]) -> bool: 

36 """Check if the firing rule is valid against the constraints.""" 

37 # Special case: Size 0 is always invalid 

38 if firing_rule.value == 0: 

39 from o2.util.logger import warn 

40 

41 warn("Warning: Size 0 rule was created!") 

42 return False 

43 

44 # Special case: If the rule is a <(=) min size, it is invalid if min size is set 

45 if self.min_size and self.min_size > 0 and (firing_rule.is_lt_or_lte): 

46 return False 

47 

48 # Special case: If the rule is a >(=) max size, it is invalid if max size is set 

49 # TODO: Think about this 

50 # if self.max_size and (firing_rule.is_gt_or_gte): 

51 # return False 

52 

53 if (self.min_size and firing_rule.value < self.min_size) or ( # noqa: SIM103 

54 self.max_size and firing_rule.value > self.max_size 

55 ): 

56 return False 

57 

58 return True 

59 

60 @property 

61 def cost_fn_lambda(self) -> Callable[[float], float]: 

62 """Get the cost function as a lambda function.""" 

63 return lambdify(Symbol("size"), self.cost_fn) 

64 

65 @property 

66 def duration_fn_lambda(self) -> Callable[[float], float]: 

67 """Get the duration function as a lambda function.""" 

68 return lambdify(Symbol("size"), self.duration_fn) 

69 

70 

71def is_size_constraint(val: BatchingConstraints) -> TypeGuard[SizeRuleConstraints]: 

72 """Check if the constraint is a size constraint.""" 

73 return val.rule_type == RULE_TYPE.SIZE