{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab579ff1",
   "metadata": {},
   "source": [
    "## Option 1: Run single flow with Prefect `ThreadPoolTaskRunner`.\n",
    "\n",
    "This method will start Prefect ThreadPoolTaskRunner by default and run tasks concurrently. Here only one single task is being run, so we won't see multiple tasks being run concurrently. See https://docs.prefect.io/v3/concepts/task-runners#available-task-runners for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa42c417-7377-4113-b4a1-2fc1d522a3e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from mlip_arena.models import REGISTRY\n",
    "from mlip_arena.tasks.combustion.flow import hydrogen_combustion\n",
    "\n",
    "model = \"MACE-MPA\"\n",
    "run_dir = Path(\".\").parent / REGISTRY[model][\"family\"]\n",
    "\n",
    "hydrogen_combustion.with_options(persist_result=True)(\n",
    "    model=model,  # can be MLIPEnum name or ASE calculator object\n",
    "    run_dir=run_dir,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1c8b50e",
   "metadata": {},
   "source": [
    "## Option 2: Run single flow with `DaskTaskRunner`.\n",
    "\n",
    "To achieve true parallelism, we can use `DaskTaskRunner` or `RayTaskRunner` to ask HPC resourcces and submit tasks to the cluster for queuing and execution. The task are automatically resubmitted if killed or running overtime. This might be more desirable option becase this task takes longer. See https://docs.prefect.io/integrations/prefect-dask/index#prefect-dask for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bec78f8-0bd0-4a5f-a3bf-fbd2be98f9db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "from dask_jobqueue import SLURMCluster\n",
    "from prefect_dask import DaskTaskRunner\n",
    "\n",
    "nodes_per_alloc = 1\n",
    "gpus_per_alloc = 1\n",
    "\n",
    "cluster_kwargs = dict(\n",
    "    cores=1,\n",
    "    memory=\"64 GB\",\n",
    "    processes=1,\n",
    "    shebang=\"#!/bin/bash\",\n",
    "    account=\"matgen\",\n",
    "    walltime=\"04:00:00\",\n",
    "    job_mem=\"0\",\n",
    "    job_script_prologue=[\n",
    "        \"source ~/.bashrc\",\n",
    "        \"module load python\",\n",
    "        \"source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\",\n",
    "    ],\n",
    "    job_directives_skip=[\"-n\", \"--cpus-per-task\", \"-J\"],\n",
    "    job_extra_directives=[\n",
    "        \"-J arena-combustion\",\n",
    "        \"-q preempt\",\n",
    "        \"--time-min=00:30:00\",\n",
    "        \"--comment=12:00:00\",\n",
    "        f\"-N {nodes_per_alloc}\",\n",
    "        \"-C gpu\",\n",
    "        f\"-G {gpus_per_alloc}\",\n",
    "    ],\n",
    ")\n",
    "cluster = SLURMCluster(**cluster_kwargs)\n",
    "print(cluster.job_script())\n",
    "cluster.adapt(minimum_jobs=1, maximum_jobs=2)\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fbb5f61-7940-4d37-a4ef-ec1b407b9c5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "hydrogen_combustion.with_options(\n",
    "    persist_result=True,\n",
    "    task_runner=DaskTaskRunner(address=client.scheduler.address),\n",
    "    log_prints=True,\n",
    ")(model=model, run_dir=run_dir)"
   ]
  }
 ],
 "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.11.13"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "22140084d8b047aa8b04a991d81c30ac": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "29b37f8fd00546e6aee3f1cd8a1a9438": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "3b9b53f3e9e640c78cd15c16042fdb24": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "5396dbda76074f06a3ab8a978a572821": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "54364b0188b042a9ac68590f8ca5c69e": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "bar_style": "danger",
       "layout": "IPY_MODEL_3b9b53f3e9e640c78cd15c16042fdb24",
       "max": 2000000,
       "style": "IPY_MODEL_b6307bc40b174c48b80d1d447d8fdaa5"
      }
     },
     "652cfcaa7e8542b5b3a33ffce36d84f6": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "layout": "IPY_MODEL_29b37f8fd00546e6aee3f1cd8a1a9438",
       "style": "IPY_MODEL_ba4005a4dcc448b2be982387c82a91be",
       "value": " 0/2000000 [02:32<?, ?it/s]"
      }
     },
     "ac2c57e840774a598cae608b15d0b1d2": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "b41147f209144946914d5955fa380dfa": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "children": [
        "IPY_MODEL_cacc853eb75e42b7a78b3e5e011b580f",
        "IPY_MODEL_54364b0188b042a9ac68590f8ca5c69e",
        "IPY_MODEL_652cfcaa7e8542b5b3a33ffce36d84f6"
       ],
       "layout": "IPY_MODEL_5396dbda76074f06a3ab8a978a572821"
      }
     },
     "b6307bc40b174c48b80d1d447d8fdaa5": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "description_width": ""
      }
     },
     "ba4005a4dcc448b2be982387c82a91be": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "cacc853eb75e42b7a78b3e5e011b580f": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "layout": "IPY_MODEL_22140084d8b047aa8b04a991d81c30ac",
       "style": "IPY_MODEL_ac2c57e840774a598cae608b15d0b1d2",
       "value": "MD H256O128:   0%"
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}