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

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Unit tests for saving and retrieving additional parameters of pending trials. 

7""" 

8from datetime import datetime 

9 

10from pytz import UTC 

11 

12from mlos_bench.storage.base_storage import Storage 

13from mlos_bench.tunables.tunable_groups import TunableGroups 

14 

15 

16def test_exp_trial_pending(exp_storage: Storage.Experiment, 

17 tunable_groups: TunableGroups) -> None: 

18 """ 

19 Schedule a trial and check that it is pending and has the right configuration. 

20 """ 

21 config = {"location": "westus2", "num_repeats": 100} 

22 trial = exp_storage.new_trial(tunable_groups, config=config) 

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

24 assert pending.trial_id == trial.trial_id 

25 assert pending.tunables == tunable_groups 

26 assert pending.config() == { 

27 "location": "westus2", 

28 "num_repeats": "100", 

29 "experiment_id": "Test-001", 

30 "trial_id": trial.trial_id, 

31 } 

32 

33 

34def test_exp_trial_configs(exp_storage: Storage.Experiment, 

35 tunable_groups: TunableGroups) -> None: 

36 """ 

37 Start multiple trials with two different configs and check that 

38 we store only two config objects in the DB. 

39 """ 

40 config1 = tunable_groups.copy().assign({'idle': 'mwait'}) 

41 trials1 = [ 

42 exp_storage.new_trial(config1), 

43 exp_storage.new_trial(config1), 

44 exp_storage.new_trial(config1.copy()), # Same values, different instance 

45 ] 

46 assert trials1[0].tunable_config_id == trials1[1].tunable_config_id 

47 assert trials1[0].tunable_config_id == trials1[2].tunable_config_id 

48 

49 config2 = tunable_groups.copy().assign({'idle': 'halt'}) 

50 trials2 = [ 

51 exp_storage.new_trial(config2), 

52 exp_storage.new_trial(config2), 

53 exp_storage.new_trial(config2.copy()), # Same values, different instance 

54 ] 

55 assert trials2[0].tunable_config_id == trials2[1].tunable_config_id 

56 assert trials2[0].tunable_config_id == trials2[2].tunable_config_id 

57 

58 assert trials1[0].tunable_config_id != trials2[0].tunable_config_id 

59 

60 pending_ids = [ 

61 pending.tunable_config_id 

62 for pending in exp_storage.pending_trials(datetime.now(UTC), running=True) 

63 ] 

64 assert len(pending_ids) == 6 

65 assert len(set(pending_ids)) == 2 

66 assert set(pending_ids) == {trials1[0].tunable_config_id, trials2[0].tunable_config_id} 

67 

68 

69def test_exp_trial_no_config(exp_no_tunables_storage: Storage.Experiment) -> None: 

70 """ 

71 Schedule a trial that has an empty tunable groups config. 

72 """ 

73 empty_config: dict = {} 

74 tunable_groups = TunableGroups(config=empty_config) 

75 trial = exp_no_tunables_storage.new_trial(tunable_groups, config=empty_config) 

76 (pending,) = exp_no_tunables_storage.pending_trials(datetime.now(UTC), running=True) 

77 assert pending.trial_id == trial.trial_id 

78 assert pending.tunables == tunable_groups 

79 assert pending.config() == { 

80 "experiment_id": "Test-003", 

81 "trial_id": trial.trial_id, 

82 }