update gitignore

.gitignore

@@ -2,7 +2,7 @@
 *.ipynb
 *.extxyz
 *.traj
-mlip_arena/tasks/*/*/*/
+mlip_arena/tasks/*/*
 lab/
 manuscripts/
 

mlip_arena/tasks/eos/run.ipynb

@@ -11,6 +11,7 @@
     "from dask.distributed import Client\n",
     "from dask_jobqueue import SLURMCluster\n",
     "from prefect import task, flow\n",
+    "from prefect.task_runners import ThreadPoolTaskRunner\n",
     "from prefect_dask import DaskTaskRunner\n",
     "from pymatgen.core.structure import Structure\n",
     "from dotenv import load_dotenv\n",
@@ -43,7 +44,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bf1668eac663453aa9a20149ecf726c4",
+       "model_id": "bb6c1969c89840888c556f8fa59b4a67",
        "version_major": 2,
        "version_minor": 0
       },
@@ -69,7 +70,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,7 +89,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -97,7 +98,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {
     "tags": []
    },
@@ -108,7 +109,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -128,32 +129,9 @@
       "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.41:45323 --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.49:36289 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n",
       "\n"
      ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Exception in thread Profile:\n",
-      "Traceback (most recent call last):\n",
-      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/threading.py\", line 1045, in _bootstrap_inner\n",
-      "    self.run()\n",
-      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ipykernel/ipkernel.py\", line 766, in run_closure\n",
-      "    _threading_Thread_run(self)\n",
-      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/threading.py\", line 982, in run\n",
-      "    self._target(*self._args, **self._kwargs)\n",
-      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/profile.py\", line 366, in _watch\n",
-      "    process(frame, None, recent, omit=omit)\n",
-      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/profile.py\", line 183, in process\n",
-      "    if any(frame.f_code.co_filename.endswith(o) for o in omit):\n",
-      "       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/profile.py\", line 183, in <genexpr>\n",
-      "    if any(frame.f_code.co_filename.endswith(o) for o in omit):\n",
-      "           ^^^^^^^^^^^^^^^^^^^^^^^^\n",
-      "AttributeError: 'dict' object has no attribute 'co_filename'\n"
-     ]
     }
    ],
    "source": [
@@ -185,13 +163,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {
     "tags": []
    },
    "outputs": [],
    "source": [
-    "\n",
+    "from prefect.concurrency.sync import concurrency\n",
+    "from prefect.runtime import flow_run, task_run\n",
     "\n",
     "def postprocess(output, model: str, formula: str):\n",
     "    row = {\n",
@@ -215,8 +194,23 @@
     "    df.drop_duplicates(subset=[\"formula\", \"method\"], keep='last', inplace=True)\n",
     "    df.to_parquet(fpath)\n",
     "\n",
-    "from prefect.runtime import flow_run, task_run\n",
-    "EOS = EOS.with_options(timeout_seconds=240, result_storage=None)\n",
+    "\n",
+    "\n",
+    "task_runner = DaskTaskRunner(address=client.scheduler.address)\n",
+    "EOS = EOS.with_options(\n",
+    "    # task_runner=task_runner, \n",
+    "    log_prints=True,\n",
+    "    timeout_seconds=120, \n",
+    "    # result_storage=None\n",
+    ")\n",
+    "\n",
+    "from prefect import get_client\n",
+    "\n",
+    "async with get_client() as client:\n",
+    "    limit_id = await client.create_concurrency_limit(\n",
+    "        tag=\"bottleneck\", \n",
+    "        concurrency_limit=2\n",
+    "    )\n",
     "\n",
     "def generate_task_run_name():\n",
     "    task_name = task_run.task_name\n",
@@ -227,66 +221,116 @@
     "    \n",
     "    return f\"{task_name}: {atoms.get_chemical_formula()}\"\n",
     "\n",
-    "\n",
-    "@task(task_run_name=generate_task_run_name, result_storage=None)\n",
-    "def fit_one(atoms: Atoms):\n",
+    "@task(task_run_name=generate_task_run_name, tags=[\"bottleneck\"], timeout_seconds=150)\n",
+    "def fit_one(atoms: Atoms, model: str):\n",
     "    \n",
-    "    outputs = []\n",
-    "    for model in MLIPEnum:\n",
-    "        try:\n",
-    "            eos = EOS(\n",
-    "                atoms=atoms,\n",
-    "                calculator_name=model.name,\n",
-    "                calculator_kwargs={},\n",
-    "                device=None,\n",
-    "                optimizer=\"QuasiNewton\",\n",
-    "                optimizer_kwargs=None,\n",
-    "                filter=\"FrechetCell\",\n",
-    "                filter_kwargs=None,\n",
-    "                criterion=dict(\n",
-    "                    fmax=0.1,\n",
-    "                ),\n",
-    "                max_abs_strain=0.1,\n",
-    "                npoints=7,\n",
-    "            )\n",
-    "            if isinstance(eos, dict):\n",
-    "                postprocess(output=eos, model=model.name, formula=atoms.get_chemical_formula())\n",
-    "                eos[\"method\"] = model.name\n",
-    "                outputs.append(eos)\n",
-    "        except:\n",
-    "            continue\n",
+    "    eos = EOS(\n",
+    "        atoms=atoms,\n",
+    "        calculator_name=model,\n",
+    "        calculator_kwargs={},\n",
+    "        device=None,\n",
+    "        optimizer=\"QuasiNewton\",\n",
+    "        optimizer_kwargs=None,\n",
+    "        filter=\"FrechetCell\",\n",
+    "        filter_kwargs=None,\n",
+    "        criterion=dict(\n",
+    "            fmax=0.1,\n",
+    "        ),\n",
+    "        max_abs_strain=0.1,\n",
+    "        npoints=7,\n",
+    "    )\n",
+    "    if isinstance(eos, dict):\n",
+    "        postprocess(output=eos, model=model, formula=atoms.get_chemical_formula())\n",
+    "        eos[\"method\"] = model\n",
     "    \n",
-    "    return outputs\n",
-    "\n",
-    "\n",
-    "@flow(task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True, result_storage=None)\n",
+    "    return eos\n",
+    "    \n",
+    "#https://docs-3.prefect.io/3.0/develop/task-runners#use-multiple-task-runners\n",
+    "# @flow(task_runner=ThreadPoolTaskRunner(max_workers=50), log_prints=True)\n",
+    "@flow(task_runner=task_runner, log_prints=True)\n",
     "def fit_all(atoms_list: list[Atoms]):\n",
     "    \n",
     "    futures = []\n",
     "    for atoms in atoms_list:\n",
-    "        future = fit_one.submit(atoms)\n",
-    "        futures.append(future)\n",
+    "        futures_per_atoms = []\n",
+    "        for model in MLIPEnum:\n",
     "            \n",
-    "    wait(futures)\n",
+    "            # with concurrency(\"bottleneck\", occupy=2):\n",
+    "            future = fit_one.submit(atoms, model.name)\n",
+    "            # if not futures_per_atoms:\n",
+    "            #     if not futures:\n",
+    "            #         future = fit_one.submit(atoms, model.name)\n",
+    "            #     else:\n",
+    "            #         future = fit_one.submit(atoms, model.name, wait_for=[futures[-1]])                    \n",
+    "            # else:\n",
+    "            #     future = fit_one.submit(atoms, model.name, wait_for=[future])\n",
+    "            futures_per_atoms.append(future)\n",
+    "            \n",
+    "        futures.extend(futures_per_atoms)\n",
+    "\n",
+    "    return [f.result() for f in futures]\n",
+    "\n",
+    "\n",
+    "# @task(task_run_name=generate_task_run_name, result_storage=None)\n",
+    "# def fit_one(atoms: Atoms):\n",
+    "    \n",
+    "#     outputs = []\n",
+    "#     for model in MLIPEnum:\n",
+    "#         try:\n",
+    "#             eos = EOS(\n",
+    "#                 atoms=atoms,\n",
+    "#                 calculator_name=model.name,\n",
+    "#                 calculator_kwargs={},\n",
+    "#                 device=None,\n",
+    "#                 optimizer=\"QuasiNewton\",\n",
+    "#                 optimizer_kwargs=None,\n",
+    "#                 filter=\"FrechetCell\",\n",
+    "#                 filter_kwargs=None,\n",
+    "#                 criterion=dict(\n",
+    "#                     fmax=0.1,\n",
+    "#                 ),\n",
+    "#                 max_abs_strain=0.1,\n",
+    "#                 npoints=7,\n",
+    "#             )\n",
+    "#             if isinstance(eos, dict):\n",
+    "#                 postprocess(output=eos, model=model.name, formula=atoms.get_chemical_formula())\n",
+    "#                 eos[\"method\"] = model.name\n",
+    "#                 outputs.append(eos)\n",
+    "#         except:\n",
+    "#             continue\n",
+    "    \n",
+    "#     return outputs\n",
+    "\n",
+    "# # https://orion-docs.prefect.io/latest/concepts/task-runners/#using-multiple-task-runners\n",
+    "# @flow(task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True, result_storage=None)\n",
+    "# def fit_all(atoms_list: list[Atoms]):\n",
     "    \n",
-    "    return [f.result(raise_on_failure=False) for f in futures]"
+    "#     futures = []\n",
+    "#     for atoms in atoms_list:\n",
+    "#         future = fit_one.submit(atoms)\n",
+    "#         futures.append(future)\n",
+    "            \n",
+    "#     wait(futures)\n",
+    "    \n",
+    "#     return [f.result(raise_on_failure=False) for f in futures]"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {
+    "scrolled": true,
     "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\">00:51:45.507 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'thistle-pronghorn'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'fit-all'</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:53:47.335 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'vengeful-malkoha'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'fit-all'</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "00:51:45.507 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'thistle-pronghorn'\u001b[0m for flow\u001b[1;35m 'fit-all'\u001b[0m\n"
+       "18:53:47.335 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'vengeful-malkoha'\u001b[0m for flow\u001b[1;35m 'fit-all'\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -295,11 +339,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\">00:51:45.509 | <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/9b293dae-4c2d-43b7-b8bb-9a4749eee74e</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:53:47.341 | <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/909d2bc4-695f-4eeb-8b7c-7660397a0692</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "00:51:45.509 | \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/9b293dae-4c2d-43b7-b8bb-9a4749eee74e\u001b[0m\n"
+       "18:53:47.341 | \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/909d2bc4-695f-4eeb-8b7c-7660397a0692\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -308,86 +352,37 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">00:51:45.818 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(5d2876ed, 'tcp://128.55.64.41:45323', workers=0, threads=0, memory=0 B)\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:53:47.654 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(df8c3d55, 'tcp://128.55.64.49:36289', workers=0, threads=0, memory=0 B)\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "00:51:45.818 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(5d2876ed, 'tcp://128.55.64.41:45323', workers=0, threads=0, memory=0 B)\n"
+       "18:53:47.654 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(df8c3d55, 'tcp://128.55.64.49:36289', 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\">00:58:09.352 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'thistle-pronghorn'</span> - Crash detected! Execution was aborted by an interrupt signal.\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "00:58:09.352 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'thistle-pronghorn'\u001b[0m - Crash detected! Execution was aborted by an interrupt signal.\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\">00:58:09.642 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'thistle-pronghorn'</span> - Finished in state <span style=\"color: #ff0000; text-decoration-color: #ff0000\">Crashed</span>('Execution was aborted by an interrupt signal.')\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "00:58:09.642 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'thistle-pronghorn'\u001b[0m - Finished in state \u001b[91mCrashed\u001b[0m('Execution was aborted by an interrupt signal.')\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "ename": "KeyboardInterrupt",
-     "evalue": "",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[6], line 21\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# import os\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# import tempfile\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m# import shutil\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     18\u001b[0m \n\u001b[1;32m     19\u001b[0m \u001b[38;5;66;03m# with twd():\u001b[39;00m\n\u001b[0;32m---> 21\u001b[0m \u001b[43mfit_all\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms_list\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: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[5], line 72\u001b[0m, in \u001b[0;36mfit_all\u001b[0;34m(atoms_list)\u001b[0m\n\u001b[1;32m     70\u001b[0m futures \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m     71\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m atoms \u001b[38;5;129;01min\u001b[39;00m atoms_list:\n\u001b[0;32m---> 72\u001b[0m     future \u001b[38;5;241m=\u001b[39m \u001b[43mfit_one\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     73\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[1;32m     75\u001b[0m wait(futures)\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/tasks.py:1163\u001b[0m, in \u001b[0;36mTask.submit\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1158\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m VisualizationUnsupportedError(\n\u001b[1;32m   1159\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`task.submit()` is not currently supported by `flow.visualize()`\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1160\u001b[0m     )\n\u001b[1;32m   1162\u001b[0m task_runner \u001b[38;5;241m=\u001b[39m flow_run_context\u001b[38;5;241m.\u001b[39mtask_runner\n\u001b[0;32m-> 1163\u001b[0m future \u001b[38;5;241m=\u001b[39m \u001b[43mtask_runner\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m return_state:\n\u001b[1;32m   1165\u001b[0m     future\u001b[38;5;241m.\u001b[39mwait()\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py:353\u001b[0m, in \u001b[0;36mDaskTaskRunner.submit\u001b[0;34m(self, task, parameters, wait_for, dependencies)\u001b[0m\n\u001b[1;32m    349\u001b[0m \u001b[38;5;66;03m# unpack the upstream call in order to cast Prefect futures to Dask futures\u001b[39;00m\n\u001b[1;32m    350\u001b[0m \u001b[38;5;66;03m# where possible to optimize Dask task scheduling\u001b[39;00m\n\u001b[1;32m    351\u001b[0m parameters \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_optimize_futures(parameters)\n\u001b[0;32m--> 353\u001b[0m future \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_client\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    354\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    355\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    356\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    357\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdependencies\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdependencies\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    358\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mstate\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    359\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    360\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m PrefectDaskFuture(wrapped_future\u001b[38;5;241m=\u001b[39mfuture, task_run_id\u001b[38;5;241m=\u001b[39mfuture\u001b[38;5;241m.\u001b[39mtask_run_id)\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/client.py:64\u001b[0m, in \u001b[0;36mPrefectDaskClient.submit\u001b[0;34m(self, func, key, workers, resources, retries, priority, fifo_timeout, allow_other_workers, actor, actors, pure, *args, **kwargs)\u001b[0m\n\u001b[1;32m     61\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     62\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m run_task_sync(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m---> 64\u001b[0m future \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     65\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwrapper_func\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     66\u001b[0m \u001b[43m    \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     67\u001b[0m \u001b[43m    \u001b[49m\u001b[43mworkers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mworkers\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     68\u001b[0m \u001b[43m    \u001b[49m\u001b[43mresources\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mresources\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     69\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretries\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mretries\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     70\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpriority\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpriority\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     71\u001b[0m \u001b[43m    \u001b[49m\u001b[43mfifo_timeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfifo_timeout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     72\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_other_workers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mallow_other_workers\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     73\u001b[0m \u001b[43m    \u001b[49m\u001b[43mactor\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mactor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     74\u001b[0m \u001b[43m    \u001b[49m\u001b[43mactors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mactors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     75\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpure\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     76\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mrun_task_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     77\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m future\u001b[38;5;241m.\u001b[39mtask_run_id \u001b[38;5;241m=\u001b[39m run_task_kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtask_run_id\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m future\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py:1931\u001b[0m, in \u001b[0;36mClient.submit\u001b[0;34m(self, func, key, workers, resources, retries, priority, fifo_timeout, allow_other_workers, actor, actors, pure, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1929\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m key \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   1930\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m pure:\n\u001b[0;32m-> 1931\u001b[0m         key \u001b[38;5;241m=\u001b[39m funcname(func) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m-\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[43mtokenize\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1932\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1933\u001b[0m         key \u001b[38;5;241m=\u001b[39m funcname(func) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m-\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(uuid\u001b[38;5;241m.\u001b[39muuid4())\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1033\u001b[0m, in \u001b[0;36mtokenize\u001b[0;34m(ensure_deterministic, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1015\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Deterministic token\u001b[39;00m\n\u001b[1;32m   1016\u001b[0m \n\u001b[1;32m   1017\u001b[0m \u001b[38;5;124;03m>>> tokenize([1, 2, '3'])\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1030\u001b[0m \u001b[38;5;124;03m    Defaults to the `tokenize.ensure-deterministic` configuration parameter.\u001b[39;00m\n\u001b[1;32m   1031\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m   1032\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m _seen_ctx(reset\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m), _ensure_deterministic_ctx(ensure_deterministic):\n\u001b[0;32m-> 1033\u001b[0m     token: \u001b[38;5;28mobject\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1034\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m kwargs:\n\u001b[1;32m   1035\u001b[0m         token \u001b[38;5;241m=\u001b[39m token, _normalize_seq_func(\u001b[38;5;28msorted\u001b[39m(kwargs\u001b[38;5;241m.\u001b[39mitems()))\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1159\u001b[0m, in \u001b[0;36m_normalize_seq_func\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1157\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1158\u001b[0m         seen[\u001b[38;5;28mid\u001b[39m(item)] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(seen), item\n\u001b[0;32m-> 1159\u001b[0m         item \u001b[38;5;241m=\u001b[39m \u001b[43mnormalize_token\u001b[49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1160\u001b[0m     out\u001b[38;5;241m.\u001b[39mappend(item)\n\u001b[1;32m   1161\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/utils.py:773\u001b[0m, in \u001b[0;36mDispatch.__call__\u001b[0;34m(self, arg, *args, **kwargs)\u001b[0m\n\u001b[1;32m    769\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    770\u001b[0m \u001b[38;5;124;03mCall the corresponding method based on type of argument.\u001b[39;00m\n\u001b[1;32m    771\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    772\u001b[0m meth \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch(\u001b[38;5;28mtype\u001b[39m(arg))\n\u001b[0;32m--> 773\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m,\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",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1129\u001b[0m, in \u001b[0;36mnormalize_dict\u001b[0;34m(d)\u001b[0m\n\u001b[1;32m   1127\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((types\u001b[38;5;241m.\u001b[39mMappingProxyType, \u001b[38;5;28mdict\u001b[39m))\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_dict\u001b[39m(d):\n\u001b[0;32m-> 1129\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdict\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43msorted\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43md\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkv\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mkv\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\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/dask/base.py:1159\u001b[0m, in \u001b[0;36m_normalize_seq_func\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1157\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1158\u001b[0m         seen[\u001b[38;5;28mid\u001b[39m(item)] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(seen), item\n\u001b[0;32m-> 1159\u001b[0m         item \u001b[38;5;241m=\u001b[39m \u001b[43mnormalize_token\u001b[49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1160\u001b[0m     out\u001b[38;5;241m.\u001b[39mappend(item)\n\u001b[1;32m   1161\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/utils.py:773\u001b[0m, in \u001b[0;36mDispatch.__call__\u001b[0;34m(self, arg, *args, **kwargs)\u001b[0m\n\u001b[1;32m    769\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    770\u001b[0m \u001b[38;5;124;03mCall the corresponding method based on type of argument.\u001b[39;00m\n\u001b[1;32m    771\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    772\u001b[0m meth \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch(\u001b[38;5;28mtype\u001b[39m(arg))\n\u001b[0;32m--> 773\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m,\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",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1166\u001b[0m, in \u001b[0;36mnormalize_seq\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((\u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mlist\u001b[39m))\n\u001b[1;32m   1165\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_seq\u001b[39m(seq):\n\u001b[0;32m-> 1166\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(seq)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mseq\u001b[49m\u001b[43m)\u001b[49m\n",
-      "    \u001b[0;31m[... skipping similar frames: Dispatch.__call__ at line 773 (1 times), _normalize_seq_func at line 1159 (1 times)]\u001b[0m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1129\u001b[0m, in \u001b[0;36mnormalize_dict\u001b[0;34m(d)\u001b[0m\n\u001b[1;32m   1127\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((types\u001b[38;5;241m.\u001b[39mMappingProxyType, \u001b[38;5;28mdict\u001b[39m))\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_dict\u001b[39m(d):\n\u001b[0;32m-> 1129\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdict\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43msorted\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43md\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkv\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mkv\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
-      "    \u001b[0;31m[... skipping similar frames: Dispatch.__call__ at line 773 (1 times), _normalize_seq_func at line 1159 (1 times)]\u001b[0m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1166\u001b[0m, in \u001b[0;36mnormalize_seq\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((\u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mlist\u001b[39m))\n\u001b[1;32m   1165\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_seq\u001b[39m(seq):\n\u001b[0;32m-> 1166\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(seq)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mseq\u001b[49m\u001b[43m)\u001b[49m\n",
-      "    \u001b[0;31m[... skipping similar frames: Dispatch.__call__ at line 773 (3 times), _normalize_seq_func at line 1159 (3 times), normalize_dict at line 1129 (1 times), normalize_seq at line 1166 (1 times)]\u001b[0m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1129\u001b[0m, in \u001b[0;36mnormalize_dict\u001b[0;34m(d)\u001b[0m\n\u001b[1;32m   1127\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((types\u001b[38;5;241m.\u001b[39mMappingProxyType, \u001b[38;5;28mdict\u001b[39m))\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_dict\u001b[39m(d):\n\u001b[0;32m-> 1129\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdict\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43msorted\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43md\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkv\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mkv\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
-      "    \u001b[0;31m[... skipping similar frames: Dispatch.__call__ at line 773 (2 times), _normalize_seq_func at line 1159 (2 times), normalize_seq at line 1166 (1 times)]\u001b[0m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1166\u001b[0m, in \u001b[0;36mnormalize_seq\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((\u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mlist\u001b[39m))\n\u001b[1;32m   1165\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_seq\u001b[39m(seq):\n\u001b[0;32m-> 1166\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(seq)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mseq\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1159\u001b[0m, in \u001b[0;36m_normalize_seq_func\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1157\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1158\u001b[0m         seen[\u001b[38;5;28mid\u001b[39m(item)] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(seen), item\n\u001b[0;32m-> 1159\u001b[0m         item \u001b[38;5;241m=\u001b[39m \u001b[43mnormalize_token\u001b[49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1160\u001b[0m     out\u001b[38;5;241m.\u001b[39mappend(item)\n\u001b[1;32m   1161\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/utils.py:773\u001b[0m, in \u001b[0;36mDispatch.__call__\u001b[0;34m(self, arg, *args, **kwargs)\u001b[0m\n\u001b[1;32m    769\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    770\u001b[0m \u001b[38;5;124;03mCall the corresponding method based on type of argument.\u001b[39;00m\n\u001b[1;32m    771\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    772\u001b[0m meth \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch(\u001b[38;5;28mtype\u001b[39m(arg))\n\u001b[0;32m--> 773\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m,\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",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1212\u001b[0m, in \u001b[0;36mnormalize_object\u001b[0;34m(o)\u001b[0m\n\u001b[1;32m   1209\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _normalize_dataclass(o)\n\u001b[1;32m   1211\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1212\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_normalize_pickle\u001b[49m\u001b[43m(\u001b[49m\u001b[43mo\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1213\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n\u001b[1;32m   1214\u001b[0m     _maybe_raise_nondeterministic(\n\u001b[1;32m   1215\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mObject \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mo\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m cannot be deterministically hashed. See \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1216\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhttps://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1217\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfor more information.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1218\u001b[0m     )\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1250\u001b[0m, in \u001b[0;36m_normalize_pickle\u001b[0;34m(o)\u001b[0m\n\u001b[1;32m   1247\u001b[0m     buffers\u001b[38;5;241m.\u001b[39mclear()\n\u001b[1;32m   1248\u001b[0m     pik \u001b[38;5;241m=\u001b[39m cloudpickle\u001b[38;5;241m.\u001b[39mdumps(o, protocol\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m5\u001b[39m, buffer_callback\u001b[38;5;241m=\u001b[39mbuffers\u001b[38;5;241m.\u001b[39mappend)\n\u001b[0;32m-> 1250\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhash_buffer_hex\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpik\u001b[49m\u001b[43m)\u001b[49m, [hash_buffer_hex(buf) \u001b[38;5;28;01mfor\u001b[39;00m buf \u001b[38;5;129;01min\u001b[39;00m buffers]\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/hashing.py:98\u001b[0m, in \u001b[0;36mhash_buffer_hex\u001b[0;34m(buf, hasher)\u001b[0m\n\u001b[1;32m     94\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mhash_buffer_hex\u001b[39m(buf, hasher\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m     95\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     96\u001b[0m \u001b[38;5;124;03m    Same as hash_buffer, but returns its result in hex-encoded form.\u001b[39;00m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 98\u001b[0m     h \u001b[38;5;241m=\u001b[39m \u001b[43mhash_buffer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbuf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhasher\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     99\u001b[0m     s \u001b[38;5;241m=\u001b[39m binascii\u001b[38;5;241m.\u001b[39mb2a_hex(h)\n\u001b[1;32m    100\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m s\u001b[38;5;241m.\u001b[39mdecode()\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/hashing.py:88\u001b[0m, in \u001b[0;36mhash_buffer\u001b[0;34m(buf, hasher)\u001b[0m\n\u001b[1;32m     86\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m hasher \u001b[38;5;129;01min\u001b[39;00m hashers:\n\u001b[1;32m     87\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 88\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhasher\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbuf\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     89\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mTypeError\u001b[39;00m, \u001b[38;5;167;01mOverflowError\u001b[39;00m):\n\u001b[1;32m     90\u001b[0m         \u001b[38;5;28;01mpass\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/hashing.py:44\u001b[0m, in \u001b[0;36m_hash_xxhash\u001b[0;34m(buf)\u001b[0m\n\u001b[1;32m     40\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_hash_xxhash\u001b[39m(buf):\n\u001b[1;32m     41\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     42\u001b[0m \u001b[38;5;124;03m    Produce a 8-bytes hash of *buf* using xxHash.\u001b[39;00m\n\u001b[1;32m     43\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 44\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mxxhash\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxxh64\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbuf\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mdigest()\n",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
-     ]
     }
    ],
+   "source": [
+    "fit_all(atoms_list)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "```\n",
+    "Note that, because the DaskTaskRunner uses multiprocessing, calls to flows in scripts must be guarded with if __name__ == \"__main__\": or you will encounter warnings and errors.\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
    "source": [
     "# import os\n",
     "# import tempfile\n",
@@ -409,12 +404,12 @@
     "\n",
     "# with twd():\n",
     "\n",
-    "fit_all(atoms_list)"
+    "# fit_all(atoms_list)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 10,
    "metadata": {
     "tags": []
    },
@@ -427,7 +422,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
@@ -460,15 +455,47 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>0</th>\n",
+       "      <th>1</th>\n",
        "      <td>Ac2O3</td>\n",
        "      <td>MACE-MP(M)</td>\n",
        "      <td>[82.36010147441682, 85.41047560309894, 88.4608...</td>\n",
-       "      <td>[-39.47665786743164, -39.65583419799805, -39.7...</td>\n",
-       "      <td>95.869141</td>\n",
+       "      <td>[-39.47541427612305, -39.65580749511719, -39.7...</td>\n",
+       "      <td>95.755459</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>1</th>\n",
+       "      <th>2</th>\n",
+       "      <td>Ac6In2</td>\n",
+       "      <td>MACE-MP(M)</td>\n",
+       "      <td>[278.3036976131417, 288.61124196918433, 298.91...</td>\n",
+       "      <td>[-31.21324348449707, -31.40914535522461, -31.5...</td>\n",
+       "      <td>33.370214</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>Ac6Tl2</td>\n",
+       "      <td>MACE-MP(M)</td>\n",
+       "      <td>[278.30267000598286, 288.6101763025008, 298.91...</td>\n",
+       "      <td>[-29.572534561157227, -29.833026885986328, -30...</td>\n",
+       "      <td>29.065081</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>Ac3Sn</td>\n",
+       "      <td>MACE-MP(M)</td>\n",
+       "      <td>[135.293532345587, 140.30440391394214, 145.315...</td>\n",
+       "      <td>[-17.135194778442383, -17.228239059448242, -17...</td>\n",
+       "      <td>30.622045</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>AcAg</td>\n",
+       "      <td>MACE-MP(M)</td>\n",
+       "      <td>[55.376437498321394, 57.4274166649259, 59.4783...</td>\n",
+       "      <td>[-7.274301528930664, -7.346108913421631, -7.39...</td>\n",
+       "      <td>40.212164</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
        "      <td>Ac4</td>\n",
        "      <td>MACE-MP(M)</td>\n",
        "      <td>[166.09086069175856, 172.2423740507126, 178.39...</td>\n",
@@ -476,12 +503,12 @@
        "      <td>25.409891</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>3</th>\n",
+       "      <th>7</th>\n",
        "      <td>Ac16S24</td>\n",
        "      <td>MACE-MP(M)</td>\n",
        "      <td>[1006.5670668063424, 1043.84732853991, 1081.12...</td>\n",
-       "      <td>[-249.42129516601562, -250.79556274414062, -25...</td>\n",
-       "      <td>61.372858</td>\n",
+       "      <td>[-249.4179229736328, -250.7970733642578, -251....</td>\n",
+       "      <td>61.734158</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -489,17 +516,25 @@
       ],
       "text/plain": [
        "   formula      method                                            volumes  \\\n",
-       "0    Ac2O3  MACE-MP(M)  [82.36010147441682, 85.41047560309894, 88.4608...   \n",
-       "1      Ac4  MACE-MP(M)  [166.09086069175856, 172.2423740507126, 178.39...   \n",
-       "3  Ac16S24  MACE-MP(M)  [1006.5670668063424, 1043.84732853991, 1081.12...   \n",
+       "1    Ac2O3  MACE-MP(M)  [82.36010147441682, 85.41047560309894, 88.4608...   \n",
+       "2   Ac6In2  MACE-MP(M)  [278.3036976131417, 288.61124196918433, 298.91...   \n",
+       "3   Ac6Tl2  MACE-MP(M)  [278.30267000598286, 288.6101763025008, 298.91...   \n",
+       "4    Ac3Sn  MACE-MP(M)  [135.293532345587, 140.30440391394214, 145.315...   \n",
+       "5     AcAg  MACE-MP(M)  [55.376437498321394, 57.4274166649259, 59.4783...   \n",
+       "6      Ac4  MACE-MP(M)  [166.09086069175856, 172.2423740507126, 178.39...   \n",
+       "7  Ac16S24  MACE-MP(M)  [1006.5670668063424, 1043.84732853991, 1081.12...   \n",
        "\n",
        "                                            energies          K  \n",
-       "0  [-39.47665786743164, -39.65583419799805, -39.7...  95.869141  \n",
-       "1  [-16.326059341430664, -16.406923294067383, -16...  25.409891  \n",
-       "3  [-249.42129516601562, -250.79556274414062, -25...  61.372858  "
+       "1  [-39.47541427612305, -39.65580749511719, -39.7...  95.755459  \n",
+       "2  [-31.21324348449707, -31.40914535522461, -31.5...  33.370214  \n",
+       "3  [-29.572534561157227, -29.833026885986328, -30...  29.065081  \n",
+       "4  [-17.135194778442383, -17.228239059448242, -17...  30.622045  \n",
+       "5  [-7.274301528930664, -7.346108913421631, -7.39...  40.212164  \n",
+       "6  [-16.326059341430664, -16.406923294067383, -16...  25.409891  \n",
+       "7  [-249.4179229736328, -250.7970733642578, -251....  61.734158  "
       ]
      },
-     "execution_count": 8,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -536,105 +571,7 @@
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {
-    "state": {
-     "1e2c24979c4e4b3b91e8f8a8e214d56f": {
-      "model_module": "@jupyter-widgets/controls",
-      "model_module_version": "2.0.0",
-      "model_name": "HTMLModel",
-      "state": {
-       "layout": "IPY_MODEL_e03d78a760b8402cad0827279f4075c1",
-       "style": "IPY_MODEL_743568418ec24e369c2ac0f9a9573a1c",
-       "value": "Retrieving SummaryDoc documents: 100%"
-      }
-     },
-     "41e363c9f8514e6a993426e30d45ec98": {
-      "model_module": "@jupyter-widgets/controls",
-      "model_module_version": "2.0.0",
-      "model_name": "HTMLStyleModel",
-      "state": {
-       "description_width": "",
-       "font_size": null,
-       "text_color": null
-      }
-     },
-     "4dde6375de954082a87b1483182e40a3": {
-      "model_module": "@jupyter-widgets/base",
-      "model_module_version": "2.0.0",
-      "model_name": "LayoutModel",
-      "state": {}
-     },
-     "607687cda1ce4276a772b626ef3e6179": {
-      "model_module": "@jupyter-widgets/controls",
-      "model_module_version": "2.0.0",
-      "model_name": "HTMLModel",
-      "state": {
-       "layout": "IPY_MODEL_93a95097ee984d95b1dfedcb2d913435",
-       "style": "IPY_MODEL_41e363c9f8514e6a993426e30d45ec98",
-       "value": " 5135/5135 [00:04&lt;00:00, 958.06it/s]"
-      }
-     },
-     "743568418ec24e369c2ac0f9a9573a1c": {
-      "model_module": "@jupyter-widgets/controls",
-      "model_module_version": "2.0.0",
-      "model_name": "HTMLStyleModel",
-      "state": {
-       "description_width": "",
-       "font_size": null,
-       "text_color": null
-      }
-     },
-     "8d602e7a8e734c50b0fad54bf76cc7d5": {
-      "model_module": "@jupyter-widgets/controls",
-      "model_module_version": "2.0.0",
-      "model_name": "ProgressStyleModel",
-      "state": {
-       "description_width": ""
-      }
-     },
-     "93a95097ee984d95b1dfedcb2d913435": {
-      "model_module": "@jupyter-widgets/base",
-      "model_module_version": "2.0.0",
-      "model_name": "LayoutModel",
-      "state": {}
-     },
-     "b09359f88f954062aff30df3bf2ba46b": {
-      "model_module": "@jupyter-widgets/base",
-      "model_module_version": "2.0.0",
-      "model_name": "LayoutModel",
-      "state": {}
-     },
-     "bf1668eac663453aa9a20149ecf726c4": {
-      "model_module": "@jupyter-widgets/controls",
-      "model_module_version": "2.0.0",
-      "model_name": "HBoxModel",
-      "state": {
-       "children": [
-        "IPY_MODEL_1e2c24979c4e4b3b91e8f8a8e214d56f",
-        "IPY_MODEL_cf1131b45bf04178ac93962259718700",
-        "IPY_MODEL_607687cda1ce4276a772b626ef3e6179"
-       ],
-       "layout": "IPY_MODEL_4dde6375de954082a87b1483182e40a3"
-      }
-     },
-     "cf1131b45bf04178ac93962259718700": {
-      "model_module": "@jupyter-widgets/controls",
-      "model_module_version": "2.0.0",
-      "model_name": "FloatProgressModel",
-      "state": {
-       "bar_style": "success",
-       "layout": "IPY_MODEL_b09359f88f954062aff30df3bf2ba46b",
-       "max": 5135,
-       "style": "IPY_MODEL_8d602e7a8e734c50b0fad54bf76cc7d5",
-       "value": 5135
-      }
-     },
-     "e03d78a760b8402cad0827279f4075c1": {
-      "model_module": "@jupyter-widgets/base",
-      "model_module_version": "2.0.0",
-      "model_name": "LayoutModel",
-      "state": {}
-     }
-    },
+    "state": {},
     "version_major": 2,
     "version_minor": 0
    }

mlip_arena/tasks/eos/run.py

@@ -47,6 +47,7 @@ def generate_task_run_name():
     return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
 
 
+# https://docs.prefect.io/3.0/develop/write-tasks#custom-retry-behavior
 # @task(task_run_name=generate_task_run_name)
 @flow(flow_run_name=generate_flow_run_name, validate_parameters=False)
 def fit(
@@ -121,8 +122,16 @@ def fit(
 
     wait(futures)
 
-    volumes = [f.result()["atoms"].get_volume() for f in futures if isinstance(f.result(), dict)]
-    energies = [f.result()["atoms"].get_potential_energy() for f in futures if isinstance(f.result(), dict)]
+    volumes = [
+        f.result()["atoms"].get_volume()
+        for f in futures
+        if isinstance(f.result(), dict)
+    ]
+    energies = [
+        f.result()["atoms"].get_potential_energy()
+        for f in futures
+        if isinstance(f.result(), dict)
+    ]
 
     volumes, energies = map(
         list,