Add equation of state task and flow

.github/workflows/test.yaml

@@ -30,5 +30,8 @@ jobs:
         pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"
 
     - name: Run tests
+      env:
+        PREFECT_API_KEY: ${{ secrets.PREFECT_API_KEY }}
+        PREFECT_API_URL: ${{ secrets.PREFECT_API_URL }}
       run: |
         pytest -v tests

mlip_arena/tasks/eos/run.py

@@ -6,7 +6,6 @@
 
 from __future__ import annotations
 
-import asyncio
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -15,9 +14,9 @@
 from ase.optimize import *  # type: ignore
 from ase.optimize.optimize import Optimizer
 from prefect import flow
-from pymatgen.analysis.eos import BirchMurnaghan
-
 from prefect.futures import wait
+from prefect.runtime import flow_run, task_run
+from pymatgen.analysis.eos import BirchMurnaghan
 
 from mlip_arena.models.utils import MLIPEnum
 from mlip_arena.tasks.optimize import run as OPT
@@ -26,13 +25,36 @@
     from ase.filters import Filter
 
 
-@flow
+def generate_flow_run_name():
+    flow_name = flow_run.flow_name
+
+    parameters = flow_run.parameters
+
+    atoms = parameters["atoms"]
+    calculator_name = parameters["calculator_name"]
+
+    return f"{flow_name}: {atoms.get_chemical_formula()} - {calculator_name}"
+
+
+def generate_task_run_name():
+    task_name = task_run.task_name
+
+    parameters = task_run.parameters
+
+    atoms = parameters["atoms"]
+    calculator_name = parameters["calculator_name"]
+
+    return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
+
+
+# @task(task_run_name=generate_task_run_name)
+@flow(flow_run_name=generate_flow_run_name, validate_parameters=False)
 def fit(
     atoms: Atoms,
     calculator_name: str | MLIPEnum,
     calculator_kwargs: dict | None,
     device: str | None = None,
-    optimizer: Optimizer | str = BFGSLineSearch, # type: ignore
+    optimizer: Optimizer | str = "BFGSLineSearch",  # type: ignore
     optimizer_kwargs: dict | None = None,
     filter: Filter | str | None = None,
     filter_kwargs: dict | None = None,
@@ -59,7 +81,7 @@ def fit(
     Returns:
         A dictionary containing the EOS data and the bulk modulus.
     """
-    result = OPT(
+    first_relax = OPT(
         atoms=atoms,
         calculator_name=calculator_name,
         calculator_kwargs=calculator_kwargs,
@@ -71,7 +93,7 @@ def fit(
         criterion=criterion,
     )
 
-    relaxed = result["atoms"]
+    relaxed = first_relax["atoms"]
 
     # p0 = relaxed.get_positions()
     c0 = relaxed.get_cell()
@@ -99,14 +121,20 @@ def fit(
 
     wait(futures)
 
-    volumes = [f.result()["atoms"].get_volume() for f in futures]
-    energies = [f.result()["atoms"].get_potential_energy() for f in 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, energies = map(
+        list,
+        zip(
+            *sorted(zip(volumes, energies, strict=True), key=lambda i: i[0]),
+            strict=True,
+        ),
+    )
 
     bm = BirchMurnaghan(volumes=volumes, energies=energies)
     bm.fit()
 
-    volumes, energies = map(list, zip(*sorted(zip(volumes, energies, strict=False), key=lambda i: i[0]), strict=False))
-
     return {
         "eos": {"volumes": volumes, "energies": energies},
         "K": bm.b0_GPa,

mlip_arena/tasks/optimize.py

@@ -41,7 +41,11 @@
 } # type: ignore
 
 
-@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(days=1))
+# @task(
+#     cache_key_fn=task_input_hash, 
+#     cache_expiration=timedelta(days=1),
+#     timeout_seconds=120)
+@task(timeout_seconds=120, result_storage=None)
 def run(
     atoms: Atoms,
     calculator_name: str | MLIPEnum,

tests/test_eos.py

@@ -12,7 +12,7 @@
 
 @pytest.mark.parametrize("model", [MLIPEnum["MACE-MP(M)"]])
 @pytest.mark.skipif(
-    sys.version_info != (3, 11), reason="requires Python 3.11 to use prefect"
+    tuple(sys.version_info)[:2] != (3, 11), reason="requires Python 3.11 to use prefect"
 )
 def test_eos(model: MLIPEnum):
     """