{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "a2486387",
      "metadata": {},
      "source": [
        "# Tierkreis and HPC\n",
        "\n",
        "Tierkreis supports execution in HPC environments through executors natively.\n",
        "Before we take at a look at the executor lets define a problem to run.\n",
        "\n",
        "For this example we are using symbolic circuits, which are supported through the `tkr_pytket_worker`.\n",
        "As before were defining inputs and outputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3c07e149",
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install tierkreis pytket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "54afa191",
      "metadata": {},
      "outputs": [],
      "source": [
        "from typing import NamedTuple\n",
        "\n",
        "from tierkreis.controller.data.models import TKR, OpaqueType\n",
        "\n",
        "\n",
        "class SymbolicCircuitsInputs(NamedTuple):\n",
        "    a: TKR[float]\n",
        "    b: TKR[float]\n",
        "    c: TKR[float]\n",
        "    ansatz: TKR[OpaqueType[\"pytket._tket.circuit.Circuit\"]]  # noqa: F821\n",
        "\n",
        "\n",
        "class SymbolicCircuitsOutputs(NamedTuple):\n",
        "    expectation: TKR[float]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ab5d8c6d",
      "metadata": {},
      "source": [
        "and defining a graph under the assumption that our circuit will have three free symbols"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cd405df6",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.builder import GraphBuilder\n",
        "from substitution_worker import substitute\n",
        "from tierkreis.pytket_worker import (\n",
        "    add_measure_all,\n",
        "    expectation,\n",
        "    optimise_phase_gadgets,\n",
        ")\n",
        "from tierkreis.aer_worker import submit_single\n",
        "\n",
        "\n",
        "def symbolic_execution() -> GraphBuilder:\n",
        "    \"\"\"A graph that substitutes 3 parameters into a circuit and gets an expectation value.\"\"\"\n",
        "    g = GraphBuilder(SymbolicCircuitsInputs, SymbolicCircuitsOutputs)\n",
        "    a = g.inputs.a\n",
        "    b = g.inputs.b\n",
        "    c = g.inputs.c\n",
        "    ansatz = g.inputs.ansatz\n",
        "    n_shots = g.const(100)\n",
        "\n",
        "    substituted_circuit = g.task(substitute(a=a, b=b, c=c, circuit=ansatz))\n",
        "    measurement_circuit = g.task(add_measure_all(circuit=substituted_circuit))\n",
        "\n",
        "    compiled_circuit = g.task(optimise_phase_gadgets(circuit=measurement_circuit))\n",
        "    backend_result = g.task(submit_single(circuit=compiled_circuit, n_shots=n_shots))\n",
        "    av = g.task(expectation(backend_result=backend_result))\n",
        "\n",
        "    g.outputs(SymbolicCircuitsOutputs(expectation=av))\n",
        "    return g"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c1f49f6d",
      "metadata": {},
      "source": [
        "defining the ansatz:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7a806d7c",
      "metadata": {},
      "outputs": [],
      "source": [
        "from pytket._tket.circuit import Circuit, fresh_symbol\n",
        "\n",
        "\n",
        "def build_ansatz() -> Circuit:\n",
        "    a = fresh_symbol(\"a\")\n",
        "    b = fresh_symbol(\"b\")\n",
        "    c = fresh_symbol(\"c\")\n",
        "    circ = Circuit(4)\n",
        "    circ.CX(0, 1)\n",
        "    circ.CX(1, 2)\n",
        "    circ.CX(2, 3)\n",
        "    circ.Rz(a, 3)\n",
        "    circ.CX(2, 3)\n",
        "    circ.CX(1, 2)\n",
        "    circ.CX(0, 1)\n",
        "    circ.Rz(b, 0)\n",
        "    circ.CX(0, 1)\n",
        "    circ.CX(1, 2)\n",
        "    circ.CX(2, 3)\n",
        "    circ.Rz(c, 3)\n",
        "    circ.CX(2, 3)\n",
        "    circ.CX(1, 2)\n",
        "    circ.CX(0, 1)\n",
        "    return circ"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "30ff75de",
      "metadata": {},
      "source": [
        "Most HPC systems use a shared file system which can be used as a storage.\n",
        "Make sure the the checkpoints live on there!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "473ecd6e",
      "metadata": {},
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "from uuid import UUID\n",
        "\n",
        "from tierkreis.controller.storage.filestorage import ControllerFileStorage\n",
        "\n",
        "storage = ControllerFileStorage(\n",
        "    UUID(int=101),\n",
        "    name=\"symbolic_circuits\",\n",
        "    do_cleanup=True,\n",
        "    tierkreis_directory=Path.home()\n",
        "    / \".tierkreis\"\n",
        "    / \"checkpoints\",  # This is the default, change this to the shared fs\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8ca8a19d",
      "metadata": {},
      "source": [
        "For everything that doesn't need HPC nodes we are going to use the regular executor based on `uv` which we will use for a composed executor."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ecfc1612",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.consts import PACKAGE_PATH\n",
        "from tierkreis.executor import UvExecutor\n",
        "\n",
        "executor = UvExecutor(PACKAGE_PATH.parent / \"tierkreis_workers\", storage.logs_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0706f0b6",
      "metadata": {},
      "source": [
        "If we wan to use HPC resources we need to provide a definition of our job.\n",
        "This includes common information like number of nodes, memory, walltime etc.\n",
        "For this we are defining a job spec which runs `uv` on a compute node.\n",
        "In this case we use a pre-build environment in case we have different architectures as the login node  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5f7a476f",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.controller.executor.hpc.job_spec import JobSpec, ResourceSpec\n",
        "\n",
        "spec = JobSpec(\n",
        "    job_name=\"tkr_symbolic_circuits\",\n",
        "    account=\"<group_name>\",  # TODO replace with an actual account / group name\n",
        "    command=\"env UV_PROJECT_ENVIRONMENT=compute_venv uv run main.py\",  # Takes care of the architecture\n",
        "    resource=ResourceSpec(nodes=1, memory_gb=None, gpus_per_node=None),\n",
        "    walltime=\"00:15:00\",\n",
        "    output_path=Path(storage.logs_path),\n",
        "    error_path=Path(storage.logs_path),\n",
        "    include_no_check_directory_flag=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "999ebb3d",
      "metadata": {},
      "source": [
        "From this we can now construct an executor.\n",
        "There are different implementations for the following schedulers:\n",
        "- slurm\n",
        "- PBS\n",
        "- PJSUB\n",
        "\n",
        "We are using `PJSUB`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eff751cf",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.controller.executor.hpc.pjsub import PJSUBExecutor\n",
        "\n",
        "custom_executor = PJSUBExecutor(\n",
        "    spec=spec,\n",
        "    registry_path=Path().parent / \"example_workers\",\n",
        "    logs_path=storage.logs_path,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "255397d7",
      "metadata": {},
      "source": [
        "Combining this into one common executor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b8db8e5b",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.controller.executor.multiple import MultipleExecutor\n",
        "\n",
        "multi_executor = MultipleExecutor(\n",
        "    executor,\n",
        "    executors={\"custom\": custom_executor},\n",
        "    assignments={\"substitution_worker\": \"custom\"},\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "14637924",
      "metadata": {},
      "source": [
        "We can now run the example after defining the inputs.\n",
        "Make sure that the code is running on an HPC system and you have selected the correct group and executor!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5b39740d",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.controller import run_graph\n",
        "from tierkreis.storage import read_outputs\n",
        "\n",
        "run_graph(\n",
        "    storage,\n",
        "    multi_executor,\n",
        "    symbolic_execution().data,\n",
        "    {\n",
        "        \"ansatz\": build_ansatz(),\n",
        "        \"a\": 0.2,\n",
        "        \"b\": 0.55,\n",
        "        \"c\": 0.75,\n",
        "    },\n",
        "    polling_interval_seconds=0.1,\n",
        ")\n",
        "output = read_outputs(symbolic_execution().data, storage)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}