Coverage for src/flag_gems/utils/models/model.py: 76%

1import inspect 

2import threading 

3from abc import abstractmethod 

4from typing import Dict, Final, Optional, Sequence, Tuple, Union, overload 

5 

6import triton 

7 

8 

9class PersistantModel(object): 

10 signature: Final[inspect.Signature] = inspect.signature(triton.Config) 

11 

12 def __init__(self, *args, **kwargs) -> None: 

13 super().__init__(*args, **kwargs) 

14 self.lock: Final[threading.Lock] = threading.Lock() 

15 

16 @staticmethod 

17 def parse_config(config: triton.Config) -> Dict[str, Union[int, float, str]]: 

18 return { 

19 k: v 

20 for k, v in config.all_kwargs().items() 

21 if isinstance(v, (int, float, str)) 

22 } 

23 

24 @abstractmethod 

25 def get_config( 

26 self, name: str, key: Sequence[Union[bool, int, float, str]] 

27 ) -> Optional[triton.Config]: 

28 ... 

29 

30 @abstractmethod 

31 def get_benchmark( 

32 self, 

33 name: str, 

34 keys: Sequence[Union[bool, int, float, str]], 

35 config: triton.Config, 

36 ) -> Optional[Tuple[float, float, float]]: 

37 ... 

38 

39 @overload 

40 def put_config( 

41 self, 

42 name: str, 

43 keys: Sequence[Union[bool, int, float, str]], 

44 config: triton.Config, 

45 ) -> None: 

46 ... 

47 

48 @overload 

49 def put_config( 

50 self, 

51 name: str, 

52 keys: Sequence[Union[bool, int, float, str]], 

53 config: Dict[str, Union[bool, int, float, str]], 

54 ) -> None: 

55 ... 

56 

57 @abstractmethod 

58 def put_config( 

59 self, 

60 name: str, 

61 keys: Sequence[Union[bool, int, float, str]], 

62 config: Union[triton.Config, Dict[str, Union[bool, int, float, str]]], 

63 ) -> None: 

64 ... 

65 

66 @overload 

67 def put_benchmark( 

68 self, 

69 name: str, 

70 keys: Sequence[Union[bool, int, float, str]], 

71 config: triton.Config, 

72 benchmark: Tuple[float, float, float], 

73 ) -> None: 

74 ... 

75 

76 @overload 

77 def put_benchmark( 

78 self, 

79 name: str, 

80 keys: Sequence[Union[bool, int, float, str]], 

81 config: Dict[str, Union[bool, int, float, str]], 

82 benchmark: Tuple[float, float, float], 

83 ) -> None: 

84 ... 

85 

86 @abstractmethod 

87 def put_benchmark( 

88 self, 

89 name: str, 

90 keys: Sequence[Union[bool, int, float, str]], 

91 config: Union[triton.Config, Dict[str, Union[bool, int, float, str]]], 

92 benchmark: Tuple[float, float, float], 

93 ) -> None: 

94 ...