refactor eos into task (#28)

* refactor eos into task

* task name and cache

* avoid duplicate test on PR branches

.github/README.md

@@ -53,7 +53,7 @@ streamlit run serve/app.py
 > The following are some tasks implemented:
 > - [Prefect structure optimization (OPT)](../mlip_arena/tasks/optimize.py)
 > - [Prefect molecular dynamics (MD)](../mlip_arena/tasks/md.py)
-> - [Prefect equation of states (EOS)](../mlip_arena/tasks/eos/run.py)
+> - [Prefect equation of states (EOS)](../mlip_arena/tasks/eos.py)
 
 1. Follow the task template to implement the task class and upload the script along with metadata to the MLIP Arena [here](../mlip_arena/tasks/README.md).
 2. Code a benchmark script to evaluate the performance of your model on the task. The script should be able to load the model and the dataset, and output the evaluation metrics.

.github/workflows/test.yaml

@@ -1,6 +1,10 @@
 name: Python Test
 
-on: [push, pull_request]
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
 
 # env:
 #   UV_SYSTEM_PYTHON: 1

mlip_arena/tasks/eos/run.py → mlip_arena/tasks/eos.py

@@ -1,5 +1,5 @@
 """
-Define equation of state flows.
+Define equation of state task.
 
 https://github.com/materialsvirtuallab/matcalc/blob/main/matcalc/eos.py
 """
@@ -9,36 +9,25 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+from prefect import task
+from prefect.futures import wait
+from prefect.runtime import task_run
+from prefect.tasks import task_input_hash
+
 from ase import Atoms
 from ase.filters import *  # type: ignore
 from ase.optimize import *  # type: ignore
 from ase.optimize.optimize import Optimizer
-from prefect import flow
-from prefect.futures import wait
-from prefect.runtime import flow_run, task_run
-from pymatgen.analysis.eos import BirchMurnaghan
-
 from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.optimize import run as OPT
+from pymatgen.analysis.eos import BirchMurnaghan
 
 if TYPE_CHECKING:
     from ase.filters import Filter
 
 
-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():
+def _generate_task_run_name():
     task_name = task_run.task_name
-
     parameters = task_run.parameters
 
     atoms = parameters["atoms"]
@@ -47,10 +36,12 @@ def generate_task_run_name():
     return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
 
 
-# https://docs.prefect.io/3.0/develop/write-tasks#custom-retry-behavior
-# @task(task_run_name=generate_task_run_name)
-@flow(flow_run_name=generate_flow_run_name, validate_parameters=False)
-def fit(
+@task(
+    name="EOS",
+    task_run_name=_generate_task_run_name,
+    cache_key_fn=task_input_hash,
+)
+def run(
     atoms: Atoms,
     calculator_name: str | MLIPEnum,
     calculator_kwargs: dict | None,