2-tutorial-remote-dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# High-compute workflow orchestration with [*Covalent*](http://covalent.xyz).\n",
    "\n",
    "Run complex workflows with ease; using cloud, on-prem, and/or local backends.\n",
    "\n",
    "## Covalent is available from PYPI\n",
    "\n",
    "```\n",
    "pip install covalent\n",
    "```\n",
    "\n",
    "## Usage\n",
    "\n",
    "👉 Add one or more [@ct.electron](https://docs.covalent.xyz/docs/user-documentation/concepts/covalent-basics#electron) decorators to designate workflow tasks (i.e. *electrons*).\n",
    "\n",
    "👉 Specify [executors](https://docs.covalent.xyz/docs/plugin) to choose electron backends.\n",
    "\n",
    "\n",
    "# 🔗 See links below for more information!\n",
    "\n",
    "- 📚 [Covalent Documentation](https://docs.covalent.xyz)\n",
    "- 🌟 [Covalent GitHub](https://github.com/AgnostiqHQ/covalent)\n",
    "- ✍️ [Covalent Blogs](https://www.covalent.xyz/blog/)\n",
    "    - [Simplifying generative AI workflows with Covalent and Streamlit](https://blog.streamlit.io/simplifying-generative-ai-workflows/)\n",
    "    - [Improving Chest X-ray Pneumonia Detection with Federated Learning and Covalent](https://medium.com/@filip_98594/improving-chest-x-ray-pneumonia-detection-with-federated-learning-and-covalent-ff60eef7946c)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment\n",
    "\n",
    "Run the following commands to create and activate the environment:\n",
    "```shell\n",
    "$ conda env create -f \"environment.yml\"\n",
    "$ conda activate covalent_pydata_2023\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Infrastructure\n",
    "\n",
    "See `terraform/tutorial-dask/README.md` for instructions on how to create the infrastructure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start Covalent\n",
    "\n",
    "Run the following command to start Covalent:\n",
    "\n",
    "```shell\n",
    "$ covalent start\n",
    "\n",
    "Covalent server has started at http://localhost:48008\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TUTORIAL: Automatic Image Captioning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dataclasses import dataclass\n",
    "from datetime import datetime\n",
    "from typing import List\n",
    "\n",
    "import covalent as ct\n",
    "import torch\n",
    "from PIL import Image\n",
    "from transformers import BlipForConditionalGeneration, BlipProcessor\n",
    "\n",
    "DASK_SCHEDULER_ADDRESS = \"52.87.169.192:8786\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROW_BLANK = \"{datetime} | {image_name:<30} | {description}\\n\"\n",
    "\n",
    "TRIGGER_DIR = os.path.abspath(\"./images-triggered\")\n",
    "CATALOG = os.path.join(TRIGGER_DIR, \"catalog.txt\")\n",
    "\n",
    "if not os.path.exists(CATALOG):\n",
    "    with open(CATALOG, \"w\") as f:\n",
    "        f.write('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model(type_: object, id_: str, label=None, **params):\n",
    "    \"\"\"Download files for the model if necessary and return the model.\"\"\"\n",
    "\n",
    "    model_id = id_.replace('/', '_')\n",
    "    if label:\n",
    "        model_id += f\"_{label}\"\n",
    "\n",
    "    model_path = os.path.abspath(f\"./{model_id}\")\n",
    "    if os.path.exists(model_path):\n",
    "        return type_.from_pretrained(model_path, **params)\n",
    "\n",
    "    model = type_.from_pretrained(id_, **params)\n",
    "    model.save_pretrained(model_path)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class InputImage:\n",
    "    file: str\n",
    "    description: str = ''\n",
    "    image: Image = None\n",
    "\n",
    "    def __post_init__(self):\n",
    "        self.file = os.path.abspath(os.path.join(TRIGGER_DIR, self.file))\n",
    "        self.image = self.image or Image.open(self.file)\n",
    "\n",
    "\n",
    "def only_image(p: str) -> bool:\n",
    "    ext = os.path.splitext(p)[1]\n",
    "    return ext in {\".jpg\", \".jpeg\", \".png\"}\n",
    "\n",
    "\n",
    "def new_image(input_image) -> bool:\n",
    "    with open(CATALOG, \"r\") as f:\n",
    "        name = os.path.basename(input_image.file)\n",
    "        line = f.readline()\n",
    "\n",
    "        while line:\n",
    "            if name in line.split('|')[1].strip():\n",
    "                return False\n",
    "            line = f.readline()\n",
    "\n",
    "        return True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now, let's define a workflow!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Electrons"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Read image files (local)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron(executor=\"local\")\n",
    "def read_new_images() -> List[InputImage]:\n",
    "    \"\"\"Read new images from the trigger directory.\"\"\"\n",
    "\n",
    "    images = []\n",
    "    for image_file in filter(only_image, os.listdir(TRIGGER_DIR)):\n",
    "        images.append(InputImage(image_file))\n",
    "\n",
    "    return list(filter(new_image, images))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Get image descriptions (Dask cluster on AWS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "aws_dask_exec = ct.executor.DaskExecutor(DASK_SCHEDULER_ADDRESS)\n",
    "\n",
    "deps_pip = [\n",
    "    \"pillow==10.1.0\",\n",
    "    \"torch==2.1.0\",\n",
    "    \"torchvision==0.16.0\",\n",
    "    \"transformers==4.33.1\",\n",
    "]\n",
    "\n",
    "@ct.electron(executor=aws_dask_exec, deps_pip=deps_pip)\n",
    "def get_descriptions(input_images: List[InputImage]) -> List[InputImage]:\n",
    "    \"\"\"Use BLIP to obtain a short description of an image.\"\"\"\n",
    "\n",
    "    processor = get_model(BlipProcessor, \"Salesforce/blip-image-captioning-base\", label=\"processor\")\n",
    "    model = get_model(BlipForConditionalGeneration, \"Salesforce/blip-image-captioning-base\", label=\"model\")\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        model = model.to(\"cuda\")\n",
    "\n",
    "    descriptions = []\n",
    "    for input_image in input_images:\n",
    "        inputs = processor(input_image.image, return_tensors=\"pt\")\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            inputs = inputs.to(\"cuda\")\n",
    "\n",
    "        out = model.generate(**inputs)\n",
    "        input_image.description = processor.decode(out[0], skip_special_tokens=True)\n",
    "        descriptions.append(input_image)\n",
    "\n",
    "    return descriptions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Record descriptions (local)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron(executor=\"local\")\n",
    "def write_descriptions(described_images: List[InputImage]) -> None:\n",
    "    \"\"\"Write descriptions to the catalog file.\"\"\"\n",
    "\n",
    "    for input_image in filter(new_image, described_images):\n",
    "        image_name = os.path.basename(input_image.file)\n",
    "        with open(CATALOG, \"a\", encoding=\"utf-8\") as f:\n",
    "            f.write(\n",
    "                ROW_BLANK.format(\n",
    "                    image_name=image_name,\n",
    "                    datetime=datetime.now(),\n",
    "                    description=input_image.description,\n",
    "                )\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='Lattice'></a>\n",
    "\n",
    "## Lattice\n",
    "\n",
    "#### Definition of the main workflow function\n",
    "\n",
    "This is a [**triggered** workflow](https://docs.covalent.xyz/docs/features/triggers). Once we execute the cell below, it will run whenever two or more images are added to `TRIGGER_DIR`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3fedd6f4-f4f0-411f-9246-c697dca17221'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trigger = ct.triggers.DirTrigger(\n",
    "    dir_path=TRIGGER_DIR,\n",
    "    event_names=[\"created\"],\n",
    "    batch_size=2,\n",
    "    recursive=False\n",
    ")\n",
    "\n",
    "@ct.lattice(triggers=[trigger])\n",
    "def describe_images():\n",
    "    images = read_new_images()\n",
    "    described_images = get_descriptions(images)\n",
    "    write_descriptions(described_images)\n",
    "\n",
    "ct.dispatch(describe_images)()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mnist",
   "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.18"
  },
  "vscode": {
   "interpreter": {
    "hash": "09f31237d8087fb717019708595cb8a07fe1ba4ed422094db14ddfd267d3815a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}