{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python Runtime\n",
    "\n",
    "In addition to TensorFlow and ONNX model runtimes, Wallaroo supports Python runtimes. A Python model is treated like other model types and can be a step in a pipeline.\n",
    "\n",
    "## Uploading\n",
    "The python should be saved to a `.py` file. Available imports include all standard python 3.8 libraries, Pandas, and Numpy.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\"\"\"Pre process adds some numbers into a tensor and specifies the field\"\"\"\n",
      "# Arbitrary standard library import. Json is loaded by generated wrapper.\n",
      "import math\n",
      "\n",
      "def wallaroo_json(data):\n",
      "    \"\"\"wallaroo_json takes a json string as input\"\"\"\n",
      "    obj = json.loads(data)\n",
      "    #>>> math.exp(1)\n",
      "    #2.718281828459045\n",
      "    #>>> math.exp(2)\n",
      "    #7.38905609893065\n",
      "    if \"query\" in obj:\n",
      "        return {\n",
      "            \"tensor_fields\": [\"foo\"],\n",
      "            \"query\": obj[\"query\"],\n",
      "            \"foo\": [[math.exp(1), math.exp(2)]]\n",
      "            }\n",
      "    else:\n",
      "        return {\n",
      "            \"tensor_fields\": [\"passthrough\"],\n",
      "            \"passthrough\": obj\n",
      "        }\n"
     ]
    }
   ],
   "source": [
    "! cat ./test_resources/pre_process.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Entrypoint\n",
    "The python must contain the global entrypoint function `wallaroo_json()`.\n",
    "\n",
    "The function will be passed a string containing valid JSON.  The function must return a Python Dict. Example input string:\n",
    "\n",
    "```json\n",
    "{ \"tensor_fields\": [\"foo\"], \n",
    "  \"otherstuff\": \"charge\",\n",
    "  \"state\": \"spin-up\",\n",
    "  \"foo\": [1, 2, 3, 4, 5] }\n",
    "```\n",
    "\n",
    "## Tensor fields\n",
    "The `tensor_fields` key may be present; its value is an array of field names which identify which fields in the structure are input tensors to the step. The default tensor field is `tensor`.\n",
    "\n",
    "## State\n",
    "The JSON may also contain a field with key `state`, which will be passed to all steps in the pipeline. If the step modifies the value of this field, the new value will be propagate to later steps in the pipeline.\n",
    "\n",
    "## Example Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import wallaroo\n",
    "import pandas as pd\n",
    "\n",
    "wl = wallaroo.Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload models\n",
    "model1 = wl.upload_model(\"preprocess\", \"./test_resources/pre_process.py\").configure('python')\n",
    "model2 = wl.upload_model(\"noopfloats\", \"./test_resources/no-op-floats.onnx\")\n",
    "model3 = wl.upload_model(\"postprocess\", \"./test_resources/post_process.py\").configure('python')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "python_pipeline = (wl.build_pipeline(\"pythonpipeline\")\n",
    "        .add_model_step(model1)\n",
    "        .add_model_step(model2)\n",
    "        .add_model_step(model3)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Waiting for deployment - this will take up to 45s ..... ok\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'name': 'pythonpipeline', 'create_time': datetime.datetime(2022, 3, 3, 23, 21, 4, 117316, tzinfo=tzutc()), 'definition': \"[{'ModelInference': {'models': [{'name': 'preprocess', 'version': 'd761cad1-d740-4fb2-9b74-2fff62741cfc', 'sha': '5553423b0d1f5863d76e74a5e852ee9a99a1e327762fb1fca78e6fa034b61257'}]}}, {'ModelInference': {'models': [{'name': 'noopfloats', 'version': '30359068-ddf0-48fc-9f5f-60000d93e949', 'sha': '4dc88d159249ccce83942ada69b919cb91455d5fd0e4bfc287de3f21d1aafb1b'}]}}, {'ModelInference': {'models': [{'name': 'postprocess', 'version': '9112aace-c7f2-4008-8602-03a5933c64ea', 'sha': '50202d07672dad4b9a4d9e07dc83d9c8a9730ff6b5f2c491522dc13ebfd6fb7c'}]}}]\"}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "python_pipeline.deploy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running inference on the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[InferenceResult({'check_failures': [],\n",
       "  'elapsed': 1201278,\n",
       "  'model_name': 'postprocess',\n",
       "  'model_version': '27f63c7a-4761-4076-b67f-3120ee784b3d',\n",
       "  'original_data': {'query': 'this will give us the first conditional branch'},\n",
       "  'outputs': [{'Json': {'data': [{'original': {'outputs': [{'Float': {'data': [2.7182817459106445,\n",
       "                                                                               7.389056205749512],\n",
       "                                                                      'dim': [1,\n",
       "                                                                              2],\n",
       "                                                                      'v': 1}}]},\n",
       "                                  'product': 20.085536603596665}],\n",
       "                        'dim': [1],\n",
       "                        'v': 1}}],\n",
       "  'pipeline_name': 'pythonpipeline',\n",
       "  'time': 1644011878036})]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "python_pipeline.infer({\"query\": \"this will give us the first conditional branch\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Under the hood\n",
    "\n",
    "### First Step: Preprocessing Python\n",
    "\n",
    "This is an example of what is happening in the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Pre process adds some numbers into a tensor and specifies the field\"\"\"\n",
    "# Arbitrary standard library import. Json is loaded by generated wrapper. It is added here so it can run in the notebook\n",
    "import math\n",
    "import json\n",
    "\n",
    "def wallaroo_json(data):\n",
    "    \"\"\"wallaroo_json takes a json string as input\"\"\"\n",
    "    obj = json.loads(data)\n",
    "    if \"query\" in obj:\n",
    "        return {\n",
    "            \"tensor_fields\": [\"foo\"],\n",
    "            \"query\": obj[\"query\"],\n",
    "            \"foo\": [[math.exp(1), math.exp(2)]],\n",
    "            \"state\": \"First branch\"\n",
    "            }\n",
    "    else:\n",
    "        return {\n",
    "            \"tensor_fields\": [\"foo\"],\n",
    "            \"query\": obj[\"query\"],\n",
    "            \"foo\": [[math.exp(2), math.exp(3)]],\n",
    "            \"state\": \"Second branch\"\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tensor_fields': ['foo'],\n",
       " 'query': 'this will give us the first conditional branch',\n",
       " 'foo': [[2.718281828459045, 7.38905609893065]],\n",
       " 'state': 'First branch'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wallaroo_json('{\"query\": \"this will give us the first conditional branch\"}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Second Step: Model Inference\n",
    "In this example, we're just running the tensor through a no op ONNX model, just so we're clear on the output.\n",
    "\n",
    "The input that the next python step will receive will be\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"outputs\" :[{\"Float\":{\"v\":1,\"dim\":[1,2],\"data\":[2.7182817459106445,7.389056205749512]}}],\n",
    "    \"state\": \"First branch\"\n",
    "}\n",
    "```\n",
    "\n",
    "### Third Step: Post Processing\n",
    "\n",
    "This step will calcluate the product of the output tensor from the previous step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Post process gets the product of the final output\"\"\"\n",
    "import math\n",
    "import json\n",
    "def wallaroo_json(data):\n",
    "    \"\"\"wallaroo_json takes a json string as input\"\"\"\n",
    "    obj = json.loads(data)\n",
    "    result = {\n",
    "        \"original\": obj\n",
    "    }\n",
    "    result[\"product\"] = math.prod(obj[\"outputs\"][0]['Float']['data'])\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'original': {'outputs': [{'Float': {'v': 1,\n",
       "     'dim': [1, 2],\n",
       "     'data': [2.7182817459106445, 7.389056205749512]}}],\n",
       "  'state': 'First branch'},\n",
       " 'product': 20.085536603596665}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wallaroo_json('{\"outputs\" :[{\"Float\":{\"v\":1,\"dim\":[1,2],\"data\":[2.7182817459106445,7.389056205749512]}}],\"state\": \"First branch\"}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The json that will come back from the pipeline inference call will be\n",
    "\n",
    "```json\n",
    "{\n",
    "  \"model_id\": \"post_process\",\n",
    "  \"model_version\": \"version\",\n",
    "  \"pipeline_id\": \"intermediate\",\n",
    "  \"outputs\": [\n",
    "    {\n",
    "      \"Json\": {\n",
    "        \"v\": 1,\n",
    "        \"dim\": [\n",
    "          1\n",
    "        ],\n",
    "        \"data\": [\n",
    "          {\n",
    "            \"original\": {\n",
    "              \"outputs\": [\n",
    "                {\n",
    "                  \"Float\": {\n",
    "                    \"v\": 1,\n",
    "                    \"dim\": [\n",
    "                      1,\n",
    "                      2\n",
    "                    ],\n",
    "                    \"data\": [\n",
    "                      2.7182817459106445,\n",
    "                      7.389056205749512\n",
    "                    ]\n",
    "                  }\n",
    "                }\n",
    "              ],\n",
    "              \"state\": \"First branch\"\n",
    "            },\n",
    "            \"product\": 20.085536603596665\n",
    "          }\n",
    "        ]\n",
    "      }\n",
    "    }\n",
    "  ],\n",
    "  \"elapsed\": 379505,\n",
    "  \"time\": 2,\n",
    "  \"original_data\": {\n",
    "    \"query\": \"The preprocessor operates conditionally on this key\"\n",
    "  },\n",
    "  \"check_failures\": []\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'pythonpipeline', 'create_time': datetime.datetime(2022, 3, 3, 23, 21, 4, 117316, tzinfo=tzutc()), 'definition': \"[{'ModelInference': {'models': [{'name': 'preprocess', 'version': 'd761cad1-d740-4fb2-9b74-2fff62741cfc', 'sha': '5553423b0d1f5863d76e74a5e852ee9a99a1e327762fb1fca78e6fa034b61257'}]}}, {'ModelInference': {'models': [{'name': 'noopfloats', 'version': '30359068-ddf0-48fc-9f5f-60000d93e949', 'sha': '4dc88d159249ccce83942ada69b919cb91455d5fd0e4bfc287de3f21d1aafb1b'}]}}, {'ModelInference': {'models': [{'name': 'postprocess', 'version': '9112aace-c7f2-4008-8602-03a5933c64ea', 'sha': '50202d07672dad4b9a4d9e07dc83d9c8a9730ff6b5f2c491522dc13ebfd6fb7c'}]}}]\"}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "python_pipeline.undeploy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
