{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "72c5f606",
   "metadata": {},
   "source": [
    "# Writing tierkreis workers\n",
    "One of the core concepts in tierkreis are workers.\n",
    "In this example, we will cover the conceptual ideas of workers and how to write one.\n",
    "For this we will write a simple worker from scratch.\n",
    "The `tkr init` cli provides a convenient set up for this which will automate most of the tasks described after.\n",
    "\n",
    "## Concepts\n",
    "\n",
    "Workers encapsulate a set of functionalities which are described in their API (stubs).\n",
    "The API contains a list of typed function interfaces which are implemented by the worker.\n",
    "Tierkreis expects the functions to be stateless and will only checkpoint inputs and outputs of the function call.\n",
    "Although, it is still possible to introduce state, but this then has to be managed by the programmer.\n",
    "\n",
    "Each worker is an independent program, can have multiple implementations (e.g. based on architecture) of these interfaces and has its own dependencies.\n",
    "There are two main ways to write workers:\n",
    "- Using the tierkreis `Worker` class for python based workers\n",
    "- Using a spec file to define the interface for non-python workers\n",
    "\n",
    "In this example, we will write a simple worker, that greets a subject.\n",
    "`greet: str -> str -> str `\n",
    "\n",
    "## Worker Setup\n",
    "\n",
    "A worker consist of four conceptual parts which we will set up in this example:=\n",
    "1. Dependency information. We recommend using `uv` and a `pyproject.toml`.\n",
    "2. A main entry point into the worker called `main.py`\n",
    "3. The api definition of the worker typically called `stubs.py` or `api.py`. For python workers this can be generated from the main file.\n",
    "4. (Optional) Library code containing the core logic.\n",
    "\n",
    "### 1.Dependencies\n",
    "\n",
    "Dependency information should be provide, such that the worker can run in its installed location.\n",
    "For python, this is canonically done through the `pyproject.toml` which tierkreis also uses.\n",
    "\n",
    "```toml\n",
    "[project]\n",
    "name = \"hello-world-worker\"\n",
    "version = \"0.1.0\"\n",
    "requires-python = \">=3.12\"\n",
    "dependencies = [\"tierkreis\"]\n",
    "\n",
    "[dependency-groups]\n",
    "# Optional dev dependency group\n",
    "dev = [\"ruff\"]\n",
    "```\n",
    "Ruff is an optional dependency to format the code later.\n",
    "\n",
    "### 2. Worker entry point\n",
    "The entrypoint contains the implementations of the interface.\n",
    "We are going to use the build in `Worker` class to define the worker.\n",
    "Each worker needs a unique name, here we chose `hello_world_worker`.\n",
    "In tierkreis the worker is defined by its name and the directory it lives in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4153a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install tierkreis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ab026a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from pathlib import Path\n",
    "\n",
    "from tierkreis import Worker\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "worker = Worker(\"hello_world_worker\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa15cdbe",
   "metadata": {},
   "source": [
    "From here we can add functions to the worker by decorating python functions with `@worker.task()`.\n",
    "To make use of the type safety in tierkreis, you should provide narrow types to the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deebc7b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "@worker.task()\n",
    "def greet(greeting: str, subject: str) -> str:\n",
    "    logger.info(\"%s %s\", greeting, subject)\n",
    "    return greeting + subject"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7167a5ee",
   "metadata": {},
   "source": [
    "Now we have defined a greet function that concatenates a two strings.\n",
    "It now lives in the namespace `hello_world_worker.greet`.\n",
    "To make sure the worker runs correctly we need to add the following lines:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3435525",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%script false --no-raise-error\n",
    "from sys import argv\n",
    "if __name__ == \"__main__\":\n",
    "    worker.app(argv)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d246abd5",
   "metadata": {},
   "source": [
    "They make sure that tierkreis can call the worker correctly.\n",
    "\n",
    "### 3. Api definition\n",
    "For python workers, tierkreis can automatically generate the api definition.\n",
    "You can either run the app as `uv run main.py --stubs-path ./api/stubs.py` or call the function directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8de17d1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker.namespace.write_stubs(Path(\"./api/stubs.py\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1996bcbd",
   "metadata": {},
   "source": [
    "This will generate a file `stubs.py` that contains the interface definition that can now be used as `task` in a tierkreis workflow.\n",
    "The code will look like this:\n",
    "```python\n",
    "\"\"\"Code generated from hello_world_worker namespace. Please do not edit.\"\"\"\n",
    "\n",
    "from typing import NamedTuple\n",
    "from tierkreis.controller.data.models import TKR\n",
    "\n",
    "\n",
    "class greet(NamedTuple):\n",
    "    greeting: TKR[str]  # noqa: F821 # fmt: skip\n",
    "    subject: TKR[str]  # noqa: F821 # fmt: skip\n",
    "\n",
    "    @staticmethod\n",
    "    def out() -> type[TKR[str]]:  # noqa: F821 # fmt: skip\n",
    "        return TKR[str]  # noqa: F821 # fmt: skip\n",
    "\n",
    "    @property\n",
    "    def namespace(self) -> str:\n",
    "        return \"hello_world_worker\"\n",
    "\n",
    "```\n",
    "Which defines a class describing the function."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tierkreis (3.12.10)",
   "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.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}