add relax task, eos workflow, and prefect test

mlip_arena/tasks/README.md

@@ -1,4 +1,3 @@
-
 ## Note on task registration
 
 1. Use `ast` to parse task classes from the uploaded script.

mlip_arena/tasks/combustion/water.ipynb

@@ -19,7 +19,17 @@
     "from pymatgen.io.packmol import PackmolBoxGen\n",
     "\n",
     "from mlip_arena.models.utils import REGISTRY, MLIPEnum\n",
-    "from mlip_arena.tasks.run import md as MD"
+    "from mlip_arena.tasks.md import run as MD"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "jp-MarkdownHeadingCollapsed": true,
+    "tags": []
+   },
+   "source": [
+    "## Create initial configuration"
    ]
   },
   {
@@ -97,6 +107,13 @@
     "write(f'{atoms.get_chemical_formula()}.extxyz', atoms)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run workflow"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 2,
@@ -119,7 +136,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -130,17 +147,15 @@
       "\n",
       "#SBATCH -A matgen\n",
       "#SBATCH --mem=0\n",
-      "#SBATCH -t 02:00:00\n",
-      "#SBATCH -J combustion-water\n",
-      "#SBATCH -q regular\n",
+      "#SBATCH -t 00:30:00\n",
       "#SBATCH -N 1\n",
+      "#SBATCH -q debug\n",
       "#SBATCH -C gpu\n",
-      "#SBATCH -G 4\n",
-      "#SBATCH --exclusive\n",
+      "#SBATCH -J combustion-water\n",
       "source ~/.bashrc\n",
       "module load python\n",
       "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\n",
-      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.18:43951 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.19:34913 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n",
       "\n"
      ]
     },
@@ -148,14 +163,10 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2024-09-27 03:54:04,812 - distributed.scheduler - ERROR - Task md-2-9b52cfa3e22647de99659385eca3d3f2-1 marked as failed because 4 workers died while trying to run it\n",
-      "2024-09-27 03:54:04,820 - distributed.scheduler - ERROR - Task md-5-2a33088343d64b66b612e4145ffaf870-1 marked as failed because 4 workers died while trying to run it\n",
-      "2024-09-27 03:54:04,824 - distributed.scheduler - ERROR - Task md-0-017eef8ab8744d7c958badc5d53fdab4-1 marked as failed because 4 workers died while trying to run it\n",
-      "2024-09-27 05:55:56,468 - distributed.scheduler - ERROR - Task md-8-bcd5e0bb82854952b2c5c430d4c2e7e8-1 marked as failed because 4 workers died while trying to run it\n",
-      "2024-09-27 05:55:56,906 - distributed.scheduler - ERROR - Task md-6-a575c0cc3b1e4fc984560ff101dbfc6c-1 marked as failed because 4 workers died while trying to run it\n",
-      "2024-09-27 05:55:56,907 - distributed.scheduler - ERROR - Task md-9-f98b1f7c61c749c8b67cd669517a49bd-1 marked as failed because 4 workers died while trying to run it\n",
-      "2024-09-27 05:55:56,910 - distributed.scheduler - ERROR - Task md-4-0e1acf2be8994ca7ae198d52d5e04d61-1 marked as failed because 4 workers died while trying to run it\n",
-      "2024-09-27 05:57:02,776 - distributed.scheduler - ERROR - Task md-7-40b7b4ede41745e29d5c5fd0cc4d9e42-1 marked as failed because 4 workers died while trying to run it\n"
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/node.py:182: UserWarning: Port 8787 is already in use.\n",
+      "Perhaps you already have a cluster running?\n",
+      "Hosting the HTTP server on port 38239 instead\n",
+      "  warnings.warn(\n"
      ]
     }
    ],
@@ -164,45 +175,62 @@
     "gpus_per_alloc = 4\n",
     "ntasks = 1\n",
     "\n",
-    "cluster_kwargs = dict(\n",
-    "    cores=1,\n",
-    "    memory=\"64 GB\",\n",
-    "    shebang=\"#!/bin/bash\",\n",
-    "    account=\"matgen\",\n",
-    "    walltime=\"02:00:00\",\n",
-    "    job_mem=\"0\",\n",
-    "    job_script_prologue=[\n",
+    "# cluster_kwargs = dict(\n",
+    "#     cores=1,\n",
+    "#     memory=\"64 GB\",\n",
+    "#     shebang=\"#!/bin/bash\",\n",
+    "#     account=\"matgen\",\n",
+    "#     walltime=\"02: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 combustion-water\",\n",
+    "#         \"-q regular\",\n",
+    "#         f\"-N {nodes_per_alloc}\",\n",
+    "#         \"-C gpu\",\n",
+    "#         f\"-G {gpus_per_alloc}\",\n",
+    "#         f\"--exclusive\",\n",
+    "#         # \"--time-min=00:30:00\",\n",
+    "#         # \"--comment=1-00:00:00\",\n",
+    "#         # \"--signal=B:USR1@60\",\n",
+    "#         # \"--requeue\",\n",
+    "#         # \"--open-mode=append\"\n",
+    "#     ],\n",
+    "#     death_timeout=86400\n",
+    "# )\n",
+    "\n",
+    "# cluster = SLURMCluster(**cluster_kwargs)\n",
+    "\n",
+    "cluster_kwargs = {\n",
+    "    \"cores\": 1,\n",
+    "    \"memory\": \"64 GB\",\n",
+    "    \"shebang\": \"#!/bin/bash\",\n",
+    "    \"account\": \"matgen\",\n",
+    "    \"walltime\": \"00:30: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 combustion-water\",\n",
-    "        \"-q regular\",\n",
-    "        f\"-N {nodes_per_alloc}\",\n",
-    "        \"-C gpu\",\n",
-    "        f\"-G {gpus_per_alloc}\",\n",
-    "        f\"--exclusive\",\n",
-    "        # \"--time-min=00:30:00\",\n",
-    "        # \"--comment=1-00:00:00\",\n",
-    "        # \"--signal=B:USR1@60\",\n",
-    "        # \"--requeue\",\n",
-    "        # \"--open-mode=append\"\n",
-    "    ],\n",
-    "    death_timeout=86400\n",
-    ")\n",
-    "\n",
-    "\n",
+    "    \"job_directives_skip\": [\"-n\", \"--cpus-per-task\", \"-J\"],\n",
+    "    \"job_extra_directives\": [f\"-N {nodes_per_alloc}\", \"-q debug\", \"-C gpu\", \"-J combustion-water\"],\n",
+    "}\n",
     "cluster = SLURMCluster(**cluster_kwargs)\n",
