workshop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Union Workflows on a Jupyter Notebook\n",
    "\n",
    "First create a Union Serverless account:\n",
    "👉 https://signup.union.ai/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First install the dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install \"flytekit==1.14.0b6\" union \"pydantic>2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, login to Union on this notebook session:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!union create login --auth device-flow --serverless"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create some data\n",
    "\n",
    "Let's create some toy data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import flytekit as fl\n",
    "import sys\n",
    "\n",
    "\n",
    "image = fl.ImageSpec(\n",
    "    name=\"jupyter-notebook-workshop\",\n",
    "    packages=[\"flytekit==1.14.0b6\", \"union==0.1.95\"],\n",
    "    python_version=f\"{sys.version_info.major}.{sys.version_info.minor}\",\n",
    ")\n",
    "\n",
    "\n",
    "task = fl.task(container_image=image)\n",
    "\n",
    "\n",
    "@task\n",
    "def create_dataset(n: int) -> list[int]:\n",
    "    return [*range(n)]\n",
    "\n",
    "\n",
    "@fl.workflow\n",
    "def dataset_wf(n: int = 100) -> list[int]:\n",
    "    return create_dataset(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running on Union\n",
    "\n",
    "Let's define a `UnionRemote` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from union.remote import UnionRemote\n",
    "\n",
    "serverless = UnionRemote()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we can execute the workflow like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_exec = serverless.execute(dataset_wf, inputs={\"n\": 50})\n",
    "print(dataset_exec.execution_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can fetch the output of the workflow back into our notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait for the task to complete\n",
    "dataset_exec = dataset_exec.wait(poll_interval=1)\n",
    "\n",
    "# Print the outputs\n",
    "print(dataset_exec.outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute the mean\n",
    "\n",
    "We can then pass the data back into another workflow that computes the mean:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@task\n",
    "def compute_mean(data: list[int]) -> float:\n",
    "    return sum(data) / len(data)\n",
    "\n",
    "\n",
    "@fl.workflow\n",
    "def mean_wf(data: list[int]) -> float:\n",
    "    return compute_mean(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the workflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute the workflow\n",
    "data, *_ = dataset_exec.outputs.values()\n",
    "mean_wf_exec = serverless.execute(mean_wf, inputs={\"data\": data})\n",
    "print(mean_wf_exec.execution_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait for the task to complete\n",
    "mean_wf_exec = mean_wf_exec.wait(poll_interval=1)\n",
    "\n",
    "# Print the outputs\n",
    "print(mean_wf_exec.outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallelizing with `map_task`\n",
    "\n",
    "We can also parallelize the computation of the mean by using the `map_task` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_compute_mean = fl.map_task(compute_mean)\n",
    "\n",
    "@fl.workflow\n",
    "def parallel_mean_wf(data_list: list[list[int]]) -> list[float]:\n",
    "    return map_compute_mean(data=data_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute the workflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parallel_mean_wf_exec = serverless.execute(parallel_mean_wf, inputs={\"data_list\": [[1, 2, 3], [2, 3, 4], [3, 4, 5]]})\n",
    "print(parallel_mean_wf_exec.execution_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait for the task to complete\n",
    "parallel_mean_wf_exec = parallel_mean_wf_exec.wait(poll_interval=1)\n",
    "\n",
    "# Print the outputs\n",
    "print(parallel_mean_wf_exec.outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's it! You can now iterate on your workflows and tasks directly inside a Jupyter Notebook 🚀"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jupyter-notebook-workshop",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}