From 073cb5b6adaa76cfda5c27797a707e367f29df52 Mon Sep 17 00:00:00 2001 From: Kelly Kostopoulou Date: Sun, 7 Jul 2024 06:02:57 +0000 Subject: [PATCH] patch for managed identity storage authentication --- .../config/cli/azure-redis-bench.jsonc | 6 +- .../environments/apps/redis/redis.jsonc | 21 ++- .../os/linux/boot/boot-ubuntu.jsonc | 27 ++- .../os/linux/runtime/setup-ubuntu.jsonc | 7 +- .../experiments/experiment_RedisBench.jsonc | 5 + .../azure-deployment-service-subschema.json | 4 + .../azure-fileshare-service-subschema.json | 8 +- .../remote/azure/service-fileshare.jsonc | 3 +- .../remote/azure/service-linux-vm-ops.jsonc | 1 + .../environments/remote/host_env.py | 6 +- .../remote/azure/azure_deployment_services.py | 176 ++++++++++++++++++ .../services/remote/azure/azure_fileshare.py | 21 ++- .../remote/azure/azure_vm_services.py | 48 ++++- .../services/types/host_provisioner_type.py | 32 ++++ 14 files changed, 335 insertions(+), 30 deletions(-) diff --git a/mlos_bench/mlos_bench/config/cli/azure-redis-bench.jsonc b/mlos_bench/mlos_bench/config/cli/azure-redis-bench.jsonc index 5a3163e102..f06f8c8b04 100644 --- a/mlos_bench/mlos_bench/config/cli/azure-redis-bench.jsonc +++ b/mlos_bench/mlos_bench/config/cli/azure-redis-bench.jsonc @@ -14,9 +14,9 @@ "environment": "environments/root/root-azure-redis.jsonc", "storage": "storage/sqlite.jsonc", - "globals": [ - "global_config_azure.jsonc" - ], + // "globals": [ + // "global_config_azure.jsonc" + // ], "teardown": false, diff --git a/mlos_bench/mlos_bench/config/environments/apps/redis/redis.jsonc b/mlos_bench/mlos_bench/config/environments/apps/redis/redis.jsonc index dc44871abd..5e3eebc269 100644 --- a/mlos_bench/mlos_bench/config/environments/apps/redis/redis.jsonc +++ b/mlos_bench/mlos_bench/config/environments/apps/redis/redis.jsonc @@ -51,19 +51,22 @@ "trial_id" ], "setup": [ - "$mountPoint/$experiment_id/$trial_id/scripts/setup-workload.sh", - "$mountPoint/$experiment_id/$trial_id/scripts/setup-app.sh" + "sudo /usr/local/bin/azure_file_share_service --function 'download' --local_path $mountPoint/$experiment_id/$trial_id/scripts/ --remote_path $experiment_id/$trial_id/scripts", + "sudo /usr/local/bin/azure_file_share_service --function 'download' --local_path $mountPoint/$experiment_id/$trial_id/input/redis.cfg --remote_path $experiment_id/$trial_id/input/redis.cfg", + "sudo bash $mountPoint/$experiment_id/$trial_id/scripts/setup-workload.sh", + "sudo bash $mountPoint/$experiment_id/$trial_id/scripts/setup-app.sh" ], "run": [ - "mkdir -p /tmp/mlos_bench/output/", - "$mountPoint/$experiment_id/$trial_id/scripts/run-workload.sh", - "mkdir -p $mountPoint/$experiment_id/$trial_id/output/", - "cp -r /tmp/mlos_bench/output/* $mountPoint/$experiment_id/$trial_id/output/" + "sudo mkdir -p /tmp/mlos_bench/output/", + "sudo bash $mountPoint/$experiment_id/$trial_id/scripts/run-workload.sh", + "sudo /usr/local/bin/azure_file_share_service --function 'upload' --local_path /tmp/mlos_bench/output/ --remote_path $experiment_id/$trial_id/output" + // "mkdir -p $mountPoint/$experiment_id/$trial_id/output/", + // "cp -r /tmp/mlos_bench/output/* $mountPoint/$experiment_id/$trial_id/output/" ], "teardown": [ - "$mountPoint/$experiment_id/$trial_id/scripts/cleanup-workload.sh", - "$mountPoint/$experiment_id/$trial_id/scripts/cleanup-app.sh", - "rm -r /tmp/mlos_bench/" + "sudo bash $mountPoint/$experiment_id/$trial_id/scripts/cleanup-workload.sh", + "sudo bash $mountPoint/$experiment_id/$trial_id/scripts/cleanup-app.sh" + // "rm -r /tmp/mlos_bench/" ] } }, diff --git a/mlos_bench/mlos_bench/config/environments/os/linux/boot/boot-ubuntu.jsonc b/mlos_bench/mlos_bench/config/environments/os/linux/boot/boot-ubuntu.jsonc index 51a5e378fc..98efb4eaae 100644 --- a/mlos_bench/mlos_bench/config/environments/os/linux/boot/boot-ubuntu.jsonc +++ b/mlos_bench/mlos_bench/config/environments/os/linux/boot/boot-ubuntu.jsonc @@ -52,34 +52,51 @@ "config": { "required_args": [ "vmName", + "tenantId", + "managedIdentityClientId", "storageAccountName", "storageFileShareName", - "storageAccountKey", "mountPoint", "experiment_id", "trial_id" ], "wait_boot": true, "shell_env_params": [ + "tenantId", + "managedIdentityClientId", "storageAccountName", "storageFileShareName", - "storageAccountKey", "mountPoint", "experiment_id", "trial_id" ], "setup": [ "sudo mkdir -p $mountPoint", - "sudo mount -t cifs //$storageAccountName.file.core.windows.net/$storageFileShareName $mountPoint -o username=$storageAccountName,password=\"$storageAccountKey\",serverino,nosharesock,actimeo=30", + "sudo chmod +x $mountPoint", + "sudo apt install software-properties-common -y", + "sudo add-apt-repository ppa:deadsnakes/ppa -y", + "sudo apt-get update -y", + "sudo apt-get install python3.8 -y", + "sudo apt install python3-pip -y", + "sudo update-alternatives --install /usr/bin/python python /usr/bin/python3.8 1", + "sudo python -m pip install --upgrade pip setuptools wheel", + "sudo python -m pip install azure-file-share-service", + "echo storageAccountName=$storageAccountName | sudo tee -a /etc/environment", + "echo storageFileShareName=$storageFileShareName | sudo tee -a /etc/environment", + "echo managedIdentityClientId=$managedIdentityClientId | sudo tee -a /etc/environment", + "echo tenantId=$tenantId | sudo tee -a /etc/environment", + "source /etc/environment" + + // "sudo mount -t cifs //$storageAccountName.file.core.windows.net/$storageFileShareName $mountPoint -o username=$storageAccountName,password=\"$storageAccountKey\",serverino,nosharesock,actimeo=30", // TODO: Apply GRUB parameters from config on shared storage at: // "$mountPoint/$experiment_id/$trial_id/input/grub.cfg" // "sudo update-grub", // "sudo shutdown -r now" // Reboot to apply the parameters - "cat $mountPoint/$experiment_id/$trial_id/input/grub.cfg" + // "cat $mountPoint/$experiment_id/$trial_id/input/grub.cfg" ], "teardown": [ // "sudo umount $mountPoint" - "echo sudo umount $mountPoint" + "echo sudo rm -rf $mountPoint" ] } } diff --git a/mlos_bench/mlos_bench/config/environments/os/linux/runtime/setup-ubuntu.jsonc b/mlos_bench/mlos_bench/config/environments/os/linux/runtime/setup-ubuntu.jsonc index 64eec5e1c7..2314198649 100644 --- a/mlos_bench/mlos_bench/config/environments/os/linux/runtime/setup-ubuntu.jsonc +++ b/mlos_bench/mlos_bench/config/environments/os/linux/runtime/setup-ubuntu.jsonc @@ -56,7 +56,6 @@ "vmName", "storageAccountName", "storageFileShareName", - "storageAccountKey", "mountPoint", "experiment_id", "trial_id" @@ -65,15 +64,15 @@ "shell_env_params": [ "storageAccountName", "storageFileShareName", - "storageAccountKey", "mountPoint", "experiment_id", "trial_id" ], "setup": [ "sudo mkdir -p $mountPoint", - "sudo mount -t cifs //$storageAccountName.file.core.windows.net/$storageFileShareName $mountPoint -o username=$storageAccountName,password=\"$storageAccountKey\",serverino,nosharesock,actimeo=30", - "sudo $mountPoint/$experiment_id/$trial_id/input/config-kernel.sh" + "sudo /usr/local/bin/azure_file_share_service --function 'download' --local_path $mountPoint/$experiment_id/$trial_id/input/ --remote_path $experiment_id/$trial_id/input", + "sudo bash $mountPoint/$experiment_id/$trial_id/input/config-kernel.sh" + // "sudo mount -t cifs //$storageAccountName.file.core.windows.net/$storageFileShareName $mountPoint -o username=$storageAccountName,password=\"$storageAccountKey\",serverino,nosharesock,actimeo=30", ] } } diff --git a/mlos_bench/mlos_bench/config/experiments/experiment_RedisBench.jsonc b/mlos_bench/mlos_bench/config/experiments/experiment_RedisBench.jsonc index 3da9e89e1c..f66bf45d17 100644 --- a/mlos_bench/mlos_bench/config/experiments/experiment_RedisBench.jsonc +++ b/mlos_bench/mlos_bench/config/experiments/experiment_RedisBench.jsonc @@ -10,6 +10,11 @@ "deploymentName": "$experiment_id", "vmName": "os-autotune-linux-vm", + "subscription": "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx", + "managedIdentityClientId": "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx", + "managedIdentityName": "mlos-managed-identity", + "tenantId": "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx", + "resourceGroup": "os-autotune", "location": "westus2", diff --git a/mlos_bench/mlos_bench/config/schemas/services/remote/azure/azure-deployment-service-subschema.json b/mlos_bench/mlos_bench/config/schemas/services/remote/azure/azure-deployment-service-subschema.json index 8abf7dd678..6152f9ae3e 100644 --- a/mlos_bench/mlos_bench/config/schemas/services/remote/azure/azure-deployment-service-subschema.json +++ b/mlos_bench/mlos_bench/config/schemas/services/remote/azure/azure-deployment-service-subschema.json @@ -20,6 +20,10 @@ "description": "The name of the resource group to place the deployment in (typically provided in the global config in order to omit from source control).", "type": "string" }, + "managedIdentityName": { + "description": "The name of the managed identity used for storage authentication (typically provided in the global config in order to omit from source control).", + "type": "string" + }, "deploymentTemplatePath": { "description": "Path to an ARM template file, or null if it should be skipped.", "type": ["string", "null"], diff --git a/mlos_bench/mlos_bench/config/schemas/services/remote/azure/azure-fileshare-service-subschema.json b/mlos_bench/mlos_bench/config/schemas/services/remote/azure/azure-fileshare-service-subschema.json index 283dc75ee5..cdf08ff9d0 100644 --- a/mlos_bench/mlos_bench/config/schemas/services/remote/azure/azure-fileshare-service-subschema.json +++ b/mlos_bench/mlos_bench/config/schemas/services/remote/azure/azure-fileshare-service-subschema.json @@ -28,8 +28,12 @@ "description": "Azure storage file share name.", "type": "string" }, - "storageAccountKey": { - "description": "Azure storage account key (typically provided in the global config in order to omit from source control).", + "managedIdentityClientId": { + "description": "Azure Managed Identity Client Id (typically provided in the global config in order to omit from source control).", + "type": "string" + }, + "tenantId": { + "description": "Azure Tenant Id (typically provided in the global config in order to omit from source control).", "type": "string" } }, diff --git a/mlos_bench/mlos_bench/config/services/remote/azure/service-fileshare.jsonc b/mlos_bench/mlos_bench/config/services/remote/azure/service-fileshare.jsonc index d800494c4c..973f24df22 100644 --- a/mlos_bench/mlos_bench/config/services/remote/azure/service-fileshare.jsonc +++ b/mlos_bench/mlos_bench/config/services/remote/azure/service-fileshare.jsonc @@ -3,6 +3,7 @@ "config": { "storageAccountName": "PLACEHOLDER; e.g.: osatsharedstorage", "storageFileShareName": "PLACEHOLDER; e.g.: os-autotune-file-share", - "storageAccountKey": "PLACEHOLDER; comes from global config" + "managedIdentityClientId": "PLACEHOLDER; comes from global config", + "tenantId": "PLACEHOLDER; comes from global config" } } diff --git a/mlos_bench/mlos_bench/config/services/remote/azure/service-linux-vm-ops.jsonc b/mlos_bench/mlos_bench/config/services/remote/azure/service-linux-vm-ops.jsonc index 31680698a1..ac9b72e5d3 100644 --- a/mlos_bench/mlos_bench/config/services/remote/azure/service-linux-vm-ops.jsonc +++ b/mlos_bench/mlos_bench/config/services/remote/azure/service-linux-vm-ops.jsonc @@ -10,6 +10,7 @@ // can be overridden by the parameters pushed from the caller Environment. "subscription": "PLACEHOLDER; AZURE SUBSCRIPTION ID", "resourceGroup": "PLACEHOLDER; e.g., os-autotune", + "managedIdentityName": "PLACEHOLDER; e.g., mlos-managed-identity", "deploymentTemplatePath": "services/remote/azure/arm-templates/azuredeploy-ubuntu-vm.jsonc", diff --git a/mlos_bench/mlos_bench/environments/remote/host_env.py b/mlos_bench/mlos_bench/environments/remote/host_env.py index 4b63e47278..1135e9adf4 100644 --- a/mlos_bench/mlos_bench/environments/remote/host_env.py +++ b/mlos_bench/mlos_bench/environments/remote/host_env.py @@ -84,7 +84,11 @@ def setup(self, tunables: TunableGroups, global_config: Optional[dict] = None) - if status.is_pending(): (status, _) = self._host_service.wait_host_deployment(params, is_setup=True) - self._is_ready = status.is_succeeded() + (status_id_assign, params) = self._host_service.assign_managed_identity(self._params) + if status_id_assign.is_pending(): + (status_id_assign, _) = self._host_service.wait_host_managed_identity_assignment(params) + + self._is_ready = status.is_succeeded() and status_id_assign.is_succeeded() return self._is_ready def teardown(self) -> None: diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_deployment_services.py b/mlos_bench/mlos_bench/services/remote/azure/azure_deployment_services.py index 187b7c055b..992295d739 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_deployment_services.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_deployment_services.py @@ -48,6 +48,15 @@ class AzureDeploymentService(Service, metaclass=abc.ABCMeta): "?api-version=2022-05-01" ) + _URL_ASSIGN_IDENTITY = ( + "https://management.azure.com" + + "/subscriptions/{subscription}" + + "/resourceGroups/{resource_group}" + + "/providers/Microsoft.Compute" + + "/virtualMachines/{vm_name}" + + "?api-version=2022-03-01" + ) + def __init__(self, config: Optional[Dict[str, Any]] = None, global_config: Optional[Dict[str, Any]] = None, @@ -73,6 +82,7 @@ def __init__(self, check_required_params(self.config, [ "subscription", "resourceGroup", + "managedIdentityName" ]) # These parameters can come from command line as strings, so conversion is needed. @@ -392,6 +402,101 @@ def _check_deployment(self, params: dict) -> Tuple[Status, dict]: # pylint: di _LOG.error("Response: %s :: %s", response, response.text) return (Status.FAILED, {}) + def _wait_host_managed_identity_assignment(self, params: dict) -> Tuple[Status, dict]: + """ + Waits for a pending operation on an Azure resource to resolve to SUCCEEDED or FAILED. + Return TIMED_OUT when timing out. + + Parameters + ---------- + params : dict + Flat dictionary of (key, value) pairs of tunable parameters. + + Returns + ------- + result : (Status, dict) + A pair of Status and result. + Status is one of {PENDING, SUCCEEDED, FAILED, TIMED_OUT} + Result is info on the operation runtime if SUCCEEDED, otherwise {}. + """ + params = self._set_default_params(params) + _LOG.info("Wait for %s to %s", self._deploy_params["vmName"], "be assigned managed identity") + return self._wait_while(self._check_managed_identity_assinment, Status.PENDING, params) + + def _check_managed_identity_assinment(self, params: dict) -> Tuple[Status, dict]: # pylint: disable=too-many-return-statements + """ + Check if the identity is assigned to the VM. + Return SUCCEEDED if true, PENDING otherwise. + + Parameters + ---------- + _params : dict + Flat dictionary of (key, value) pairs of tunable parameters. + This parameter is not used; we need it for compatibility with + other polling functions used in `_wait_while()`. + + Returns + ------- + result : (Status, dict={}) + A pair of Status and result. The result is always {}. + Status is one of {SUCCEEDED, PENDING, FAILED} + """ + params = self._set_default_params(params) + config = merge_parameters( + dest=self.config.copy(), + source=params, + required_keys=[ + "subscription", + "resourceGroup", + "managedIdentityName", + ] + ) + + vmName = self._deploy_params["vmName"] + subscriptionId = config["subscription"] + resourceGroup = config["resourceGroup"] + managedIdentityName = config["managedIdentityName"] + + _LOG.info("Check identity assignment: %s", vmName) + + url = self._URL_ASSIGN_IDENTITY.format( + subscription=subscriptionId, + resource_group=resourceGroup, + vm_name=vmName, + ) + + session = self._get_session(params) + try: + response = session.get(url, timeout=self._request_timeout) + except requests.exceptions.ReadTimeout: + _LOG.warning("Request timed out after %.2f s: %s", self._request_timeout, url) + return Status.RUNNING, {} + except requests.exceptions.RequestException as ex: + _LOG.exception("Error in request checking deployment", exc_info=ex) + return (Status.FAILED, {}) + + _LOG.debug("Response: %s", response) + + if response.status_code == 200: + output = response.json() + state = output.get("properties", {}).get("provisioningState", "") + print("state", state) + + print(output) + mid = f"/subscriptions/{subscriptionId}/resourceGroups/{resourceGroup}/providers/Microsoft.ManagedIdentity/userAssignedIdentities/{managedIdentityName}" + if state == "Succeeded" and "identity" in output and mid in output["identity"]["userAssignedIdentities"]: + return (Status.SUCCEEDED, {}) + elif state in {"Accepted", "Creating", "Deleting", "Running", "Updating"}: + return (Status.PENDING, {}) + else: + _LOG.error("Response: %s :: %s", response, json.dumps(output, indent=2)) + return (Status.FAILED, {}) + elif response.status_code == 404: + return (Status.PENDING, {}) + + _LOG.error("Response: %s :: %s", response, response.text) + return (Status.FAILED, {}) + def _provision_resource(self, params: dict) -> Tuple[Status, dict]: """ Attempts to (re)deploy a resource. @@ -465,3 +570,74 @@ def _provision_resource(self, params: dict) -> Tuple[Status, dict]: _LOG.error("Response: %s :: %s", response, response.text) # _LOG.error("Bad Request:\n%s", response.request.body) return (Status.FAILED, {}) + + def _assign_managed_identity(self, params: dict) -> Tuple[Status, dict]: + """ + Attempts to assign a managed identity to resource. + + Parameters + ---------- + params : dict + Flat dictionary of (key, value) pairs of tunable parameters. + Tunables are variable parameters that, together with the + Environment configuration, are sufficient to provision the resource. + + Returns + ------- + result : (Status, dict={}) + A pair of Status and result. The result is the input `params` plus the + parameters extracted from the response JSON, or {} if the status is FAILED. + Status is one of {PENDING, SUCCEEDED, FAILED} + """ + params = self._set_default_params(params) + config = merge_parameters(dest=self.config.copy(), source=params) + + vmName = self._deploy_params["vmName"] + subscriptionId = config["subscription"] + resourceGroup = config["resourceGroup"] + managedIdentityName = config["managedIdentityName"] + + _LOG.info("Assign managed identity: %s :: %s", vmName, params) + + url = self._URL_ASSIGN_IDENTITY.format( + subscription=subscriptionId, + resource_group=resourceGroup, + vm_name=vmName, + ) + + json_req = { + "identity": { + "type": "UserAssigned", + "userAssignedIdentities": { + f"/subscriptions/{subscriptionId}/resourceGroups/{resourceGroup}/providers/Microsoft.ManagedIdentity/userAssignedIdentities/{managedIdentityName}": {} + } + } + } + + if _LOG.isEnabledFor(logging.DEBUG): + _LOG.debug("Request: PATCH %s\n%s", url, json.dumps(json_req, indent=2)) + + response = requests.patch(url, json=json_req, + headers=self._get_headers(), timeout=self._request_timeout) + + if _LOG.isEnabledFor(logging.DEBUG): + _LOG.debug("Response: %s\n%s", response, + json.dumps(response.json(), indent=2) + if response.content else "") + else: + _LOG.info("Response: %s", response) + + if response.status_code == 200: + return (Status.PENDING, config) + elif response.status_code == 201: + output = self._extract_arm_parameters(response.json()) + if _LOG.isEnabledFor(logging.DEBUG): + _LOG.debug("Extracted parameters:\n%s", json.dumps(output, indent=2)) + params.update(output) + params.setdefault("asyncResultsUrl", url) + params.setdefault("deploymentName", config["deploymentName"]) + return (Status.PENDING, params) + else: + _LOG.error("Response: %s :: %s", response, response.text) + # _LOG.error("Bad Request:\n%s", response.request.body) + return (Status.FAILED, {}) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index af09f4c723..a1e8594727 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -12,6 +12,7 @@ from typing import Any, Callable, Dict, List, Optional, Set, Union from azure.storage.fileshare import ShareClient +from azure.identity import DefaultAzureCredential from azure.core.exceptions import ResourceNotFoundError from mlos_bench.services.base_service import Service @@ -58,16 +59,32 @@ def __init__(self, self.config, { "storageAccountName", "storageFileShareName", - "storageAccountKey", + "managedIdentityClientId", + "tenantId" } ) + os.environ["AZURE_CLIENT_ID"] = self.config["managedIdentityClientId"] + os.environ["AZURE_TENANT_ID"] = self.config["tenantId"] + + args = { + "exclude_workload_identity_credential": True, + "exclude_developer_cli_credential": True, + "exclude_cli_credential": True, + "exclude_powershell_credential": True, + "exclude_shared_token_cache_credential": True, + "exclude_interactive_browser_credential": True, + "exclude_visual_studio_code_credential": True, + "exclude_environment_credential": True, + } + self._share_client = ShareClient.from_share_url( AzureFileShareService._SHARE_URL.format( account_name=self.config["storageAccountName"], fs_name=self.config["storageFileShareName"], ), - credential=self.config["storageAccountKey"], + credential=DefaultAzureCredential(**args), + token_intent="backup", ) def download(self, params: dict, remote_path: str, local_path: str, recursive: bool = True) -> None: diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_vm_services.py b/mlos_bench/mlos_bench/services/remote/azure/azure_vm_services.py index 7fdbdc18df..a3d9bcfdfd 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_vm_services.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_vm_services.py @@ -131,6 +131,8 @@ def __init__(self, config, global_config, parent, self.merge_methods(methods, [ # SupportsHostProvisioning + self.assign_managed_identity, + self.wait_host_managed_identity_assignment, self.provision_host, self.deprovision_host, self.deallocate_host, @@ -164,10 +166,11 @@ def __init__(self, def _set_default_params(self, params: dict) -> dict: # pylint: disable=no-self-use # Try and provide a semi sane default for the deploymentName if not provided - # since this is a common way to set the deploymentName and can same some + # since this is a common way to set the deploymentName and can save some # config work for the caller. - if "vmName" in params and "deploymentName" not in params: - params["deploymentName"] = f"{params['vmName']}-deployment" + deploy_params = super().deploy_params + if "vmName" in deploy_params and "deploymentName" not in params: + params["deploymentName"] = f"{deploy_params['vmName']}-deployment" _LOG.info("deploymentName missing from params. Defaulting to '%s'.", params["deploymentName"]) return params @@ -192,6 +195,25 @@ def wait_host_deployment(self, params: dict, *, is_setup: bool) -> Tuple[Status, """ return self._wait_deployment(params, is_setup=is_setup) + def wait_host_managed_identity_assignment(self, params: dict) -> Tuple[Status, dict]: + """ + Waits for a pending identity assignment on an Azure VM to resolve to SUCCEEDED or FAILED. + Return TIMED_OUT when timing out. + + Parameters + ---------- + params : dict + Flat dictionary of (key, value) pairs of tunable parameters. + + Returns + ------- + result : (Status, dict) + A pair of Status and result. + Status is one of {PENDING, SUCCEEDED, FAILED, TIMED_OUT} + Result is info on the operation runtime if SUCCEEDED, otherwise {}. + """ + return self._wait_host_managed_identity_assignment(params) + def wait_host_operation(self, params: dict) -> Tuple[Status, dict]: """ Waits for a pending operation on an Azure VM to resolve to SUCCEEDED or FAILED. @@ -274,6 +296,26 @@ def deprovision_host(self, params: dict) -> Tuple[Status, dict]: vm_name=config["vmName"], )) + def assign_managed_identity(self, params: dict) -> Tuple[Status, dict]: + """ + Assign managed identity to Azure VM + + Parameters + ---------- + params : dict + Flat dictionary of (key, value) pairs of tunable parameters. + HostEnv tunables are variable parameters that, together with the + HostEnv configuration, are sufficient to provision a VM. + + Returns + ------- + result : (Status, dict={}) + A pair of Status and result. The result is the input `params` plus the + parameters extracted from the response JSON, or {} if the status is FAILED. + Status is one of {PENDING, SUCCEEDED, FAILED} + """ + return self._assign_managed_identity(params) + def deallocate_host(self, params: dict) -> Tuple[Status, dict]: """ Deallocates the VM on Azure by shutting it down then releasing the compute resources. diff --git a/mlos_bench/mlos_bench/services/types/host_provisioner_type.py b/mlos_bench/mlos_bench/services/types/host_provisioner_type.py index b3560783fc..2ef07cc167 100644 --- a/mlos_bench/mlos_bench/services/types/host_provisioner_type.py +++ b/mlos_bench/mlos_bench/services/types/host_provisioner_type.py @@ -90,3 +90,35 @@ def deallocate_host(self, params: dict) -> Tuple["Status", dict]: A pair of Status and result. The result is always {}. Status is one of {PENDING, SUCCEEDED, FAILED} """ + + def assign_managed_identity(self, params: dict) -> Tuple["Status", dict]: + """ + Assigns a managed identity to the Host/VM. + + Parameters + ---------- + params : dict + Flat dictionary of (key, value) pairs of tunable parameters. + + Returns + ------- + result : (Status, dict={}) + A pair of Status and result. The result is always {}. + Status is one of {PENDING, SUCCEEDED, FAILED} + """ + + def wait_host_managed_identity_assignment(self, params: dict) -> Tuple["Status", dict]: + """ + Waits for a pending managed identity assignment operation to resolve to SUCCEEDED or FAILED. + + Parameters + ---------- + params : dict + Flat dictionary of (key, value) pairs of tunable parameters. + + Returns + ------- + result : (Status, dict={}) + A pair of Status and result. The result is always {}. + Status is one of {PENDING, SUCCEEDED, FAILED} + """