Coverage for mlos_bench/mlos_bench/environments/remote/saas_env.py: 44%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Cloud-based (configurable) SaaS 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_ops_type import SupportsHostOps 

16from mlos_bench.services.types.remote_config_type import SupportsRemoteConfig 

17from mlos_bench.tunables.tunable_groups import TunableGroups 

18 

19_LOG = logging.getLogger(__name__) 

20 

21 

22class SaaSEnv(Environment): 

23 """ 

24 Cloud-based (configurable) SaaS environment. 

25 """ 

26 

27 def __init__(self, 

28 *, 

29 name: str, 

30 config: dict, 

31 global_config: Optional[dict] = None, 

32 tunables: Optional[TunableGroups] = None, 

33 service: Optional[Service] = None): 

34 """ 

35 Create a new environment for (configurable) cloud-based SaaS instance. 

36 

37 Parameters 

38 ---------- 

39 name: str 

40 Human-readable name of the environment. 

41 config : dict 

42 Free-format dictionary that contains the benchmark environment 

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

44 and the "const_args" sections. 

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 

52 (e.g., providing methods to configure the remote service). 

53 """ 

54 super().__init__(name=name, config=config, global_config=global_config, 

55 tunables=tunables, service=service) 

56 

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

58 "RemoteEnv requires a service that supports host operations" 

59 self._host_service: SupportsHostOps = self._service 

60 

61 assert self._service is not None and isinstance(self._service, SupportsRemoteConfig), \ 

62 "SaaSEnv requires a service that supports remote host configuration API" 

63 self._config_service: SupportsRemoteConfig = self._service 

64 

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

66 """ 

67 Update the configuration of a remote SaaS instance. 

68 

69 Parameters 

70 ---------- 

71 tunables : TunableGroups 

72 A collection of groups of tunable parameters along with the 

73 parameters' values. 

74 global_config : dict 

75 Free-format dictionary of global parameters of the environment 

76 that are not used in the optimization process. 

77 

78 Returns 

79 ------- 

80 is_success : bool 

81 True if operation is successful, false otherwise. 

82 """ 

83 _LOG.info("SaaS set up: %s :: %s", self, tunables) 

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

85 return False 

86 

87 (status, _) = self._config_service.configure( 

88 self._params, self._tunable_params.get_param_values()) 

89 if not status.is_succeeded(): 

90 return False 

91 

92 (status, res) = self._config_service.is_config_pending(self._params) 

93 if not status.is_succeeded(): 

94 return False 

95 

96 # Azure Flex DB instances currently require a VM reboot after reconfiguration. 

97 if res.get('isConfigPendingRestart') or res.get('isConfigPendingReboot'): 

98 _LOG.info("Restarting: %s", self) 

99 (status, params) = self._host_service.restart_host(self._params) 

100 if status.is_pending(): 

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

102 if not status.is_succeeded(): 

103 return False 

104 

105 _LOG.info("Wait to restart: %s", self) 

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

107 if status.is_pending(): 

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

109 

110 self._is_ready = status.is_succeeded() 

111 return self._is_ready