sdk/python/endpoints/online/managed/online-endpoints-keyvault.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Access Keyvault secrets from a Managed Online Endpoint\n",
    "\n",
    "In this example we create a Keyvault, set a secret, and then retrieve the secret from a Managed Online Endpoint using the endpoint's system-assigned managed identity. By using the managed identity, the need to pass secrets as well as any other credentials in the image or deployment is avoided.\n",
    "\n",
    "## Prerequisites: \n",
    "* The following additional Python packages should be installed: \n",
    "    * [azure-mgmt-keyvault](https://pypi.org/project/azure-mgmt-keyvault/) - Used to create a keyvault\n",
    "    * [azure-keyvault](https://pypi.org/project/azure-keyvault/)- Used to set the secret and permissions\n",
    "\n",
    "Install the prerequisites with the following code: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install azure-mgmt-keyvault\n",
    "%pip install azure-keyvault"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Connect to Azure Machine Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-import-required-libraries"
   },
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient\n",
    "from azure.ai.ml.entities import (\n",
    "    ManagedOnlineEndpoint,\n",
    "    ManagedOnlineDeployment,\n",
    "    Model,\n",
    "    CodeConfiguration,\n",
    "    Environment,\n",
    ")\n",
    "from azure.mgmt.keyvault import KeyVaultManagementClient\n",
    "from azure.keyvault.secrets import SecretClient\n",
    "from azure.mgmt.keyvault.models import (\n",
    "    VaultCreateOrUpdateParameters,\n",
    "    VaultProperties,\n",
    "    Sku,\n",
    ")\n",
    "from azure.mgmt.keyvault.models import AccessPolicyEntry, Permissions, SecretPermissions\n",
    "from azure.identity import DefaultAzureCredential, AzureCliCredential\n",
    "import random, os, json, base64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Set workspace details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-set-workspace-details"
   },
   "outputs": [],
   "source": [
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace_name = \"<AML_WORKSPACE_NAME>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Set variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-set-variables"
   },
   "outputs": [],
   "source": [
    "rand = random.randint(0, 10000)\n",
    "\n",
    "endpoint_name = f\"endpt-moe-{rand}\"\n",
    "keyvault_name = f\"kvexample{rand}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 Create an MLClient instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-create-mlclient-instance"
   },
   "outputs": [],
   "source": [
    "credential = DefaultAzureCredential()\n",
    "ml_client = MLClient(\n",
    "    credential,\n",
    "    subscription_id=subscription_id,\n",
    "    resource_group_name=resource_group,\n",
    "    workspace_name=workspace_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.5 Create a Keyvault Management client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "1-create-keyvault-management"
   },
   "outputs": [],
   "source": [
    "keyvault_mgmt_client = KeyVaultManagementClient(\n",
    "    credential=credential, subscription_id=subscription_id\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create an endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "2-create-endpoint"
   },
   "outputs": [],
   "source": [
    "endpoint = ManagedOnlineEndpoint(name=endpoint_name)\n",
    "endpoint = ml_client.online_endpoints.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create a Keyvault\n",
    "\n",
    "To set and retrieve secrets or other assets in the Keyvault, access policies need to be defined for each role. We need the Object (Principal) ID for the user or SP/managed identity that will read or write secrets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Get your Object ID\n",
    "The `oid` in your JWT [access token](https://learn.microsoft.com/en-us/azure/active-directory/develop/access-tokens) represents the Object ID of the current user or SP logged into the Azure CLI. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "3-get-object-id"
   },
   "outputs": [],
   "source": [
    "cli_credential = AzureCliCredential()\n",
    "token = cli_credential.get_token(\"https://management.azure.com\").token\n",
    "user_or_sp_object_id = json.loads(base64.b64decode(token.split(\".\")[1] + \"===\")).get(\n",
    "    \"oid\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Define an AccessPolicy for the Endpoint\n",
    "Allow the endpoint to get secrets in the Keyvault."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "3-define-accesspolicy-endpoint"
   },
   "outputs": [],
   "source": [
    "endpoint_access_policy = AccessPolicyEntry(\n",
    "    tenant_id=endpoint.identity.tenant_id,\n",
    "    object_id=endpoint.identity.principal_id,\n",
    "    permissions=Permissions(secrets=[SecretPermissions.GET]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Define an AccessPolicy for the current user\n",
    "Allow all secret permissions for the current user or Service Principal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "3-define-accesspolicy-current"
   },
   "outputs": [],
   "source": [
    "user_or_sp_access_policy = AccessPolicyEntry(\n",
    "    tenant_id=endpoint.identity.tenant_id,\n",
    "    object_id=user_or_sp_object_id,\n",
    "    permissions=Permissions(secrets=[SecretPermissions.ALL]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.5 Create the Keyvault"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "3-create-keyvault"
   },
   "outputs": [],
   "source": [
    "keyvault = keyvault_mgmt_client.vaults.begin_create_or_update(\n",
    "    vault_name=keyvault_name,\n",
    "    resource_group_name=resource_group,\n",
    "    parameters=VaultCreateOrUpdateParameters(\n",
    "        location=endpoint.location,\n",
    "        properties=VaultProperties(\n",
    "            tenant_id=endpoint.identity.tenant_id,\n",
    "            sku=Sku(name=\"Standard\", family=\"A\"),\n",
    "            access_policies=[endpoint_access_policy, user_or_sp_access_policy],\n",
    "        ),\n",
    "    ),\n",
    ").result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Set a Keyvault secret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "4-set-keyvault-secret"
   },
   "outputs": [],
   "source": [
    "secret_client = SecretClient(\n",
    "    credential=credential, vault_url=f\"https://{keyvault_name}.vault.azure.net\"\n",
    ")\n",
    "secret = secret_client.set_secret(name=\"multiplier\", value=str(7))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Create a Deployment\n",
    "\n",
    "The [scoring script](keyvault/code/score.py) uses a `ManagedIdentityCredential` to authenticate itself to the Keyvault via a `SecretClient` from the `azure-keyvault` package. No arguments are needed to instantiate the credential object when this code is executed in a deployment, because it reads the environment variables `MSI_SECRET` and `MSI_ENDPOINT` which are already present.\n",
    "\n",
    "As part of the deployment, we will pass an environment variable called `KV_SECRET_MULTIPLIER` and give it the value `multiplier@https://<VAULT_NAME>.vault.azure.net`. The convenience function `load_secrets` looks for environment variables with `KV_SECRET` and replaces their values with the actual value of the secret from the keyvault. \n",
    "\n",
    "When a request is received, `input` is multiplied by our secret. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 Define the deployment\n",
    "The environment variable `KV_SECRET_MULTIPLIER` is set to `multiplier@<KEYVAULT_URL>`. In the scoring script, this value is parsed and passed to a SecretClient to retrieve the secret from the Keyvault."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "5-define-deployment"
   },
   "outputs": [],
   "source": [
    "deployment = ManagedOnlineDeployment(\n",
    "    name=\"kvdep\",\n",
    "    endpoint_name=endpoint_name,\n",
    "    model=Model(path=\"keyvault\"),\n",
    "    code_configuration=CodeConfiguration(\n",
    "        code=\"keyvault/code\", scoring_script=\"score.py\"\n",
    "    ),\n",
    "    environment=Environment(\n",
    "        conda_file=\"keyvault/env.yml\",\n",
    "        image=\"mcr.microsoft.com/azureml/minimal-ubuntu22.04-py39-cpu-inference:latest\",\n",
    "    ),\n",
    "    environment_variables={\n",
    "        \"KV_SECRET_MULTIPLIER\": f\"multiplier@https://{keyvault_name}.vault.azure.net\"\n",
    "    },\n",
    "    instance_type=\"Standard_DS3_v2\",\n",
    "    instance_count=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.2 Create the deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "5-create-deployment"
   },
   "outputs": [],
   "source": [
    "deployment = ml_client.online_deployments.begin_create_or_update(deployment).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.3 Update endpoint traffic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "5-update-endpoint-traffic"
   },
   "outputs": [],
   "source": [
    "endpoint.traffic = {\"kvdep\": 100}\n",
    "endpoint = ml_client.online_endpoints.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Test the endpoint\n",
    "The endpoint returns the value of `input` multiplied by the secret."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "6-test-endpoint"
   },
   "outputs": [],
   "source": [
    "ml_client.online_endpoints.invoke(\n",
    "    endpoint_name=endpoint_name, request_file=\"keyvault/sample_request.json\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Delete assets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.1 Delete the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "7-delete-endpoint"
   },
   "outputs": [],
   "source": [
    "ml_client.online_endpoints.begin_delete(name=endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.2 Delete the keyvault"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "name": "7-delete-keyvault"
   },
   "outputs": [],
   "source": [
    "keyvault_mgmt_client.vaults.delete(\n",
    "    resource_group_name=resource_group, vault_name=keyvault_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"
  },
  "vscode": {
   "interpreter": {
    "hash": "103e4dffdc87d5e761c21b3875987ce70e93823f5f2e969a3724eab733263b67"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}