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

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Cloud-based (configurable) SaaS environment.""" 

6 

7import logging 

8from typing import Optional 

9 

10from mlos_bench.environments.base_environment import Environment 

11from mlos_bench.services.base_service import Service 

12from mlos_bench.services.types.host_ops_type import SupportsHostOps 

13from mlos_bench.services.types.remote_config_type import SupportsRemoteConfig 

14from mlos_bench.tunables.tunable_groups import TunableGroups 

15 

16_LOG = logging.getLogger(__name__) 

17 

18 

19class SaaSEnv(Environment): 

20 """Cloud-based (configurable) SaaS environment.""" 

21 

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

23 self, 

24 *, 

25 name: str, 

26 config: dict, 

27 global_config: Optional[dict] = None, 

28 tunables: Optional[TunableGroups] = None, 

29 service: Optional[Service] = None, 

30 ): 

31 """ 

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

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 global_config : dict 

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

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

45 tunables : TunableGroups 

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

47 service: Service 

48 An optional service object 

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

50 """ 

51 super().__init__( 

52 name=name, 

53 config=config, 

54 global_config=global_config, 

55 tunables=tunables, 

56 service=service, 

57 ) 

58 

59 assert self._service is not None and isinstance( 

60 self._service, SupportsHostOps 

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

62 self._host_service: SupportsHostOps = self._service 

63 

64 assert self._service is not None and isinstance( 

65 self._service, SupportsRemoteConfig 

66 ), "SaaSEnv requires a service that supports remote host configuration API" 

67 self._config_service: SupportsRemoteConfig = self._service 

68 

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

70 """ 

71 Update the configuration of a remote SaaS instance. 

72 

73 Parameters 

74 ---------- 

75 tunables : TunableGroups 

76 A collection of groups of tunable parameters along with the 

77 parameters' values. 

78 global_config : dict 

79 Free-format dictionary of global parameters of the environment 

80 that are not used in the optimization process. 

81 

82 Returns 

83 ------- 

84 is_success : bool 

85 True if operation is successful, false otherwise. 

86 """ 

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

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

89 return False 

90 

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

92 self._params, 

93 self._tunable_params.get_param_values(), 

94 ) 

95 if not status.is_succeeded(): 

96 return False 

97 

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

99 if not status.is_succeeded(): 

100 return False 

101 

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

103 if res.get("isConfigPendingRestart") or res.get("isConfigPendingReboot"): 

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

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

106 if status.is_pending(): 

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

108 if not status.is_succeeded(): 

109 return False 

110 

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

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

113 if status.is_pending(): 

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

115 

116 self._is_ready = status.is_succeeded() 

117 return self._is_ready