use cpu_affinity to avoid painfully slow MDs when num_workers_per_blo…

…ck >> 1
molmod · Nov 1, 2023 · acbfa51 · acbfa51
1 parent e159a94
commit acbfa51
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/psiflow/execution.py b/psiflow/execution.py
@@ -342,6 +342,7 @@ def load(
                     cores_per_worker=definition.cores_per_worker,
                     max_workers=max_workers,
                     provider=definition.parsl_provider,
+                    cpu_affinity="block",
                 )
             else:
                 executor = ThreadPoolExecutor(

diff --git a/tests/test_reference.py b/tests/test_reference.py
@@ -541,9 +541,18 @@ def test_pyscf_timeout(context):
     basis = "cc-pvqz"
     spin = 0
     reference = PySCFReference(routine, basis, spin)
-    atoms = FlowAtoms(
-        numbers=np.ones(4),
-        positions=np.array([[0, 0, 0], [0.74, 0, 0], [0, 3, 0], [0.74, 3, 0]]),
+    atoms = FlowAtoms(  # 3 x H2
+        numbers=np.ones(6),
+        positions=np.array(
+            [
+                [0, 0, 0],
+                [0.74, 0, 0],
+                [0, 3, 0],
+                [0.74, 3, 0],
+                [0, 6, 0],
+                [0.74, 6, 0],
+            ]
+        ),
     )
     reference.evaluate(atoms).result()
     assert not atoms.reference_status