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

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

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

6 

7import logging 

8 

9from mlos_bench.environments.base_environment import Environment 

10from mlos_bench.environments.status import Status 

11from mlos_bench.services.base_service import Service 

12from mlos_bench.services.types.host_ops_type import SupportsHostOps 

13from mlos_bench.services.types.os_ops_type import SupportsOSOps 

14from mlos_bench.tunables.tunable_groups import TunableGroups 

15 

16_LOG = logging.getLogger(__name__) 

17 

18 

19class OSEnv(Environment): 

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

21 

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

23 self, 

24 *, 

25 name: str, 

26 config: dict, 

27 global_config: dict | None = None, 

28 tunables: TunableGroups | None = None, 

29 service: Service | None = None, 

30 ): 

31 """ 

32 Create a new environment for remote execution. 

33 

34 Parameters 

35 ---------- 

36 name: str 

37 Human-readable name of the environment. 

38 config : dict 

39 Free-format dictionary that contains the benchmark environment 

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

41 and the "const_args" sections. 

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

43 {setup, run, teardown, wait_boot} 

44 global_config : dict 

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

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

47 tunables : TunableGroups 

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

49 service: Service 

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

51 deploy or reboot a VM, etc.). 

52 """ 

53 super().__init__( 

54 name=name, 

55 config=config, 

56 global_config=global_config, 

57 tunables=tunables, 

58 service=service, 

59 ) 

60 

61 assert self._service is not None and isinstance( 

62 self._service, SupportsHostOps 

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

64 self._host_service: SupportsHostOps = self._service 

65 

66 assert self._service is not None and isinstance( 

67 self._service, SupportsOSOps 

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

69 self._os_service: SupportsOSOps = self._service 

70 

71 def setup(self, tunables: TunableGroups, global_config: dict | None = None) -> bool: 

72 """ 

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

74 

75 Parameters 

76 ---------- 

77 tunables : TunableGroups 

78 A collection of groups of tunable parameters along with the 

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

80 together with the HostEnv configuration, are sufficient to provision 

81 and start a host. 

82 global_config : dict 

83 Free-format dictionary of global parameters of the environment 

84 that are not used in the optimization process. 

85 

86 Returns 

87 ------- 

88 is_success : bool 

89 True if operation is successful, false otherwise. 

90 """ 

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

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

93 return False 

94 

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

96 if status.is_pending(): 

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

98 

99 # TODO: configure OS settings here? 

100 

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

102 return self._is_ready 

103 

104 def teardown(self) -> None: 

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

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

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

108 if status.is_pending(): 

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

110 

111 super().teardown() 

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