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