sample_workflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo workflow\n",
    "\n",
    "Let's look at how covalent can be used when you already have certain tasks and want to bring them together into one workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import covalent as ct\n",
    "\n",
    "@ct.electron\n",
    "def add_large_num(num_1, num_2):\n",
    "    # Add two large numbers\n",
    "    \n",
    "    print(f\"Adding {num_1} and {num_2}\")\n",
    "    \n",
    "    return num_1 + num_2\n",
    "\n",
    "\n",
    "@ct.electron\n",
    "def multiply_large_num(num_1, num_2):\n",
    "    # Multiply two large numbers\n",
    "\n",
    "    print(f\"Multiplying {num_1} and {num_2}\")\n",
    "\n",
    "    return num_1 * num_2\n",
    "\n",
    "@ct.electron\n",
    "def square_large_num(num):\n",
    "    # Square a large number\n",
    "\n",
    "    print(f\"Squaring {num}\")\n",
    "\n",
    "    return num ** 2\n",
    "\n",
    "\n",
    "@ct.lattice\n",
    "def sample_workflow(large_num_1, large_num_2, large_num_3):\n",
    "    # Sample workflow to do 3 independent jobs and 1 dependent job\n",
    "\n",
    "    added_num = add_large_num(large_num_1, large_num_2)\n",
    "    multiplied_num = multiply_large_num(large_num_1, large_num_2)\n",
    "\n",
    "    square_large_num(large_num_3)\n",
    "\n",
    "    return add_large_num(added_num, multiplied_num)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execution without covalent\n",
    "\n",
    "Even if you've decorated these functions, they are still executable as normal python functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding 1234567 and 7654321\n",
      "Multiplying 1234567 and 7654321\n",
      "Squaring 12321\n",
      "Adding 8888888 and 9449772114007\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "9449781002895"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l1 = 1234567\n",
    "l2 = 7654321\n",
    "l3 = 12321\n",
    "\n",
    "sample_workflow(l1, l2, l3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Drawing the graph\n",
    "\n",
    "Suppose you want to see how the workflow graph looks like without executing the tasks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_workflow.draw(l1, l2, l3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can go over to [Covalent UI](localhost:48008) and click on the preview tab on your left to check out the workflow graph."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dispatching the workflow\n",
    "\n",
    "Let's dispatch the workflow to covalent server and check its realtime status in the UI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dispatch_id = ct.dispatch(sample_workflow)(l1, l2, l3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting the result back\n",
    "\n",
    "Using the `dispatch_id` you received you can get the result of any execution and see information ranging from its status, each task's output, start and end times to error if any occured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lattice Result\n",
      "==============\n",
      "status: COMPLETED\n",
      "result: 9449781002895\n",
      "inputs: {'args': [1234567, 7654321, 12321], 'kwargs': {}}\n",
      "error: None\n",
      "\n",
      "start_time: 2022-06-10 13:47:40.173737+00:00\n",
      "end_time: 2022-06-10 13:47:40.451817+00:00\n",
      "\n",
      "results_dir: /Users/sankalpsanand/dev/tutorials_covalent_mlops_2022/results\n",
      "dispatch_id: e387be5b-d731-447c-893e-f520c5dc71dc\n",
      "\n",
      "Node Outputs\n",
      "------------\n",
      "add_large_num(0): 8888888\n",
      ":parameter:1234567(1): 1234567\n",
      ":parameter:7654321(2): 7654321\n",
      "multiply_large_num(3): 9449772114007\n",
      ":parameter:1234567(4): 1234567\n",
      ":parameter:7654321(5): 7654321\n",
      "square_large_num(6): 151807041\n",
      ":parameter:12321(7): 12321\n",
      "add_large_num(8): 9449781002895\n",
      "\n"
     ]
    }
   ],
   "source": [
    "sample_workflow_result = ct.get_result(dispatch_id)\n",
    "print(sample_workflow_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Executors\n",
    "\n",
    "Executors are responsible for running your task on a particular backend of your choice. Let's start with a simple `DaskExecutor` which will run your task on a Dask cluster. You can use your custom executors by just overriding one function - [How to create your own executor](https://covalent.readthedocs.io/en/latest/how_to/execution/creating_custom_executors.html).\n",
    "\n",
    "### Workflow level assignment\n",
    "In this case we'll use the executor for the whole workflow (`lattice`) so each electron inside it will be submitted to the running Dask cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dask_executor = ct.executor.DaskExecutor(scheduler_address='tcp://127.0.0.1:55201')\n",
    "\n",
    "@ct.lattice(executor=dask_executor)\n",
    "def sample_workflow(large_num_1, large_num_2, large_num_3):\n",
    "    # Sample workflow to do 2 independent jobs and 1 dependent job\n",
    "\n",
    "    added_num = add_large_num(large_num_1, large_num_2)\n",
    "    multiplied_num = multiply_large_num(large_num_1, large_num_2)\n",
    "\n",
    "    square_large_num(large_num_3)\n",
    "\n",
    "    return add_large_num(added_num, multiplied_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48463d1f-f18f-4fa2-8289-6d2292ab0ee8\n"
     ]
    }
   ],
   "source": [
    "dispatch_id = ct.dispatch(sample_workflow)(l1, l2, l3)\n",
    "print(dispatch_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task level assignment\n",
    "\n",
    "Now let's suppose we want to run all of these different tasks on different hardware backends. We can do that by creating the respective `Executor` objects. To start with we currently provide `SlurmExecutor` - task will be run as a SLURM job and `SSHExecutor` - task will be run on a remote machine connected through ssh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "slurm_executor = ct.executor.SlurmExecutor(\n",
    "    username=\"sankalp\",\n",
    "    address=\"beehive.agnostiq.ai\",\n",
    "    remote_workdir=\"/federation/sankalp/workdir\",\n",
    "    ssh_key_file=\"/Users/sankalpsanand/.ssh/beehive\",\n",
    "    poll_freq=10,\n",
    "    options={\n",
    "        \"partition\": \"debug\", \"ntasks\": 1, \"cpus-per-task\": 1,\n",
    "        \"chdir\": \"/federation/sankalp/workdir\", \"nodelist\": \"beehive-debug-st-t2medium-1\"},\n",
    ")\n",
    "\n",
    "ssh_executor = ct.executor.SSHExecutor(\n",
    "    username=\"sankalp\",\n",
    "    hostname=\"beehive.agnostiq.ai\",\n",
    "    remote_dir=\"/federation/sankalp/workdir\",\n",
    "    ssh_key_file=\"/Users/sankalpsanand/.ssh/beehive\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's assign these executors to the different electrons as desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron(executor=ssh_executor)\n",
    "def add_large_num(num_1, num_2):\n",
    "    # Add two large numbers\n",
    "    \n",
    "    print(f\"Adding {num_1} and {num_2}\")\n",
    "    \n",
    "    return num_1 + num_2\n",
    "\n",
    "\n",
    "@ct.electron(executor=dask_executor)\n",
    "def multiply_large_num(num_1, num_2):\n",
    "    # Multiply two large numbers\n",
    "\n",
    "    print(f\"Multiplying {num_1} and {num_2}\")\n",
    "\n",
    "    return num_1 * num_2\n",
    "\n",
    "@ct.electron(executor=slurm_executor)\n",
    "def square_large_num(num):\n",
    "    # Square a large number\n",
    "\n",
    "    print(f\"Squaring {num}\")\n",
    "\n",
    "    return num ** 2\n",
    "\n",
    "\n",
    "@ct.lattice\n",
    "def sample_workflow(large_num_1, large_num_2, large_num_3):\n",
    "    # Sample workflow to do 2 independent jobs and 1 dependent job\n",
    "\n",
    "    added_num = add_large_num(large_num_1, large_num_2)\n",
    "    multiplied_num = multiply_large_num(large_num_1, large_num_2)\n",
    "\n",
    "    square_large_num(large_num_3)\n",
    "\n",
    "    return add_large_num(added_num, multiplied_num)\n",
    "\n",
    "l1 = 1234567\n",
    "l2 = 7654321\n",
    "l3 = 12321\n",
    "\n",
    "dispatch_id = ct.dispatch(sample_workflow)(l1, l2, l3)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "197454157c080c19c876ac9bd28ff3f683250ebc7e59376e497fe430017b1105"
  },
  "kernelspec": {
   "display_name": "Python 3.8.13 ('cova-qa-1')",
   "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.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}