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"""Base interface for accessing the stored benchmark trial data.""" 

6from abc import ABCMeta, abstractmethod 

7from datetime import datetime 

8from typing import TYPE_CHECKING, Any, Dict, Optional 

9 

10import pandas 

11from pytz import UTC 

12 

13from mlos_bench.environments.status import Status 

14from mlos_bench.storage.base_tunable_config_data import TunableConfigData 

15from mlos_bench.storage.util import kv_df_to_dict 

16from mlos_bench.tunables.tunable import TunableValue 

17 

18if TYPE_CHECKING: 

19 from mlos_bench.storage.base_tunable_config_trial_group_data import ( 

20 TunableConfigTrialGroupData, 

21 ) 

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 of 

29 tunable parameters). 

30 """ 

31 

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

33 self, 

34 *, 

35 experiment_id: str, 

36 trial_id: int, 

37 tunable_config_id: int, 

38 ts_start: datetime, 

39 ts_end: Optional[datetime], 

40 status: Status, 

41 ): 

42 self._experiment_id = experiment_id 

43 self._trial_id = trial_id 

44 self._tunable_config_id = tunable_config_id 

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

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

47 self._ts_start = ts_start 

48 self._ts_end = ts_end 

49 self._status = status 

50 

51 def __repr__(self) -> str: 

52 return ( 

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

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

55 ) 

56 

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

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

59 return False 

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

61 

62 @property 

63 def experiment_id(self) -> str: 

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

65 return self._experiment_id 

66 

67 @property 

68 def trial_id(self) -> int: 

69 """ID of the trial.""" 

70 return self._trial_id 

71 

72 @property 

73 def ts_start(self) -> datetime: 

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

75 return self._ts_start 

76 

77 @property 

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

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

80 return self._ts_end 

81 

82 @property 

83 def status(self) -> Status: 

84 """Status of the trial.""" 

85 return self._status 

86 

87 @property 

88 def tunable_config_id(self) -> int: 

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

90 return self._tunable_config_id 

91 

92 @property 

93 @abstractmethod 

94 def tunable_config(self) -> TunableConfigData: 

95 """ 

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

97 

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

99 

100 Returns 

101 ------- 

102 tunable_config : TunableConfigData 

103 A TunableConfigData object. 

104 """ 

105 

106 @property 

107 @abstractmethod 

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

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

110 

111 @property 

112 @abstractmethod 

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

114 """ 

115 Retrieve the trials' results from the storage. 

116 

117 Returns 

118 ------- 

119 results : pandas.DataFrame 

120 A dataframe with the trial results. 

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

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

123 """ 

124 

125 @property 

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

127 """ 

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

129 

130 Returns 

131 ------- 

132 results : dict 

133 """ 

134 return kv_df_to_dict(self.results_df) 

135 

136 @property 

137 @abstractmethod 

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

139 """ 

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

141 

142 Returns 

143 ------- 

144 config : pandas.DataFrame 

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

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

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

148 the dataframe is empty. 

149 """ 

150 

151 @property 

152 @abstractmethod 

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

154 """ 

155 Retrieve the trials' metadata parameters as a dataframe. 

156 

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

158 

159 Returns 

160 ------- 

161 metadata : pandas.DataFrame 

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

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

164 Returns an empty dataframe if there is no metadata. 

165 """ 

166 

167 @property 

168 def metadata_dict(self) -> dict: 

169 """ 

170 Retrieve the trials' metadata parameters as a dict. 

171 

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

173 

174 Returns 

175 ------- 

176 metadata : dict 

177 """ 

178 return kv_df_to_dict(self.metadata_df)