+    "\n",
     "print(cluster.job_script())\n",
-    "cluster.adapt(minimum_jobs=10, maximum_jobs=11)\n",
+    "cluster.adapt(minimum_jobs=2, maximum_jobs=2)\n",
     "client = Client(cluster)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -236,7 +264,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 7,
    "metadata": {
     "tags": []
    },
@@ -244,115 +272,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:58.587 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'combustion'</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:19:58.587 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'tangerine-mussel'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:58.592 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - View at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/ea339a90-929f-491d-bf23-e257c925ad6b</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:19:58.592 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/ea339a90-929f-491d-bf23-e257c925ad6b\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:58.593 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.18:43951\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:19:58.593 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.18:43951\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:58.599 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - The Dask dashboard is available at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">http://128.55.64.18:8787/status</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:19:58.599 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - The Dask dashboard is available at \u001b[94mhttp://128.55.64.18:8787/status\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:59.454 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-2' for task 'md'\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:19:59.454 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-2' for task 'md'\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:59.484 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-6' for task 'md'\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:19:59.484 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-6' for task 'md'\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:59.488 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-3' for task 'md'\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:19:59.488 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-3' for task 'md'\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:59.491 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-5' for task 'md'\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:19:59.491 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-5' for task 'md'\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:59.494 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-9' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.449 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'liberal-beetle'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'combustion'</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "23:19:59.494 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-9' for task 'md'\n"
+       "14:50:51.449 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'liberal-beetle'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -361,11 +285,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:59.497 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-1' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.453 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'liberal-beetle'</span> - View at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/4803afc0-0472-4cf3-8d76-794fc218e2c3</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "23:19:59.497 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-1' for task 'md'\n"
+       "14:50:51.453 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'liberal-beetle'\u001b[0m - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/4803afc0-0472-4cf3-8d76-794fc218e2c3\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -374,11 +298,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:59.646 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-0' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.454 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.19:34913\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "23:19:59.646 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-0' for task 'md'\n"
+       "14:50:51.454 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.19:34913\n"
       ]
      },
      "metadata": {},
@@ -387,11 +311,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:19:59.938 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-7' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.460 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - The Dask dashboard is available at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">http://128.55.64.19:38239/status</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "23:19:59.938 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-7' for task 'md'\n"
+       "14:50:51.460 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - The Dask dashboard is available at \u001b[94mhttp://128.55.64.19:38239/status\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -400,163 +324,20 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.111 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-5' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.111 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-5' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.119 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-8' for task 'md'\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.119 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-8' for task 'md'\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.125 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-8' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.125 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-8' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.131 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Created task run 'md-4' for task 'md'\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.131 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Created task run 'md-4' for task 'md'\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.137 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-4' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.137 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-4' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.154 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-6' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.154 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-6' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.205 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-0' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.205 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-0' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.212 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-1' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.212 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-1' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.283 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-2' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.283 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-2' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.291 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-9' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.291 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-9' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.400 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-3' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.400 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-3' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">23:20:00.554 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Submitted task run 'md-7' for execution.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "23:20:00.554 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Submitted task run 'md-7' for execution.\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:54:04.823 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Encountered exception during execution:\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.876 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'liberal-beetle'</span> - Encountered exception during execution:\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 50, in get_call_parameters\n",
+       "    bound_signature = inspect.signature(fn).bind(*call_args, **call_kwargs)\n",
+       "                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py\", line 3212, in bind\n",
+       "    return self._bind(args, kwargs)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py\", line 3201, in _bind\n",
+       "    raise TypeError(\n",
+       "TypeError: got an unexpected keyword argument 'mb_velocity_seed'\n",
+       "\n",
+       "During handling of the above exception, another exception occurred:\n",
+       "\n",
        "Traceback (most recent call last):\n",
        "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py\", line 867, in orchestrate_flow_run\n",
        "    result = await flow_call.aresult()\n",
@@ -567,41 +348,33 @@
        "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 352, in _run_sync\n",
        "    result = self.fn(*self.args, **self.kwargs)\n",
        "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_2266116/1920202683.py\", line 26, in combustion\n",
-       "    return [future.result() for future in futures]\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_2266116/1920202683.py\", line 26, in &lt;listcomp&gt;\n",
-       "    return [future.result() for future in futures]\n",
-       "            ^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py\", line 228, in result\n",
-       "    return from_sync.call_soon_in_loop_thread(result).result()\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 318, in result\n",
-       "    return self.future.result(timeout=timeout)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 188, in result\n",
-       "    return self.__get_result()\n",
-       "           ^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py\", line 401, in __get_result\n",
-       "    raise self._exception\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 389, in _run_async\n",
-       "    result = await coro\n",
+       "  File \"/tmp/ipykernel_997512/1920202683.py\", line 6, in combustion\n",
+       "    future = MD.submit(\n",
        "             ^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py\", line 237, in _result\n",
