change from future to state to avoid OOM error for single task

examples/eos_alloy/run_Fe-Ni-Cr.ipynb

@@ -96,7 +96,7 @@
      "text": [
       "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/node.py:187: UserWarning: Port 8787 is already in use.\n",
       "Perhaps you already have a cluster running?\n",
-      "Hosting the HTTP server on port 40791 instead\n",
+      "Hosting the HTTP server on port 36753 instead\n",
       "  warnings.warn(\n"
      ]
     },
@@ -114,21 +114,22 @@
       "#SBATCH -N 1\n",
       "#SBATCH -C gpu\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.21:32827 --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.21:42119 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n",
       "\n"
      ]
     },
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">19:45:15.706 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'small-orca'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'run-from-db'</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">01:28:01.483 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'maroon-seagull'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'run-from-db'</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "19:45:15.706 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'small-orca'\u001b[0m for flow\u001b[1;35m 'run-from-db'\u001b[0m\n"
+       "01:28:01.483 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'maroon-seagull'\u001b[0m for flow\u001b[1;35m 'run-from-db'\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -137,11 +138,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\">19:45:15.711 | <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/405009c6-b3c3-49fd-8cbe-dd7771281f0c</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">01:28:01.487 | <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/3e9f3df7-6054-4f2e-b81c-8b7735c168fe</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "19:45:15.711 | \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/405009c6-b3c3-49fd-8cbe-dd7771281f0c\u001b[0m\n"
+       "01:28:01.487 | \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/3e9f3df7-6054-4f2e-b81c-8b7735c168fe\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -150,11 +151,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\">19:45:16.218 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(293f6d25, 'tcp://128.55.64.21:32827', 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\">01:28:02.091 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(ae948a05, 'tcp://128.55.64.21:42119', workers=0, threads=0, memory=0 B)\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "19:45:16.218 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(293f6d25, 'tcp://128.55.64.21:32827', workers=0, threads=0, memory=0 B)\n"
+       "01:28:02.091 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(ae948a05, 'tcp://128.55.64.21:42119', workers=0, threads=0, memory=0 B)\n"
       ]
      },
      "metadata": {},
@@ -163,24 +164,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\">19:45:57.865 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'get_atoms_from_db-936' - Created task run 'get_atoms_from_db-936' for task 'get_atoms_from_db'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">01:29:21.579 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'get_atoms_from_db-3a9' - Created task run 'get_atoms_from_db-3a9' for task 'get_atoms_from_db'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "19:45:57.865 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-936' - Created task run 'get_atoms_from_db-936' for task 'get_atoms_from_db'\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\">19:46:27.860 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'get_atoms_from_db-936' - Finished in state <span style=\"color: #008000; text-decoration-color: #008000\">Completed</span>()\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "19:46:27.860 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-936' - Finished in state \u001b[32mCompleted\u001b[0m()\n"
+       "01:29:21.579 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-3a9' - Created task run 'get_atoms_from_db-3a9' for task 'get_atoms_from_db'\n"
       ]
      },
      "metadata": {},
@@ -214,6 +202,7 @@
     "        f\"-N {nodes_per_alloc}\",\n",
     "        \"-C gpu\",\n",
     "        f\"-G {gpus_per_alloc}\",\n",
+    "        \"--exclusive\"\n",
     "    ],\n",
     ")\n",
     "\n",

mlip_arena/tasks/eos_alloy/flow.py

@@ -105,7 +105,7 @@ def run_from_db(
         on_completion=[partial(save_to_hdf, fpath=out_path, table_name=table_name)]
     )
 
-    futures = []
+    states = []
     for atoms in get_atoms_from_db(db_path):
         for mlip in MLIPEnum:
             if not REGISTRY[mlip.name]["npt"]:
@@ -115,7 +115,7 @@ def run_from_db(
                 + REGISTRY[mlip.name].get("gpu-tasks", [])
             ):
                 continue
-            future = EOS_.submit(
+            state = EOS_.submit(
                 atoms=atoms,
                 calculator_name=mlip.name,
                 calculator_kwargs=dict(),
@@ -127,13 +127,14 @@ def run_from_db(
                 max_abs_strain=max_abs_strain,
                 concurrent=concurrent,
                 cache_opt=False,
+                return_state=True
             )
-            futures.append(future)
+            states.append(state)
 
-    wait(futures)
+    wait(states)
 
     return [
-        f.result(timeout=None, raise_on_failure=False)
-        for f in futures
-        if f.state.is_completed()
+        s.result(timeout=None, raise_on_failure=False)
+        for s in states
+        if s.is_completed()
     ]