fix dflow dftd3 and grace oom

lambench/models/ase_models.py

@@ -84,6 +84,9 @@ def __init__(self, *args, **kwargs):
     @property
     def calc(self) -> Calculator:
         """ASE Calculator with the model loaded."""
+        if self._calc is not None:
+            return self._calc
+
         calculator_dispatch = {
             "MACE": self._init_mace_calculator,
             "ORB": self._init_orb_calculator,

lambench/tasks/calculator/calculator_tasks.yml

@@ -10,10 +10,11 @@ phonon_mdr:
     distance: 0.01
 inference_efficiency:
   test_data: /bohr/lambench-efficiency-rg7a/v3/efficiency
+  machine_type: 1 * NVIDIA L20_48g
   calculator_params:
     warmup_ratio: 0.1
     natoms_upper_limit:
-      dpa3_250211_v3_0_0_7M: 850
+      dpa3_250211_v3_0_0_7M: 1275
 torsionnet:
   test_data: /bohr/lambench-torsionnet-e4sc/v3/torsionnet500_ccsdt
   calculator_params: null

lambench/tasks/calculator/inference_efficiency/efficiency_utils.py

@@ -24,18 +24,31 @@ def get_efv(atoms: Atoms) -> tuple[float, np.ndarray, np.ndarray]:
     return e, f, v
 
 
+_OOM_MARKERS = (
+    "out of memory",
+    "oom",
+    "dst tensor is not initialized",
+    "resource_exhausted",
+)
+
+
 def catch_oom_error(atoms: Atoms) -> bool:
     """
     Catch OOM error when running inference.
+
+    TensorFlow reports several messages for the same failure: a genuine BFC OOM
+    often contains "ran out of memory", but a fully exhausted pool can instead
+    raise "Dst tensor is not initialized" when the destination buffer cannot be
+    allocated. Treat both as OOM so binary search can back off.
     """
     try:
         get_efv(atoms)
         return False
     except Exception as e:
-        if "out of memory" in str(e) or "OOM" in str(e):
+        if isinstance(e, MemoryError):
             return True
-        else:
-            return False
+        msg = str(e).lower()
+        return any(marker in msg for marker in _OOM_MARKERS)
 
 
 def get_divisors(num: int) -> list[int]:

lambench/tasks/calculator/inference_efficiency/inference_efficiency.py

@@ -19,7 +19,7 @@ def run_inference(
     model: ASEModel,
     test_data: Path,
     warmup_ratio: float,
-    natoms_upper_limit: int = 1000,
+    natoms_upper_limit: int = 1500,
 ) -> dict[str, dict[str, float]]:
     """
     Inference for all trajectories, return average time and success rate for each system.

lambench/workflow/dflow.py

@@ -18,7 +18,27 @@
 import lambench
 from lambench.models.basemodel import BaseLargeAtomModel
 from lambench.tasks.base_task import BaseTask
-from lambench.workflow.entrypoint import job_list
+
+# ASEModel imports dftd3 at module load, so the worker must install it before
+# `from lambench.workflow.dflow import run_task_op`. PythonOPTemplate prepends
+# this block to `script`. Do not use str.format placeholders here: dflow calls
+# pre_script.format(tmp_root=...). Worker images often pin a Tsinghua PyPI
+# mirror that 403s on dftd3 wheels; force indexes that Bohrium can reach.
+_DFTD3_PRE_SCRIPT = """\
+import importlib.util
+import os
+import subprocess
+import sys
+
+os.environ.pop("PIP_INDEX_URL", None)
+os.environ.pop("PIP_EXTRA_INDEX_URL", None)
+if importlib.util.find_spec("dftd3") is None:
+    install = [sys.executable, "-m", "pip", "install", "dftd3"]
+    try:
+        subprocess.check_call(install + ["--index-url", "https://mirrors.aliyun.com/pypi/simple", "--trusted-host", "mirrors.aliyun.com"])
+    except subprocess.CalledProcessError:
+        subprocess.check_call(install + ["--index-url", "https://pypi.org/simple", "--trusted-host", "pypi.org"])
+"""
 
 
 @OP.function
@@ -40,7 +60,7 @@ def get_dataset(paths: list[Optional[Path]]) -> Optional[list[BohriumDatasetsArt
 
 
 def submit_tasks_dflow(
-    jobs: job_list,
+    jobs: list[tuple[BaseTask, BaseLargeAtomModel]],
     name="lambench",
 ):
     job_group_id: int = create_job_group(name)
@@ -73,7 +93,7 @@ def submit_tasks_dflow(
                 python_packages=[
                     Path(package.__path__[0]) for package in [lambench, dpdata]
                 ],
-                pre_script="import os\nos.system('pip install dftd3')\n",
+                pre_script=_DFTD3_PRE_SCRIPT,
             ),
             parameters={
                 "task": task,