Allow fine grained intermediate cache for optimization in EOS task (#35)

* add eos cache

* remove unused import

* add uv into test workflow

* avoid bash inside uv

* disable concurrent eos test due to gh action instability

.github/workflows/test.yaml

@@ -6,8 +6,8 @@ on:
   pull_request:
     branches: [main]
 
-# env:
-#   UV_SYSTEM_PYTHON: 1
+env:
+  UV_SYSTEM_PYTHON: 1
 
 jobs:
   test:
@@ -23,30 +23,30 @@ jobs:
     - name: Checkout repository
       uses: actions/checkout@v4
 
-    # - name: Install uv
-    #   uses: astral-sh/setup-uv@v3
-    #   with:
-    #     enable-cache: true
-    #     cache-dependency-glob: "pyproject.toml"
-
-    # - name: Set up Python ${{ matrix.python-version }}
-    #   run: uv python install ${{ matrix.python-version }}
+    - name: Install uv
+      uses: astral-sh/setup-uv@v4
+      with:
+        enable-cahce: true
+        cache-dependency-glob: "pyproject.toml"
 
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
+        python-version-file: "pyproject.toml"
         cache: 'pip'
 
     - name: Install dependencies
       run: |
-        python -m pip install --upgrade pip
-        pip install torch==2.2.0
-        bash scripts/install-pyg.sh
-        bash scripts/install-dgl.sh
-        pip install -e .[test]
-        pip install -e .[mace]
-        pip install -e .[deepmd]
+        TORCH=2.2
+        CUDA=cu121
+        uv pip install torch==${TORCH}.0
+        uv pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
+        uv pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
+        uv pip install dgl -f https://data.dgl.ai/wheels/torch-${TORCH}/${CUDA}/repo.html
+        uv pip install -e .[test]
+        uv pip install -e .[mace]
+        uv pip install -e .[deepmd]
 
     - name: List dependencies
       run: pip list

mlip_arena/tasks/eos.py

@@ -56,6 +56,7 @@ def run(
     max_abs_strain: float = 0.1,
     npoints: int = 11,
     concurrent: bool = True,
+    cache_opt: bool = True,
 ) -> dict[str, Any] | State:
     """
     Compute the equation of state (EOS) for the given atoms and calculator.
@@ -73,11 +74,17 @@ 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.
+        cache_opt: Whether to cache the intermediate optimization results.
 
     Returns:
         A dictionary containing the EOS data, bulk modulus, equilibrium volume, and equilibrium energy if successful. Otherwise, a prefect state object.
     """
-    state = OPT(
+
+    OPT_ = OPT.with_options(
+        refresh_cache=not cache_opt
+    )
+
+    state = OPT_(
         atoms=atoms,
         calculator_name=calculator_name,
         calculator_kwargs=calculator_kwargs,
@@ -92,8 +99,11 @@ def run(
 
     if state.is_failed():
         return state
-
-    first_relax = state.result(raise_on_failure=False)
+    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"]
 
@@ -108,7 +118,7 @@ def run(
             atoms = relaxed.copy()
             atoms.set_cell(c0 * f, scale_atoms=True)
 
-            future = OPT.submit(
+            future = OPT_.submit(
                 atoms=atoms,
                 calculator_name=calculator_name,
                 calculator_kwargs=calculator_kwargs,
@@ -134,7 +144,7 @@ def run(
             atoms = relaxed.copy()
             atoms.set_cell(c0 * f, scale_atoms=True)
 
-            state = OPT(
+            state = OPT_(
                 atoms=atoms,
                 calculator_name=calculator_name,
                 calculator_kwargs=calculator_kwargs,

tests/test_eos.py

@@ -10,10 +10,14 @@ from mlip_arena.tasks.eos import run as EOS
 
 
 @flow
-def single_eos_flow(calculator_name):
+def single_eos_flow(calculator_name, concurrent=True, cache=False):
     atoms = bulk("Cu", "fcc", a=3.6)
 
-    return EOS(
+    EOS_ = EOS.with_options(
+        refresh_cache=not cache,
+    )
+    
+    return EOS_(
         atoms=atoms,
         calculator_name=calculator_name,
         calculator_kwargs={},
@@ -27,7 +31,8 @@ def single_eos_flow(calculator_name):
         ),
         max_abs_strain=0.1,
         npoints=6,
-        concurrent=True
+        concurrent=concurrent,
+        cache_opt=cache
     )
 
 
@@ -35,8 +40,9 @@ def single_eos_flow(calculator_name):
     sys.version_info[:2] != (3, 11),
     reason="avoid prefect race condition on concurrent tasks",
 )
+@pytest.mark.parametrize("concurrent", [False])
 @pytest.mark.parametrize("model", [MLIPEnum["MACE-MP(M)"]])
-def test_eos(model: MLIPEnum):
+def test_eos(model: MLIPEnum, concurrent: bool):
     """
     Test EOS prefect workflow with a simple cubic lattice.
     """
@@ -44,6 +50,19 @@ def test_eos(model: MLIPEnum):
     with prefect_test_harness():
         result = single_eos_flow(
             calculator_name=model.name,
+            concurrent=concurrent,
+            cache=False,
         )
+        assert isinstance(b0_scratch := result["b0"], float)
+
+        # @pytest.mark.dependency(depends=["test_eos"])
+        # @pytest.mark.parametrize("model", [MLIPEnum["MACE-MP(M)"]])
+        # def test_eos_cache(model: MLIPEnum):
 
-        assert isinstance(result["b0"], float)
+        result = single_eos_flow(
+            calculator_name=model.name,
+            concurrent=concurrent,
+            cache=True,
+        )
+        assert isinstance(b0_cache := result["b0"], float)
+        assert b0_scratch == pytest.approx(b0_cache, rel=1e-6)