Coverage for mlos_bench/mlos_bench/tests/storage/exp_load_test.py: 100%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Unit tests for the storage subsystem.""" 

6from datetime import datetime, tzinfo 

7 

8import pytest 

9from pytz import UTC 

10 

11from mlos_bench.environments.status import Status 

12from mlos_bench.storage.base_storage import Storage 

13from mlos_bench.tests import ZONE_INFO 

14from mlos_bench.tunables.tunable_groups import TunableGroups 

15 

16 

17def test_exp_load_empty(exp_storage: Storage.Experiment) -> None: 

18 """Try to retrieve old experimental data from the empty storage.""" 

19 (trial_ids, configs, scores, status) = exp_storage.load() 

20 assert not trial_ids 

21 assert not configs 

22 assert not scores 

23 assert not status 

24 

25 

26def test_exp_pending_empty(exp_storage: Storage.Experiment) -> None: 

27 """Try to retrieve pending experiments from the empty storage.""" 

28 trials = list(exp_storage.pending_trials(datetime.now(UTC), running=True)) 

29 assert not trials 

30 

31 

32@pytest.mark.parametrize(("zone_info"), ZONE_INFO) 

33def test_exp_trial_pending( 

34 exp_storage: Storage.Experiment, 

35 tunable_groups: TunableGroups, 

36 zone_info: tzinfo | None, 

37) -> None: 

38 """Start a trial and check that it is pending.""" 

39 trial = exp_storage.new_trial(tunable_groups) 

40 (pending,) = list(exp_storage.pending_trials(datetime.now(zone_info), running=True)) 

41 assert pending.trial_id == trial.trial_id 

42 assert pending.tunables == tunable_groups 

43 

44 

45@pytest.mark.parametrize(("zone_info"), ZONE_INFO) 

46def test_exp_trial_pending_many( 

47 exp_storage: Storage.Experiment, 

48 tunable_groups: TunableGroups, 

49 zone_info: tzinfo | None, 

50) -> None: 

51 """Start THREE trials and check that both are pending.""" 

52 config1 = tunable_groups.copy().assign({"idle": "mwait"}) 

53 config2 = tunable_groups.copy().assign({"idle": "noidle"}) 

54 trial_ids = { 

55 exp_storage.new_trial(config1).trial_id, 

56 exp_storage.new_trial(config2).trial_id, 

57 exp_storage.new_trial(config2).trial_id, # Submit same config twice 

58 } 

59 pending_ids = { 

60 pending.trial_id 

61 for pending in exp_storage.pending_trials(datetime.now(zone_info), running=True) 

62 } 

63 assert len(pending_ids) == 3 

64 assert trial_ids == pending_ids 

65 

66 

67@pytest.mark.parametrize(("zone_info"), ZONE_INFO) 

68def test_exp_trial_pending_fail( 

69 exp_storage: Storage.Experiment, 

70 tunable_groups: TunableGroups, 

71 zone_info: tzinfo | None, 

72) -> None: 

73 """Start a trial, fail it, and and check that it is NOT pending.""" 

74 trial = exp_storage.new_trial(tunable_groups) 

75 trial.update(Status.FAILED, datetime.now(zone_info)) 

76 trials = list(exp_storage.pending_trials(datetime.now(zone_info), running=True)) 

77 assert not trials 

78 

79 

80@pytest.mark.parametrize(("zone_info"), ZONE_INFO) 

81def test_exp_trial_success( 

82 exp_storage: Storage.Experiment, 

83 tunable_groups: TunableGroups, 

84 zone_info: tzinfo | None, 

85) -> None: 

86 """Start a trial, finish it successfully, and and check that it is NOT pending.""" 

87 trial = exp_storage.new_trial(tunable_groups) 

88 trial.update(Status.SUCCEEDED, datetime.now(zone_info), {"score": 99.9}) 

89 trials = list(exp_storage.pending_trials(datetime.now(zone_info), running=True)) 

90 assert not trials 

91 

92 

93@pytest.mark.parametrize(("zone_info"), ZONE_INFO) 

94def test_exp_trial_update_categ( 

95 exp_storage: Storage.Experiment, 

96 tunable_groups: TunableGroups, 

97 zone_info: tzinfo | None, 

98) -> None: 

99 """Update the trial with multiple metrics, some of which are categorical.""" 

100 trial = exp_storage.new_trial(tunable_groups) 

101 trial.update(Status.SUCCEEDED, datetime.now(zone_info), {"score": 99.9, "benchmark": "test"}) 

102 assert exp_storage.load() == ( 

103 [trial.trial_id], 

104 [ 

105 { 

106 "idle": "halt", 

107 "kernel_sched_latency_ns": "2000000", 

108 "kernel_sched_migration_cost_ns": "-1", 

109 "vmSize": "Standard_B4ms", 

110 } 

111 ], 

112 [{"score": "99.9", "benchmark": "test"}], 

113 [Status.SUCCEEDED], 

114 ) 

115 

116 

117@pytest.mark.parametrize(("zone_info"), ZONE_INFO) 

118def test_exp_trial_update_twice( 

119 exp_storage: Storage.Experiment, 

120 tunable_groups: TunableGroups, 

121 zone_info: tzinfo | None, 

122) -> None: 

123 """Update the trial status twice and receive an error.""" 

124 trial = exp_storage.new_trial(tunable_groups) 

125 trial.update(Status.FAILED, datetime.now(zone_info)) 

126 with pytest.raises(RuntimeError): 

127 trial.update(Status.SUCCEEDED, datetime.now(UTC), {"score": 99.9}) 

128 

129 

130@pytest.mark.parametrize(("zone_info"), ZONE_INFO) 

131def test_exp_trial_pending_3( 

132 exp_storage: Storage.Experiment, 

133 tunable_groups: TunableGroups, 

134 zone_info: tzinfo | None, 

135) -> None: 

136 """ 

137 Start THREE trials, let one succeed, another one fail and keep one not updated. 

138 

139 Check that one is still pending another one can be loaded into the optimizer. 

140 """ 

141 score = 99.9 

142 

143 trial_fail = exp_storage.new_trial(tunable_groups) 

144 trial_succ = exp_storage.new_trial(tunable_groups) 

145 trial_pend = exp_storage.new_trial(tunable_groups) 

146 

147 trial_fail.update(Status.FAILED, datetime.now(zone_info)) 

148 trial_succ.update(Status.SUCCEEDED, datetime.now(zone_info), {"score": score}) 

149 

150 (pending,) = list(exp_storage.pending_trials(datetime.now(UTC), running=True)) 

151 assert pending.trial_id == trial_pend.trial_id 

152 

153 (trial_ids, configs, scores, status) = exp_storage.load() 

154 assert trial_ids == [trial_fail.trial_id, trial_succ.trial_id] 

155 assert len(configs) == 2 

156 assert scores == [None, {"score": f"{score}"}] 

157 assert status == [Status.FAILED, Status.SUCCEEDED] 

158 assert tunable_groups.copy().assign(configs[0]).reset() == trial_fail.tunables 

159 assert tunable_groups.copy().assign(configs[1]).reset() == trial_succ.tunables