Coverage for mlos_bench/mlos_bench/services/remote/azure/azure_network_services.py: 94%
34 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"""A collection Service functions for managing virtual networks on Azure."""
7import logging
8from collections.abc import Callable
9from typing import Any
11from mlos_bench.environments.status import Status
12from mlos_bench.services.base_service import Service
13from mlos_bench.services.remote.azure.azure_deployment_services import (
14 AzureDeploymentService,
15)
16from mlos_bench.services.types.network_provisioner_type import (
17 SupportsNetworkProvisioning,
18)
19from mlos_bench.util import merge_parameters
21_LOG = logging.getLogger(__name__)
24class AzureNetworkService(AzureDeploymentService, SupportsNetworkProvisioning):
25 """Helper methods to manage Virtual Networks on Azure."""
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 # pylint: disable=line-too-long # noqa
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__(
42 self,
43 config: dict[str, Any] | None = None,
44 global_config: dict[str, Any] | None = None,
45 parent: Service | None = None,
46 methods: dict[str, Callable] | list[Callable] | None = None,
47 ):
48 """
49 Create a new instance of Azure Network services proxy.
51 Parameters
52 ----------
53 config : dict
54 Free-format dictionary that contains the benchmark environment
55 configuration.
56 global_config : dict
57 Free-format dictionary of global parameters.
58 parent : Service
59 Parent service that can provide mixin functions.
60 methods : Union[dict[str, Callable], list[Callable], None]
61 New methods to register with the service.
62 """
63 super().__init__(
64 config,
65 global_config,
66 parent,
67 self.merge_methods(
68 methods,
69 [
70 # SupportsNetworkProvisioning
71 self.provision_network,
72 self.deprovision_network,
73 self.wait_network_deployment,
74 ],
75 ),
76 )
77 if not self._deploy_template:
78 raise ValueError(
79 "AzureNetworkService requires a deployment template:\n"
80 + f"config={config}\nglobal_config={global_config}"
81 )
83 def _set_default_params(self, params: dict) -> dict: # pylint: disable=no-self-use
84 # Try and provide a semi sane default for the deploymentName if not provided
85 # since this is a common way to set the deploymentName and can same some
86 # config work for the caller.
87 if "vnetName" in params and "deploymentName" not in params:
88 params["deploymentName"] = f"""{params["vnetName"]}-deployment"""
89 _LOG.info(
90 "deploymentName missing from params. Defaulting to '%s'.",
91 params["deploymentName"],
92 )
93 return params
95 def wait_network_deployment(self, params: dict, *, is_setup: bool) -> tuple[Status, dict]:
96 """
97 Waits for a pending operation on an Azure VM to resolve to SUCCEEDED or FAILED.
98 Return TIMED_OUT when timing out.
100 Parameters
101 ----------
102 params : dict
103 Flat dictionary of (key, value) pairs of tunable parameters.
104 is_setup : bool
105 If True, wait for VM being deployed; otherwise, wait for successful deprovisioning.
107 Returns
108 -------
109 result : (Status, dict)
110 A pair of Status and result.
111 Status is one of {PENDING, SUCCEEDED, FAILED, TIMED_OUT}
112 Result is info on the operation runtime if SUCCEEDED, otherwise {}.
113 """
114 return self._wait_deployment(params, is_setup=is_setup)
116 def provision_network(self, params: dict) -> tuple[Status, dict]:
117 """
118 Deploy a virtual network, if necessary.
120 Parameters
121 ----------
122 params : dict
123 Flat dictionary of (key, value) pairs of tunable parameters.
124 NetworkEnv tunables are variable parameters that, together with the
125 NetworkEnv configuration, are sufficient to provision a virtual network.
127 Returns
128 -------
129 result : (Status, dict)
130 A pair of Status and result. The result is the input `params` plus the
131 parameters extracted from the response JSON, or {} if the status is FAILED.
132 Status is one of {PENDING, SUCCEEDED, FAILED}
133 """
134 return self._provision_resource(params)
136 def deprovision_network(self, params: dict, ignore_errors: bool = True) -> tuple[Status, dict]:
137 """
138 Deprovisions the virtual network on Azure by deleting it.
140 Parameters
141 ----------
142 params : dict
143 Flat dictionary of (key, value) pairs of tunable parameters.
144 ignore_errors : bool
145 Whether to ignore errors (default) encountered during the operation
146 (e.g., due to dependent resources still in use).
148 Returns
149 -------
150 result : (Status, dict)
151 A pair of Status and result. The result is always {}.
152 Status is one of {PENDING, SUCCEEDED, FAILED}
153 """
154 params = self._set_default_params(params)
155 config = merge_parameters(
156 dest=self.config.copy(),
157 source=params,
158 required_keys=[
159 "subscription",
160 "resourceGroup",
161 "deploymentName",
162 "vnetName",
163 ],
164 )
165 _LOG.info("Deprovision Network: %s", config["vnetName"])
166 _LOG.info("Deprovision deployment: %s", config["deploymentName"])
167 (status, results) = self._azure_rest_api_post_helper(
168 config,
169 self._URL_DEPROVISION.format(
170 subscription=config["subscription"],
171 resource_group=config["resourceGroup"],
172 vnet_name=config["vnetName"],
173 ),
174 )
175 if ignore_errors and status == Status.FAILED:
176 _LOG.warning("Ignoring error: %s", results)
177 status = Status.SUCCEEDED
178 return (status, results)