Coverage for mlos_bench/mlos_bench/environments/remote/os_env.py: 58%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6OS-level remote Environment on Azure. 

7""" 

8 

9from typing import Optional 

10 

11import logging 

12 

13from mlos_bench.environments.base_environment import Environment 

14from mlos_bench.environments.status import Status 

15from mlos_bench.services.base_service import Service 

16from mlos_bench.services.types.host_ops_type import SupportsHostOps 

17from mlos_bench.services.types.os_ops_type import SupportsOSOps 

18from mlos_bench.tunables.tunable_groups import TunableGroups 

19 

20_LOG = logging.getLogger(__name__) 

21 

22 

23class OSEnv(Environment): 

24 """ 

25 OS Level Environment for a host. 

26 """ 

27 

28 def __init__(self, 

29 *, 

30 name: str, 

31 config: dict, 

32 global_config: Optional[dict] = None, 

33 tunables: Optional[TunableGroups] = None, 

34 service: Optional[Service] = None): 

35 """ 

36 Create a new environment for remote execution. 

37 

38 Parameters 

39 ---------- 

40 name: str 

41 Human-readable name of the environment. 

42 config : dict 

43 Free-format dictionary that contains the benchmark environment 

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

45 and the "const_args" sections. 

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

47 {setup, run, teardown, wait_boot} 

48 global_config : dict 

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

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

51 tunables : TunableGroups 

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

53 service: Service 

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

55 deploy or reboot a VM, etc.). 

56 """ 

57 super().__init__(name=name, config=config, global_config=global_config, tunables=tunables, service=service) 

58 

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

60 "RemoteEnv requires a service that supports host operations" 

61 self._host_service: SupportsHostOps = self._service 

62 

63 assert self._service is not None and isinstance(self._service, SupportsOSOps), \ 

64 "RemoteEnv requires a service that supports host operations" 

65 self._os_service: SupportsOSOps = self._service 

66 

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

68 """ 

69 Check if the host is up and running; boot it, if necessary. 

70 

71 Parameters 

72 ---------- 

73 tunables : TunableGroups 

74 A collection of groups of tunable parameters along with the 

75 parameters' values. HostEnv tunables are variable parameters that, 

76 together with the HostEnv configuration, are sufficient to provision 

77 and start a host. 

78 global_config : dict 

79 Free-format dictionary of global parameters of the environment 

80 that are not used in the optimization process. 

81 

82 Returns 

83 ------- 

84 is_success : bool 

85 True if operation is successful, false otherwise. 

86 """ 

87 _LOG.info("OS set up: %s :: %s", self, tunables) 

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

89 return False 

90 

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

92 if status.is_pending(): 

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

94 

95 # TODO: configure OS settings here? 

96 

97 self._is_ready = status in {Status.SUCCEEDED, Status.READY} 

98 return self._is_ready 

99 

100 def teardown(self) -> None: 

101 """ 

102 Clean up and shut down the host without deprovisioning it. 

103 """ 

104 _LOG.info("OS tear down: %s", self) 

105 (status, params) = self._os_service.shutdown(self._params) 

106 if status.is_pending(): 

107 (status, _) = self._os_service.wait_os_operation(params) 

108 

109 super().teardown() 

110 _LOG.debug("Final status of OS stopping: %s :: %s", self, status)