Coverage for mlos_bench/mlos_bench/environments/remote/host_env.py: 52%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Remote host Environment. 

7""" 

8 

9from typing import Optional 

10 

11import logging 

12 

13from mlos_bench.environments.base_environment import Environment 

14from mlos_bench.services.base_service import Service 

15from mlos_bench.services.types.host_provisioner_type import SupportsHostProvisioning 

16from mlos_bench.tunables.tunable_groups import TunableGroups 

17 

18_LOG = logging.getLogger(__name__) 

19 

20 

21class HostEnv(Environment): 

22 """ 

23 Remote host environment. 

24 """ 

25 

26 def __init__(self, 

27 *, 

28 name: str, 

29 config: dict, 

30 global_config: Optional[dict] = None, 

31 tunables: Optional[TunableGroups] = None, 

32 service: Optional[Service] = None): 

33 """ 

34 Create a new environment for host operations. 

35 

36 Parameters 

37 ---------- 

38 name: str 

39 Human-readable name of the environment. 

40 config : dict 

41 Free-format dictionary that contains the benchmark environment 

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

43 and the "const_args" sections. 

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/host, etc.). 

52 """ 

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

54 

55 assert self._service is not None and isinstance(self._service, SupportsHostProvisioning), \ 

56 "HostEnv requires a service that supports host provisioning operations" 

57 self._host_service: SupportsHostProvisioning = self._service 

58 

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

60 """ 

61 Check if host is ready. (Re)provision and start it, if necessary. 

62 

63 Parameters 

64 ---------- 

65 tunables : TunableGroups 

66 A collection of groups of tunable parameters along with the 

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

68 together with the HostEnv configuration, are sufficient to provision 

69 and start a Host. 

70 global_config : dict 

71 Free-format dictionary of global parameters of the environment 

72 that are not used in the optimization process. 

73 

74 Returns 

75 ------- 

76 is_success : bool 

77 True if operation is successful, false otherwise. 

78 """ 

79 _LOG.info("Host set up: %s :: %s", self, tunables) 

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

81 return False 

82 

83 (status, params) = self._host_service.provision_host(self._params) 

84 if status.is_pending(): 

85 (status, _) = self._host_service.wait_host_deployment(params, is_setup=True) 

86 

87 self._is_ready = status.is_succeeded() 

88 return self._is_ready 

89 

90 def teardown(self) -> None: 

91 """ 

92 Shut down the Host and release it. 

93 """ 

94 _LOG.info("Host tear down: %s", self) 

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

96 if status.is_pending(): 

97 (status, _) = self._host_service.wait_host_deployment(params, is_setup=False) 

98 

99 super().teardown() 

100 _LOG.debug("Final status of Host deprovisioning: %s :: %s", self, status)