Coverage for mlos_bench/mlos_bench/services/remote/azure/azure_network_services.py: 94%
33 statements
« prev ^ index » next coverage.py v7.6.7, created at 2024-11-22 01:18 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""A collection Service functions for managing virtual networks on Azure."""
7import logging
8from typing import Any, Callable, Dict, List, Optional, Tuple, Union
10from mlos_bench.environments.status import Status
11from mlos_bench.services.base_service import Service
12from mlos_bench.services.remote.azure.azure_deployment_services import (
13 AzureDeploymentService,
14)
15from mlos_bench.services.types.network_provisioner_type import (
16 SupportsNetworkProvisioning,
17)
18from mlos_bench.util import merge_parameters
20_LOG = logging.getLogger(__name__)
23class AzureNetworkService(AzureDeploymentService, SupportsNetworkProvisioning):
24 """Helper methods to manage Virtual Networks on Azure."""
26 # Azure Compute REST API calls as described in
27 # https://learn.microsoft.com/en-us/rest/api/virtualnetwork/virtual-networks?view=rest-virtualnetwork-2023-05-01
29 # From: https://learn.microsoft.com/en-us/rest/api/virtualnetwork/virtual-networks?view=rest-virtualnetwork-2023-05-01 # pylint: disable=line-too-long # noqa
30 _URL_DEPROVISION = (
31 "https://management.azure.com"
32 "/subscriptions/{subscription}"
33 "/resourceGroups/{resource_group}"
34 "/providers/Microsoft.Network"
35 "/virtualNetwork/{vnet_name}"
36 "/delete"
37 "?api-version=2023-05-01"
38 )
40 def __init__(
41 self,
42 config: Optional[Dict[str, Any]] = None,
43 global_config: Optional[Dict[str, Any]] = None,
44 parent: Optional[Service] = None,
45 methods: Union[Dict[str, Callable], List[Callable], None] = None,
46 ):
47 """
48 Create a new instance of Azure Network services proxy.
50 Parameters
51 ----------
52 config : dict
53 Free-format dictionary that contains the benchmark environment
54 configuration.
55 global_config : dict
56 Free-format dictionary of global parameters.
57 parent : Service
58 Parent service that can provide mixin functions.
59 methods : Union[Dict[str, Callable], List[Callable], None]
60 New methods to register with the service.
61 """
62 super().__init__(
63 config,
64 global_config,
65 parent,
66 self.merge_methods(
67 methods,
68 [
69 # SupportsNetworkProvisioning
70 self.provision_network,
71 self.deprovision_network,
72 self.wait_network_deployment,
73 ],
74 ),
75 )
76 if not self._deploy_template:
77 raise ValueError(
78 "AzureNetworkService requires a deployment template:\n"
79 + f"config={config}\nglobal_config={global_config}"
80 )
82 def _set_default_params(self, params: dict) -> dict: # pylint: disable=no-self-use
83 # Try and provide a semi sane default for the deploymentName if not provided
84 # since this is a common way to set the deploymentName and can same some
85 # config work for the caller.
86 if "vnetName" in params and "deploymentName" not in params:
87 params["deploymentName"] = f"{params['vnetName']}-deployment"
88 _LOG.info(
89 "deploymentName missing from params. Defaulting to '%s'.",
90 params["deploymentName"],
91 )
92 return params
94 def wait_network_deployment(self, params: dict, *, is_setup: bool) -> Tuple[Status, dict]:
95 """
96 Waits for a pending operation on an Azure VM to resolve to SUCCEEDED or FAILED.
97 Return TIMED_OUT when timing out.
99 Parameters
100 ----------
101 params : dict
102 Flat dictionary of (key, value) pairs of tunable parameters.
103 is_setup : bool
104 If True, wait for VM being deployed; otherwise, wait for successful deprovisioning.
106 Returns
107 -------
108 result : (Status, dict)
109 A pair of Status and result.
110 Status is one of {PENDING, SUCCEEDED, FAILED, TIMED_OUT}
111 Result is info on the operation runtime if SUCCEEDED, otherwise {}.
112 """
113 return self._wait_deployment(params, is_setup=is_setup)
115 def provision_network(self, params: dict) -> Tuple[Status, dict]:
116 """
117 Deploy a virtual network, if necessary.
119 Parameters
120 ----------
121 params : dict
122 Flat dictionary of (key, value) pairs of tunable parameters.
123 NetworkEnv tunables are variable parameters that, together with the
124 NetworkEnv configuration, are sufficient to provision a virtual network.
126 Returns
127 -------
128 result : (Status, dict)
129 A pair of Status and result. The result is the input `params` plus the
130 parameters extracted from the response JSON, or {} if the status is FAILED.
131 Status is one of {PENDING, SUCCEEDED, FAILED}
132 """
133 return self._provision_resource(params)
135 def deprovision_network(self, params: dict, ignore_errors: bool = True) -> Tuple[Status, dict]:
136 """
137 Deprovisions the virtual network on Azure by deleting it.
139 Parameters
140 ----------
141 params : dict
142 Flat dictionary of (key, value) pairs of tunable parameters.
143 ignore_errors : bool
144 Whether to ignore errors (default) encountered during the operation
145 (e.g., due to dependent resources still in use).
147 Returns
148 -------
149 result : (Status, dict)
150 A pair of Status and result. The result is always {}.
151 Status is one of {PENDING, SUCCEEDED, FAILED}
152 """
153 params = self._set_default_params(params)
154 config = merge_parameters(
155 dest=self.config.copy(),
156 source=params,
157 required_keys=[
158 "subscription",
159 "resourceGroup",
160 "deploymentName",
161 "vnetName",
162 ],
163 )
164 _LOG.info("Deprovision Network: %s", config["vnetName"])
165 _LOG.info("Deprovision deployment: %s", config["deploymentName"])
166 (status, results) = self._azure_rest_api_post_helper(
167 config,
168 self._URL_DEPROVISION.format(
169 subscription=config["subscription"],
170 resource_group=config["resourceGroup"],
171 vnet_name=config["vnetName"],
172 ),
173 )
174 if ignore_errors and status == Status.FAILED:
175 _LOG.warning("Ignoring error: %s", results)
176 status = Status.SUCCEEDED
177 return (status, results)