comment out unsupported models from registry

mlip_arena/models/registry.yaml

@@ -290,40 +290,40 @@ ORBv2:
   npt: true
   license: Apache-2.0
 
-ORBv2(MPTrj):
-  module: externals
-  class: ORBv2
-  family: orb
-  package: 
-  checkpoint: orb-mptraj-only-v2-20241014.ckpt
-  username: 
-  last-update: 2024-10-29T00:00:00
-  datetime: 2024-10-29T00:00:00 # TODO: Fake datetime
-  datasets:
-    - MPTrj
-  github: https://github.com/orbital-materials/orb-models
-  doi: 
-  date: 2024-10-15
-  prediction: EFS
-  nvt: true
-  npt: true
-  license: Apache-2.0
+# ORBv2(MPTrj):
+#   module: externals
+#   class: ORBv2
+#   family: orb
+#   package: 
+#   checkpoint: orb-mptraj-only-v2-20241014.ckpt
+#   username: 
+#   last-update: 2024-10-29T00:00:00
+#   datetime: 2024-10-29T00:00:00 # TODO: Fake datetime
+#   datasets:
+#     - MPTrj
+#   github: https://github.com/orbital-materials/orb-models
+#   doi: 
+#   date: 2024-10-15
+#   prediction: EFS
+#   nvt: true
+#   npt: true
+#   license: Apache-2.0
 
-eqV2(MPTrj-S):
-  module: externals
-  class: eqV2
-  family: fairchem
-  package: fairchem-core==1.2.0
-  checkpoint: eqV2_31M_mp.pt
-  username: fairchem # HF handle
-  last-update: 2024-10-29T00:00:00
-  datetime: 2024-10-29T00:00:00
-  datasets:
-    - MPTrj
-  prediction: EFS
-  nvt: true
-  npt: true
-  date: 2024-10-18
-  github: https://github.com/FAIR-Chem/fairchem
-  doi: https://arxiv.org/abs/2410.12771
-  license: Modified Apache-2.0 (Meta)
+# eqV2(MPTrj-S):
+#   module: externals
+#   class: eqV2
+#   family: fairchem
+#   package: fairchem-core==1.2.0
+#   checkpoint: eqV2_31M_mp.pt
+#   username: fairchem # HF handle
+#   last-update: 2024-10-29T00:00:00
+#   datetime: 2024-10-29T00:00:00
+#   datasets:
+#     - MPTrj
+#   prediction: EFS
+#   nvt: true
+#   npt: true
+#   date: 2024-10-18
+#   github: https://github.com/FAIR-Chem/fairchem
+#   doi: https://arxiv.org/abs/2410.12771
+#   license: Modified Apache-2.0 (Meta)

mlip_arena/tasks/combustion/water.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -18,18 +18,18 @@
     "from pymatgen.core import Molecule\n",
     "from pymatgen.io.packmol import PackmolBoxGen\n",
     "\n",
