convert `ResultRecord` to `result` if fetched from cache (#39)

* convert `ResultRecord` to `result` if fetched from cache

* format

* hotfix chgnet

* move warnings filter before chgnet imports

* use simplefilter instead

* replace `raise_deadline_warning` as hotfix

* exclude chgnet for now

* install pymatgen from repo

* disable setup-python pip cache

* cleanup

.github/workflows/test.yaml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       # max-parallel: 2
       matrix:
-        python-version: ["3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12"]
         
 
     steps:
@@ -33,8 +33,6 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
-        python-version-file: "pyproject.toml"
-        cache: 'pip'
 
     - name: Install dependencies
       run: |

mlip_arena/tasks/__init__.py

@@ -6,10 +6,6 @@ from huggingface_hub import HfApi, HfFileSystem, hf_hub_download
 from mlip_arena.models import MLIP
 from mlip_arena.models import REGISTRY as MODEL_REGISTRY
 
-# from .run import md as MD
-
-# __all__ = ["MD"]
-
 with open(Path(__file__).parent / "registry.yaml") as f:
     REGISTRY = yaml.safe_load(f)
 

mlip_arena/tasks/eos.py

@@ -12,6 +12,7 @@ import numpy as np
 from prefect import task
 from prefect.cache_policies import INPUTS, TASK_SOURCE
 from prefect.futures import wait
+from prefect.results import ResultRecord
 from prefect.runtime import task_run
 from prefect.states import State
 
@@ -38,10 +39,7 @@ def _generate_task_run_name():
 
 
 @task(
-    name="EOS",
-    task_run_name=_generate_task_run_name,
-    cache_policy=TASK_SOURCE + INPUTS
-    # cache_key_fn=task_input_hash,
+    name="EOS", task_run_name=_generate_task_run_name, cache_policy=TASK_SOURCE + INPUTS
 )
 def run(
     atoms: Atoms,
@@ -75,6 +73,7 @@ def run(
         max_abs_strain: The maximum absolute strain to use.
         npoints: The number of points to sample.
         concurrent: Whether to relax multiple structures concurrently.
+        persist_opt: Whether to persist the optimization results.
         cache_opt: Whether to cache the intermediate optimization results.
 
     Returns:
@@ -83,7 +82,7 @@ def run(
 
     OPT_ = OPT.with_options(
         refresh_cache=not cache_opt,
-        persist_result=persist_opt,     
+        persist_result=persist_opt,
     )
 
     state = OPT_(
@@ -101,13 +100,13 @@ def run(
 
     if state.is_failed():
         return state
-    elif state.is_completed() and state.name in ["Completed", "Cached"]:
-        first_relax = state.result(raise_on_failure=False)
-    elif state.is_completed() and state.name in ["Rollback"]:
-        first_relax = state.result(raise_on_failure=False)
-    
-    assert isinstance(first_relax, dict)
-    relaxed = first_relax["atoms"]
+
+    first_relax = state.result(raise_on_failure=False)
+
+    if isinstance(first_relax, ResultRecord):
+        relaxed = first_relax.result["atoms"]
+    else:
+        relaxed = first_relax["atoms"]
 
     # p0 = relaxed.get_positions()
     c0 = relaxed.get_cell()
@@ -159,12 +158,13 @@ def run(
                 return_state=True,
             )
             states.append(state)
-        results = [
-            s.result(raise_on_failure=False) for s in states if s.is_completed()
-        ]
 
-    volumes = [f["atoms"].get_volume() for f in results]
-    energies = [f["atoms"].get_potential_energy() for f in results]
+        results = [s.result(raise_on_failure=False) for s in states if s.is_completed()]
+
+    results = [r.result if isinstance(r, ResultRecord) else r for r in results]
+
+    volumes = [r["atoms"].get_volume() for r in results]
+    energies = [r["atoms"].get_potential_energy() for r in results]
 
     volumes, energies = map(
         list,

pyproject.toml

@@ -66,6 +66,7 @@ test = [
     "pytest",
     "pytest-xdist",
     "prefect>=3.0.4,<3.1.0",
+    "pymatgen@git+https://github.com/chiang-yuan/pymatgen" # hotfix for deprecated pymatgen property https://github.com/materialsproject/pymatgen/pull/4242
 ]
 mace = [
     "mace-torch==0.3.4",