Coverage for mlos_bench/mlos_bench/environments/remote/network_env.py: 45%
31 statements
« prev ^ index » next coverage.py v7.6.10, created at 2025-01-21 01:50 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""Network Environment."""
7import logging
9from mlos_bench.environments.base_environment import Environment
10from mlos_bench.services.base_service import Service
11from mlos_bench.services.types.network_provisioner_type import (
12 SupportsNetworkProvisioning,
13)
14from mlos_bench.tunables.tunable_groups import TunableGroups
16_LOG = logging.getLogger(__name__)
19class NetworkEnv(Environment):
20 """
21 Network Environment.
23 Used to model creating a virtual network (and network security group),
24 but no real tuning is expected for it ... yet.
25 """
27 def __init__( # pylint: disable=too-many-arguments
28 self,
29 *,
30 name: str,
31 config: dict,
32 global_config: dict | None = None,
33 tunables: TunableGroups | None = None,
34 service: Service | None = None,
35 ):
36 """
37 Create a new environment for network operations.
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__(
57 name=name,
58 config=config,
59 global_config=global_config,
60 tunables=tunables,
61 service=service,
62 )
64 # Virtual networks can be used for more than one experiment, so by default
65 # we don't attempt to deprovision them.
66 self._deprovision_on_teardown = config.get("deprovision_on_teardown", False)
68 assert self._service is not None and isinstance(
69 self._service, SupportsNetworkProvisioning
70 ), "NetworkEnv requires a service that supports network provisioning"
71 self._network_service: SupportsNetworkProvisioning = self._service
73 def setup(self, tunables: TunableGroups, global_config: dict | None = None) -> bool:
74 """
75 Check if network is ready. Provision, if necessary.
77 Parameters
78 ----------
79 tunables : TunableGroups
80 A collection of groups of tunable parameters along with the
81 parameters' values. NetworkEnv tunables are variable parameters that,
82 together with the NetworkEnv configuration, are sufficient to provision
83 and start a set of network resources (e.g., virtual network and network
84 security group).
85 global_config : dict
86 Free-format dictionary of global parameters of the environment
87 that are not used in the optimization process.
89 Returns
90 -------
91 is_success : bool
92 True if operation is successful, false otherwise.
93 """
94 _LOG.info("Network set up: %s :: %s", self, tunables)
95 if not super().setup(tunables, global_config):
96 return False
98 (status, params) = self._network_service.provision_network(self._params)
99 if status.is_pending():
100 (status, _) = self._network_service.wait_network_deployment(params, is_setup=True)
102 self._is_ready = status.is_succeeded()
103 return self._is_ready
105 def teardown(self) -> None:
106 """Shut down the Network and releases it."""
107 if not self._deprovision_on_teardown:
108 _LOG.info("Skipping Network deprovision: %s", self)
109 return
110 # Else
111 _LOG.info("Network tear down: %s", self)
112 (status, params) = self._network_service.deprovision_network(
113 self._params,
114 ignore_errors=True,
115 )
116 if status.is_pending():
117 (status, _) = self._network_service.wait_network_deployment(params, is_setup=False)
119 super().teardown()
120 _LOG.debug("Final status of Network deprovisioning: %s :: %s", self, status)