From 3a20637cf06ef3b6deac7f2347dd73807ce907d4 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 26 Aug 2026 16:56:40 +0800
Subject: [PATCH] fix: revert proxy

---
 lambench/models/ase_models.py |  3 ++-
 lambench/workflow/dflow.py    | 23 +----------------------
 2 files changed, 3 insertions(+), 23 deletions(-)

diff --git a/lambench/models/ase_models.py b/lambench/models/ase_models.py
index 149007f..cf4cfc0 100644
--- a/lambench/models/ase_models.py
+++ b/lambench/models/ase_models.py
@@ -17,7 +17,6 @@
 from ase.io import write
 from ase.optimize import FIRE
 from ase.calculators.emt import EMT
-from dftd3.ase import DFTD3
 from tqdm import tqdm
 
 from lambench.models.basemodel import BaseLargeAtomModel
@@ -379,6 +378,8 @@ def run_ase_dptest(
 
         calc = model.calc
         if dispersion_correction is not None:
+            from dftd3.ase import DFTD3
+
             calc = SumCalculator(
                 [calc, DFTD3(method="PBE", damping=dispersion_correction)]
             )
diff --git a/lambench/workflow/dflow.py b/lambench/workflow/dflow.py
index dbb66c1..9da2db4 100644
--- a/lambench/workflow/dflow.py
+++ b/lambench/workflow/dflow.py
@@ -19,27 +19,6 @@
 from lambench.models.basemodel import BaseLargeAtomModel
 from lambench.tasks.base_task import BaseTask
 
-# 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
 def run_task_op(
@@ -93,7 +72,7 @@ def submit_tasks_dflow(
                 python_packages=[
                     Path(package.__path__[0]) for package in [lambench, dpdata]
                 ],
-                pre_script=_DFTD3_PRE_SCRIPT,
+                pre_script="import os\nos.system('pip install dftd3')\n",
             ),
             parameters={
                 "task": task,