-    "from mlip_arena.models.utils import REGISTRY, MLIPEnum\n",
+    "from mlip_arena.models import REGISTRY, MLIPEnum\n",
+    "\n",
     "from mlip_arena.tasks.md import run as MD"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {
-    "jp-MarkdownHeadingCollapsed": true,
     "tags": []
    },
    "source": [
-    "## Create initial configuration"
+    "# Intial configuration"
    ]
   },
   {
@@ -124,7 +124,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 4,
    "metadata": {
     "tags": []
    },
@@ -144,9 +144,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/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 38185 instead\n",
+      "  warnings.warn(\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
@@ -155,29 +165,19 @@
       "\n",
       "#SBATCH -A matgen\n",
       "#SBATCH --mem=0\n",
-      "#SBATCH -t 00:30:00\n",
+      "#SBATCH -t 02:00:00\n",
+      "#SBATCH -J combustion-water\n",
+      "#SBATCH -q regular\n",
       "#SBATCH -N 1\n",
-      "#SBATCH -q debug\n",
       "#SBATCH -C gpu\n",
-      "#SBATCH -J combustion-water\n",
+      "#SBATCH -G 4\n",
+      "#SBATCH --exclusive\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.19:39737 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n",
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.23:40945 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
       "\n"
      ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/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 38693 instead\n",
-      "  warnings.warn(\n",
-      "2024-10-08 03:30:05,696 - distributed.scheduler - ERROR - Task run-b8e0881eb1e8959f855ada9515884ada marked as failed because 4 workers died while trying to run it\n",
-      "2024-10-08 03:30:05,723 - distributed.scheduler - ERROR - Task run-3307095fc41590db704148cc499bad14 marked as failed because 4 workers died while trying to run it\n"
-     ]
     }
    ],
    "source": [
@@ -185,54 +185,49 @@
     "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",
-    "#         \"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",
+    "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\": [f\"-N {nodes_per_alloc}\", \"-q debug\", \"-C gpu\", \"-J combustion-water\"],\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",
+    "    ],\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\": [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=2, maximum_jobs=2)\n",
     "client = Client(cluster)"
@@ -240,7 +235,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -274,153 +269,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {
     "tags": []
    },
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">01:25:41.097 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'green-turaco'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'combustion'</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "01:25:41.097 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'green-turaco'\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\">01:25:41.101 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - 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/runs/flow-run/784dc0e0-9f4f-4320-8bfa-4c5e8d3b35fb</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "01:25:41.101 | \u001b[36mINFO\u001b[0m    | prefect.engine - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/784dc0e0-9f4f-4320-8bfa-4c5e8d3b35fb\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\">01:25:41.409 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(91d5c261, 'tcp://128.55.64.19:39737', workers=0, threads=0, memory=0 B)\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "01:25:41.409 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(91d5c261, 'tcp://128.55.64.19:39737', workers=0, threads=0, memory=0 B)\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:30:05.739 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'green-turaco'</span> - Encountered exception during execution: KilledWorker('run-3307095fc41590db704148cc499bad14', &lt;WorkerState 'tcp://128.55.65.2:37007', name: SLURMCluster-1, status: closed, memory: 0, processing: 0&gt;, 3)\n",
-       "Traceback (most recent call last):\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 652, in run_context\n",
-       "    yield self\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 696, in run_flow_sync\n",
-       "    engine.call_flow_fn()\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 675, in call_flow_fn\n",
-       "    result = call_with_parameters(self.flow.fn, self.parameters)\n",
-       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
-       "    return fn(*args, **kwargs)\n",
-       "           ^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_892374/2043615938.py\", line 26, in combustion\n",
-       "    return [future.result() for future in futures]\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_892374/2043615938.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_dask/task_runners.py\", line 132, in result\n",
-       "    future_result = self._wrapped_future.result(timeout=timeout)\n",
-       "                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 328, in result\n",
-       "    return self.client.sync(self._result, callback_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 'run-3307095fc41590db704148cc499bad14' 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.2:37007. 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": [
-       "03:30:05.739 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'green-turaco'\u001b[0m - Encountered exception during execution: KilledWorker('run-3307095fc41590db704148cc499bad14', <WorkerState 'tcp://128.55.65.2:37007', name: SLURMCluster-1, status: closed, memory: 0, processing: 0>, 3)\n",
-       "Traceback (most recent call last):\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 652, in run_context\n",
-       "    yield self\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 696, in run_flow_sync\n",
-       "    engine.call_flow_fn()\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 675, in call_flow_fn\n",
-       "    result = call_with_parameters(self.flow.fn, self.parameters)\n",
-       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
-       "    return fn(*args, **kwargs)\n",
-       "           ^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_892374/2043615938.py\", line 26, in combustion\n",
-       "    return [future.result() for future in futures]\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_892374/2043615938.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_dask/task_runners.py\", line 132, in result\n",
-       "    future_result = self._wrapped_future.result(timeout=timeout)\n",
-       "                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 328, in result\n",
-       "    return self.client.sync(self._result, callback_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 'run-3307095fc41590db704148cc499bad14' 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.2:37007. 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\">03:30:06.504 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'green-turaco'</span> - Finished in state <span style=\"color: #d70000; text-decoration-color: #d70000\">Failed</span>(\"Flow run encountered an exception: KilledWorker: Attempted to run task 'run-3307095fc41590db704148cc499bad14' 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.2:37007. 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": [
-       "03:30:06.504 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'green-turaco'\u001b[0m - Finished in state \u001b[38;5;160mFailed\u001b[0m(\"Flow run encountered an exception: KilledWorker: Attempted to run task 'run-3307095fc41590db704148cc499bad14' 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.2:37007. 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 'run-3307095fc41590db704148cc499bad14' 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.2:37007. 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.",
-     "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",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flows.py:1345\u001b[0m, in \u001b[0;36mFlow.__call__\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1341\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[1;32m   1343\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;01mflow_engine\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m run_flow\n\u001b[0;32m-> 1345\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1346\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1347\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1348\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   1349\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   1350\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/flow_engine.py:818\u001b[0m, in \u001b[0;36mrun_flow\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    816\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m run_flow_async(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    817\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 818\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow_sync\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:698\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[1;32m    696\u001b[0m             engine\u001b[38;5;241m.\u001b[39mcall_flow_fn()\n\u001b[0;32m--> 698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m 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;01melse\u001b[39;00m \u001b[43mengine\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/flow_engine.py:255\u001b[0m, in \u001b[0;36mFlowRunEngine.result\u001b[0;34m(self, raise_on_failure)\u001b[0m\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m NotSet:\n\u001b[1;32m    254\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure:\n\u001b[0;32m--> 255\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised\n\u001b[1;32m    256\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised\n\u001b[1;32m    258\u001b[0m \u001b[38;5;66;03m# This is a fall through case which leans on the existing state result mechanics to get the\u001b[39;00m\n\u001b[1;32m    259\u001b[0m \u001b[38;5;66;03m# return value. This is necessary because we currently will return a State object if the\u001b[39;00m\n\u001b[1;32m    260\u001b[0m \u001b[38;5;66;03m# the State was Prefect-created.\u001b[39;00m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;66;03m# TODO: Remove the need to get the result from a State except in cases where the return value\u001b[39;00m\n\u001b[1;32m    262\u001b[0m \u001b[38;5;66;03m# is a State object.\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:652\u001b[0m, in \u001b[0;36mFlowRunEngine.run_context\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    645\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m timeout_context(\n\u001b[1;32m    646\u001b[0m         seconds\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow\u001b[38;5;241m.\u001b[39mtimeout_seconds,\n\u001b[1;32m    647\u001b[0m         timeout_exc_type\u001b[38;5;241m=\u001b[39mFlowRunTimeoutError,\n\u001b[1;32m    648\u001b[0m     ):\n\u001b[1;32m    649\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\n\u001b[1;32m    650\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExecuting flow \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m for flow run \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow_run\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m...\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    651\u001b[0m         )\n\u001b[0;32m--> 652\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n\u001b[1;32m    653\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    654\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_timeout(exc)\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:696\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    694\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mis_running():\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[0;32m--> 696\u001b[0m             \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_flow_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m 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;01melse\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mresult()\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:675\u001b[0m, in \u001b[0;36mFlowRunEngine.call_flow_fn\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    673\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _call_flow_fn()\n\u001b[1;32m    674\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 675\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[43mcall_with_parameters\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    676\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_success(result)\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:206\u001b[0m, in \u001b[0;36mcall_with_parameters\u001b[0;34m(fn, parameters)\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;124;03mCall a function with parameters extracted with `get_call_parameters`\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    203\u001b[0m \u001b[38;5;124;03mthe args/kwargs using `parameters_to_positional_and_keyword` directly\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    205\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m parameters_to_args_kwargs(fn, parameters)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[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[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\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_dask/task_runners.py:132\u001b[0m, in \u001b[0;36mPrefectDaskFuture.result\u001b[0;34m(self, timeout, raise_on_failure)\u001b[0m\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_final_state:\n\u001b[1;32m    131\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 132\u001b[0m         future_result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_wrapped_future\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\u001b[1;32m    133\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m distributed\u001b[38;5;241m.\u001b[39mTimeoutError \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    134\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m(\n\u001b[1;32m    135\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTask run \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtask_run_id\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m did not complete within \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtimeout\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m seconds\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    136\u001b[0m         ) \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/distributed/client.py:328\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    326\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_verify_initialized()\n\u001b[1;32m    327\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m shorten_traceback():\n\u001b[0;32m--> 328\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[43mclient\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msync\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_result\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback_timeout\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/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 'run-3307095fc41590db704148cc499bad14' 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.2:37007. 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."
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "results = combustion(atoms)"
    ]
@@ -430,7 +283,9 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "mlip_arena"
+   ]
   }
  ],
  "metadata": {