Merge pull request #97 from pdobbelaere/serialization

Serialization refactor

Author: pdobbelaere

psiflow/__init__.py

@@ -5,8 +5,8 @@
 from .serialization import (  # noqa: F401
     _DataFuture,
     deserialize,
-    serializable,
     serialize,
+    register_serializable,
 )
 
 
@@ -30,4 +30,3 @@ def resolve_and_check(path: Path) -> Path:
 load = ExecutionContextLoader.load
 context = ExecutionContextLoader.context
 wait = ExecutionContextLoader.wait
-

psiflow/data/dataset.py

@@ -36,7 +36,7 @@
 )
 
 
-@psiflow.serializable
+@psiflow.register_serializable
 class Dataset:
     """
     A class representing a dataset of atomic structures.

psiflow/execution.py

@@ -287,6 +287,7 @@ def __init__(
         # TODO: how to handle env variables?
         # disable thread affinity and busy-idling until we can isolate task resources
         default_env_vars = {
+            "PYTHONUNBUFFERED": "TRUE",
             "OMP_PROC_BIND": "FALSE",
             "OMP_WAIT_POLICY": "PASSIVE",
             "OMP_DISPLAY_ENV": "VERBOSE",  # verbose OMP log

psiflow/functions.py

@@ -297,7 +297,7 @@ def __post_init__(self):
         # OMP_NUM_THREADS for parallel evaluation does not work..
         # https://github.com/dftd3/simple-dftd3/issues/49
         # TODO: check whether this is still the case
-        os.environ["OMP_NUM_THREADS"] = str(self.num_threads * 10)
+        os.environ["OMP_NUM_THREADS"] = str(self.num_threads)
 
         from dftd3.ase import DFTD3
 

psiflow/hamiltonians.py

@@ -1,7 +1,7 @@
 import urllib
 from functools import partial
 from pathlib import Path
-from typing import ClassVar, Optional, Union, Callable, Sequence
+from typing import Optional, Union, Callable, Sequence
 
 import numpy as np
 from parsl.app.app import python_app
@@ -26,13 +26,16 @@
 from psiflow.utils.io import dump_json
 
 
-@psiflow.serializable
+apply_threads = python_app(_apply, executors=["default_threads"])
+apply_htex = python_app(_apply, executors=["default_htex"])
+apply_modelevaluation = python_app(_apply, executors=["ModelEvaluation"])
+
+
 class Hamiltonian(Computable):
-    # TODO: app is actually an instance variable, but serialization complains..
-    outputs: ClassVar[tuple] = ("energy", "forces", "stress")
+    outputs: tuple = ("energy", "forces", "stress")
     batch_size = 1000
-    app: ClassVar[Callable]
-    function_name: ClassVar[str]
+    app: Callable
+    function_name: str
 
     def compute(
         self,
@@ -44,12 +47,7 @@ def compute(
             outputs = tuple(self.__class__.outputs)
         if batch_size == -1:
             batch_size = self.__class__.batch_size
-        return compute(
-            arg,
-            self.app,
-            outputs_=outputs,
-            batch_size=batch_size,
-        )
+        return compute(arg, self.get_app(), outputs_=outputs, batch_size=batch_size)
 
     def __eq__(self, hamiltonian: "Hamiltonian") -> bool:
         raise NotImplementedError
@@ -80,32 +78,13 @@ def serialize_function(self, **kwargs) -> DataFuture:
         ).outputs[0]
 
     def parameters(self) -> dict:
-        return {}
-
-
-@psiflow.serializable
-class Zero(Hamiltonian):
-
-    def __init__(self):
-        apply_zero = python_app(_apply, executors=["default_threads"])
-        self.app = partial(apply_zero, function_cls=ZeroFunction)
-
-    def __eq__(self, hamiltonian: Hamiltonian) -> bool:
-        if type(hamiltonian) is Zero:
-            return True
-        return False
-
-    def __mul__(self, a: float) -> "Zero":
-        return Zero()
-
-    def __add__(self, hamiltonian: Hamiltonian) -> Hamiltonian:
-        # (Zero + Hamiltonian) is different from (Hamiltonian + Zero)
-        return hamiltonian
+        raise NotImplementedError
 
-    __rmul__ = __mul__  # handle float * Zero
+    def get_app(self) -> Callable:
+        raise NotImplementedError
 
 
-@psiflow.serializable
+@psiflow.register_serializable
 class MixtureHamiltonian(Hamiltonian):
     hamiltonians: list[Hamiltonian]
     coefficients: list[float]
@@ -127,10 +106,9 @@ def compute(  # override compute for efficient batching
     ) -> Union[list[AppFuture], AppFuture]:
         if outputs is None:
             outputs = list(self.__class__.outputs)
-        apply_apps = [h.app for h in self.hamiltonians]
+        apply_apps = [h.get_app() for h in self.hamiltonians]
         reduce_func = partial(
-            aggregate_multiple,
-            coefficients=np.array(self.coefficients),
+            aggregate_multiple, coefficients=np.array(self.coefficients)
         )
         return compute(
             arg,
@@ -231,25 +209,47 @@ def serialize(self, **kwargs) -> list[DataFuture]:
         return [h.serialize_function(**kwargs) for h in self.hamiltonians]
 
 
-@psiflow.serializable
+@psiflow.register_serializable
+class Zero(Hamiltonian):
+    function_name: str = "ZeroFunction"
+
+    def __init__(self):
+        pass
+
+    def get_app(self) -> Callable:
+        return partial(apply_threads, function_cls=ZeroFunction)
+
+    def __eq__(self, hamiltonian: Hamiltonian) -> bool:
+        if type(hamiltonian) is Zero:
+            return True
+        return False
+
+    def __mul__(self, a: float) -> "Zero":
+        return Zero()
+
+    def __add__(self, hamiltonian: Hamiltonian) -> Hamiltonian:
+        # (Zero + Hamiltonian) is different from (Hamiltonian + Zero)
+        return hamiltonian
+
+    __rmul__ = __mul__  # handle float * Zero
+
+
+@psiflow.register_serializable
 class EinsteinCrystal(Hamiltonian):
-    reference_geometry: Union[Geometry, AppFuture]
+    # TODO: logic not consistent depending on Geometry | AppFuture
+    reference_geometry: Geometry | AppFuture
     force_constant: float
-    function_name: ClassVar[str] = "EinsteinCrystalFunction"
+    function_name: str = "EinsteinCrystalFunction"
 
-    def __init__(
-        self, geometry: Union[Geometry, AppFuture[Geometry]], force_constant: float
-    ):
+    def __init__(self, geometry: Union[Geometry, AppFuture], force_constant: float):
         super().__init__()
         self.reference_geometry = copy_app_future(geometry)
         self.force_constant = force_constant
         self.external = None  # needed
-        self._create_apps()
 
-    def _create_apps(self):
-        apply_app = python_app(_apply, executors=["default_threads"])
-        self.app = partial(
-            apply_app, function_cls=EinsteinCrystalFunction, **self.parameters()
+    def get_app(self) -> Callable:
+        return partial(
+            apply_threads, function_cls=EinsteinCrystalFunction, **self.parameters()
         )
 
     def parameters(self) -> dict:
@@ -269,34 +269,27 @@ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
         return True
 
 
-@psiflow.serializable
+@psiflow.register_serializable
 class PlumedHamiltonian(Hamiltonian):
     plumed_input: str  # TODO: or future?
     external: Optional[psiflow._DataFuture]
-    function_name: ClassVar[str] = "PlumedFunction"
+    function_name: str = "PlumedFunction"
 
     def __init__(
         self,
         plumed_input: str,
         external: Union[None, str, Path, File, DataFuture] = None,
     ):
         super().__init__()
-
         self.plumed_input = remove_comments_printflush(plumed_input)
         if type(external) in [str, Path]:
             external = File(str(external))
         if external is not None:
             assert external.filepath in self.plumed_input
         self.external = external
-        self._create_apps()
 
-    def _create_apps(self):
-        apply_app = python_app(_apply, executors=["default_htex"])
-        self.app = partial(
-            apply_app,
-            function_cls=PlumedFunction,
-            **self.parameters(),
-        )
+    def get_app(self) -> Callable:
+        return partial(apply_htex, function_cls=PlumedFunction, **self.parameters())
 
     def parameters(self) -> dict:
         if self.external is not None:  # ensure parameters depends on self.external
@@ -314,28 +307,24 @@ def __eq__(self, other: Hamiltonian) -> bool:
         return True
 
 
-@psiflow.serializable
+@psiflow.register_serializable
 class Harmonic(Hamiltonian):
-    reference_geometry: Union[Geometry, AppFuture[Geometry]]
-    hessian: Union[np.ndarray, AppFuture[np.ndarray]]
-    function_name: ClassVar[str] = "HarmonicFunction"
+    reference_geometry: Geometry | AppFuture
+    hessian: np.ndarray | AppFuture
+    function_name: str = "HarmonicFunction"
 
     def __init__(
         self,
-        reference_geometry: Union[Geometry, AppFuture[Geometry]],
-        hessian: Union[np.ndarray, AppFuture[np.ndarray]],
+        reference_geometry: Geometry | AppFuture,
+        hessian: np.ndarray | AppFuture,
     ):
         # TODO: why not copy_app_future(geometry) like others?
         self.reference_geometry = reference_geometry
         self.hessian = hessian
-        self._create_apps()
 
-    def _create_apps(self):
-        apply_app = python_app(_apply, executors=["default_threads"])
-        self.app = partial(
-            apply_app,
-            function_cls=HarmonicFunction,
-            **self.parameters(),
+    def get_app(self) -> Callable:
+        return partial(
+            apply_threads, function_cls=HarmonicFunction, **self.parameters()
         )
 
     def parameters(self) -> dict:
@@ -367,30 +356,26 @@ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
         return True
 
 
-@psiflow.serializable
+@psiflow.register_serializable
 class D3Hamiltonian(Hamiltonian):
     method: str
     damping: str
-    function_name: ClassVar[str] = "DispersionFunction"
+    function_name: str = "DispersionFunction"
 
     def __init__(self, method: str, damping: str = "d3bj"):
         self.method, self.damping = method, damping
-        self._create_apps()
 
-    def _create_apps(self):
-        # TODO: does this make sense? GPU settings are useless for example
+    def get_app(self) -> Callable:
+        # execution-side parameters of function are not included in self.parameters()
         evaluation = psiflow.context().definitions["ModelEvaluation"]
-        apply_app = python_app(_apply, executors=["ModelEvaluation"])
         resources = evaluation.wq_resources(1)
-        resources.pop("gpus", None)
-
-        # execution-side parameters of function are not included in self.parameters()
-        self.app = partial(
-            apply_app,
+        resources.pop("gpus", None)  # do not request GPU
+        return partial(
+            apply_modelevaluation,
             function_cls=DispersionFunction,
             parsl_resource_specification=resources,
             **self.parameters(),
-            num_threads=resources.get("cores", 1),  # TODO: sloppy
+            num_threads=resources.get("cores"),
         )
 
     def parameters(self) -> dict:
@@ -406,11 +391,11 @@ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
         return True
 
 
-@psiflow.serializable
+@psiflow.register_serializable
 class MACEHamiltonian(Hamiltonian):
     external: psiflow._DataFuture
     atomic_energies: dict[str, float]
-    function_name: ClassVar[str] = "MACEFunction"
+    function_name: str = "MACEFunction"
 
     def __init__(
         self,
@@ -422,16 +407,13 @@ def __init__(
             self.external = File(external)
         else:
             self.external = external
-        self._create_apps()
 
-    def _create_apps(self):
+    def get_app(self) -> Callable:
+        # execution-side parameters of function are not included in self.parameters()
         evaluation = psiflow.context().definitions["ModelEvaluation"]
-        apply_app = python_app(_apply, executors=["ModelEvaluation"])
         resources = evaluation.wq_resources(1)
-
-        # execution-side parameters of function are not included in self.parameters()
-        self.app = partial(
-            apply_app,
+        return partial(
+            apply_modelevaluation,
             function_cls=MACEFunction,
             parsl_resource_specification=resources,
             **self.parameters(),
@@ -469,7 +451,7 @@ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
     # TODO: the methods below are outdated..
 
     @classmethod
-    def mace_mp0(cls, size: str = "small") -> 'MACEHamiltonian':
+    def mace_mp0(cls, size: str = "small") -> "MACEHamiltonian":
         urls = dict(
             small="https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model",  # 2023-12-10-mace-128-L0_energy_epoch-249.model
             large="https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-03-mace-128-L1_epoch-199.model",
@@ -483,7 +465,7 @@ def mace_mp0(cls, size: str = "small") -> 'MACEHamiltonian':
         return cls(parsl_file, {})
 
     @classmethod
-    def mace_cc(cls) -> 'MACEHamiltonian':
+    def mace_cc(cls) -> "MACEHamiltonian":
         url = "https://github.com/molmod/psiflow/raw/main/examples/data/ani500k_cc_cpu.model"
         parsl_file = psiflow.context().new_file("mace_mp_", ".pth")
         urllib.request.urlretrieve(

psiflow/learning.py

@@ -100,7 +100,7 @@ def evaluate_outputs(
 
 
 @typeguard.typechecked
-@psiflow.serializable
+# @psiflow.serializable
 class Learning:
     reference: Reference
     path_output: str

psiflow/metrics.py

@@ -454,7 +454,7 @@ def _initialize_wandb(
 
 
 @typeguard.typechecked
-@psiflow.serializable
+# @psiflow.serializable
 class Metrics:
     wandb_group: Optional[str]
     wandb_project: Optional[str]