-       "    return await final_state.result(raise_on_failure=raise_on_failure, fetch=True)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py\", line 91, in _get_state_result\n",
-       "    raise await get_state_exception(state)\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 311, in wait\n",
-       "    return await future.result(timeout=timeout)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 336, in _result\n",
-       "    raise exc.with_traceback(tb)\n",
-       "distributed.scheduler.KilledWorker: Attempted to run task 'md-2-9b52cfa3e22647de99659385eca3d3f2-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/tasks.py\", line 811, in submit\n",
+       "    parameters = get_call_parameters(self.fn, args, kwargs)\n",
+       "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 52, in get_call_parameters\n",
+       "    raise ParameterBindError.from_bind_failure(fn, exc, call_args, call_kwargs)\n",
+       "prefect.exceptions.ParameterBindError: Error binding parameters for function 'run': got an unexpected keyword argument 'mb_velocity_seed'.\n",
+       "Function 'run' has signature 'atoms: 'Atoms', calculator_name: 'str | MLIPEnum', calculator_kwargs: 'dict | None', dispersion: 'str | None' = None, dispersion_kwargs: 'dict | None' = None, device: 'str | None' = None, ensemble: \"Literal['nve', 'nvt', 'npt']\" = 'nvt', dynamics: 'str | MolecularDynamics' = 'langevin', time_step: 'float | None' = None, total_time: 'float' = 1000, temperature: 'float | Sequence | np.ndarray | None' = 300.0, pressure: 'float | Sequence | np.ndarray | None' = None, ase_md_kwargs: 'dict | None' = None, md_velocity_seed: 'int | None' = None, zero_linear_momentum: 'bool' = True, zero_angular_momentum: 'bool' = True, traj_file: 'str | Path | None' = None, traj_interval: 'int' = 1, restart: 'bool' = True' but received args: () and kwargs: ['atoms', 'calculator_name', 'calculator_kwargs', 'ensemble', 'dynamics', 'time_step', 'ase_md_kwargs', 'total_time', 'temperature', 'pressure', 'mb_velocity_seed', 'traj_file', 'traj_interval', 'restart'].\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "03:54:04.823 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Encountered exception during execution:\n",
+       "14:50:51.876 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'liberal-beetle'\u001b[0m - Encountered exception during execution:\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 50, in get_call_parameters\n",
+       "    bound_signature = inspect.signature(fn).bind(*call_args, **call_kwargs)\n",
+       "                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py\", line 3212, in bind\n",
+       "    return self._bind(args, kwargs)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py\", line 3201, in _bind\n",
+       "    raise TypeError(\n",
+       "TypeError: got an unexpected keyword argument 'mb_velocity_seed'\n",
+       "\n",
+       "During handling of the above exception, another exception occurred:\n",
+       "\n",
        "Traceback (most recent call last):\n",
        "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py\", line 867, in orchestrate_flow_run\n",
        "    result = await flow_call.aresult()\n",
@@ -612,89 +385,16 @@
        "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 352, in _run_sync\n",
        "    result = self.fn(*self.args, **self.kwargs)\n",
        "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_2266116/1920202683.py\", line 26, in combustion\n",
-       "    return [future.result() for future in futures]\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_2266116/1920202683.py\", line 26, in <listcomp>\n",
-       "    return [future.result() for future in futures]\n",
-       "            ^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py\", line 228, in result\n",
-       "    return from_sync.call_soon_in_loop_thread(result).result()\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 318, in result\n",
-       "    return self.future.result(timeout=timeout)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 188, in result\n",
-       "    return self.__get_result()\n",
-       "           ^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py\", line 401, in __get_result\n",
-       "    raise self._exception\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 389, in _run_async\n",
-       "    result = await coro\n",
+       "  File \"/tmp/ipykernel_997512/1920202683.py\", line 6, in combustion\n",
+       "    future = MD.submit(\n",
        "             ^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py\", line 237, in _result\n",
-       "    return await final_state.result(raise_on_failure=raise_on_failure, fetch=True)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py\", line 91, in _get_state_result\n",
-       "    raise await get_state_exception(state)\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 311, in wait\n",
-       "    return await future.result(timeout=timeout)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 336, in _result\n",
-       "    raise exc.with_traceback(tb)\n",
-       "distributed.scheduler.KilledWorker: Attempted to run task 'md-2-9b52cfa3e22647de99659385eca3d3f2-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">05:57:03.100 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-2' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-2-9b52cfa3e22647de99659385eca3d3f2-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "05:57:03.100 | \u001b[36mINFO\u001b[0m    | Task run 'md-2' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-2-9b52cfa3e22647de99659385eca3d3f2-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">05:57:03.472 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-5' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-5-2a33088343d64b66b612e4145ffaf870-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "05:57:03.472 | \u001b[36mINFO\u001b[0m    | Task run 'md-5' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-5-2a33088343d64b66b612e4145ffaf870-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">05:57:03.991 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-6' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-6-a575c0cc3b1e4fc984560ff101dbfc6c-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.69.100:45983. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "05:57:03.991 | \u001b[36mINFO\u001b[0m    | Task run 'md-6' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-6-a575c0cc3b1e4fc984560ff101dbfc6c-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.69.100:45983. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">05:57:04.333 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-7' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-7-40b7b4ede41745e29d5c5fd0cc4d9e42-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.68.6:39739. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "05:57:04.333 | \u001b[36mINFO\u001b[0m    | Task run 'md-7' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-7-40b7b4ede41745e29d5c5fd0cc4d9e42-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.68.6:39739. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/tasks.py\", line 811, in submit\n",
+       "    parameters = get_call_parameters(self.fn, args, kwargs)\n",
+       "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 52, in get_call_parameters\n",
+       "    raise ParameterBindError.from_bind_failure(fn, exc, call_args, call_kwargs)\n",
+       "prefect.exceptions.ParameterBindError: Error binding parameters for function 'run': got an unexpected keyword argument 'mb_velocity_seed'.\n",
+       "Function 'run' has signature 'atoms: 'Atoms', calculator_name: 'str | MLIPEnum', calculator_kwargs: 'dict | None', dispersion: 'str | None' = None, dispersion_kwargs: 'dict | None' = None, device: 'str | None' = None, ensemble: \"Literal['nve', 'nvt', 'npt']\" = 'nvt', dynamics: 'str | MolecularDynamics' = 'langevin', time_step: 'float | None' = None, total_time: 'float' = 1000, temperature: 'float | Sequence | np.ndarray | None' = 300.0, pressure: 'float | Sequence | np.ndarray | None' = None, ase_md_kwargs: 'dict | None' = None, md_velocity_seed: 'int | None' = None, zero_linear_momentum: 'bool' = True, zero_angular_momentum: 'bool' = True, traj_file: 'str | Path | None' = None, traj_interval: 'int' = 1, restart: 'bool' = True' but received args: () and kwargs: ['atoms', 'calculator_name', 'calculator_kwargs', 'ensemble', 'dynamics', 'time_step', 'ase_md_kwargs', 'total_time', 'temperature', 'pressure', 'mb_velocity_seed', 'traj_file', 'traj_interval', 'restart'].\n"
       ]
      },
      "metadata": {},
