demo script for switching function execution between solveit @ Modal

run_on.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f6249c8",
   "metadata": {},
   "source": [
    "# run_on() utility (`SolveIt` or `Modal`)\n",
    "\n",
    "Here are the beginnings of a simple utility for running things generically, either locally on Solveit or sending them to `Modal`.\n",
    "\n",
    "It's based on a couple things I learned while reading [Chris Thomas's gist on calling Modal from SolveIt](https://gist.github.com/chris-thomas/1d03d6099307bcc1c63f23236c4f00f4). Check there for more context and additional options, commands, etc. But I wanted something with which I could somewhat seamlessly switch back and forth between where a function was being executed. Hence this.\n",
    "\n",
    "We'll define a function `mymain()` that has a dependency `square()` from another cell. \n",
    "\n",
    "We'll set both of those cells as exportable and then we'll run our `run_on` code block (below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9e24740",
   "metadata": {
    "time_run": "4:18:27a"
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "def square(x):\n",
    "    return x**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e55cacc5",
   "metadata": {
    "time_run": "4:18:29a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# non-exported testing code: \n",
    "square(3) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b236589e",
   "metadata": {
    "time_run": "4:18:31a"
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "def mymain(a=6):\n",
    "    b = square(a) \n",
    "    print(f\"a was {a}, now b is {b}\") # maybe it even has print statements .;-) \n",
    "    return b "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4a917b3",
   "metadata": {
    "time_run": "4:18:32a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a was 6, now b is 36\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "36"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mymain() "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b49acd9c",
   "metadata": {},
   "source": [
    "Say we want to call mymain() and run it locally or remotely.   We need to somehow know what else it's going to call. \n",
    "\n",
    "# Our Utility Code \n",
    "\n",
    "The simple way we do that is we just export all the exportable code cells in the notebook and send those. Period. Obviously, additional imports and things can be managed according to the modal documentation. I was just trying to integrate this workflow into the rapid testing environment and so forth.\n",
    "\n",
    "*warning: janky work in progress. new to `modal`*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1751e70",
   "metadata": {},
   "source": [
    "First a little util to export the code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d090300",
   "metadata": {
    "time_run": "4:18:36a"
   },
   "outputs": [],
   "source": [
    "from dialoghelper import find_msgs\n",
    "\n",
    "def export_code_cells(filename):\n",
    "    exported_msgs = [m for m in find_msgs(msg_type='code') if m.get('is_exported')] \n",
    "    with open(filename, 'w') as f:\n",
    "        for msg in exported_msgs: f.write(msg['content'] + '\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "179cfeb8",
   "metadata": {
    "time_run": "4:18:36a"
   },
   "outputs": [],
   "source": [
    "# run the routine, for testing purposes\n",
    "export_code_cells('sendtomodal.py') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4875b32",
   "metadata": {
    "time_run": "4:18:37a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def square(x):\r\n",
      "    return x**2\r\n",
      "\r\n",
      "def mymain(a=6):\r\n",
      "    b = square(a) \r\n",
      "    print(f\"a was {a}, now b is {b}\") # maybe it even has print statements .;-) \r\n",
      "    return b \r\n",
      "\r\n",
      "# NOTE: for this to work, any cell you want to run on Modal has to be \r\n",
      "# marked as exported, even the function you're calling. \r\n",
      "def test_stderr(x=5):\r\n",
      "    print(f\"Testing with x={x}\")\r\n",
      "    assert x > 10, f\"x must be greater than 10, but got {x}\"\r\n",
      "    return x\r\n",
      "\r\n"
     ]
    }
   ],
   "source": [
    "# check that it worked\n",
    "!cat sendtomodal.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4568631e",
   "metadata": {},
   "source": [
    "That looks good.  \n",
    "\n",
    "Now on to our main code... "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeaf49c1",
   "metadata": {
    "time_run": "4:02:09a"
   },
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "import modal\n",
    "from io import StringIO\n",
    "import sys\n",
    "\n",
    "app = modal.App(\"my-app\")\n",
    "\n",
    "EXPORT_FILE = 'sendtomodal.py' # i hate using a static global name but tmpfile gave me issues\n",
    "open(EXPORT_FILE, 'a').close() # create a blank version, overwrite it below before calling\n",
    "\n",
    "IMAGE = modal.Image.debian_slim().add_local_file(EXPORT_FILE, remote_path=\"/root/\"+EXPORT_FILE)\n",
    "\n",
    "@app.function(image=IMAGE)\n",
    "def modal_wrapper(func_name, *args, **kwargs):\n",
    "    import sys\n",
    "    from io import StringIO\n",
    "    sys.path.insert(0, '/root')\n",
    "    import sendtomodal\n",
    "\n",
    "    stdout_capture = StringIO()  # Modal captures stderr for us automatically btw\n",
    "    sys.stdout = stdout_capture\n",
    "    func = getattr(sendtomodal, func_name)\n",
    "    result = func(*args, **kwargs)\n",
    "    sys.stdout = sys.__stdout__\n",
    "    output = \"=== MODAL ===\\n\" + stdout_capture.getvalue()\n",
    "    return {'result': result, 'output': output}\n",
    "\n",
    "\n",
    "def run_on(func, *args, dest: Literal['solveit', 'modal'] = 'solveit', **kwargs):\n",
    "    global IMAGE\n",
    "    if dest == 'solveit':\n",
    "        return func(*args, **kwargs)\n",
    "\n",
    "    elif dest == 'modal':\n",
    "        export_code_cells(EXPORT_FILE)\n",
    "        IMAGE = modal.Image.debian_slim().add_local_file(EXPORT_FILE, remote_path=\"/root/\"+EXPORT_FILE) # rebuild the image\n",
    "        with app.run():\n",
    "            result = modal_wrapper.remote(func.__name__, *args, **kwargs)\n",
    "            print(result['output']) # make prints happen just like local run \n",
    "            return result['result']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ccf8e87",
   "metadata": {
    "time_run": "4:02:10a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a was 6, now b is 36\n",
      "Result: 36\n"
     ]
    }
   ],
   "source": [
    "result = run_on(mymain, dest='solveit')\n",
    "print(f\"Result: {result}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08f9d3e0",
   "metadata": {
    "time_run": "4:02:13a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== MODAL ===\n",
      "a was 6, now b is 36\n",
      "\n",
      "Result: 36\n"
     ]
    }
   ],
   "source": [
    "result = run_on(mymain, dest='modal')\n",
    "print(f\"Result: {result}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7165e3a2",
   "metadata": {},
   "source": [
    "## What about error messages? \n",
    "Let's test `stderr` capture by triggering a failed assert... "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75041e9c",
   "metadata": {
    "time_run": "4:02:23a"
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "# NOTE: for this to work, any cell you want to run on Modal has to be \n",
    "# marked as exported, even the function you're calling. \n",
    "def test_stderr(x=5):\n",
    "    print(f\"Testing with x={x}\")\n",
    "    assert x > 10, f\"x must be greater than 10, but got {x}\"\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed24e57d",
   "metadata": {},
   "source": [
    "First run it locally... "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da295f4e",
   "metadata": {
    "time_run": "4:02:25a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing with x=5\n"
     ]
    },
    {
     "ename": "AssertionError",
     "evalue": "x must be greater than 10, but got 5",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mAssertionError\u001b[39m                            Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[79]\u001b[39m\u001b[32m, line 1\u001b[39m",
      "\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43mrun_on\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtest_stderr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdest\u001b[49m\u001b[43m=\u001b[49m\u001b[33;43m'\u001b[39;49m\u001b[33;43msolveit\u001b[39;49m\u001b[33;43m'\u001b[39;49m\u001b[43m)\u001b[49m",
      "",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[69]\u001b[39m\u001b[32m, line 32\u001b[39m, in \u001b[36mrun_on\u001b[39m\u001b[34m(func, dest, *args, **kwargs)\u001b[39m",
      "\u001b[32m     30\u001b[39m \u001b[38;5;28;01mglobal\u001b[39;00m IMAGE",
      "\u001b[32m     31\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m dest == \u001b[33m'\u001b[39m\u001b[33msolveit\u001b[39m\u001b[33m'\u001b[39m:",
      "\u001b[32m---> \u001b[39m\u001b[32m32\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m",
      "\u001b[32m     34\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m dest == \u001b[33m'\u001b[39m\u001b[33mmodal\u001b[39m\u001b[33m'\u001b[39m:",
      "\u001b[32m     35\u001b[39m     export_code_cells(EXPORT_FILE)",
      "",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[76]\u001b[39m\u001b[32m, line 5\u001b[39m, in \u001b[36mtest_stderr\u001b[39m\u001b[34m(x)\u001b[39m",
      "\u001b[32m      3\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mtest_stderr\u001b[39m(x=\u001b[32m5\u001b[39m):",
      "\u001b[32m      4\u001b[39m     \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mTesting with x=\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mx\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m)",
      "\u001b[32m----> \u001b[39m\u001b[32m5\u001b[39m     \u001b[38;5;28;01massert\u001b[39;00m x > \u001b[32m10\u001b[39m, \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mx must be greater than 10, but got \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mx\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m",
      "\u001b[32m      6\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m x",
      "",
      "\u001b[31mAssertionError\u001b[39m: x must be greater than 10, but got 5"
     ]
    }
   ],
   "source": [
    "run_on(test_stderr, dest='solveit')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "675942a3",
   "metadata": {},
   "source": [
    "And now on Modal..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5bed574",
   "metadata": {
    "time_run": "4:02:33a"
   },
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "x must be greater than 10, but got 5",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mAssertionError\u001b[39m                            Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[82]\u001b[39m\u001b[32m, line 1\u001b[39m",
      "\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43mrun_on\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtest_stderr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdest\u001b[49m\u001b[43m=\u001b[49m\u001b[33;43m'\u001b[39;49m\u001b[33;43mmodal\u001b[39;49m\u001b[33;43m'\u001b[39;49m\u001b[43m)\u001b[49m",
      "",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[69]\u001b[39m\u001b[32m, line 38\u001b[39m, in \u001b[36mrun_on\u001b[39m\u001b[34m(func, dest, *args, **kwargs)\u001b[39m",
      "\u001b[32m     36\u001b[39m IMAGE = modal.Image.debian_slim().add_local_file(EXPORT_FILE, remote_path=\u001b[33m\"\u001b[39m\u001b[33m/root/\u001b[39m\u001b[33m\"\u001b[39m+EXPORT_FILE) \u001b[38;5;66;03m# rebuild the image\u001b[39;00m",
      "\u001b[32m     37\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m app.run():",
      "\u001b[32m---> \u001b[39m\u001b[32m38\u001b[39m     result = \u001b[43mmodal_wrapper\u001b[49m\u001b[43m.\u001b[49m\u001b[43mremote\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m.\u001b[49m\u001b[34;43m__name__\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m",
      "\u001b[32m     39\u001b[39m     \u001b[38;5;28mprint\u001b[39m(result[\u001b[33m'\u001b[39m\u001b[33moutput\u001b[39m\u001b[33m'\u001b[39m]) \u001b[38;5;66;03m# make prints happen just like local run \u001b[39;00m",
      "\u001b[32m     40\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m result[\u001b[33m'\u001b[39m\u001b[33mresult\u001b[39m\u001b[33m'\u001b[39m]",
      "",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.local/lib/python3.12/site-packages/modal/_object.py:304\u001b[39m, in \u001b[36mlive_method.<locals>.wrapped\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m",
      "\u001b[32m    301\u001b[39m \u001b[38;5;129m@wraps\u001b[39m(method)",
      "\u001b[32m    302\u001b[39m \u001b[38;5;28;01masync\u001b[39;00m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mwrapped\u001b[39m(\u001b[38;5;28mself\u001b[39m, *args, **kwargs):",
      "\u001b[32m    303\u001b[39m     \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m.hydrate()",
      "\u001b[32m--> \u001b[39m\u001b[32m304\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m method(\u001b[38;5;28mself\u001b[39m, *args, **kwargs)",
      "",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.local/lib/python3.12/site-packages/modal/_functions.py:1722\u001b[39m, in \u001b[36m_Function.remote\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m",
      "\u001b[32m   1717\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m._is_generator:",
      "\u001b[32m   1718\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m InvalidError(",
      "\u001b[32m   1719\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mA generator function cannot be called with `.remote(...)`. Use `.remote_gen(...)` instead.\u001b[39m\u001b[33m\"\u001b[39m",
      "\u001b[32m   1720\u001b[39m     )",
      "\u001b[32m-> \u001b[39m\u001b[32m1722\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m._call_function(args, kwargs)",
      "",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.local/lib/python3.12/site-packages/modal/_functions.py:1666\u001b[39m, in \u001b[36m_Function._call_function\u001b[39m\u001b[34m(self, args, kwargs)\u001b[39m",
      "\u001b[32m   1657\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:",
      "\u001b[32m   1658\u001b[39m     invocation = \u001b[38;5;28;01mawait\u001b[39;00m _Invocation.create(",
      "\u001b[32m   1659\u001b[39m         \u001b[38;5;28mself\u001b[39m,",
      "\u001b[32m   1660\u001b[39m         args,",
      "\u001b[32m   (...)\u001b[39m\u001b[32m   1663\u001b[39m         function_call_invocation_type=api_pb2.FUNCTION_CALL_INVOCATION_TYPE_SYNC,",
      "\u001b[32m   1664\u001b[39m     )",
      "\u001b[32m-> \u001b[39m\u001b[32m1666\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m invocation.run_function()",
      "",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.local/lib/python3.12/site-packages/modal/_functions.py:300\u001b[39m, in \u001b[36m_Invocation.run_function\u001b[39m\u001b[34m(self)\u001b[39m",
      "\u001b[32m    292\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m (",
      "\u001b[32m    293\u001b[39m     \u001b[38;5;129;01mnot\u001b[39;00m ctx",
      "\u001b[32m    294\u001b[39m     \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m ctx.retry_policy",
      "\u001b[32m   (...)\u001b[39m\u001b[32m    297\u001b[39m     \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m ctx.sync_client_retries_enabled",
      "\u001b[32m    298\u001b[39m ):",
      "\u001b[32m    299\u001b[39m     item = \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m._get_single_output()",
      "\u001b[32m--> \u001b[39m\u001b[32m300\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m _process_result(item.result, item.data_format, \u001b[38;5;28mself\u001b[39m.stub, \u001b[38;5;28mself\u001b[39m.client)",
      "\u001b[32m    302\u001b[39m \u001b[38;5;66;03m# User errors including timeouts are managed by the user specified retry policy.\u001b[39;00m",
      "\u001b[32m    303\u001b[39m user_retry_manager = RetryManager(ctx.retry_policy)",
      "",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.local/lib/python3.12/site-packages/modal/_utils/function_utils.py:518\u001b[39m, in \u001b[36m_process_result\u001b[39m\u001b[34m(result, data_format, stub, client)\u001b[39m",
      "\u001b[32m    515\u001b[39m             \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:",
      "\u001b[32m    516\u001b[39m                 \u001b[38;5;28;01mpass\u001b[39;00m",
      "\u001b[32m--> \u001b[39m\u001b[32m518\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m exc_with_hints(exc)",
      "\u001b[32m    520\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m RemoteError(result.exception)",
      "\u001b[32m    522\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:",
      "",
      "\u001b[36mFile \u001b[39m\u001b[32m<ta-01K6VTYV85KDHGB6J544PFGSB7>:/tmp/ipykernel_554/3601753661.py:23\u001b[39m, in \u001b[36mmodal_wrapper\u001b[39m\u001b[34m()\u001b[39m",
      "",
      "\u001b[36mFile \u001b[39m\u001b[32m<ta-01K6VTYV85KDHGB6J544PFGSB7>:/root/sendtomodal.py:13\u001b[39m, in \u001b[36mtest_stderr\u001b[39m\u001b[34m()\u001b[39m",
      "",
      "\u001b[31mAssertionError\u001b[39m: x must be greater than 10, but got 5"
     ]
    }
   ],
   "source": [
    "run_on(test_stderr, dest='modal')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bb1945d",
   "metadata": {},
   "source": [
    "Notice that the file containing the offense is `/root/sendtomodal.py`, i.e. the file on Modal! :-) "
   ]
  }
 ],
 "metadata": {
  "solveit_dialog_mode": "learning",
  "solveit_ver": 2
 },
 "nbformat": 4,
 "nbformat_minor": 5
}