Papermill Flyte Inputs

papermill.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "910fa854-eb78-4330-9753-b6fa3d3be52a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77bad3f35a6f4f6fb777e9a14481babc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Executing:   0%|          | 0/7 [00:00<?, ?cell/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from flytekit.models.literals import Literal\n",
    "from flytekit import StructuredDataset\n",
    "from flytekit.types.file import FlyteFile\n",
    "from flytekit.core.context_manager import FlyteContext \n",
    "from flytekit.core.type_engine import TypeEngine\n",
    "from flytekit.core import context_manager\n",
    "from flyteidl.core import literals_pb2\n",
    "from flytekit.core import utils\n",
    "\n",
    "import tempfile\n",
    "\n",
    "import papermill as pm\n",
    "\n",
    "\n",
    "### I papermill task execution\n",
    "def convert_inputs_to_lit_and_save(**kwargs) -> str:\n",
    "    \"\"\"\n",
    "    Use this method to record outputs from a notebook.\n",
    "    It will convert all outputs to a Flyte understandable format. For Files, Directories, please use FlyteFile or\n",
    "    FlyteDirectory, or wrap up your paths in these decorators.\n",
    "    \"\"\"\n",
    "    if kwargs is None:\n",
    "        return \"\"\n",
    "\n",
    "    m = {}\n",
    "    ctx = FlyteContext.current_context()\n",
    "    for k, v in kwargs.items():\n",
    "        expected = TypeEngine.to_literal_type(type(v))\n",
    "        lit = TypeEngine.to_literal(ctx, python_type=type(v), python_val=v, expected=expected)\n",
    "        \n",
    "        tmp = tempfile.mktemp(suffix=\"bin\")\n",
    "        utils.write_proto_to_file(lit.to_flyte_idl(), tmp)\n",
    "        m[k] = tmp\n",
    "\n",
    "    return m\n",
    "\n",
    "df = pd.DataFrame({\"a\": [1, 2,3]})\n",
    "ff = FlyteFile(path=\"./env.yaml\")\n",
    "\n",
    "inputs = convert_inputs_to_lit_and_save(df=df, ff=ff, a=1)\n",
    "\n",
    "_ = pm.execute_notebook(\n",
    "   './dummy_notebook.ipynb',\n",
    "   './dummy_notebook-rendered.ipynb',\n",
    "   parameters=inputs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2c6b5f5-0592-49a4-95eb-5b97b2ff693f",
   "metadata": {
    "tags": []
   },
   "source": [
    "### NOTEBOOK\n",
    "```python\n",
    "from flytekit.core import context_manager\n",
    "from flyteidl.core import literals_pb2\n",
    "from flytekit.core import utils\n",
    "from flytekit.core.type_engine import TypeEngine\n",
    "from flytekit.models.literals import Literal\n",
    "from flytekit.types.file import FlyteFile\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "def read_input(path: str, dtype):\n",
    "    proto = utils.load_proto_from_file(literals_pb2.Literal, path)\n",
    "    lit = Literal.from_flyte_idl(proto)\n",
    "    ctx = context_manager.FlyteContext.current_context()\n",
    "    python_value = TypeEngine.to_python_value(ctx, lit, dtype)    \n",
    "    return python_value\n",
    "\n",
    "### PAPERMILL FILLS IN \n",
    "df = \"tmp/path\"\n",
    "ff = \"tmp/path\"\n",
    "a = \"tmp/path\"\n",
    "\n",
    "### USERS need to do this\n",
    "#### possibly only required for non json serializable inputs\n",
    "df = read_input(df, pd.DataFrame)\n",
    "ff = read_input(ff, FlyteFile)\n",
    "a = read_input(a, int)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fefb8f07-ca3e-4c68-998e-8e3458ac9613",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}