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 

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

9tunable parameters). 

10""" 

11from abc import ABCMeta, abstractmethod 

12from datetime import datetime 

13from typing import TYPE_CHECKING, Any, Dict, Optional 

14 

15import pandas 

16from pytz import UTC 

17 

18from mlos_bench.environments.status import Status 

19from mlos_bench.storage.base_tunable_config_data import TunableConfigData 

20from mlos_bench.storage.util import kv_df_to_dict 

21from mlos_bench.tunables.tunable import TunableValue 

22 

23if TYPE_CHECKING: 

24 from mlos_bench.storage.base_tunable_config_trial_group_data import ( 

25 TunableConfigTrialGroupData, 

26 ) 

27 

28 

29class TrialData(metaclass=ABCMeta): 

30 """ 

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

32 

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

34 tunable parameters). 

35 """ 

36 

37 def __init__( # pylint: disable=too-many-arguments 

38 self, 

39 *, 

40 experiment_id: str, 

41 trial_id: int, 

42 tunable_config_id: int, 

43 ts_start: datetime, 

44 ts_end: Optional[datetime], 

45 status: Status, 

46 ): 

47 self._experiment_id = experiment_id 

48 self._trial_id = trial_id 

49 self._tunable_config_id = tunable_config_id 

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

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

52 self._ts_start = ts_start 

53 self._ts_end = ts_end 

54 self._status = status 

55 

56 def __repr__(self) -> str: 

57 return ( 

58 f"Trial :: {self._experiment_id}:{self._trial_id} " 

59 f"cid:{self._tunable_config_id} {self._status.name}" 

60 ) 

61 

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

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

64 return False 

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

66 

67 @property 

68 def experiment_id(self) -> str: 

69 """ID of the experiment this trial belongs to.""" 

70 return self._experiment_id 

71 

72 @property 

73 def trial_id(self) -> int: 

74 """ID of the trial.""" 

75 return self._trial_id 

76 

77 @property 

78 def ts_start(self) -> datetime: 

79 """Start timestamp of the trial (UTC).""" 

80 return self._ts_start 

81 

82 @property 

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

84 """End timestamp of the trial (UTC).""" 

85 return self._ts_end 

86 

87 @property 

88 def status(self) -> Status: 

89 """Status of the trial.""" 

90 return self._status 

91 

92 @property 

93 def tunable_config_id(self) -> int: 

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

95 return self._tunable_config_id 

96 

97 @property 

98 @abstractmethod 

99 def tunable_config(self) -> TunableConfigData: 

100 """ 

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

102 

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

104 

105 Returns 

106 ------- 

107 tunable_config : TunableConfigData 

108 A TunableConfigData object. 

109 """ 

110 

111 @property 

112 @abstractmethod 

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

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

115 

116 @property 

117 @abstractmethod 

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

119 """ 

120 Retrieve the trials' results from the storage. 

121 

122 Returns 

123 ------- 

124 results : pandas.DataFrame 

125 A dataframe with the trial results. 

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

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

128 """ 

129 

130 @property 

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

132 """ 

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

134 

135 Returns 

136 ------- 

137 results : dict 

138 """ 

139 return kv_df_to_dict(self.results_df) 

140 

141 @property 

142 @abstractmethod 

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

144 """ 

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

146 

147 Returns 

148 ------- 

149 config : pandas.DataFrame 

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

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

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

153 the dataframe is empty. 

154 """ 

155 

156 @property 

157 @abstractmethod 

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

159 """ 

160 Retrieve the trials' metadata parameters as a dataframe. 

161 

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

163 

164 Returns 

165 ------- 

166 metadata : pandas.DataFrame 

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

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

169 Returns an empty dataframe if there is no metadata. 

170 """ 

171 

172 @property 

173 def metadata_dict(self) -> dict: 

174 """ 

175 Retrieve the trials' metadata parameters as a dict. 

176 

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

178 

179 Returns 

180 ------- 

181 metadata : dict 

182 """ 

183 return kv_df_to_dict(self.metadata_df)