notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modal demo\n",
    "\n",
    "This notebook uses Modal to run code remotely. Before running this notebook, you need to authenticate:\n",
    "\n",
    "```bash\n",
    "uv run modal setup\n",
    "```\n",
    "\n",
    "Then restart the notebook kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local metrics\n",
    "\n",
    "Let's get some metrics displayed right here in the notebook! We'll define a function to draw a loss chart. This function will be called several times during training, and it should update the chart each time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from IPython.display import display\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from utils.modal import SyncHandler\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Metrics:\n",
    "    epoch: int\n",
    "    loss: float\n",
    "\n",
    "\n",
    "def plot_history(history: list[Metrics]):\n",
    "    xs = [h.epoch for h in history]\n",
    "    ys = [h.loss for h in history]\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(8, 4))\n",
    "    ax.set_title('Training progress')\n",
    "    ax.set_xlabel('Epoch')\n",
    "    ax.set_ybound(0, 1)\n",
    "    ax.plot(xs, ys, label='Loss')\n",
    "    ax.legend()\n",
    "    plt.close(fig)\n",
    "    return fig\n",
    "\n",
    "\n",
    "def progress() -> SyncHandler[Metrics]:\n",
    "    history: list[Metrics] = []\n",
    "    display_id = display(plot_history(history), display_id=True)\n",
    "\n",
    "    def receive(metrics: list[Metrics]):\n",
    "        history.extend(metrics)\n",
    "        display_id.update(plot_history(history))\n",
    "\n",
    "    return receive"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remote function\n",
    "\n",
    "Here we define a mock training function that will run remotely. It just loops a few times and returns a stub model function.\n",
    "\n",
    "We specify the exact packages that we'll need in the image to keep it small. Version specifiers are needed (see [`freeze`](src/dair/requirements.py)), so that:\n",
    "- The remote function behaves exactly how it would locally\n",
    "- Objects can be pickled and sent back and forth.\n",
    "\n",
    "We'll add `modal` itself as a dependency, because it's used by `emit_metrics` (see _Training_ below).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import sleep\n",
    "import modal\n",
    "\n",
    "from utils.requirements import freeze\n",
    "from utils.modal import SyncHandler\n",
    "\n",
    "image=(\n",
    "    modal.Image\n",
    "    .debian_slim()\n",
    "    .pip_install(freeze('modal', 'matplotlib'))\n",
    "    .add_local_python_source('utils')\n",
    ")\n",
    "\n",
    "app = modal.App()\n",
    "\n",
    "@app.function(image=image, gpu=None)\n",
    "def train(epochs: int, emit_metrics: SyncHandler[Metrics]):\n",
    "    print(\"Training...\")\n",
    "\n",
    "    for i in range(epochs):\n",
    "        emit_metrics([Metrics(epoch=i+1, loss=1/(i+1))])\n",
    "        sleep(0.2)\n",
    "\n",
    "    def stub_model(x):\n",
    "        if x == \"What is your quest?\":\n",
    "            return \"To seek the Holy Grail.\"\n",
    "        elif x == \"What is the air-speed velocity of an unladen swallow?\":\n",
    "            return \"What do you mean? An African or European swallow?\"\n",
    "        else:\n",
    "            return \"I don't know that!\"\n",
    "\n",
    "    print(\"Training complete\")\n",
    "    return stub_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prevent build logs from showing in the next step. Currently Modal doesn't provide a way to separate the build logs from container stdout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@app.function(image=image)\n",
    "def prebuild(): pass\n",
    "with app.run(): prebuild.remote()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "Now let's run the training code remotely.\n",
    "\n",
    "A [distributed `Queue`](https://modal.com/docs/reference/modal.Queue) is used to send progress information back during training. You can push rich data onto the queue (like actual Matplotlib figures), and it transparently handles serialization - but in this example, a simple dataclass is emitted. The progress function is wrapped in [`send_to`](src/dair/comms.py): a context manager that provides a simple interface over the queue.\n",
    "\n",
    "If we only cared about the final result, or if we were happy just printing progress to stdout, we could call `train` synchronously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.modal import run, send_to\n",
    "\n",
    "async with run(app), send_to.batch(progress()) as emit_metrics:\n",
    "    model = await train.remote.aio(20, emit_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing\n",
    "\n",
    "The model was created remotely, serialized, and sent back. Now we can run it locally!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from textwrap import dedent\n",
    "\n",
    "x = \"What is your quest?\"\n",
    "print(dedent(f\"\"\"\n",
    "    {x}\n",
    "    {model(x)}\n",
    "    \"\"\").strip())\n",
    "\n",
    "x = \"What is the air-speed velocity of an unladen swallow?\"\n",
    "print(dedent(f\"\"\"\n",
    "    {x}\n",
    "    {model(x)}\n",
    "    {model(model(x))}\n",
    "    \"\"\").strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}