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

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Network 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.network_provisioner_type import SupportsNetworkProvisioning 

16from mlos_bench.tunables.tunable_groups import TunableGroups 

17 

18_LOG = logging.getLogger(__name__) 

19 

20 

21class NetworkEnv(Environment): 

22 """ 

23 Network Environment. 

24 

25 Used to model creating a virtual network (and network security group), 

26 but no real tuning is expected for it ... yet. 

27 """ 

28 

29 def __init__(self, 

30 *, 

31 name: str, 

32 config: dict, 

33 global_config: Optional[dict] = None, 

34 tunables: Optional[TunableGroups] = None, 

35 service: Optional[Service] = None): 

36 """ 

37 Create a new environment for network operations. 

38 

39 Parameters 

40 ---------- 

41 name: str 

42 Human-readable name of the environment. 

43 config : dict 

44 Free-format dictionary that contains the benchmark environment 

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

46 and the "const_args" sections. 

47 global_config : dict 

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

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

50 tunables : TunableGroups 

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

52 service: Service 

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

54 deploy a network, etc.). 

55 """ 

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

57 

58 # Virtual networks can be used for more than one experiment, so by default 

59 # we don't attempt to deprovision them. 

60 self._deprovision_on_teardown = config.get("deprovision_on_teardown", False) 

61 

62 assert self._service is not None and isinstance(self._service, SupportsNetworkProvisioning), \ 

63 "NetworkEnv requires a service that supports network provisioning" 

64 self._network_service: SupportsNetworkProvisioning = self._service 

65 

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

67 """ 

68 Check if network is ready. Provision, if necessary. 

69 

70 Parameters 

71 ---------- 

72 tunables : TunableGroups 

73 A collection of groups of tunable parameters along with the 

74 parameters' values. NetworkEnv tunables are variable parameters that, 

75 together with the NetworkEnv configuration, are sufficient to provision 

76 and start a set of network resources (e.g., virtual network and network 

77 security group). 

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("Network set up: %s :: %s", self, tunables) 

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

89 return False 

90 

91 (status, params) = self._network_service.provision_network(self._params) 

92 if status.is_pending(): 

93 (status, _) = self._network_service.wait_network_deployment(params, is_setup=True) 

94 

95 self._is_ready = status.is_succeeded() 

96 return self._is_ready 

97 

98 def teardown(self) -> None: 

99 """ 

100 Shut down the Network and releases it. 

101 """ 

102 if not self._deprovision_on_teardown: 

103 _LOG.info("Skipping Network deprovision: %s", self) 

104 return 

105 # Else 

106 _LOG.info("Network tear down: %s", self) 

107 (status, params) = self._network_service.deprovision_network(self._params, ignore_errors=True) 

108 if status.is_pending(): 

109 (status, _) = self._network_service.wait_network_deployment(params, is_setup=False) 

110 

111 super().teardown() 

112 _LOG.debug("Final status of Network deprovisioning: %s :: %s", self, status)