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

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

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

12from mlos_bench.services.types.remote_config_type import SupportsRemoteConfig 

13from mlos_bench.tunables.tunable_groups import TunableGroups 

14 

15_LOG = logging.getLogger(__name__) 

16 

17 

18class SaaSEnv(Environment): 

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

20 

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

22 self, 

23 *, 

24 name: str, 

25 config: dict, 

26 global_config: dict | None = None, 

27 tunables: TunableGroups | None = None, 

28 service: Service | None = None, 

29 ): 

30 """ 

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

32 

33 Parameters 

34 ---------- 

35 name: str 

36 Human-readable name of the environment. 

37 config : dict 

38 Free-format dictionary that contains the benchmark environment 

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

40 and the "const_args" sections. 

41 global_config : dict 

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

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

44 tunables : TunableGroups 

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

46 service: Service 

47 An optional service object 

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

49 """ 

50 super().__init__( 

51 name=name, 

52 config=config, 

53 global_config=global_config, 

54 tunables=tunables, 

55 service=service, 

56 ) 

57 

58 assert self._service is not None and isinstance( 

59 self._service, SupportsHostOps 

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

61 self._host_service: SupportsHostOps = self._service 

62 

63 assert self._service is not None and isinstance( 

64 self._service, SupportsRemoteConfig 

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

66 self._config_service: SupportsRemoteConfig = self._service 

67 

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

69 """ 

70 Update the configuration of a remote SaaS instance. 

71 

72 Parameters 

73 ---------- 

74 tunables : TunableGroups 

75 A collection of groups of tunable parameters along with the 

76 parameters' values. 

77 global_config : dict 

78 Free-format dictionary of global parameters of the environment 

79 that are not used in the optimization process. 

80 

81 Returns 

82 ------- 

83 is_success : bool 

84 True if operation is successful, false otherwise. 

85 """ 

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

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

88 return False 

89 

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

91 self._params, 

92 self._tunable_params.get_param_values(), 

93 ) 

94 if not status.is_succeeded(): 

95 return False 

96 

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

98 if not status.is_succeeded(): 

99 return False 

100 

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

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

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

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

105 if status.is_pending(): 

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

107 if not status.is_succeeded(): 

108 return False 

109 

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

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

112 if status.is_pending(): 

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

114 

115 self._is_ready = status.is_succeeded() 

116 return self._is_ready