{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-31T17:54:21.337343Z",
     "iopub.status.busy": "2023-07-31T17:54:21.336919Z",
     "iopub.status.idle": "2023-07-31T17:54:21.495527Z",
     "shell.execute_reply": "2023-07-31T17:54:21.494981Z",
     "shell.execute_reply.started": "2023-07-31T17:54:21.337320Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import wonderwords\n",
    "from tqdm import tqdm\n",
    "from tritonclient.utils import *\n",
    "from random import choice, randrange\n",
    "import tritonclient.http as httpclient\n",
    "from multiprocessing.pool import ThreadPool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-31T17:54:22.003332Z",
     "iopub.status.busy": "2023-07-31T17:54:22.002684Z",
     "iopub.status.idle": "2023-07-31T17:54:22.008703Z",
     "shell.execute_reply": "2023-07-31T17:54:22.007697Z",
     "shell.execute_reply.started": "2023-07-31T17:54:22.003294Z"
    }
   },
   "outputs": [],
   "source": [
    "shape = [1]\n",
    "MIN_WORDS, MAX_WORDS = 4, 20\n",
    "model_name = \"ssmt_pipeline\"\n",
    "rs = wonderwords.RandomWord()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Hit the tritonserver with a random sentence to a random model\n",
    "* See https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/metrics.html for metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-31T17:54:23.090209Z",
     "iopub.status.busy": "2023-07-31T17:54:23.089654Z",
     "iopub.status.idle": "2023-07-31T17:54:23.100644Z",
     "shell.execute_reply": "2023-07-31T17:54:23.100011Z",
     "shell.execute_reply.started": "2023-07-31T17:54:23.090177Z"
    }
   },
   "outputs": [],
   "source": [
    "def task(x):\n",
    "    lang_pair_map = list({'en-hi': 1, 'hi-en': 2, 'te-en': 4, 'hi-te': 6, 'te-hi': 7, 'en-gu': 8, 'gu-en': 9}.keys())\n",
    "    with httpclient.InferenceServerClient(\"localhost:8000\") as client:\n",
    "        async_responses = []\n",
    "        for i in range(10):\n",
    "            s = ' '.join(rs.random_words(randrange(MIN_WORDS, MAX_WORDS)) + ['.']) # 'this is a sentence.' Use a constant sentence if you want to hit the cache\n",
    "            source_data = np.array([[s]], dtype='object')\n",
    "            inputs = [httpclient.InferInput(\"INPUT_TEXT\", source_data.shape, np_to_triton_dtype(source_data.dtype)), httpclient.InferInput(\"INPUT_LANGUAGE_ID\", source_data.shape, np_to_triton_dtype(source_data.dtype)), httpclient.InferInput(\"OUTPUT_LANGUAGE_ID\", source_data.shape, np_to_triton_dtype(source_data.dtype))]\n",
    "            inputs[0].set_data_from_numpy(np.array([[s]], dtype='object'))\n",
    "            langpair = choice(lang_pair_map)\n",
    "            inputs[1].set_data_from_numpy(np.array([[langpair.split('-')[0].strip()]], dtype='object'))\n",
    "            inputs[2].set_data_from_numpy(np.array([[langpair.split('-')[1].strip()]], dtype='object'))\n",
    "            outputs = [httpclient.InferRequestedOutput(\"OUTPUT_TEXT\")]\n",
    "            async_responses.append(client.async_infer(model_name, inputs, request_id=str(1), outputs=outputs))\n",
    "        for r in async_responses: r.get_result(timeout=10).get_response()\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-31T17:55:24.235964Z",
     "iopub.status.busy": "2023-07-31T17:55:24.235574Z",
     "iopub.status.idle": "2023-07-31T17:58:30.757911Z",
     "shell.execute_reply": "2023-07-31T17:58:30.756271Z",
     "shell.execute_reply.started": "2023-07-31T17:55:24.235935Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [03:06<00:00,  5.36it/s]\n"
     ]
    }
   ],
   "source": [
    "with ThreadPool(100) as pool:\n",
    "    for output in tqdm(pool.imap_unordered(task, range(1000), chunksize=1), total=1000): pass"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mt-model-deploy-dhruva",
   "language": "python",
   "name": "mt-model-deploy-dhruva"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}