Coverage for mlos_bench/mlos_bench/storage/base_trial_data.py: 98%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Base interface for accessing the stored benchmark trial data. 

7""" 

8from abc import ABCMeta, abstractmethod 

9from datetime import datetime 

10from typing import Any, Dict, Optional, TYPE_CHECKING 

11 

12import pandas 

13from pytz import UTC 

14 

15from mlos_bench.environments.status import Status 

16from mlos_bench.tunables.tunable import TunableValue 

17from mlos_bench.storage.base_tunable_config_data import TunableConfigData 

18from mlos_bench.storage.util import kv_df_to_dict 

19 

20if TYPE_CHECKING: 

21 from mlos_bench.storage.base_tunable_config_trial_group_data import TunableConfigTrialGroupData 

22 

23 

24class TrialData(metaclass=ABCMeta): 

25 """ 

26 Base interface for accessing the stored experiment benchmark trial data. 

27 

28 A trial is a single run of an experiment with a given configuration (e.g., set 

29 of tunable parameters). 

30 """ 

31 

32 def __init__(self, *, 

33 experiment_id: str, 

34 trial_id: int, 

35 tunable_config_id: int, 

36 ts_start: datetime, 

37 ts_end: Optional[datetime], 

38 status: Status): 

39 self._experiment_id = experiment_id 

40 self._trial_id = trial_id 

41 self._tunable_config_id = tunable_config_id 

42 assert ts_start.tzinfo == UTC, "ts_start must be in UTC" 

43 assert ts_end is None or ts_end.tzinfo == UTC, "ts_end must be in UTC if not None" 

44 self._ts_start = ts_start 

45 self._ts_end = ts_end 

46 self._status = status 

47 

48 def __repr__(self) -> str: 

49 return f"Trial :: {self._experiment_id}:{self._trial_id} cid:{self._tunable_config_id} {self._status.name}" 

50 

51 def __eq__(self, other: Any) -> bool: 

52 if not isinstance(other, self.__class__): 

53 return False 

54 return self._experiment_id == other._experiment_id and self._trial_id == other._trial_id 

55 

56 @property 

57 def experiment_id(self) -> str: 

58 """ 

59 ID of the experiment this trial belongs to. 

60 """ 

61 return self._experiment_id 

62 

63 @property 

64 def trial_id(self) -> int: 

65 """ 

66 ID of the trial. 

67 """ 

68 return self._trial_id 

69 

70 @property 

71 def ts_start(self) -> datetime: 

72 """ 

73 Start timestamp of the trial (UTC). 

74 """ 

75 return self._ts_start 

76 

77 @property 

78 def ts_end(self) -> Optional[datetime]: 

79 """ 

80 End timestamp of the trial (UTC). 

81 """ 

82 return self._ts_end 

83 

84 @property 

85 def status(self) -> Status: 

86 """ 

87 Status of the trial. 

88 """ 

89 return self._status 

90 

91 @property 

92 def tunable_config_id(self) -> int: 

93 """ 

94 ID of the (tunable) configuration of the trial. 

95 """ 

96 return self._tunable_config_id 

97 

98 @property 

99 @abstractmethod 

100 def tunable_config(self) -> TunableConfigData: 

101 """ 

102 Retrieve the trials' tunable configuration data from the storage. 

103 

104 Note: this corresponds to the Trial object's "tunables" property. 

105 

106 Returns 

107 ------- 

108 tunable_config : TunableConfigData 

109 A TunableConfigData object. 

110 """ 

111 

112 @property 

113 @abstractmethod 

114 def tunable_config_trial_group(self) -> "TunableConfigTrialGroupData": 

115 """ 

116 Retrieve the trial's (tunable) config trial group data from the storage. 

117 """ 

118 

119 @property 

120 @abstractmethod 

121 def results_df(self) -> pandas.DataFrame: 

122 """ 

123 Retrieve the trials' results from the storage. 

124 

125 Returns 

126 ------- 

127 results : pandas.DataFrame 

128 A dataframe with the trial results. 

129 It has two `str` columns, "metric" and "value". 

130 If the trial status is not SUCCEEDED, the dataframe is empty. 

131 """ 

132 

133 @property 

134 def results_dict(self) -> Dict[str, Optional[TunableValue]]: 

135 """ 

136 Retrieve the trials' results from the storage as a dict. 

137 

138 Returns 

139 ------- 

140 results : dict 

141 """ 

142 return kv_df_to_dict(self.results_df) 

143 

144 @property 

145 @abstractmethod 

146 def telemetry_df(self) -> pandas.DataFrame: 

147 """ 

148 Retrieve the trials' telemetry from the storage as a dataframe. 

149 

150 Returns 

151 ------- 

152 config : pandas.DataFrame 

153 A dataframe with the trial telemetry, if there is any. 

154 It has one `datetime` column, "ts", and two `str` columns, "metric" and "value". 

155 If the trial status is not SUCCEEDED, or there is no telemetry data, 

156 the dataframe is empty. 

157 """ 

158 

159 @property 

160 @abstractmethod 

161 def metadata_df(self) -> pandas.DataFrame: 

162 """ 

163 Retrieve the trials' metadata parameters as a dataframe. 

164 

165 Note: this corresponds to the Trial object's "config" property. 

166 

167 Returns 

168 ------- 

169 metadata : pandas.DataFrame 

170 An optional dataframe with the metadata associated with the trial. 

171 It has two `str` columns, "parameter" and "value". 

172 Returns an empty dataframe if there is no metadata. 

173 """ 

174 

175 @property 

176 def metadata_dict(self) -> dict: 

177 """ 

178 Retrieve the trials' metadata parameters as a dict. 

179 

180 Note: this corresponds to the Trial object's "config" property. 

181 

182 Returns 

183 ------- 

184 metadata : dict 

185 """ 

186 return kv_df_to_dict(self.metadata_df)