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` : 

14 The base storage module for mlos_bench, which includes some basic examples in 

15 the documentation. 

16""" 

17from abc import ABCMeta, abstractmethod 

18from datetime import datetime 

19from typing import TYPE_CHECKING, Any 

20 

21import pandas 

22from pytz import UTC 

23 

24from mlos_bench.environments.status import Status 

25from mlos_bench.storage.base_tunable_config_data import TunableConfigData 

26from mlos_bench.storage.util import kv_df_to_dict 

27from mlos_bench.tunables.tunable_types import TunableValue 

28 

29if TYPE_CHECKING: 

30 from mlos_bench.storage.base_tunable_config_trial_group_data import ( 

31 TunableConfigTrialGroupData, 

32 ) 

33 

34 

35class TrialData(metaclass=ABCMeta): 

36 """ 

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

38 

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

40 tunable parameters). 

41 """ 

42 

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

44 self, 

45 *, 

46 experiment_id: str, 

47 trial_id: int, 

48 tunable_config_id: int, 

49 ts_start: datetime, 

50 ts_end: datetime | None, 

51 status: Status, 

52 trial_runner_id: int | None = None, 

53 ): 

54 self._experiment_id = experiment_id 

55 self._trial_id = trial_id 

56 self._tunable_config_id = tunable_config_id 

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

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

59 self._ts_start = ts_start 

60 self._ts_end = ts_end 

61 self._status = status 

62 self._trial_runner_id = trial_runner_id 

63 

64 def __repr__(self) -> str: 

65 return ( 

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

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

68 ) 

69 

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

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

72 return False 

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

74 

75 @property 

76 def experiment_id(self) -> str: 

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

78 return self._experiment_id 

79 

80 @property 

81 def trial_id(self) -> int: 

82 """ID of the trial.""" 

83 return self._trial_id 

84 

85 @property 

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

87 """ID of the TrialRunner.""" 

88 return self._trial_runner_id 

89 

90 @property 

91 def ts_start(self) -> datetime: 

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

93 return self._ts_start 

94 

95 @property 

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

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

98 return self._ts_end 

99 

100 @property 

101 def status(self) -> Status: 

102 """Status of the trial.""" 

103 return self._status 

104 

105 @property 

106 def tunable_config_id(self) -> int: 

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

108 return self._tunable_config_id 

109 

110 @property 

111 @abstractmethod 

112 def tunable_config(self) -> TunableConfigData: 

113 """ 

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

115 

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

117 

118 Returns 

119 ------- 

120 tunable_config : TunableConfigData 

121 A TunableConfigData object. 

122 """ 

123 

124 @property 

125 @abstractmethod 

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

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

128 

129 @property 

130 @abstractmethod 

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

132 """ 

133 Retrieve the trials' results from the storage. 

134 

135 Returns 

136 ------- 

137 results : pandas.DataFrame 

138 A dataframe with the trial results. 

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

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

141 """ 

142 

143 @property 

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

145 """ 

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

147 

148 Returns 

149 ------- 

150 results : dict 

151 """ 

152 return kv_df_to_dict(self.results_df) 

153 

154 @property 

155 @abstractmethod 

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

157 """ 

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

159 

160 Returns 

161 ------- 

162 config : pandas.DataFrame 

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

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

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

166 the dataframe is empty. 

167 """ 

168 

169 @property 

170 @abstractmethod 

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

172 """ 

173 Retrieve the trials' metadata parameters as a dataframe. 

174 

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

176 

177 Returns 

178 ------- 

179 metadata : pandas.DataFrame 

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

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

182 Returns an empty dataframe if there is no metadata. 

183 """ 

184 

185 @property 

186 def metadata_dict(self) -> dict: 

187 """ 

188 Retrieve the trials' metadata parameters as a dict. 

189 

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

191 

192 Returns 

193 ------- 

194 metadata : dict 

195 """ 

196 return kv_df_to_dict(self.metadata_df)