Coverage for mlos_bench/mlos_bench/run.py: 87%

1#!/usr/bin/env python3 

2# 

3# Copyright (c) Microsoft Corporation. 

4# Licensed under the MIT License. 

5# 

6""" 

7OS Autotune main optimization loop. 

8 

9Note: this script is also available as a CLI tool via pip under the name "mlos_bench". 

10 

11See `--help` output for details. 

12""" 

13 

14import logging 

15import sys 

16from typing import Dict, List, Optional, Tuple 

17 

18import numpy as np 

19 

20from mlos_bench.launcher import Launcher 

21from mlos_bench.tunables.tunable_groups import TunableGroups 

22 

23_LOG = logging.getLogger(__name__) 

24 

25 

26def _sanity_check_results(launcher: Launcher) -> None: 

27 """Do some sanity checking on the results and throw an exception if it looks like 

28 something went wrong. 

29 """ 

30 basic_err_msg = "Check configuration, scripts, and logs for details." 

31 

32 # Check if the scheduler has any trials. 

33 if not launcher.scheduler.trial_count: 

34 raise RuntimeError(f"No trials were run. {basic_err_msg}") 

35 

36 # Check if the scheduler ran the expected number of trials. 

37 expected_trial_count = min( 

38 launcher.scheduler.max_trials if launcher.scheduler.max_trials > 0 else np.inf, 

39 launcher.scheduler.trial_config_repeat_count * launcher.optimizer.max_suggestions, 

40 ) 

41 if launcher.scheduler.trial_count < expected_trial_count: 

42 raise RuntimeError( 

43 f"Expected {expected_trial_count} trials, " 

44 f"but only {launcher.scheduler.trial_count} were run. {basic_err_msg}" 

45 ) 

46 

47 # Check to see if "too many" trials seem to have failed (#523). 

48 unsuccessful_trials = [t for t in launcher.scheduler.ran_trials if not t.status.is_succeeded()] 

49 if len(unsuccessful_trials) > 0.2 * launcher.scheduler.trial_count: 

50 raise RuntimeWarning( 

51 "Too many trials failed: " 

52 f"{len(unsuccessful_trials)} out of {launcher.scheduler.trial_count}. " 

53 f"{basic_err_msg}" 

54 ) 

55 

56 

57def _main( 

58 argv: Optional[List[str]] = None, 

59) -> Tuple[Optional[Dict[str, float]], Optional[TunableGroups]]: 

60 launcher = Launcher("mlos_bench", "Systems autotuning and benchmarking tool", argv=argv) 

61 

62 with launcher.scheduler as scheduler_context: 

63 scheduler_context.start() 

64 scheduler_context.teardown() 

65 

66 _sanity_check_results(launcher) 

67 

68 (score, _config) = result = launcher.scheduler.get_best_observation() 

69 # NOTE: This log line is used in test_launch_main_app_* unit tests: 

70 _LOG.info("Final score: %s", score) 

71 return result 

72 

73 

74def _shell_main( 

75 argv: Optional[List[str]] = None, 

76) -> int: 

77 (best_score, best_config) = _main(argv) 

78 # Exit zero if it looks like the overall operation was successful. 

79 # TODO: Improve this sanity check to be more robust. 

80 if ( 

81 best_score 

82 and best_config 

83 and all(isinstance(score_value, float) for score_value in best_score.values()) 

84 ): 

85 return 0 

86 else: 

87 raise ValueError(f"Unexpected result: {best_score=}, {best_config=}") 

88 

89 

90if __name__ == "__main__": 

91 sys.exit(_shell_main())