test_reference tweaks

Author: pdobbelaere

configs/local_test.yaml

@@ -15,18 +15,18 @@ ModelTraining:
   max_workers: 1  # suppress assertion for multigpu training
 
 CP2K:
-  cores_per_worker: 2
-  max_evaluation_time: 1
+  cores_per_worker: 1
+  max_evaluation_time: 0.1
   container_uri: 'oras://ghcr.io/molmod/cp2k:2024.1'
 
 GPAW:
-  cores_per_worker: 2
-  max_evaluation_time: 1
+  cores_per_worker: 1
+  max_evaluation_time: 0.1
   container_uri: 'oras://ghcr.io/molmod/gpaw:24.1'
 
 ORCA:
-  cores_per_worker: 2
-  max_evaluation_time: 1
+  cores_per_worker: 1
+  max_evaluation_time: 0.1
 
 
 ...

psiflow/__init__.py

@@ -56,6 +56,7 @@ def resolve_and_check(path: Path) -> Path:
 #  - see chatgpt convo for process memory limits and such
 #  - make /tmp for app workdirs an option?
 #  - what is scaling_cores_per_worker in WQ
+#  - can we clean up psiflow_internal slightly?
 #  -
 # TODO: REFERENCE
 #  - reference MPI args not really checked

psiflow/execution.py

@@ -35,6 +35,8 @@
 logger = logging.getLogger(__name__)  # logging per module
 
 
+PSIFLOW_INTERNAL = "psiflow_internal"
+
 EXECUTION_KWARGS = (
     "container_uri",
     "container_engine",
@@ -552,10 +554,10 @@ def from_config(
         default_threads: int = 4,
         htex_address: str = "127.0.0.1",
         zip_staging: Optional[bool] = None,
-        make_symlinks: bool = True,
+        make_symlinks: bool = False,
         **kwargs,
     ) -> ExecutionContext:
-        path = Path.cwd().resolve() / "psiflow_internal"
+        path = Path.cwd().resolve() / PSIFLOW_INTERNAL
         psiflow.resolve_and_check(path)
         if path.exists():
             shutil.rmtree(path)

psiflow/reference/dummy.py

@@ -42,7 +42,6 @@ def parse_output(self, stdout: str) -> dict:
         geom = Geometry.from_string(stdout)
         data = {
             "status": Status.SUCCESS,
-            "runtime": np.nan,
             "positions": geom.per_atom.positions,
             "natoms": len(geom),
             "energy": np.random.uniform(),

psiflow/reference/reference.py

@@ -1,6 +1,5 @@
 from __future__ import annotations  # necessary for type-guarding class methods
 
-import logging
 import warnings
 from typing import ClassVar, Optional, Union, Callable, Sequence
 from pathlib import Path
@@ -30,14 +29,27 @@ class Status(Enum):
     INCONSISTENT = 2
 
 
+class SinglePointResult:
+    """All dict keys update_geometry understands"""
+
+    status: Status
+    natoms: int  # optional
+    positions: np.ndarray
+    energy: float
+    forces: np.ndarray  # optional
+    stdout: Path
+    stderr: Path  # optional
+    runtime: float  # optional
+
+
 def update_geometry(geom: Geometry, data: dict) -> Geometry:
     """"""
     _, task_id, task_name = data["stdout"].stem.split("_", maxsplit=2)
     logger.info(f'Task "{task_name}" (ID {task_id}): {data["status"].name}')
 
     geom = geom.copy()
     geom.reset()
-    geom.order['status'], geom.order['task_id'] = data["status"].name, task_id
+    geom.order["status"], geom.order["task_id"] = data["status"].name, task_id
     if data["status"] != Status.SUCCESS:
         return geom
     geom.order["runtime"] = data.get("runtime")

psiflow/utils/io.py

@@ -82,6 +82,7 @@ def _save_txt(data: str, outputs: list[File] = []) -> None:
 
 @typeguard.typechecked
 def _load_metrics(inputs: list = []) -> np.recarray:
+    # TODO: stop using recarrays
     return np.load(inputs[0], allow_pickle=True)
 
 
@@ -90,6 +91,7 @@ def _load_metrics(inputs: list = []) -> np.recarray:
 
 @typeguard.typechecked
 def _save_metrics(data: np.recarray, outputs: list = []) -> None:
+    # TODO: stop using recarrays
     with open(outputs[0], "wb") as f:
         data.dump(f)
 

psiflow/utils/parse.py

@@ -1,10 +1,14 @@
 import datetime
+from pathlib import Path
 
 import numpy as np
 
+from psiflow.execution import PSIFLOW_INTERNAL
+
 
 class LineNotFoundError(Exception):
     """Call to find_line failed"""
+
     pass
 
 
@@ -27,7 +31,7 @@ def find_line(
                 return idx_start + i
             else:
                 return idx_start - i
-    raise LineNotFoundError('Could not find line starting with \"{}\".'.format(line))
+    raise LineNotFoundError('Could not find line starting with "{}".'.format(line))
 
 
 def lines_to_array(
@@ -38,8 +42,17 @@ def lines_to_array(
 
 
 def string_to_timedelta(timedelta: str) -> datetime.timedelta:
+    """"""
     allowed_units = "weeks", "days", "hours", "minutes", "seconds"
     time_list = timedelta.split()
     values, units = time_list[:-1:2], time_list[1::2]
     kwargs = {u: float(v) for u, v in zip(units, values) if u in allowed_units}
     return datetime.timedelta(**kwargs)
+
+
+def get_task_logs(task_id: int) -> tuple[Path, Path]:
+    """"""
+    path = Path.cwd().resolve() / PSIFLOW_INTERNAL / "000/task_logs"  # TODO
+    stdout = next(path.rglob(f"task_{task_id}_*.stdout"))
+    stderr = next(path.rglob(f"task_{task_id}_*.stderr"))
+    return stdout, stderr

tests/conftest.py

@@ -128,9 +128,7 @@ def dataset_h2(context):
     )
     data = [h2.copy() for i in range(20)]
     for atoms in data:
-        atoms.set_positions(
-            atoms.get_positions() + np.random.uniform(-0.05, 0.05, size=(2, 3))
-        )
+        atoms.positions += np.random.uniform(-0.05, 0.05, size=(2, 3))
     return Dataset([Geometry.from_atoms(a) for a in data])