patch for managed identity storage authentication

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,
 

mlos_bench/mlos_bench/config/environments/apps/redis/redis.jsonc

@@ -51,18 +51,21 @@
                         "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",
+                        "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/"
                     ]
                 }

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"
                     ]
                 }
             }

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",
                     ]
                 }
             }

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",
 

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"
                         }
                     },

mlos_bench/mlos_bench/config/services/remote/azure/arm-templates/azuredeploy-ubuntu-vm.jsonc

@@ -20,6 +20,12 @@
         "description": "OS Autotune Linux VM"
       }
     },
+    "managedIdentityName": {
+        "type": "string",
+        "metadata": {
+            "description": "User managed identity name"
+        }
+    },
     "customData": {
       "type": "string",
       "defaultValue": "",
@@ -187,6 +193,12 @@
       "apiVersion": "2021-11-01",
       "name": "[parameters('vmName')]",
       "location": "[parameters('location')]",
+      "identity": {
+        "type": "userAssigned",
+        "userAssignedIdentities": {
+            "[resourceID('Microsoft.ManagedIdentity/userAssignedIdentities/', parameters('managedIdentityName'))]": {}
+        }
+      },
       "properties": {
         "hardwareProfile": {
           "vmSize": "[parameters('vmSize')]"

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"
     }
 }

mlos_bench/mlos_bench/config/services/remote/azure/service-linux-vm-ops.jsonc

@@ -15,7 +15,7 @@
 
                 // Make sure to list all ARM template parameters that can be overridden by the caller.
                 "deploymentTemplateParameters": {
-
+                    "managedIdentityName": "PLACEHOLDER; e.g., mlos-managed-identity",
                     "storageAccountName": "PLACEHOLDER; e.g., osatsharedstorage",
                     "storageFileShareName": "PLACEHOLDER; e.g., os-autotune-file-share",
                     "location": "PLACEHOLDER; e.g., westus2",

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:

mlos_bench/mlos_bench/services/remote/azure/azure_vm_services.py

@@ -164,7 +164,7 @@ 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"