diff --git a/lambench/models/ase_models.py b/lambench/models/ase_models.py index bcfdd39f..6c317f8d 100644 --- a/lambench/models/ase_models.py +++ b/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, diff --git a/lambench/tasks/calculator/calculator_tasks.yml b/lambench/tasks/calculator/calculator_tasks.yml index 39272d99..c8a18f45 100644 --- a/lambench/tasks/calculator/calculator_tasks.yml +++ b/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 diff --git a/lambench/tasks/calculator/inference_efficiency/efficiency_utils.py b/lambench/tasks/calculator/inference_efficiency/efficiency_utils.py index afce54f4..a444c29a 100644 --- a/lambench/tasks/calculator/inference_efficiency/efficiency_utils.py +++ b/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]: diff --git a/lambench/tasks/calculator/inference_efficiency/inference_efficiency.py b/lambench/tasks/calculator/inference_efficiency/inference_efficiency.py index c1318d15..aa2800c8 100644 --- a/lambench/tasks/calculator/inference_efficiency/inference_efficiency.py +++ b/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. diff --git a/lambench/workflow/dflow.py b/lambench/workflow/dflow.py index dfde9bc4..dbb66c11 100644 --- a/lambench/workflow/dflow.py +++ b/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,