Coverage for mlos_bench/mlos_bench/environments/remote/remote_env.py: 88%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Remotely executed benchmark/script environment. 

7 

8e.g. Application Environment 

9""" 

10 

11import logging 

12from datetime import datetime 

13from typing import Dict, Iterable, Optional, Tuple 

14 

15from pytz import UTC 

16 

17from mlos_bench.environments.status import Status 

18from mlos_bench.environments.script_env import ScriptEnv 

19from mlos_bench.services.base_service import Service 

20from mlos_bench.services.types.remote_exec_type import SupportsRemoteExec 

21from mlos_bench.services.types.host_ops_type import SupportsHostOps 

22from mlos_bench.tunables.tunable import TunableValue 

23from mlos_bench.tunables.tunable_groups import TunableGroups 

24 

25_LOG = logging.getLogger(__name__) 

26 

27 

28class RemoteEnv(ScriptEnv): 

29 """ 

30 Environment to run benchmarks and scripts on a remote host OS. 

31 

32 e.g. Application Environment 

33 """ 

34 

35 def __init__(self, 

36 *, 

37 name: str, 

38 config: dict, 

39 global_config: Optional[dict] = None, 

40 tunables: Optional[TunableGroups] = None, 

41 service: Optional[Service] = None): 

42 """ 

43 Create a new environment for remote execution. 

44 

45 Parameters 

46 ---------- 

47 name: str 

48 Human-readable name of the environment. 

49 config : dict 

50 Free-format dictionary that contains the benchmark environment 

51 configuration. Each config must have at least the "tunable_params" 

52 and the "const_args" sections. 

53 `RemoteEnv` must also have at least some of the following parameters: 

54 {setup, run, teardown, wait_boot} 

55 global_config : dict 

56 Free-format dictionary of global parameters (e.g., security credentials) 

57 to be mixed in into the "const_args" section of the local config. 

58 tunables : TunableGroups 

59 A collection of tunable parameters for *all* environments. 

60 service: Service 

61 An optional service object (e.g., providing methods to 

62 deploy or reboot a Host, VM, OS, etc.). 

63 """ 

64 super().__init__(name=name, config=config, global_config=global_config, 

65 tunables=tunables, service=service) 

66 

67 self._wait_boot = self.config.get("wait_boot", False) 

68 

69 assert self._service is not None and isinstance(self._service, SupportsRemoteExec), \ 

70 "RemoteEnv requires a service that supports remote execution operations" 

71 self._remote_exec_service: SupportsRemoteExec = self._service 

72 

73 if self._wait_boot: 

74 assert self._service is not None and isinstance(self._service, SupportsHostOps), \ 

75 "RemoteEnv requires a service that supports host operations" 

76 self._host_service: SupportsHostOps = self._service 

77 

78 def setup(self, tunables: TunableGroups, global_config: Optional[dict] = None) -> bool: 

79 """ 

80 Check if the environment is ready and set up the application 

81 and benchmarks on a remote host. 

82 

83 Parameters 

84 ---------- 

85 tunables : TunableGroups 

86 A collection of tunable OS and application parameters along with their 

87 values. Setting these parameters should not require an OS reboot. 

88 global_config : dict 

89 Free-format dictionary of global parameters of the environment 

90 that are not used in the optimization process. 

91 

92 Returns 

93 ------- 

94 is_success : bool 

95 True if operation is successful, false otherwise. 

96 """ 

97 if not super().setup(tunables, global_config): 

98 return False 

99 

100 if self._wait_boot: 

101 _LOG.info("Wait for the remote environment to start: %s", self) 

102 (status, params) = self._host_service.start_host(self._params) 

103 if status.is_pending(): 

104 (status, _) = self._host_service.wait_host_operation(params) 

105 if not status.is_succeeded(): 

106 return False 

107 

108 if self._script_setup: 

109 _LOG.info("Set up the remote environment: %s", self) 

110 (status, _timestamp, _output) = self._remote_exec(self._script_setup) 

111 _LOG.info("Remote set up complete: %s :: %s", self, status) 

112 self._is_ready = status.is_succeeded() 

113 else: 

114 self._is_ready = True 

115 

116 return self._is_ready 

117 

118 def run(self) -> Tuple[Status, datetime, Optional[Dict[str, TunableValue]]]: 

119 """ 

120 Runs the run script on the remote environment. 

121 

122 This can be used to, for instance, submit a new experiment to the 

123 remote application environment by (re)configuring an application and 

124 launching the benchmark, or run a script that collects the results. 

125 

126 Returns 

127 ------- 

128 (status, timestamp, output) : (Status, datetime, dict) 

129 3-tuple of (Status, timestamp, output) values, where `output` is a dict 

130 with the results or None if the status is not COMPLETED. 

131 If run script is a benchmark, then the score is usually expected to 

132 be in the `score` field. 

133 """ 

134 _LOG.info("Run script remotely on: %s", self) 

135 (status, timestamp, _) = result = super().run() 

136 if not (status.is_ready() and self._script_run): 

137 return result 

138 

139 (status, timestamp, output) = self._remote_exec(self._script_run) 

140 if status.is_succeeded() and output is not None: 

141 output = self._extract_stdout_results(output.get("stdout", "")) 

142 _LOG.info("Remote run complete: %s :: %s = %s", self, status, output) 

143 return (status, timestamp, output) 

144 

145 def teardown(self) -> None: 

146 """ 

147 Clean up and shut down the remote environment. 

148 """ 

149 if self._script_teardown: 

150 _LOG.info("Remote teardown: %s", self) 

151 (status, _timestamp, _output) = self._remote_exec(self._script_teardown) 

152 _LOG.info("Remote teardown complete: %s :: %s", self, status) 

153 super().teardown() 

154 

155 def _remote_exec(self, script: Iterable[str]) -> Tuple[Status, datetime, Optional[dict]]: 

156 """ 

157 Run a script on the remote host. 

158 

159 Parameters 

160 ---------- 

161 script : [str] 

162 List of commands to be executed on the remote host. 

163 

164 Returns 

165 ------- 

166 result : (Status, datetime, dict) 

167 3-tuple of Status, timestamp, and dict with the benchmark/script results. 

168 Status is one of {PENDING, SUCCEEDED, FAILED, TIMED_OUT} 

169 """ 

170 env_params = self._get_env_params() 

171 _LOG.debug("Submit script: %s with %s", self, env_params) 

172 (status, output) = self._remote_exec_service.remote_exec( 

173 script, config=self._params, env_params=env_params) 

174 _LOG.debug("Script submitted: %s %s :: %s", self, status, output) 

175 if status in {Status.PENDING, Status.SUCCEEDED}: 

176 (status, output) = self._remote_exec_service.get_remote_exec_results(output) 

177 _LOG.debug("Status: %s :: %s", status, output) 

178 # FIXME: get the timestamp from the remote environment! 

179 timestamp = datetime.now(UTC) 

180 return (status, timestamp, output)