docs/hosting/cas-on-colab.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "cas-on-colab.ipynb",
   "provenance": [],
   "collapsed_sections": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Hosting CLIP-as-service on Google Colab with TPU/GPU support\n",
    "\n",
    "This tutorial guides you on how to implement the following architecture:\n",
    "\n",
    "[![](https://mermaid.ink/img/pako:eNp1kEFrwzAMhf-K0bkh99xGVwpjh9Ctp7oMxVYTM8cOttwy2v732fMGgzFd9Hjvk0C6gvKaoIMx4DKJ5510IldMQzW23o-WxNpbHMRBlXasSCmF8S1SOFNommbb79vXfl9TcrqKh0OBlDXk-Cgydht3_bqdmJf2QkP06p349mtTHXvCM0YVzMJfMwX_C6kU7L8xrGCmMKPR-bprcSTwRDNJ6LLUdMJkWYJ094ymRSPTRhv2AboT2kgrwMT-5cMp6Dgk-oEeDebfzN_U_RP7v2yd)](https://mermaid.live/edit#pako:eNp1kEFrwzAMhf-K0bkh99xGVwpjh9Ctp7oMxVYTM8cOttwy2v732fMGgzFd9Hjvk0C6gvKaoIMx4DKJ5510IldMQzW23o-WxNpbHMRBlXasSCmF8S1SOFNommbb79vXfl9TcrqKh0OBlDXk-Cgydht3_bqdmJf2QkP06p349mtTHXvCM0YVzMJfMwX_C6kU7L8xrGCmMKPR-bprcSTwRDNJ6LLUdMJkWYJ094ymRSPTRhv2AboT2kgrwMT-5cMp6Dgk-oEeDebfzN_U_RP7v2yd)\n",
    "\n",
    "CLIP-as-service is powered by Jina, [there is another tutorial showing you how to host Jina service on Colab in general](https://colab.research.google.com/github/jina-ai/jina/blob/master/docs/Using_Jina_on_Colab.ipynb). Highly recommended!\n",
    "\n",
    "\n",
    "## 1. Change runtime type\n",
    "\n",
    "Go to menu `Runtime -> Change run time type -> GPU/TPU`\n",
    "\n",
    "\n",
    "## 2. Install Packages\n",
    "\n",
    "As we will run the client locally, we only need to install `clip_server` package on Colab.\n",
    "\n",
    "\n",
    "**⚠️ You will be asked to \"Restart Runtime\" after this step, please click the button and restart the runtime.**"
   ],
   "metadata": {
    "id": "lbUpcvs1p1CF",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MRrB2If6kDfX",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!pip install clip_server pyngrok"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 3. Config Flow YAML\n",
    "\n",
    "\n",
    "Unlike classic entrypoint from CLI, here we need to start the Flow in Python. Let's load use Pytorch backend and write a Flow YAML. Note that we need to load the torch Python file from `clip_server` installation, hence you see `cas_path` below. More available options [can be found here](https://github.com/jina-ai/clip-as-service/tree/main/server/clip_server/executors)."
   ],
   "metadata": {
    "id": "q3bmGKIvx5S-",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "import clip_server\n",
    "cas_path = clip_server.__path__[0]"
   ],
   "metadata": {
    "id": "nypR4g9EmgOj",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "This YAML is directly [taken from this file](https://github.com/jina-ai/clip-as-service/blob/main/server/clip_server/torch-flow.yml). You can also customize it as you wish, [please check CLIP-as-service docs](https://clip-as-service.jina.ai/user-guides/server/#yaml-config)."
   ],
   "metadata": {
    "id": "5RVA1OD8ywOo",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "flow_yaml = f'''\n",
    "jtype: Flow\n",
    "with:\n",
    "  port: 51000\n",
    "executors:\n",
    "  - name: clip_t\n",
    "    uses:\n",
    "      jtype: CLIPEncoder\n",
    "      metas:\n",
    "        py_modules:\n",
    "          - {cas_path}/executors/clip_torch.py\n",
    "'''"
   ],
   "metadata": {
    "id": "q1BXWnXVkIZ8",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "flow_yaml"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "id": "Fb1PKf992rLj",
    "outputId": "a06b634a-5021-4b24-f3dc-a2c6b1d87524",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "'\\njtype: Flow\\nwith:\\n  port: 51000\\nexecutors:\\n  - name: clip_t\\n    uses:\\n      jtype: CLIPEncoder\\n      metas:\\n        py_modules:\\n          - /usr/local/lib/python3.7/dist-packages/clip_server/executors/clip_torch.py\\n'"
      ],
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      }
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 4. Start the Flow\n",
    "\n",
    "It may take a minute or so on the first start, as it will download the pretrained models. To select different pretrained models, [please check CLIP-as-service docs](https://clip-as-service.jina.ai/user-guides/server/#yaml-config)."
   ],
   "metadata": {
    "id": "GvAeaUf4y88e",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "from jina import Flow\n",
    "\n",
    "f = Flow.load_config(flow_yaml)\n",
    "f.start()"
   ],
   "metadata": {
    "id": "4UubypFpl8-K",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Remember to close it via `f.close()` when you don't use it. But let's keep it open for now."
   ],
   "metadata": {
    "id": "2BOYxmpd8YSE",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 5. Set up forwarding\n",
    "\n",
    "By default Flow uses gRPC protocol, it is highly-efficient and feature-rich. So in this tutorial, we will use gRPC protocol and use `ngrok` for forwarding. It is possible and in fact slighly easier to set up when using `Flow(protocol='http')`, [please read the turorial here](https://colab.research.google.com/github/jina-ai/jina/blob/master/docs/Using_Jina_on_Colab.ipynb#scrollTo=0ASjGLBhXono) here I won't repeat again.\n",
    "\n",
    "\n",
    "You will need to first sign up at https://dashboard.ngrok.com/signup (http do not need register, that's why I said it is easier)\n",
    "\n",
    "After signing up, you can get a token. Then simply add your token via (replacing `YOUR_TOKEN_HERE`)"
   ],
   "metadata": {
    "id": "1lTqYEwezDTP",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "!pip install pyngrok\n",
    "\n",
    "# remember to replace to your token! otherwise i can see your service, i mean i dont really have time to see it but nonetheless\n",
    "!ngrok authtoken 2ARsKtGKj47h7y4uXMQPrIeOinS_47Mkh6jkzNjFEJWuZYNEX"
   ],
   "metadata": {
    "id": "PYQPKek-oG1a",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "!ngrok tcp 51000 --log \"stdout\""
   ],
   "metadata": {
    "id": "2Hacpj4qn9nx",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "At the last line, you should see something like: \n",
    "\n",
    "```\n",
    "t=2022-06-11T20:29:11+0000 lvl=info msg=\"started tunnel\" obj=tunnels name=command_line addr=//localhost:54321 url=tcp://6.tcp.ngrok.io:18096\n",
    "```\n",
    "\n",
    "Grab the text after `url=tcp://` in my case it is `6.tcp.ngrok.io:18096`.\n",
    "\n",
    "Now build a client using this address from your local laptop/Python environment.\n",
    "\n",
    "Copy paste the code below to your local Python, remmeber to change your address.\n",
    "\n",
    "**Remember, if your last line is `url=tcp://6.tcp.ngrok.io:18096` then you should set `Client('grpc://6.tcp.ngrok.io:18096')`**\n",
    "\n",
    "### Try Embedding Task from Local\n",
    "\n",
    "```python\n",
    "# pip install clip-client\n",
    "from clip_client import Client\n",
    "\n",
    "c = Client('grpc://6.tcp.ngrok.io:18096')\n",
    "\n",
    "r = c.encode(\n",
    "    [\n",
    "        'First do it',\n",
    "        'then do it right',\n",
    "        'then do it better',\n",
    "        'https://picsum.photos/200',\n",
    "    ]\n",
    ")\n",
    "print(r)\n",
    "```\n",
    "\n",
    "And you will get \n",
    "\n",
    "```text\n",
    "[[ 0.03494263 -0.23510742  0.0104599  ... -0.5229492  -0.10021973\n",
    "  -0.08685303]\n",
    " [-0.06793213 -0.0032444   0.01506805 ... -0.50341797 -0.06143188\n",
    "  -0.08520508]\n",
    " [ 0.15063477 -0.07922363 -0.06530762 ... -0.46484375 -0.08526611\n",
    "   0.04324341]\n",
    " [-0.16088867  0.10552979 -0.20581055 ... -0.41381836  0.19543457\n",
    "   0.05718994]]\n",
    "```\n",
    "\n",
    "Showing the connection is success!\n",
    "\n",
    "\n",
    "### Try Ranking Task from Local\n",
    "\n",
    "```python\n",
    "from docarray import Document\n",
    "\n",
    "from clip_client import Client\n",
    "\n",
    "c = Client(server='grpc://6.tcp.ngrok.io:18096/rank')\n",
    "\n",
    "r = c.rank(\n",
    "    [\n",
    "        Document(\n",
    "            uri='https://picsum.photos/id/1/300/300',\n",
    "            matches=[\n",
    "                Document(text=f'a photo of a {p}')\n",
    "                for p in (\n",
    "                    'man',\n",
    "                    'woman',\n",
    "                )\n",
    "            ],\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "print(r['@m', ['text', 'scores']])\n",
    "```\n",
    "\n",
    "```\n",
    "[['a photo of a man', 'a photo of a woman'], [defaultdict(<class 'docarray.score.NamedScore'>, {'clip_score': {'value': 0.5806832313537598, 'op_name': 'softmax'}, 'clip_score_cosine': {'value': 0.2178003191947937, 'op_name': 'cosine'}}), defaultdict(<class 'docarray.score.NamedScore'>, {'clip_score': {'value': 0.41931676864624023, 'op_name': 'softmax'}, 'clip_score_cosine': {'value': 0.21454453468322754, 'op_name': 'cosine'}})]]\n",
    "```\n",
    "\n",
    "\n",
    "Now enjoy the free GPU/TPU to build your awesome CAS applications!"
   ],
   "metadata": {
    "id": "Fzxt8j3Bz9Nu",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "f.close()"
   ],
   "metadata": {
    "id": "wzj0pb7qo56c",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Push to the Limit\n",
    "\n",
    "Now let's use the biggest `ViT-L/14-336px` and fully leverage all VRAM with 4 replicas, lets see if it works.\t"
   ],
   "metadata": {
    "id": "c6yNVg69-vaw",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "flow_yaml = f'''\n",
    "jtype: Flow\n",
    "with:\n",
    "  port: 51000\n",
    "executors:\n",
    "  - name: clip_t\n",
    "    uses:\n",
    "      jtype: CLIPEncoder\n",
    "      metas:\n",
    "        py_modules:\n",
    "          - {cas_path}/executors/clip_torch.py\n",
    "    replicas: 4\n",
    "'''"
   ],
   "metadata": {
    "id": "uHHWk3WF_DaO",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 12,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "from jina import Flow\n",
    "\n",
    "f = Flow.load_config(flow_yaml)\n",
    "f.start()"
   ],
   "metadata": {
    "id": "0AGcGasu_JIv",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "!ngrok tcp 51000 --log \"stdout\""
   ],
   "metadata": {
    "id": "DQzvwOF3_K6U",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Yay it works!"
   ],
   "metadata": {
    "id": "8T2z6HXd_hKB",
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "id": "4-y_vbHW_acV",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}