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

27 statements  

« prev     ^ index     » next       coverage.py v7.6.10, created at 2025-01-21 01:50 +0000

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Remote host Environment.""" 

6 

7import logging 

8 

9from mlos_bench.environments.base_environment import Environment 

10from mlos_bench.services.base_service import Service 

11from mlos_bench.services.types.host_provisioner_type import SupportsHostProvisioning 

12from mlos_bench.tunables.tunable_groups import TunableGroups 

13 

14_LOG = logging.getLogger(__name__) 

15 

16 

17class HostEnv(Environment): 

18 """Remote host environment.""" 

19 

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

21 self, 

22 *, 

23 name: str, 

24 config: dict, 

25 global_config: dict | None = None, 

26 tunables: TunableGroups | None = None, 

27 service: Service | None = None, 

28 ): 

29 """ 

30 Create a new environment for host operations. 

31 

32 Parameters 

33 ---------- 

34 name: str 

35 Human-readable name of the environment. 

36 config : dict 

37 Free-format dictionary that contains the benchmark environment 

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

39 and the "const_args" sections. 

40 global_config : dict 

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

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

43 tunables : TunableGroups 

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

45 service: Service 

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

47 deploy or reboot a VM/host, etc.). 

48 """ 

49 super().__init__( 

50 name=name, 

51 config=config, 

52 global_config=global_config, 

53 tunables=tunables, 

54 service=service, 

55 ) 

56 

57 assert self._service is not None and isinstance( 

58 self._service, SupportsHostProvisioning 

59 ), "HostEnv requires a service that supports host provisioning operations" 

60 self._host_service: SupportsHostProvisioning = self._service 

61 

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

63 """ 

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

65 

66 Parameters 

67 ---------- 

68 tunables : TunableGroups 

69 A collection of groups of tunable parameters along with the 

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

71 together with the HostEnv configuration, are sufficient to provision 

72 and start a Host. 

73 global_config : dict 

74 Free-format dictionary of global parameters of the environment 

75 that are not used in the optimization process. 

76 

77 Returns 

78 ------- 

79 is_success : bool 

80 True if operation is successful, false otherwise. 

81 """ 

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

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

84 return False 

85 

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

87 if status.is_pending(): 

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

89 

90 self._is_ready = status.is_succeeded() 

91 return self._is_ready 

92 

93 def teardown(self) -> None: 

94 """Shut down the Host and release it.""" 

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

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

97 if status.is_pending(): 

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

99 

100 super().teardown() 

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