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

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""OS-level remote Environment on Azure.""" 

6 

7import logging 

8from typing import Optional 

9 

10from mlos_bench.environments.base_environment import Environment 

11from mlos_bench.environments.status import Status 

12from mlos_bench.services.base_service import Service 

13from mlos_bench.services.types.host_ops_type import SupportsHostOps 

14from mlos_bench.services.types.os_ops_type import SupportsOSOps 

15from mlos_bench.tunables.tunable_groups import TunableGroups 

16 

17_LOG = logging.getLogger(__name__) 

18 

19 

20class OSEnv(Environment): 

21 """OS Level Environment for a host.""" 

22 

23 def __init__( # pylint: disable=too-many-arguments 

24 self, 

25 *, 

26 name: str, 

27 config: dict, 

28 global_config: Optional[dict] = None, 

29 tunables: Optional[TunableGroups] = None, 

30 service: Optional[Service] = None, 

31 ): 

32 """ 

33 Create a new environment for remote execution. 

34 

35 Parameters 

36 ---------- 

37 name: str 

38 Human-readable name of the environment. 

39 config : dict 

40 Free-format dictionary that contains the benchmark environment 

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

42 and the "const_args" sections. 

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

44 {setup, run, teardown, wait_boot} 

45 global_config : dict 

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

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

48 tunables : TunableGroups 

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

50 service: Service 

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

52 deploy or reboot a VM, etc.). 

53 """ 

54 super().__init__( 

55 name=name, 

56 config=config, 

57 global_config=global_config, 

58 tunables=tunables, 

59 service=service, 

60 ) 

61 

62 assert self._service is not None and isinstance( 

63 self._service, SupportsHostOps 

64 ), "RemoteEnv requires a service that supports host operations" 

65 self._host_service: SupportsHostOps = self._service 

66 

67 assert self._service is not None and isinstance( 

68 self._service, SupportsOSOps 

69 ), "RemoteEnv requires a service that supports host operations" 

70 self._os_service: SupportsOSOps = self._service 

71 

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

73 """ 

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

75 

76 Parameters 

77 ---------- 

78 tunables : TunableGroups 

79 A collection of groups of tunable parameters along with the 

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

81 together with the HostEnv configuration, are sufficient to provision 

82 and start a host. 

83 global_config : dict 

84 Free-format dictionary of global parameters of the environment 

85 that are not used in the optimization process. 

86 

87 Returns 

88 ------- 

89 is_success : bool 

90 True if operation is successful, false otherwise. 

91 """ 

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

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

94 return False 

95 

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

97 if status.is_pending(): 

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

99 

100 # TODO: configure OS settings here? 

101 

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

103 return self._is_ready 

104 

105 def teardown(self) -> None: 

106 """Clean up and shut down the host without deprovisioning it.""" 

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

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

109 if status.is_pending(): 

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

111 

112 super().teardown() 

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