sdk/python/endpoints/online/managed/managed-identities/online-endpoints-managed-identity-uai.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Access Azure resources from an online endpoint with a user-assigned managed identity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additional packages are required for this example:\n",
    "\n",
    "* Microsoft Azure Msi Management Client \n",
    "* Microsoft Azure Storage Client\n",
    "* Microsoft Azure Authorization Management Client\n",
    "\n",
    "Install them with the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "install-packages"
   },
   "outputs": [],
   "source": [
    "%pip install --pre azure-mgmt-msi\n",
    "%pip install --pre azure-mgmt-storage\n",
    "%pip install --pre azure-mgmt-authorization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Configure variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Assign variables for the workspace and deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-assign-variables"
   },
   "outputs": [],
   "source": [
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace_name = \"<AML_WORKSPACE_NAME>\"\n",
    "endpoint_name = \"<ENDPOINT_NAME>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Specify the storage details and file to access\n",
    "A text file is located in the `managed-identities` folder called `hello.txt`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-specify-storage-details"
   },
   "outputs": [],
   "source": [
    "storage_account_name = \"<STORAGE_ACCOUNT_NAME>\"\n",
    "storage_container_name = \"<CONTAINER_TO_ACCESS>\"\n",
    "file_name = \"<FILE_TO_ACCESS>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Decide on the name of your user identity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-decide-name-user-identity"
   },
   "outputs": [],
   "source": [
    "uai_name = \"<USER_ASSIGNED_IDENTITY_NAME>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "rand = random.randint(0, 10000)\n",
    "\n",
    "endpoint_name = f\"endpt-moe-{rand}\"\n",
    "storage_account_name = f\"oepstorage{rand}\"\n",
    "storage_container_name = f\"hellocontainer\"\n",
    "file_name = \"hello.txt\"\n",
    "uai_name = f\"oep-user-identity-{rand}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 Retrieve the workspace location:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-retrieve-workspace-location"
   },
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient\n",
    "from azure.identity import AzureCliCredential\n",
    "from azure.ai.ml.entities import (\n",
    "    ManagedOnlineDeployment,\n",
    "    ManagedOnlineEndpoint,\n",
    "    Model,\n",
    "    CodeConfiguration,\n",
    "    Environment,\n",
    ")\n",
    "\n",
    "credential = AzureCliCredential()\n",
    "ml_client = MLClient(credential, subscription_id, resource_group, workspace_name)\n",
    "\n",
    "workspace_location = ml_client.workspaces.get(workspace_name).location"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Configure deployment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Define a deployment configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "2-define-deployment-configuration"
   },
   "outputs": [],
   "source": [
    "deployment = ManagedOnlineDeployment(\n",
    "    name=\"blue\",\n",
    "    endpoint_name=endpoint_name,\n",
    "    model=Model(path=\"../../model-1/model/\"),\n",
    "    code_configuration=CodeConfiguration(\n",
    "        code=\"../../model-1/onlinescoring/\", scoring_script=\"score_managedidentity.py\"\n",
    "    ),\n",
    "    environment=Environment(\n",
    "        conda_file=\"../../model-1/environment/conda-managedidentity.yml\",\n",
    "        image=\"mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04:latest\",\n",
    "    ),\n",
    "    instance_type=\"Standard_DS3_v2\",\n",
    "    instance_count=1,\n",
    "    environment_variables={\n",
    "        \"STORAGE_ACCOUNT_NAME\": storage_account_name,\n",
    "        \"STORAGE_CONTAINER_NAME\": storage_container_name,\n",
    "        \"FILE_NAME\": file_name,\n",
    "        # We will update this after creating an identity\n",
    "        \"UAI_CLIENT_ID\": \"uai_client_id_place_holder\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create the managed identity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Get a handle to the `ManagedServiceIdentityClient`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "3-get-handle"
   },
   "outputs": [],
   "source": [
    "from azure.mgmt.msi import ManagedServiceIdentityClient\n",
    "from azure.mgmt.msi.models import Identity\n",
    "\n",
    "credential = AzureCliCredential()\n",
    "msi_client = ManagedServiceIdentityClient(\n",
    "    subscription_id=subscription_id,\n",
    "    credential=credential,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Create the identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "3-create-identity"
   },
   "outputs": [],
   "source": [
    "msi_client.user_assigned_identities.create_or_update(\n",
    "    resource_group_name=resource_group,\n",
    "    resource_name=uai_name,\n",
    "    parameters=Identity(location=workspace_location),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Retrieve the identity object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "3-retrieve-identity-object"
   },
   "outputs": [],
   "source": [
    "uai_identity = msi_client.user_assigned_identities.get(\n",
    "    resource_group_name=resource_group,\n",
    "    resource_name=uai_name,\n",
    ")\n",
    "uai_identity.as_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Create storage account and container"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Get a handle to the `StorageManagementclient`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "4-get-handle"
   },
   "outputs": [],
   "source": [
    "from azure.mgmt.storage import StorageManagementClient\n",
    "from azure.storage.blob import ContainerClient\n",
    "from azure.mgmt.storage.models import Sku, StorageAccountCreateParameters, BlobContainer\n",
    "\n",
    "credential = AzureCliCredential()\n",
    "storage_client = StorageManagementClient(\n",
    "    credential=credential, subscription_id=subscription_id\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Define and create the storage account"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "4-create-storage-account"
   },
   "outputs": [],
   "source": [
    "storage_account_parameters = StorageAccountCreateParameters(\n",
    "    sku=Sku(name=\"Standard_LRS\"), kind=\"Storage\", location=workspace_location\n",
    ")\n",
    "\n",
    "storage_account = storage_client.storage_accounts.begin_create(\n",
    "    resource_group_name=resource_group,\n",
    "    account_name=storage_account_name,\n",
    "    parameters=storage_account_parameters,\n",
    ").result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 Create a blob container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "4-create-blob-container"
   },
   "outputs": [],
   "source": [
    "blob_container = storage_client.blob_containers.create(\n",
    "    resource_group_name=resource_group,\n",
    "    account_name=storage_account_name,\n",
    "    container_name=storage_container_name,\n",
    "    blob_container=BlobContainer(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.4 Get the storage account key and create a `ContainerClient`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "4-create-container-client"
   },
   "outputs": [],
   "source": [
    "res = storage_client.storage_accounts.list_keys(\n",
    "    resource_group_name=resource_group,\n",
    "    account_name=storage_account_name,\n",
    ")\n",
    "key = res.keys[0].value\n",
    "\n",
    "container_client = ContainerClient(\n",
    "    account_url=storage_account.primary_endpoints.blob,\n",
    "    container_name=storage_container_name,\n",
    "    credential=key,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.5 Upload a blob to the container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "4-upload-blob"
   },
   "outputs": [],
   "source": [
    "with open(file_name, \"rb\") as f:\n",
    "    container_client.upload_blob(name=file_name, data=f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Create an online endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 Create the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "5-create-online-endpoint"
   },
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import ManagedIdentityConfiguration, IdentityConfiguration\n",
    "\n",
    "endpoint = ManagedOnlineEndpoint(\n",
    "    name=endpoint_name,\n",
    "    auth_mode=\"key\",\n",
    "    identity=IdentityConfiguration(\n",
    "        type=\"user_assigned\",\n",
    "        user_assigned_identities=[\n",
    "            ManagedIdentityConfiguration(resource_id=uai_identity.id)\n",
    "        ],\n",
    "    ),\n",
    ")\n",
    "\n",
    "ml_client.online_endpoints.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.2 Confirm identity details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "5-get-details"
   },
   "outputs": [],
   "source": [
    "endpoint = ml_client.online_endpoints.get(endpoint_name)\n",
    "print(endpoint.identity.type)\n",
    "print(endpoint.identity.user_assigned_identities)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Give access permission to the managed identity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.1 Get an `AuthorizationManagementClient` to list Role Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "6-get-role-definitions-client"
   },
   "outputs": [],
   "source": [
    "from azure.mgmt.authorization import AuthorizationManagementClient\n",
    "from azure.mgmt.authorization.v2018_01_01_preview.models import RoleDefinition\n",
    "import uuid\n",
    "\n",
    "role_definition_client = AuthorizationManagementClient(\n",
    "    credential=credential,\n",
    "    subscription_id=subscription_id,\n",
    "    api_version=\"2018-01-01-preview\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.2 Get an `AuthorizationManagementClient` to make Role Assignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "6-get-role-assignments-client"
   },
   "outputs": [],
   "source": [
    "from azure.mgmt.authorization.v2020_10_01_preview.models import (\n",
    "    RoleAssignment,\n",
    "    RoleAssignmentCreateParameters,\n",
    ")\n",
    "\n",
    "role_assignment_client = AuthorizationManagementClient(\n",
    "    credential=credential,\n",
    "    subscription_id=subscription_id,\n",
    "    api_version=\"2020-10-01-preview\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.3 Get User-assigned identity details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "6-get-uai-details"
   },
   "outputs": [],
   "source": [
    "uai_identity = msi_client.user_assigned_identities.get(\n",
    "    resource_group_name=resource_group, resource_name=uai_name\n",
    ")\n",
    "uai_principal_id = uai_identity.principal_id\n",
    "uai_client_id = uai_identity.client_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.4 Assign `Storage Blob Data Reader` at the new storage account scope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "6-give-permission-user-storage-account"
   },
   "outputs": [],
   "source": [
    "role_name = \"Storage Blob Data Reader\"\n",
    "scope = storage_account.id\n",
    "\n",
    "role_defs = role_definition_client.role_definitions.list(scope=scope)\n",
    "role_def = next((r for r in role_defs if r.role_name == role_name))\n",
    "\n",
    "role_assignment_client.role_assignments.create(\n",
    "    scope=scope,\n",
    "    role_assignment_name=str(uuid.uuid4()),\n",
    "    parameters=RoleAssignmentCreateParameters(\n",
    "        role_definition_id=role_def.id,\n",
    "        principal_id=uai_principal_id,\n",
    "        principal_type=\"ServicePrincipal\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.5 Retrieve the workspace and container registry objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "6-retrieve-workspace-acr"
   },
   "outputs": [],
   "source": [
    "workspace = ml_client.workspaces.get(workspace_name)\n",
    "container_registry = workspace.container_registry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.6 Assign `AcrPull` at the workspace container registry scope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "6-give-permission-container-registry"
   },
   "outputs": [],
   "source": [
    "role_name = \"AcrPull\"\n",
    "scope = container_registry\n",
    "\n",
    "role_defs = role_definition_client.role_definitions.list(scope=scope)\n",
    "role_def = next((r for r in role_defs if r.role_name == role_name))\n",
    "\n",
    "role_assignment_client.role_assignments.create(\n",
    "    scope=scope,\n",
    "    role_assignment_name=str(uuid.uuid4()),\n",
    "    parameters=RoleAssignmentCreateParameters(\n",
    "        role_definition_id=role_def.id,\n",
    "        principal_id=uai_principal_id,\n",
    "        principal_type=\"ServicePrincipal\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.7 Assign `Storage Blob Data Reader` at the workspace storage account scope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "6-give-permission-workspace-storage"
   },
   "outputs": [],
   "source": [
    "role_name = \"Storage Blob Data Reader\"\n",
    "scope = workspace.storage_account\n",
    "\n",
    "role_defs = role_definition_client.role_definitions.list(scope=scope)\n",
    "role_def = next((r for r in role_defs if r.role_name == role_name))\n",
    "\n",
    "role_assignment_client.role_assignments.create(\n",
    "    scope=scope,\n",
    "    role_assignment_name=str(uuid.uuid4()),\n",
    "    parameters=RoleAssignmentCreateParameters(\n",
    "        role_definition_id=role_def.id,\n",
    "        principal_id=uai_principal_id,\n",
    "        principal_type=\"ServicePrincipal\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Create a deployment with your configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.1 Update the deployment configuration with the UAI Client ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "7-update-uai-client-id"
   },
   "outputs": [],
   "source": [
    "deployment.environment_variables[\"UAI_CLIENT_ID\"] = uai_client_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.2 Create the deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "7-create-deployment"
   },
   "outputs": [],
   "source": [
    "ml_client.online_deployments.begin_create_or_update(deployment).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.3 Check the status of the deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "7-check-deployment-status"
   },
   "outputs": [],
   "source": [
    "deployment = ml_client.online_deployments.get(\n",
    "    endpoint_name=endpoint_name, name=deployment.name\n",
    ")\n",
    "print(deployment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.4 Get the deployment logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "7-get-deployment-logs"
   },
   "outputs": [],
   "source": [
    "ml_client.online_deployments.get_logs(deployment.name, deployment.endpoint_name, 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.5 Set traffic to 100% for deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "7-set-traffic"
   },
   "outputs": [],
   "source": [
    "endpoint.traffic = {str(deployment.name): 100}\n",
    "ml_client.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Test the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "8-confirm-endpoint-deployed-successfully"
   },
   "outputs": [],
   "source": [
    "sample_data = \"../../model-1/sample-request.json\"\n",
    "ml_client.online_endpoints.invoke(endpoint_name=endpoint_name, request_file=sample_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Delete the endpoint and storage account"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 9.1 Delete the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "9-delete-endpoint"
   },
   "outputs": [],
   "source": [
    "ml_client.online_endpoints.begin_delete(endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 9.2 Delete the storage account"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "9-delete-storage-account"
   },
   "outputs": [],
   "source": [
    "storage_client.storage_accounts.delete(\n",
    "    resource_group_name=resource_group, account_name=storage_account_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 9.3 Delete the User-assigned managed identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "9-delete-uai"
   },
   "outputs": [],
   "source": [
    "msi_client.user_assigned_identities.delete(\n",
    "    resource_group_name=resource_group, resource_name=uai_name\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10 - SDK V2",
   "language": "python",
   "name": "python310-sdkv2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4 (main, Mar 31 2022, 08:41:55) [GCC 7.5.0]"
  },
  "vscode": {
   "interpreter": {
    "hash": "c54d4b4f21f908d21f1064b6d031502c08620e465e849bef5aa76d1f6a474870"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}