Coverage for src/flag_gems/runtime/backend/backend_utils.py: 59%

1import functools 

2import os 

3from dataclasses import dataclass 

4 

5import yaml 

6 

7 

8# Metadata template, Each vendor needs to specialize instances of this template 

9@dataclass 

10class VendorInfoBase: 

11 vendor_name: str 

12 device_name: str 

13 device_query_cmd: str 

14 dispatch_key: str = None 

15 triton_extra_name: str = None 

16 

17 

18def get_tune_config(vendor_name=None, file_mode="r", file_path=None): 

19 BACKEND_EVENT = file_path is not None 

20 config = None 

21 try: 

22 if not file_path: 

23 vendor_name = "_" + vendor_name 

24 script_path = os.path.abspath(__file__) 

25 base_dir = os.path.dirname(script_path) 

26 file_path = os.path.join(base_dir, vendor_name, "tune_configs.yaml") 

27 else: 

28 file_path = os.path.join(file_path, "tune_configs.yaml") 

29 with open(file_path, file_mode) as file: 

30 config = yaml.safe_load(file) 

31 except FileNotFoundError: 

32 if not BACKEND_EVENT: 

33 raise FileNotFoundError(f"Configuration file not found: {file_path}") 

34 except yaml.YAMLError as e: 

35 raise ValueError(f"Failed to parse YAML file: {e}") 

36 except Exception as e: 

37 raise RuntimeError(f"An unexpected error occurred: {e}") 

38 

39 return config 

40 

41 

42class BackendEventBase: 

43 def __init__(self): 

44 ... 

45 

46 def get_ops(self): 

47 ... 

48 

49 def is_available(self): 

50 ... 

51 

52 

53@functools.lru_cache(maxsize=None) 

54def _load_expand_config(file_path, file_mode="r"): 

55 with open(file_path, file_mode) as file: 

56 return yaml.safe_load(file) or {} 

57 

58 

59def get_expand_config(op_name=None, file_mode="r", file_path=None): 

60 if not file_path: 

61 raise ValueError("expand config file path is required") 

62 try: 

63 config = _load_expand_config(file_path, file_mode) 

64 except FileNotFoundError: 

65 raise FileNotFoundError(f"Configuration file not found: {file_path}") 

66 except yaml.YAMLError as e: 

67 raise ValueError(f"Failed to parse YAML file: {e}") 

68 except Exception as e: 

69 raise RuntimeError(f"An unexpected error occurred: {e}") 

70 if op_name is None: 

71 return config 

72 return config.get(op_name)