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 

11See Also 

12-------- 

13:py:mod:`mlos_bench.storage` : The base storage module for mlos_bench, which 

14 includes some basic examples in the documentation. 

15""" 

16from abc import ABCMeta, abstractmethod 

17from datetime import datetime 

18from typing import TYPE_CHECKING, Any 

19 

20import pandas 

21from pytz import UTC 

22 

23from mlos_bench.environments.status import Status 

24from mlos_bench.storage.base_tunable_config_data import TunableConfigData 

25from mlos_bench.storage.util import kv_df_to_dict 

26from mlos_bench.tunables.tunable import TunableValue 

27 

28if TYPE_CHECKING: 

29 from mlos_bench.storage.base_tunable_config_trial_group_data import ( 

30 TunableConfigTrialGroupData, 

31 ) 

32 

33 

34class TrialData(metaclass=ABCMeta): 

35 """ 

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

37 

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

39 tunable parameters). 

40 """ 

41 

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

43 self, 

44 *, 

45 experiment_id: str, 

46 trial_id: int, 

47 tunable_config_id: int, 

48 ts_start: datetime, 

49 ts_end: datetime | None, 

50 status: Status, 

51 trial_runner_id: int | None = None, 

52 ): 

53 self._experiment_id = experiment_id 

54 self._trial_id = trial_id 

55 self._tunable_config_id = tunable_config_id 

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

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

58 self._ts_start = ts_start 

59 self._ts_end = ts_end 

60 self._status = status 

61 self._trial_runner_id = trial_runner_id 

62 

63 def __repr__(self) -> str: 

64 return ( 

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

66 f"cid:{self._tunable_config_id} rid:{self._trial_runner_id} {self._status.name}" 

67 ) 

68 

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

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

71 return False 

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

73 

74 @property 

75 def experiment_id(self) -> str: 

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

77 return self._experiment_id 

78 

79 @property 

80 def trial_id(self) -> int: 

81 """ID of the trial.""" 

82 return self._trial_id 

83 

84 @property 

85 def trial_runner_id(self) -> int | None: 

86 """ID of the TrialRunner.""" 

87 return self._trial_runner_id 

88 

89 @property 

90 def ts_start(self) -> datetime: 

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

92 return self._ts_start 

93 

94 @property 

95 def ts_end(self) -> datetime | None: 

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

97 return self._ts_end 

98 

99 @property 

100 def status(self) -> Status: 

101 """Status of the trial.""" 

102 return self._status 

103 

104 @property 

105 def tunable_config_id(self) -> int: 

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

107 return self._tunable_config_id 

108 

109 @property 

110 @abstractmethod 

111 def tunable_config(self) -> TunableConfigData: 

112 """ 

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

114 

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

116 

117 Returns 

118 ------- 

119 tunable_config : TunableConfigData 

120 A TunableConfigData object. 

121 """ 

122 

123 @property 

124 @abstractmethod 

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

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

127 

128 @property 

129 @abstractmethod 

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

131 """ 

132 Retrieve the trials' results from the storage. 

133 

134 Returns 

135 ------- 

136 results : pandas.DataFrame 

137 A dataframe with the trial results. 

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

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

140 """ 

141 

142 @property 

143 def results_dict(self) -> dict[str, TunableValue | None]: 

144 """ 

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

146 

147 Returns 

148 ------- 

149 results : dict 

150 """ 

151 return kv_df_to_dict(self.results_df) 

152 

153 @property 

154 @abstractmethod 

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

156 """ 

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

158 

159 Returns 

160 ------- 

161 config : pandas.DataFrame 

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

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

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

165 the dataframe is empty. 

166 """ 

167 

168 @property 

169 @abstractmethod 

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

171 """ 

172 Retrieve the trials' metadata parameters as a dataframe. 

173 

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

175 

176 Returns 

177 ------- 

178 metadata : pandas.DataFrame 

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

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

181 Returns an empty dataframe if there is no metadata. 

182 """ 

183 

184 @property 

185 def metadata_dict(self) -> dict: 

186 """ 

187 Retrieve the trials' metadata parameters as a dict. 

188 

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

190 

191 Returns 

192 ------- 

193 metadata : dict 

194 """ 

195 return kv_df_to_dict(self.metadata_df)