setup_azureml.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "70d5c9c8",
   "metadata": {},
   "source": [
    "Setup an existing AzureML Workspace. To create a new Workspace follow instructions at: https://learn.microsoft.com/en-us/azure/machine-learning/how-to-manage-workspace?tabs=azure-portal#create-a-workspace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "159f46aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install azure-ai-ml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "7490cc0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime, time\n",
    "\n",
    "from azure.ai.ml import MLClient\n",
    "from azure.ai.ml.entities import (\n",
    "    ManagedOnlineEndpoint,\n",
    "    ManagedOnlineDeployment,\n",
    "    BatchDeployment,\n",
    "    Model,\n",
    "    Environment,\n",
    "    BuildContext,\n",
    "    CodeConfiguration,\n",
    "    OnlineRequestSettings,\n",
    "    ProbeSettings\n",
    ")\n",
    "from azure.core.exceptions import ResourceNotFoundError\n",
    "from azure.identity import DefaultAzureCredential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5ef71240",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found the config file in: .\\config.json\n",
      "Class RegistryOperations: This is an experimental class, and may change at any time. Please see https://aka.ms/azuremlexperimental for more information.\n"
     ]
    }
   ],
   "source": [
    "credential = DefaultAzureCredential()\n",
    "ml_client = MLClient.from_config(credential)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "14b460fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The below code downloads from an unofficial source.\n",
    "# You can also download the model manually from the official source: https://huggingface.co/CompVis/stable-diffusion-v-1-4-original\n",
    "\n",
    "import requests\n",
    "\n",
    "model_download_url = \"https://drive.yerf.org/wl/?id=EBfTrmcCCUAGaQBXVIj5lJmEhjoP1tgl&mode=grid&download=1\"\n",
    "model_file_path = \"./stable-diffusion-v1.ckpt\"\n",
    "\n",
    "response = requests.get(model_download_url)\n",
    "response.raise_for_status()\n",
    "with open(model_file_path, \"wb\") as f:\n",
    "    f.write(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "4c82c62a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"stable-diffusion-v1\"\n",
    "\n",
    "try:\n",
    "    model = ml_client.models.get(model_name, version=\"1\")\n",
    "except ResourceNotFoundError:\n",
    "    model = Model(\n",
    "        path=model_file_path,\n",
    "        name=model_name,\n",
    "    )\n",
    "    model = ml_client.create_or_update(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0149178a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!(git submodule init && git submodule update)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "daaba9e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: build=BuildContext(path=\"./stable-diffusion-api\"),\n",
    "\n",
    "environment_name = \"stable-diffusion-v1-new\"\n",
    "\n",
    "try:\n",
    "    env = ml_client.environments.get(environment_name, version=\"1\")\n",
    "except ResourceNotFoundError:\n",
    "    env = Environment(\n",
    "        name=environment_name,\n",
    "        image=\"docker.io/shivammg/stable-diffusion:v1\",\n",
    "        inference_config={\n",
    "            \"liveness_route\": {\"port\": 5000, \"path\": \"/api/health\"},\n",
    "            \"readiness_route\": {\"port\": 5000, \"path\": \"/api/health\"},\n",
    "            \"scoring_route\": {\"port\": 5000, \"path\": \"/api/txt2img\"},\n",
    "        },\n",
    "    )\n",
    "    env = ml_client.environments.create_or_update(env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "f53a7d3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name_prefix = \"stable-diffusion-\"\n",
    "endpoint_name = None\n",
    "endpoint = None\n",
    "\n",
    "for e in ml_client.online_endpoints.list():\n",
    "    if e.name.startswith(endpoint_name_prefix):\n",
    "        endpoint = e\n",
    "        endpoint_name = e.name\n",
    "        break\n",
    "\n",
    "if not endpoint:\n",
    "    endpoint_name = endpoint_name_prefix + datetime.datetime.now().strftime(\"%Y%m%d%H%M\")  # Endpoint name needs to be unique within a region\n",
    "    endpoint = ManagedOnlineEndpoint(\n",
    "        name=endpoint_name,\n",
    "        auth_mode=\"key\",\n",
    "    )\n",
    "    ml_client.online_endpoints.begin_create_or_update(endpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "c4ccbfb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(100):\n",
    "    endpoint = ml_client.online_endpoints.get(endpoint_name)\n",
    "    if endpoint.provisioning_state == 'Succeeded':\n",
    "        break\n",
    "    time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "fe139e60",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Check: endpoint stable-diffusion-202301061139 exists\n",
      "data_collector is not a known attribute of class <class 'azure.ai.ml._restclient.v2022_02_01_preview.models._models_py3.ManagedOnlineDeployment'> and will be ignored\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<azure.core.polling._poller.LROPoller at 0x1b82b2d9930>"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ".........................................."
     ]
    }
   ],
   "source": [
    "deployment_name = \"stable-diffusion-deployment\"\n",
    "deployment = ManagedOnlineDeployment(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=endpoint_name,\n",
    "    model=model,\n",
    "    environment=env,\n",
    "    environment_variables={\n",
    "        \"MODEL_DIR\": f\"/var/azureml-app/azureml-models/{model_name}/1\",\n",
    "    },\n",
    "    instance_type=\"Standard_NC6s_v3\",\n",
    "    instance_count=1,\n",
    "    request_settings=OnlineRequestSettings(\n",
    "        request_timeout_ms=60*1000,\n",
    "        max_concurrent_requests_per_instance=2,\n",
    "        max_queue_wait_ms=1*1000,\n",
    "    ),\n",
    "    readiness_probe=ProbeSettings(\n",
    "        failure_threshold=30,\n",
    "        timeout=10,\n",
    "        period=20,\n",
    "        initial_delay=20,\n",
    "    ),\n",
    "    liveness_probe=ProbeSettings(\n",
    "        failure_threshold=15,\n",
    "        timeout=5,\n",
    "        period=10,\n",
    "        initial_delay=10,\n",
    "    ),\n",
    ")\n",
    "ml_client.online_deployments.begin_create_or_update(deployment)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "4622ba65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "........................................................"
     ]
    }
   ],
   "source": [
    "for _ in range(100):\n",
    "    deployment = ml_client.online_deployments.get(deployment_name, endpoint_name=endpoint_name)\n",
    "    if deployment.provisioning_state == 'Succeeded':\n",
    "        break\n",
    "    time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "7e5619f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Readonly attribute principal_id will be ignored in class <class 'azure.ai.ml._restclient.v2022_05_01.models._models_py3.ManagedServiceIdentity'>\n",
      "Readonly attribute tenant_id will be ignored in class <class 'azure.ai.ml._restclient.v2022_05_01.models._models_py3.ManagedServiceIdentity'>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<azure.core.polling._poller.LROPoller at 0x1b82b6bd780>"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "endpoint.traffic = {deployment_name: 100}\n",
    "ml_client.online_endpoints.begin_create_or_update(endpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "01c0fab2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "endpoint_host = endpoint.scoring_uri.removesuffix('/score')\n",
    "keys = ml_client.online_endpoints.get_keys(endpoint_name)\n",
    "endpoint_details = {\n",
    "    'host': endpoint_host,\n",
    "    'key': keys.primary_key\n",
    "}\n",
    "\n",
    "with open('online-endpoint.json', 'w') as f:\n",
    "    f.write(json.dumps(endpoint_details))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce3a7465",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}