Add CPUs per task metadata default and port

edan-bainglass · Feb 18, 2024 · d38f89c · d38f89c
1 parent 0ab36f4
commit d38f89c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/examples/coulomb_blockade/tester.ipynb b/examples/coulomb_blockade/tester.ipynb
@@ -182,16 +182,19 @@
    "source": [
     "h, m, s = 1, 0, 0\n",
     "\n",
-    "NPROCS = 1\n",
-    "OMP_NUM_THREADS = 1\n",
-    "NUMBA_NUM_THREADS = NPROCS\n",
+    "NODES = 1\n",
+    "TASKS_PER_NODE = 1\n",
+    "CPUS_PER_TASK = 1\n",
+    "OMP_NUM_THREADS = CPUS_PER_TASK\n",
+    "NUMBA_NUM_THREADS = CPUS_PER_TASK * TASKS_PER_NODE\n",
     "\n",
     "metadata = {\n",
     "    \"options\": {\n",
-    "        \"withmpi\": False,\n",
+    "        \"withmpi\": True,\n",
     "        \"resources\": {\n",
-    "            \"num_machines\": 1,\n",
-    "            \"num_mpiprocs_per_machine\": NPROCS,\n",
+    "            \"num_machines\": NODES,\n",
+    "            \"num_mpiprocs_per_machine\": TASKS_PER_NODE,\n",
+    "            \"num_cores_per_mpiprocs\": CPUS_PER_TASK,\n",
     "        },\n",
     "        \"max_wallclock_seconds\": h * 3600 + m * 60 + s,\n",
     "        \"environment_variables\": {\n",

diff --git a/src/aiida_quantum_transport/calculations/base.py b/src/aiida_quantum_transport/calculations/base.py
@@ -26,6 +26,7 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
             "metadata.options.resources": lambda: {
                 "num_machines": 1,
                 "num_mpiprocs_per_machine": 1,
+                "num_cores_per_mpiprocs": 1,
             },
         }