sdk/python/foundation-models/system/inference/text-generation/aacs-prepare-uai.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare UAI account for Azure Content Safety enabled Llama 2 model deployment\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.1 Assign variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uai_name = f\"aacs-uai\"\n",
    "\n",
    "# The name of the AACS created in deploy_llama notebook\n",
    "# Leave it blank will assign all AACS resource under resources group to the UAI\n",
    "aacs_name = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.2 Get credential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential\n",
    "\n",
    "try:\n",
    "    credential = DefaultAzureCredential()\n",
    "    # Check if given credential can get token successfully.\n",
    "    credential.get_token(\"https://management.azure.com/.default\")\n",
    "except Exception as ex:\n",
    "    # Fall back to InteractiveBrowserCredential in case DefaultAzureCredential not work\n",
    "    credential = InteractiveBrowserCredential()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.3 Configure workspace "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient\n",
    "\n",
    "try:\n",
    "    ml_client = MLClient.from_config(credential=credential)\n",
    "except Exception as ex:\n",
    "    # enter details of your AML workspace\n",
    "    subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "    resource_group = \"<RESOURCE_GROUP>\"\n",
    "    workspace_name = \"<AML_WORKSPACE_NAME>\"\n",
    "\n",
    "    # get a handle to the workspace\n",
    "    ml_client = MLClient(credential, subscription_id, resource_group, workspace_name)\n",
    "\n",
    "\n",
    "subscription_id = ml_client.subscription_id\n",
    "resource_group = ml_client.resource_group_name\n",
    "workspace_name = ml_client.workspace_name\n",
    "workspace_resource = ml_client.workspaces.get(workspace_name)\n",
    "workspace_location = workspace_resource.location\n",
    "\n",
    "print(f\"Connected to workspace {workspace_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Create a Managed Identity for the Azure AI Content Safety enabled Llama 2 endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.1 Get a handle to the ManagedServiceIdentityClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.mgmt.msi import ManagedServiceIdentityClient\n",
    "from azure.mgmt.msi.models import Identity\n",
    "\n",
    "msi_client = ManagedServiceIdentityClient(\n",
    "    subscription_id=subscription_id,\n",
    "    credential=credential,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.2 Create the User Assigned Identity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": [
    "##### 2.3 Retrieve the identity object\n",
    "we need to retrieve the identity object so that we can use it to deploy the Azure AI Content Safety enabled Llama 2 online endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uai_identity = msi_client.user_assigned_identities.get(\n",
    "    resource_group_name=resource_group,\n",
    "    resource_name=uai_name,\n",
    ")\n",
    "uai_principal_id = uai_identity.principal_id\n",
    "uai_client_id = uai_identity.client_id\n",
    "uai_id = uai_identity.id\n",
    "print(f\"UAI principal id: {uai_principal_id}\")\n",
    "print(f\"UAI id: {uai_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.4 Grant appropriate roles to the UAI we created above.\n",
    "Note: In order to successfully run scripts in current step, your must have owner permission on the AACS resource and the Llama 2 endpoint, which we created in the previous steps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.4.1 Get an AuthorizationManagementClient to list Role Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.mgmt.authorization import AuthorizationManagementClient\n",
    "from azure.mgmt.authorization.v2022_04_01.models import (\n",
    "    RoleAssignmentCreateParameters,\n",
    ")\n",
    "import uuid\n",
    "\n",
    "role_definition_client = AuthorizationManagementClient(\n",
    "    credential=credential,\n",
    "    subscription_id=subscription_id,\n",
    "    api_version=\"2022-04-01\",\n",
    ")\n",
    "role_assignment_client = AuthorizationManagementClient(\n",
    "    credential=credential,\n",
    "    subscription_id=subscription_id,\n",
    "    api_version=\"2022-04-01\",\n",
    ")\n",
    "\n",
    "uai_role_check_list = {\n",
    "    \"Cognitive Services Contributor\": {\n",
    "        \"step\": \"4.2.2\",\n",
    "        \"description\": \"assigne the role Cognitive Services User to the UAI on the Azure AI Content Safety resource.\",\n",
    "    },\n",
    "    \"AcrPull\": {\n",
    "        \"step\": \"4.2.3\",\n",
    "        \"description\": \"assigne the role AcrPull to the UAI on the Azure Container Registry.\",\n",
    "    },\n",
    "    \"Storage Blob Data Reader\": {\n",
    "        \"step\": \"4.2.4\",\n",
    "        \"description\": \"assigne the role Storage Blob Data Reader to the UAI on the Azure Storage account.\",\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.4.2 Define util function to assign access for UAI\n",
    "Cognitive Services User role is required to access the Azure Content Safety resource."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "role_name = \"Cognitive Services Contributor\"  # minimum role required for listing AACS access keys\n",
    "\n",
    "\n",
    "def assign_access_to_acs(scope):\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",
    "    from azure.core.exceptions import ResourceExistsError\n",
    "\n",
    "    try:\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",
    "        )\n",
    "    except ResourceExistsError as ex:\n",
    "        pass\n",
    "    except Exception as ex:\n",
    "        print(ex)\n",
    "        raise ex\n",
    "\n",
    "    if role_name in uai_role_check_list:\n",
    "        del uai_role_check_list[role_name]\n",
    "    print(\n",
    "        f\"Role assignment for {role_name} at the Azure AI Content Safety resource level completed.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.4.3 Grant the user identity access to the Azure Content Safety resource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient\n",
    "\n",
    "acs_client = CognitiveServicesManagementClient(credential, subscription_id)\n",
    "\n",
    "\n",
    "def find_acs(accounts):\n",
    "    return next(x for x in accounts if x.kind == \"ContentSafety\")\n",
    "\n",
    "\n",
    "if aacs_name == \"\":\n",
    "    for acs_resource in acs_client.accounts.list_by_resource_group(resource_group):\n",
    "        assign_access_to_acs(acs_resource.id)\n",
    "else:\n",
    "    acs_resource = acs_client.accounts.get(resource_group, aacs_name)\n",
    "    assign_access_to_acs(acs_resource.id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.4.4 Assign AcrPull at the workspace container registry scope\n",
    "Since we will create the Azure AI Content Safety enabled Llama 2 endpoint with User Assigned Identity, the user managed identity must have Storage Blob Data Reader permission on the storage account for the workspace, and AcrPull permission on the Azure Container Registry (ACR) for the workspace. Make sure your User Assigned Identity has the right permission."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "container_registry = workspace_resource.container_registry\n",
    "\n",
    "role_name = \"AcrPull\"\n",
    "\n",
    "role_defs = role_definition_client.role_definitions.list(scope=container_registry)\n",
    "role_def = next((r for r in role_defs if r.role_name == role_name))\n",
    "\n",
    "from azure.core.exceptions import ResourceExistsError\n",
    "\n",
    "try:\n",
    "    role_assignment_client.role_assignments.create(\n",
    "        scope=container_registry,\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",
    "    )\n",
    "except ResourceExistsError as ex:\n",
    "    pass\n",
    "except Exception as ex:\n",
    "    print(ex)\n",
    "    raise ex\n",
    "\n",
    "if role_name in uai_role_check_list:\n",
    "    del uai_role_check_list[role_name]\n",
    "print(\"Role assignment for AcrPull at the workspace container registry completed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.4.6 Assign Storage Blob Data Reader at the workspace storage account scope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "role_name = \"Storage Blob Data Reader\"\n",
    "blob_scope = workspace_resource.storage_account\n",
    "\n",
    "role_defs = role_definition_client.role_definitions.list(scope=blob_scope)\n",
    "role_def = next((r for r in role_defs if r.role_name == role_name))\n",
    "\n",
    "from azure.core.exceptions import ResourceExistsError\n",
    "\n",
    "try:\n",
    "    role_assignment_client.role_assignments.create(\n",
    "        scope=blob_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",
    "    )\n",
    "except ResourceExistsError as ex:\n",
    "    pass\n",
    "except Exception as ex:\n",
    "    print(ex)\n",
    "    raise ex\n",
    "\n",
    "if role_name in uai_role_check_list:\n",
    "    del uai_role_check_list[role_name]\n",
    "print(\n",
    "    \"Role assignment for `Storage Blob Data Reader` at the workspace storage account completed.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.4.7 Let's make sure we didn't miss anything in the previous steps, please execute the following script to check on that:__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check everything is properly done before creating the Azure AI Content Safety Enabled Llama 2 online endpoint\n",
    "missing_steps = []\n",
    "print(\n",
    "    \"You selected UAI to deploy the Azure AI Content Safety Enabled Llama 2 online endpoint, checking if the UAI has the required roles assigned...\"\n",
    ")\n",
    "if uai_role_check_list:\n",
    "    for key, value in uai_role_check_list.items():\n",
    "        missing_steps.append(\n",
    "            f'Please go to step {value[\"step\"]} to {value[\"description\"]}'\n",
    "        )\n",
    "\n",
    "if missing_steps:\n",
    "    print(\"Seems you missed some step above.\")\n",
    "    steps = \"\\n\".join(missing_steps)\n",
    "    raise Exception(f\"Please complete the missing steps before proceeding:\\n{steps}\")\n",
    "else:\n",
    "    print(\n",
    "        \"All steps are completed, proceeding to create the Azure AI Content Safety Enabled Llama 2 online endpoint...\\n\"\n",
    "        f\"uai_id: {uai_id}\\n\"\n",
    "        f\"uai_client_id: {uai_client_id}\"\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}