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

1import os 

2from dataclasses import dataclass 

3 

4import yaml 

5 

6 

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

8@dataclass 

9class VendorInfoBase: 

10 vendor_name: str 

11 device_name: str 

12 device_query_cmd: str 

13 dispatch_key: str = None 

14 triton_extra_name: str = None 

15 

16 

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

18 BACKEND_EVENT = file_path is not None 

19 config = None 

20 try: 

21 if not file_path: 

22 vendor_name = "_" + vendor_name 

23 script_path = os.path.abspath(__file__) 

24 base_dir = os.path.dirname(script_path) 

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

26 else: 

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

28 with open(file_path, file_mode) as file: 

29 config = yaml.safe_load(file) 

30 except FileNotFoundError: 

31 if not BACKEND_EVENT: 

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

33 except yaml.YAMLError as e: 

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

35 except Exception as e: 

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

37 

38 return config