fine tune eos alloy flow

examples/eos_alloy/run_Fe-Ni-Cr.ipynb

@@ -80,7 +80,7 @@
     "from prefect_dask import DaskTaskRunner\n",
     "\n",
     "from mlip_arena.models import REGISTRY\n",
-    "from mlip_arena.tasks.eos_alloy.flow import run_from_db"
+    "from mlip_arena.tasks.eos_alloy.flow import run as EOS_ALLOY"
    ]
   },
   {
@@ -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 36753 instead\n",
+      "Hosting the HTTP server on port 36141 instead\n",
       "  warnings.warn(\n"
      ]
     },
@@ -110,7 +110,7 @@
       "#SBATCH --mem=0\n",
       "#SBATCH -t 00:30:00\n",
       "#SBATCH -J eos\n",
-      "#SBATCH -q regular\n",
+      "#SBATCH -q debug\n",
       "#SBATCH -N 1\n",
       "#SBATCH -C gpu\n",
       "#SBATCH -G 4\n",
@@ -118,18 +118,18 @@
       "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:42119 --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:41671 --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\">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 style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">16:10:48.738 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'azure-roadrunner'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'EOS Alloy'</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "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"
+       "16:10:48.738 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'azure-roadrunner'\u001b[0m for flow\u001b[1;35m 'EOS Alloy'\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -138,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\">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 style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">16:10:48.742 | <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/3c90a68c-fc1b-473a-a0e2-42212e6b1925</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "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"
+       "16:10:48.742 | \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/3c90a68c-fc1b-473a-a0e2-42212e6b1925\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -151,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\">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 style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">16:10:49.296 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(edf7e3ce, 'tcp://128.55.64.21:41671', workers=0, threads=0, memory=0 B)\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "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"
+       "16:10:49.296 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(edf7e3ce, 'tcp://128.55.64.21:41671', workers=0, threads=0, memory=0 B)\n"
       ]
      },
      "metadata": {},
@@ -164,11 +164,24 @@
     {
      "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: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 style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">16:10:49.329 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'get_atoms_from_db-656' - Created task run 'get_atoms_from_db-656' for task 'get_atoms_from_db'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "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"
+       "16:10:49.329 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-656' - Created task run 'get_atoms_from_db-656' 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\">16:11:17.219 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'get_atoms_from_db-656' - Finished in state <span style=\"color: #008000; text-decoration-color: #008000\">Completed</span>()\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "16:11:17.219 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-656' - Finished in state \u001b[32mCompleted\u001b[0m()\n"
       ]
      },
      "metadata": {},
@@ -198,7 +211,7 @@
     "    job_directives_skip=[\"-n\", \"--cpus-per-task\", \"-J\"],\n",
     "    job_extra_directives=[\n",
     "        \"-J eos\",\n",
-    "        \"-q regular\",\n",
+    "        \"-q debug\",\n",
     "        f\"-N {nodes_per_alloc}\",\n",
     "        \"-C gpu\",\n",
     "        f\"-G {gpus_per_alloc}\",\n",
@@ -208,18 +221,19 @@
     "\n",
     "cluster = SLURMCluster(**cluster_kwargs)\n",
     "print(cluster.job_script())\n",
-    "cluster.adapt(minimum_jobs=50, maximum_jobs=50)\n",
+    "cluster.adapt(minimum_jobs=2, maximum_jobs=2)\n",
     "client = Client(cluster)\n",
     "\n",
-    "run_from_db_ = run_from_db.with_options(\n",
-    "    task_runner=DaskTaskRunner(address=client.scheduler.address),\n",
-    "    # log_prints=True,\n",
-    ")\n",
-    "\n",
     "# Run the workflow\n",
     "\n",
-    "results = run_from_db_(\n",
-    "    db_path=\"sqs_Fe-Ni-Cr.db\", out_path=\"eos.h5\", table_name=\"Fe-Ni-Cr\"\n",
+    "results = EOS_ALLOY.with_options(\n",
+    "    task_runner=DaskTaskRunner(address=client.scheduler.address),\n",
+    "    # log_prints=True,\n",
+    ")(\n",
+    "    db_path=\"sqs_Fe-Ni-Cr.db\", \n",
+    "    out_path=\"eos.h5\", \n",
+    "    table_name=\"Fe-Ni-Cr\",\n",
+    "    cache=False,\n",
     ")"
    ]
   },

mlip_arena/tasks/eos.py

@@ -56,6 +56,7 @@ def run(
     max_abs_strain: float = 0.1,
     npoints: int = 11,
     concurrent: bool = True,
+    persist_opt: bool = True,
     cache_opt: bool = True,
 ) -> dict[str, Any] | State:
     """
@@ -81,7 +82,8 @@ def run(
     """
 
     OPT_ = OPT.with_options(
-        refresh_cache=not cache_opt
+        refresh_cache=not cache_opt,
+        persist_result=persist_opt,     
     )
 
     state = OPT_(

mlip_arena/tasks/eos_alloy/flow.py

@@ -88,8 +88,10 @@ def save_to_hdf(
         print(e)
 
 
-@flow
-def run_from_db(
+@flow(
+    name="EOS Alloy"
+)
+def run(
     db_path: Path | str,
     out_path: Path | str,
     table_name: str,
@@ -100,12 +102,14 @@ def run_from_db(
     criterion=dict(fmax=0.1, steps=1000),
     max_abs_strain=0.20,
     concurrent=False,
+    cache=True,
 ):
     EOS_ = EOS.with_options(
-        on_completion=[partial(save_to_hdf, fpath=out_path, table_name=table_name)]
+        on_completion=[partial(save_to_hdf, fpath=out_path, table_name=table_name)],
+        refresh_cache=not cache,
     )
 
-    states = []
+    futures = []
     for atoms in get_atoms_from_db(db_path):
         for mlip in MLIPEnum:
             if not REGISTRY[mlip.name]["npt"]:
@@ -115,7 +119,7 @@ def run_from_db(
                 + REGISTRY[mlip.name].get("gpu-tasks", [])
             ):
                 continue
-            state = EOS_.submit(
+            future = EOS_.submit(
                 atoms=atoms,
                 calculator_name=mlip.name,
                 calculator_kwargs=dict(),
@@ -126,15 +130,16 @@ def run_from_db(
                 criterion=criterion,
                 max_abs_strain=max_abs_strain,
                 concurrent=concurrent,
-                cache_opt=False,
-                return_state=True
+                persist_opt=cache,
+                cache_opt=cache,
+                # return_state=True
             )
-            states.append(state)
+            futures.append(future)
 
-    wait(states)
+    wait(futures)
 
     return [
-        s.result(timeout=None, raise_on_failure=False)
-        for s in states
-        if s.is_completed()
+        f.result(timeout=None, raise_on_failure=False)
+        for f in futures
+        if f.state.is_completed()
     ]