Coverage for mlos_bench/mlos_bench/environments/remote/network_env.py: 47%
32 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-10-07 01:52 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""Network Environment."""
7import logging
8from typing import Optional
10from mlos_bench.environments.base_environment import Environment
11from mlos_bench.services.base_service import Service
12from mlos_bench.services.types.network_provisioner_type import (
13 SupportsNetworkProvisioning,
14)
15from mlos_bench.tunables.tunable_groups import TunableGroups
17_LOG = logging.getLogger(__name__)
20class NetworkEnv(Environment):
21 """
22 Network Environment.
24 Used to model creating a virtual network (and network security group),
25 but no real tuning is expected for it ... yet.
26 """
28 def __init__( # pylint: disable=too-many-arguments
29 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 """
38 Create a new environment for network operations.
40 Parameters
41 ----------
42 name: str
43 Human-readable name of the environment.
44 config : dict
45 Free-format dictionary that contains the benchmark environment
46 configuration. Each config must have at least the "tunable_params"
47 and the "const_args" sections.
48 global_config : dict
49 Free-format dictionary of global parameters (e.g., security credentials)
50 to be mixed in into the "const_args" section of the local config.
51 tunables : TunableGroups
52 A collection of tunable parameters for *all* environments.
53 service: Service
54 An optional service object (e.g., providing methods to
55 deploy a network, etc.).
56 """
57 super().__init__(
58 name=name,
59 config=config,
60 global_config=global_config,
61 tunables=tunables,
62 service=service,
63 )
65 # Virtual networks can be used for more than one experiment, so by default
66 # we don't attempt to deprovision them.
67 self._deprovision_on_teardown = config.get("deprovision_on_teardown", False)
69 assert self._service is not None and isinstance(
70 self._service, SupportsNetworkProvisioning
71 ), "NetworkEnv requires a service that supports network provisioning"
72 self._network_service: SupportsNetworkProvisioning = self._service
74 def setup(self, tunables: TunableGroups, global_config: Optional[dict] = None) -> bool:
75 """
76 Check if network is ready. Provision, if necessary.
78 Parameters
79 ----------
80 tunables : TunableGroups
81 A collection of groups of tunable parameters along with the
82 parameters' values. NetworkEnv tunables are variable parameters that,
83 together with the NetworkEnv configuration, are sufficient to provision
84 and start a set of network resources (e.g., virtual network and network
85 security group).
86 global_config : dict
87 Free-format dictionary of global parameters of the environment
88 that are not used in the optimization process.
90 Returns
91 -------
92 is_success : bool
93 True if operation is successful, false otherwise.
94 """
95 _LOG.info("Network set up: %s :: %s", self, tunables)
96 if not super().setup(tunables, global_config):
97 return False
99 (status, params) = self._network_service.provision_network(self._params)
100 if status.is_pending():
101 (status, _) = self._network_service.wait_network_deployment(params, is_setup=True)
103 self._is_ready = status.is_succeeded()
104 return self._is_ready
106 def teardown(self) -> None:
107 """Shut down the Network and releases it."""
108 if not self._deprovision_on_teardown:
109 _LOG.info("Skipping Network deprovision: %s", self)
110 return
111 # Else
112 _LOG.info("Network tear down: %s", self)
113 (status, params) = self._network_service.deprovision_network(
114 self._params,
115 ignore_errors=True,
116 )
117 if status.is_pending():
118 (status, _) = self._network_service.wait_network_deployment(params, is_setup=False)
120 super().teardown()
121 _LOG.debug("Final status of Network deprovisioning: %s :: %s", self, status)