@@ -703,50 +403,30 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">05:57:04.976 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-8' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-8-bcd5e0bb82854952b2c5c430d4c2e7e8-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.238:37065. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:52.106 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'liberal-beetle'</span> - Finished in state <span style=\"color: #d70000; text-decoration-color: #d70000\">Failed</span>('Flow run encountered an exception. ParameterBindError: Error binding parameters for function \\'run\\': got an unexpected keyword argument \\'mb_velocity_seed\\'.\\nFunction \\'run\\' has signature \\'atoms: \\'Atoms\\', calculator_name: \\'str | MLIPEnum\\', calculator_kwargs: \\'dict | None\\', dispersion: \\'str | None\\' = None, dispersion_kwargs: \\'dict | None\\' = None, device: \\'str | None\\' = None, ensemble: \"Literal[\\'nve\\', \\'nvt\\', \\'npt\\']\" = \\'nvt\\', dynamics: \\'str | MolecularDynamics\\' = \\'langevin\\', time_step: \\'float | None\\' = None, total_time: \\'float\\' = 1000, temperature: \\'float | Sequence | np.ndarray | None\\' = 300.0, pressure: \\'float | Sequence | np.ndarray | None\\' = None, ase_md_kwargs: \\'dict | None\\' = None, md_velocity_seed: \\'int | None\\' = None, zero_linear_momentum: \\'bool\\' = True, zero_angular_momentum: \\'bool\\' = True, traj_file: \\'str | Path | None\\' = None, traj_interval: \\'int\\' = 1, restart: \\'bool\\' = True\\' but received args: () and kwargs: [\\'atoms\\', \\'calculator_name\\', \\'calculator_kwargs\\', \\'ensemble\\', \\'dynamics\\', \\'time_step\\', \\'ase_md_kwargs\\', \\'total_time\\', \\'temperature\\', \\'pressure\\', \\'mb_velocity_seed\\', \\'traj_file\\', \\'traj_interval\\', \\'restart\\'].')\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "05:57:04.976 | \u001b[36mINFO\u001b[0m    | Task run 'md-8' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-8-bcd5e0bb82854952b2c5c430d4c2e7e8-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.238:37065. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
+       "14:50:52.106 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'liberal-beetle'\u001b[0m - Finished in state \u001b[38;5;160mFailed\u001b[0m('Flow run encountered an exception. ParameterBindError: Error binding parameters for function \\'run\\': got an unexpected keyword argument \\'mb_velocity_seed\\'.\\nFunction \\'run\\' has signature \\'atoms: \\'Atoms\\', calculator_name: \\'str | MLIPEnum\\', calculator_kwargs: \\'dict | None\\', dispersion: \\'str | None\\' = None, dispersion_kwargs: \\'dict | None\\' = None, device: \\'str | None\\' = None, ensemble: \"Literal[\\'nve\\', \\'nvt\\', \\'npt\\']\" = \\'nvt\\', dynamics: \\'str | MolecularDynamics\\' = \\'langevin\\', time_step: \\'float | None\\' = None, total_time: \\'float\\' = 1000, temperature: \\'float | Sequence | np.ndarray | None\\' = 300.0, pressure: \\'float | Sequence | np.ndarray | None\\' = None, ase_md_kwargs: \\'dict | None\\' = None, md_velocity_seed: \\'int | None\\' = None, zero_linear_momentum: \\'bool\\' = True, zero_angular_momentum: \\'bool\\' = True, traj_file: \\'str | Path | None\\' = None, traj_interval: \\'int\\' = 1, restart: \\'bool\\' = True\\' but received args: () and kwargs: [\\'atoms\\', \\'calculator_name\\', \\'calculator_kwargs\\', \\'ensemble\\', \\'dynamics\\', \\'time_step\\', \\'ase_md_kwargs\\', \\'total_time\\', \\'temperature\\', \\'pressure\\', \\'mb_velocity_seed\\', \\'traj_file\\', \\'traj_interval\\', \\'restart\\'].')\n"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
     {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">05:57:05.435 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-9' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-9-f98b1f7c61c749c8b67cd669517a49bd-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.65.221:34843. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "05:57:05.435 | \u001b[36mINFO\u001b[0m    | Task run 'md-9' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-9-f98b1f7c61c749c8b67cd669517a49bd-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.65.221:34843. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">05:57:05.840 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'tangerine-mussel'</span> - Finished in state <span style=\"color: #d70000; text-decoration-color: #d70000\">Failed</span>(\"Flow run encountered an exception. KilledWorker: Attempted to run task 'md-2-9b52cfa3e22647de99659385eca3d3f2-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\")\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "05:57:05.840 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'tangerine-mussel'\u001b[0m - Finished in state \u001b[38;5;160mFailed\u001b[0m(\"Flow run encountered an exception. KilledWorker: Attempted to run task 'md-2-9b52cfa3e22647de99659385eca3d3f2-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\")\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "ename": "KilledWorker",
-     "evalue": "Attempted to run task 'md-2-9b52cfa3e22647de99659385eca3d3f2-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see https://distributed.dask.org/en/stable/killed.html.",
+     "ename": "ParameterBindError",
+     "evalue": "Error binding parameters for function 'run': got an unexpected keyword argument 'mb_velocity_seed'.\nFunction 'run' has signature 'atoms: 'Atoms', calculator_name: 'str | MLIPEnum', calculator_kwargs: 'dict | None', dispersion: 'str | None' = None, dispersion_kwargs: 'dict | None' = None, device: 'str | None' = None, ensemble: \"Literal['nve', 'nvt', 'npt']\" = 'nvt', dynamics: 'str | MolecularDynamics' = 'langevin', time_step: 'float | None' = None, total_time: 'float' = 1000, temperature: 'float | Sequence | np.ndarray | None' = 300.0, pressure: 'float | Sequence | np.ndarray | None' = None, ase_md_kwargs: 'dict | None' = None, md_velocity_seed: 'int | None' = None, zero_linear_momentum: 'bool' = True, zero_angular_momentum: 'bool' = True, traj_file: 'str | Path | None' = None, traj_interval: 'int' = 1, restart: 'bool' = True' but received args: () and kwargs: ['atoms', 'calculator_name', 'calculator_kwargs', 'ensemble', 'dynamics', 'time_step', 'ase_md_kwargs', 'total_time', 'temperature', 'pressure', 'mb_velocity_seed', 'traj_file', 'traj_interval', 'restart'].",
      "output_type": "error",
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKilledWorker\u001b[0m                              Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[43mcombustion\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:50\u001b[0m, in \u001b[0;36mget_call_parameters\u001b[0;34m(fn, call_args, call_kwargs, apply_defaults)\u001b[0m\n\u001b[1;32m     49\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 50\u001b[0m     bound_signature \u001b[38;5;241m=\u001b[39m \u001b[43minspect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msignature\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcall_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcall_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     51\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py:3212\u001b[0m, in \u001b[0;36mSignature.bind\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   3208\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Get a BoundArguments object, that maps the passed `args`\u001b[39;00m\n\u001b[1;32m   3209\u001b[0m \u001b[38;5;124;03mand `kwargs` to the function's signature.  Raises `TypeError`\u001b[39;00m\n\u001b[1;32m   3210\u001b[0m \u001b[38;5;124;03mif the passed arguments can not be bound.\u001b[39;00m\n\u001b[1;32m   3211\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m-> 3212\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py:3201\u001b[0m, in \u001b[0;36mSignature._bind\u001b[0;34m(self, args, kwargs, partial)\u001b[0m\n\u001b[1;32m   3200\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 3201\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\n\u001b[1;32m   3202\u001b[0m             \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mgot an unexpected keyword argument \u001b[39m\u001b[38;5;132;01m{arg!r}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m   3203\u001b[0m                 arg\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mnext\u001b[39m(\u001b[38;5;28miter\u001b[39m(kwargs))))\n\u001b[1;32m   3205\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_bound_arguments_cls(\u001b[38;5;28mself\u001b[39m, arguments)\n",
+      "\u001b[0;31mTypeError\u001b[0m: got an unexpected keyword argument 'mb_velocity_seed'",
+      "\nDuring handling of the above exception, another exception occurred:\n",
+      "\u001b[0;31mParameterBindError\u001b[0m                        Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[7], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[43mcombustion\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms\u001b[49m\u001b[43m)\u001b[49m\n",
       "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flows.py:1224\u001b[0m, in \u001b[0;36mFlow.__call__\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1219\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m task_viz_tracker:\n\u001b[1;32m   1220\u001b[0m     \u001b[38;5;66;03m# this is a subflow, for now return a single task and do not go further\u001b[39;00m\n\u001b[1;32m   1221\u001b[0m     \u001b[38;5;66;03m# we can add support for exploring subflows for tasks in the future.\u001b[39;00m\n\u001b[1;32m   1222\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m track_viz_task(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39misasync, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mname, parameters)\n\u001b[0;32m-> 1224\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43menter_flow_run_engine_from_flow_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1225\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1226\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1227\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_for\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1228\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1229\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
       "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py:297\u001b[0m, in \u001b[0;36menter_flow_run_engine_from_flow_call\u001b[0;34m(flow, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    290\u001b[0m     retval \u001b[38;5;241m=\u001b[39m from_async\u001b[38;5;241m.\u001b[39mwait_for_call_in_loop_thread(\n\u001b[1;32m    291\u001b[0m         begin_run,\n\u001b[1;32m    292\u001b[0m         done_callbacks\u001b[38;5;241m=\u001b[39mdone_callbacks,\n\u001b[1;32m    293\u001b[0m         contexts\u001b[38;5;241m=\u001b[39mcontexts,\n\u001b[1;32m    294\u001b[0m     )\n\u001b[1;32m    296\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 297\u001b[0m     retval \u001b[38;5;241m=\u001b[39m \u001b[43mfrom_sync\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwait_for_call_in_loop_thread\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    298\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbegin_run\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    299\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdone_callbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdone_callbacks\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    300\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcontexts\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcontexts\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    301\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    303\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m retval\n",
       "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/api.py:243\u001b[0m, in \u001b[0;36mfrom_sync.wait_for_call_in_loop_thread\u001b[0;34m(_from_sync__call, timeout, done_callbacks, contexts)\u001b[0m\n\u001b[1;32m    241\u001b[0m     stack\u001b[38;5;241m.\u001b[39menter_context(context)\n\u001b[1;32m    242\u001b[0m waiter\u001b[38;5;241m.\u001b[39mwait()\n\u001b[0;32m--> 243\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
@@ -760,18 +440,10 @@
       "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py:867\u001b[0m, in \u001b[0;36morchestrate_flow_run\u001b[0;34m(flow, flow_run, parameters, wait_for, interruptible, client, partial_flow_run_context, user_thread)\u001b[0m\n\u001b[1;32m    862\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    863\u001b[0m             from_async\u001b[38;5;241m.\u001b[39mcall_soon_in_new_thread(\n\u001b[1;32m    864\u001b[0m                 flow_call, timeout\u001b[38;5;241m=\u001b[39mflow\u001b[38;5;241m.\u001b[39mtimeout_seconds\n\u001b[1;32m    865\u001b[0m             )\n\u001b[0;32m--> 867\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m flow_call\u001b[38;5;241m.\u001b[39maresult()\n\u001b[1;32m    869\u001b[0m         waited_for_task_runs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m wait_for_task_runs_and_report_crashes(\n\u001b[1;32m    870\u001b[0m             flow_run_context\u001b[38;5;241m.\u001b[39mtask_run_futures, client\u001b[38;5;241m=\u001b[39mclient\n\u001b[1;32m    871\u001b[0m         )\n\u001b[1;32m    872\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m PausedRun \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    873\u001b[0m     \u001b[38;5;66;03m# could get raised either via utility or by returning Paused from a task run\u001b[39;00m\n\u001b[1;32m    874\u001b[0m     \u001b[38;5;66;03m# if a task run pauses, we set its state as the flow's state\u001b[39;00m\n\u001b[1;32m    875\u001b[0m     \u001b[38;5;66;03m# to preserve reschedule and timeout behavior\u001b[39;00m\n",
       "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:327\u001b[0m, in \u001b[0;36mCall.aresult\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    321\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    322\u001b[0m \u001b[38;5;124;03mWait for the result of the call.\u001b[39;00m\n\u001b[1;32m    323\u001b[0m \n\u001b[1;32m    324\u001b[0m \u001b[38;5;124;03mFor use from asynchronous contexts.\u001b[39;00m\n\u001b[1;32m    325\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    326\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 327\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m asyncio\u001b[38;5;241m.\u001b[39mwrap_future(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture)\n\u001b[1;32m    328\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m asyncio\u001b[38;5;241m.\u001b[39mCancelledError \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    329\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m CancelledError() \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mexc\u001b[39;00m\n",
       "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:352\u001b[0m, in \u001b[0;36mCall._run_sync\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    350\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39menforce_sync_deadline() \u001b[38;5;28;01mas\u001b[39;00m cancel_scope:\n\u001b[1;32m    351\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 352\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    353\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    354\u001b[0m         \u001b[38;5;66;03m# Forget this call's arguments in order to free up any memory\u001b[39;00m\n\u001b[1;32m    355\u001b[0m         \u001b[38;5;66;03m# that may be referenced by them; after a call has happened,\u001b[39;00m\n\u001b[1;32m    356\u001b[0m         \u001b[38;5;66;03m# there's no need to keep a reference to them\u001b[39;00m\n\u001b[1;32m    357\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "Cell \u001b[0;32mIn[4], line 26\u001b[0m, in \u001b[0;36mcombustion\u001b[0;34m(atoms)\u001b[0m\n\u001b[1;32m      6\u001b[0m     future \u001b[38;5;241m=\u001b[39m MD\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m      7\u001b[0m         atoms\u001b[38;5;241m=\u001b[39matoms,\n\u001b[1;32m      8\u001b[0m         calculator_name\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m         restart\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m[\u001b[49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfuture\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfutures\u001b[49m\u001b[43m]\u001b[49m\n",
-      "Cell \u001b[0;32mIn[4], line 26\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      6\u001b[0m     future \u001b[38;5;241m=\u001b[39m MD\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m      7\u001b[0m         atoms\u001b[38;5;241m=\u001b[39matoms,\n\u001b[1;32m      8\u001b[0m         calculator_name\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m         restart\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m [\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m future \u001b[38;5;129;01min\u001b[39;00m futures]\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py:228\u001b[0m, in \u001b[0;36mPrefectFuture.result\u001b[0;34m(self, timeout, raise_on_failure)\u001b[0m\n\u001b[1;32m    226\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m from_async\u001b[38;5;241m.\u001b[39mcall_soon_in_loop_thread(result)\u001b[38;5;241m.\u001b[39maresult()\n\u001b[1;32m    227\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 228\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfrom_sync\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_soon_in_loop_thread\u001b[49m\u001b[43m(\u001b[49m\u001b[43mresult\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:318\u001b[0m, in \u001b[0;36mCall.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mresult\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout: Optional[\u001b[38;5;28mfloat\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m    313\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    314\u001b[0m \u001b[38;5;124;03m    Wait for the result of the call.\u001b[39;00m\n\u001b[1;32m    315\u001b[0m \n\u001b[1;32m    316\u001b[0m \u001b[38;5;124;03m    Not safe for use from asynchronous contexts.\u001b[39;00m\n\u001b[1;32m    317\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 318\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:188\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    186\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m CancelledError()\n\u001b[1;32m    187\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;241m==\u001b[39m FINISHED:\n\u001b[0;32m--> 188\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__get_result\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    189\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    190\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m()\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py:401\u001b[0m, in \u001b[0;36mFuture.__get_result\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    399\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception:\n\u001b[1;32m    400\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 401\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    403\u001b[0m         \u001b[38;5;66;03m# Break a reference cycle with the exception in self._exception\u001b[39;00m\n\u001b[1;32m    404\u001b[0m         \u001b[38;5;28mself\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:389\u001b[0m, in \u001b[0;36mCall._run_async\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    387\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39menforce_async_deadline() \u001b[38;5;28;01mas\u001b[39;00m cancel_scope:\n\u001b[1;32m    388\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 389\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m coro\n\u001b[1;32m    390\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    391\u001b[0m         \u001b[38;5;66;03m# Forget this call's arguments in order to free up any memory\u001b[39;00m\n\u001b[1;32m    392\u001b[0m         \u001b[38;5;66;03m# that may be referenced by them; after a call has happened,\u001b[39;00m\n\u001b[1;32m    393\u001b[0m         \u001b[38;5;66;03m# there's no need to keep a reference to them\u001b[39;00m\n\u001b[1;32m    394\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py:237\u001b[0m, in \u001b[0;36mPrefectFuture._result\u001b[0;34m(self, timeout, raise_on_failure)\u001b[0m\n\u001b[1;32m    235\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m final_state:\n\u001b[1;32m    236\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCall timed out before task finished.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 237\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m final_state\u001b[38;5;241m.\u001b[39mresult(raise_on_failure\u001b[38;5;241m=\u001b[39mraise_on_failure, fetch\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py:91\u001b[0m, in \u001b[0;36m_get_state_result\u001b[0;34m(state, raise_on_failure)\u001b[0m\n\u001b[1;32m     84\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m UnfinishedRun(\n\u001b[1;32m     85\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRun is in \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstate\u001b[38;5;241m.\u001b[39mtype\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m state, its result is not available.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     86\u001b[0m     )\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure \u001b[38;5;129;01mand\u001b[39;00m (\n\u001b[1;32m     89\u001b[0m     state\u001b[38;5;241m.\u001b[39mis_crashed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_failed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_cancelled()\n\u001b[1;32m     90\u001b[0m ):\n\u001b[0;32m---> 91\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m get_state_exception(state)\n\u001b[1;32m     93\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(state\u001b[38;5;241m.\u001b[39mdata, DataDocument):\n\u001b[1;32m     94\u001b[0m     result \u001b[38;5;241m=\u001b[39m result_from_state_with_data_document(\n\u001b[1;32m     95\u001b[0m         state, raise_on_failure\u001b[38;5;241m=\u001b[39mraise_on_failure\n\u001b[1;32m     96\u001b[0m     )\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py:311\u001b[0m, in \u001b[0;36mDaskTaskRunner.wait\u001b[0;34m(self, key, timeout)\u001b[0m\n\u001b[1;32m    309\u001b[0m future \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_dask_future(key)\n\u001b[1;32m    310\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 311\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m future\u001b[38;5;241m.\u001b[39mresult(timeout\u001b[38;5;241m=\u001b[39mtimeout)\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m distributed\u001b[38;5;241m.\u001b[39mTimeoutError:\n\u001b[1;32m    313\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py:336\u001b[0m, in \u001b[0;36mFuture._result\u001b[0;34m(self, raiseit)\u001b[0m\n\u001b[1;32m    334\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m raiseit:\n\u001b[1;32m    335\u001b[0m     typ, exc, tb \u001b[38;5;241m=\u001b[39m exc\n\u001b[0;32m--> 336\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mwith_traceback(tb)\n\u001b[1;32m    337\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    338\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m exc\n",
-      "\u001b[0;31mKilledWorker\u001b[0m: Attempted to run task 'md-2-9b52cfa3e22647de99659385eca3d3f2-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.32:43067. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see https://distributed.dask.org/en/stable/killed.html."
+      "Cell \u001b[0;32mIn[6], line 6\u001b[0m, in \u001b[0;36mcombustion\u001b[0;34m(atoms)\u001b[0m\n\u001b[1;32m      3\u001b[0m futures \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m model \u001b[38;5;129;01min\u001b[39;00m MLIPEnum:\n\u001b[0;32m----> 6\u001b[0m     future \u001b[38;5;241m=\u001b[39m \u001b[43mMD\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m        \u001b[49m\u001b[43matoms\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43matoms\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcalculator_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcalculator_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[43m        \u001b[49m\u001b[43mensemble\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mnvt\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdynamics\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mnose-hoover\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     12\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtime_step\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     13\u001b[0m \u001b[43m        \u001b[49m\u001b[43mase_md_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mttime\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m25\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43munits\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpfactor\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     14\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtotal_time\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1000_000\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtemperature\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m300\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3000\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3000\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m300\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpressure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     17\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmb_velocity_seed\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtraj_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mPath\u001b[49m\u001b[43m(\u001b[49m\u001b[43mREGISTRY\u001b[49m\u001b[43m[\u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m]\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfamily\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mname\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m_\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43matoms\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_chemical_formula\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m.traj\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     20\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtraj_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1000\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     21\u001b[0m \u001b[43m        \u001b[49m\u001b[43mrestart\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     22\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m [future\u001b[38;5;241m.\u001b[39mresult() \u001b[38;5;28;01mfor\u001b[39;00m future \u001b[38;5;129;01min\u001b[39;00m futures]\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/tasks.py:811\u001b[0m, in \u001b[0;36mTask.submit\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m    808\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mprefect\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mengine\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m create_autonomous_task_run, enter_task_run_engine\n\u001b[1;32m    810\u001b[0m \u001b[38;5;66;03m# Convert the call args/kwargs to a parameter dict\u001b[39;00m\n\u001b[0;32m--> 811\u001b[0m parameters \u001b[38;5;241m=\u001b[39m \u001b[43mget_call_parameters\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    812\u001b[0m return_type \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m return_state \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfuture\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    814\u001b[0m task_viz_tracker \u001b[38;5;241m=\u001b[39m get_task_viz_tracker()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:52\u001b[0m, in \u001b[0;36mget_call_parameters\u001b[0;34m(fn, call_args, call_kwargs, apply_defaults)\u001b[0m\n\u001b[1;32m     50\u001b[0m     bound_signature \u001b[38;5;241m=\u001b[39m inspect\u001b[38;5;241m.\u001b[39msignature(fn)\u001b[38;5;241m.\u001b[39mbind(\u001b[38;5;241m*\u001b[39mcall_args, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mcall_kwargs)\n\u001b[1;32m     51\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[0;32m---> 52\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m ParameterBindError\u001b[38;5;241m.\u001b[39mfrom_bind_failure(fn, exc, call_args, call_kwargs)\n\u001b[1;32m     54\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m apply_defaults:\n\u001b[1;32m     55\u001b[0m     bound_signature\u001b[38;5;241m.\u001b[39mapply_defaults()\n",
+      "\u001b[0;31mParameterBindError\u001b[0m: Error binding parameters for function 'run': got an unexpected keyword argument 'mb_velocity_seed'.\nFunction 'run' has signature 'atoms: 'Atoms', calculator_name: 'str | MLIPEnum', calculator_kwargs: 'dict | None', dispersion: 'str | None' = None, dispersion_kwargs: 'dict | None' = None, device: 'str | None' = None, ensemble: \"Literal['nve', 'nvt', 'npt']\" = 'nvt', dynamics: 'str | MolecularDynamics' = 'langevin', time_step: 'float | None' = None, total_time: 'float' = 1000, temperature: 'float | Sequence | np.ndarray | None' = 300.0, pressure: 'float | Sequence | np.ndarray | None' = None, ase_md_kwargs: 'dict | None' = None, md_velocity_seed: 'int | None' = None, zero_linear_momentum: 'bool' = True, zero_angular_momentum: 'bool' = True, traj_file: 'str | Path | None' = None, traj_interval: 'int' = 1, restart: 'bool' = True' but received args: () and kwargs: ['atoms', 'calculator_name', 'calculator_kwargs', 'ensemble', 'dynamics', 'time_step', 'ase_md_kwargs', 'total_time', 'temperature', 'pressure', 'mb_velocity_seed', 'traj_file', 'traj_interval', 'restart']."
      ]
     }
    ],

mlip_arena/tasks/md.py

@@ -1,4 +1,6 @@
 """
+Define molecular dynamics tasks.
+
 This script has been adapted from Atomate2 MLFF MD workflow written by Aaron Kaplan and Yuan Chiang
 https://github.com/materialsproject/atomate2/blob/main/src/atomate2/forcefields/md.py
 
@@ -52,9 +54,10 @@ works thereof, in binary and source code form.
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from datetime import datetime, timedelta
 from pathlib import Path
-from typing import Literal, Sequence, Tuple
+from typing import Literal
 
 import numpy as np
 from ase import Atoms, units
@@ -83,8 +86,6 @@ from tqdm.auto import tqdm
 
 from mlip_arena.models.utils import MLIPEnum, get_freer_device
 
-# from mlip_arena.models.utils import EXTMLIPEnum, MLIPMap, external_ase_calculator
-
 _valid_dynamics: dict[str, tuple[str, ...]] = {
     "nve": ("velocityverlet",),
     "nvt": ("nose-hoover", "langevin", "andersen", "berendsen"),
@@ -117,7 +118,7 @@ def _get_ensemble_schedule(
     n_steps: int = 1000,
     temperature: float | Sequence | np.ndarray | None = 300.0,
     pressure: float | Sequence | np.ndarray | None = None,
-) -> Tuple[np.ndarray, np.ndarray]:
+) -> tuple[np.ndarray, np.ndarray]:
     if ensemble == "nve":
         # Disable thermostat and barostat
         temperature = np.nan

mlip_arena/tasks/optimize.py

@@ -0,0 +1,124 @@
+""" 
+Define structure optimization tasks.
+"""
+
+from __future__ import annotations
+
+from datetime import timedelta
+
+from ase import Atoms
+from ase.calculators.calculator import Calculator
+from ase.calculators.mixing import SumCalculator
+from ase.filters import *  # type: ignore
+from ase.filters import Filter
+from ase.optimize import *  # type: ignore
+from ase.optimize.optimize import Optimizer
+from prefect import task
+from prefect.tasks import task_input_hash
+from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
+
+from mlip_arena.models.utils import MLIPEnum, get_freer_device
+
+_valid_filters: dict[str, Filter] = {
+    "Filter": Filter,
+    "UnitCell": UnitCellFilter,
+    "ExpCell": ExpCellFilter,
+    "Strain": StrainFilter,
+    "FrechetCell": FrechetCellFilter,
+} # type: ignore
+
+_valid_optimizers: dict[str, Optimizer] = {
+    "MDMin": MDMin,
+    "FIRE": FIRE,
+    "LBFGS": LBFGS,
+    "LBFGSLineSearch": LBFGSLineSearch,
+    "BFGS": BFGS,
+    "BFGSLineSearch": BFGSLineSearch,
+    "QuasiNewton": QuasiNewton,
+    "GPMin": GPMin,
+    "CellAwareBFGS": CellAwareBFGS,
+    "ODE12r": ODE12r,
+} # type: ignore
+
+
+@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(days=1))
+def run(
+    atoms: Atoms,
+    calculator_name: str | MLIPEnum,
+    calculator_kwargs: dict | None,
+    dispersion: str | None = None,
+    dispersion_kwargs: dict | None = None,
+    device: str | None = None,
+    optimizer: Optimizer | str = BFGSLineSearch,
+    optimizer_kwargs: dict | None = None,
+    filter: Filter | str | None = None,
+    filter_kwargs: dict | None = None,
+    criterion: dict | None = None,
+    # TODO: fix symmetry
+):
+    device = device or str(get_freer_device())
+
+    print(f"Using device: {device}")
+
+    calculator_kwargs = calculator_kwargs or {}
+
+    if isinstance(calculator_name, MLIPEnum) and calculator_name in MLIPEnum:
+        assert issubclass(calculator_name.value, Calculator)
+        calc = calculator_name.value(**calculator_kwargs)
+    elif (
+        isinstance(calculator_name, str) and calculator_name in MLIPEnum._member_names_
+    ):
+        calc = MLIPEnum[calculator_name].value(**calculator_kwargs)
+    else:
+        raise ValueError(f"Invalid calculator: {calculator_name}")
+
+    print(f"Using calculator: {calc}")
+
+    dispersion_kwargs = dispersion_kwargs or {}
+
+    dispersion_kwargs.update({"device": device})
+
+    if dispersion is not None:
+        disp_calc = TorchDFTD3Calculator(
+            **dispersion_kwargs,
+        )
+        calc = SumCalculator([calc, disp_calc])
+
+        print(f"Using dispersion: {dispersion}")
+
+    atoms.calc = calc
+
+    if isinstance(filter, str):
+        if filter not in _valid_filters:
+            raise ValueError(f"Invalid filter: {filter}")
+        filter = _valid_filters[filter]
+
+    if isinstance(optimizer, str):
+        if optimizer not in _valid_optimizers:
+            raise ValueError(f"Invalid optimizer: {optimizer}")
+        optimizer = _valid_optimizers[optimizer]
+
+
+    filter_kwargs = filter_kwargs or {}
+    optimizer_kwargs = optimizer_kwargs or {}
+    criterion = criterion or {}
+
+    if isinstance(filter, type) and issubclass(filter, Filter):
+        filter_instance = filter(atoms, **filter_kwargs)
+        print(f"Using filter: {filter_instance}")
+
+        optimizer_instance = optimizer(atoms, **optimizer_kwargs)
+        print(f"Using optimizer: {optimizer_instance}")
+
+        optimizer_instance.run(**criterion)
+
+    elif filter is None:
+        optimizer_instance = optimizer(atoms, **optimizer_kwargs)
+        print(f"Using optimizer: {optimizer_instance}")
+
+        optimizer_instance.run(**criterion)
+
+    return {
+        "atoms": atoms,
+    }
+

pyproject.toml

@@ -159,6 +159,7 @@ ignore = [
     "E741",  # ambigous variable naming, i.e. one letter
     "FBT003",  # boolean positional variable in function call
     "PERF203",  # `try`-`except` within a loop incurs performance overhead (no overhead in Py 3.11+)
+    "F405",  # 'module' may be undefined, or defined from star imports
 ]
 fixable = ["ALL"]
 pydocstyle.convention = "google"

tests/test_eos.py

@@ -0,0 +1,32 @@
+import numpy as np
+from ase.build import bulk
+
+from mlip_arena.tasks.eos.run import fit as EOS
+
+atoms = bulk("Cu", "fcc", a=3.6)
+
+def test_eos():
+    """
+    Test EOS prefect workflow with a simple cubic lattice.
+    """
+
+    result = EOS(
+        atoms=atoms,
+        calculator_name="MACE-MP(M)",
+        calculator_kwargs={},
+        device=None,
+        optimizer="BFGSLineSearch",
+        optimizer_kwargs=None,
+        filter="FrechetCell",
+        filter_kwargs=None,
+        criterion=dict(
+            fmax=0.01,
+        ),
+        max_abs_strain=0.1,
+        npoints=11,
+    )
+
+    print(result)
+
+    assert isinstance(result["K"], float)
+    assert np.isclose(result["K"], 143.2353241040571, rtol=5e-2)