From bac49aac45181b0d6533302afde23b6007b09e77 Mon Sep 17 00:00:00 2001
From: Miki Bonacci <46074008+mikibonacci@users.noreply.github.com>
Date: Thu, 19 Dec 2024 18:21:07 +0100
Subject: [PATCH] Fixing the estimation of available cpus for a machine
 (localhost) (#971)

* Fixing the estimation of available cpus estimation

(now under the `machine_cpus` variable, no more `localhost_cpus`)
This is done by using the `get_default_mpiprocs_per_machine` method of the
computer node.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fixing the max cpus for localhost.

I take the information from the computer setup:

```python
default_mpiprocs = computer.get_default_mpiprocs_per_machine()
```

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* disabling (again) num_nodes if localhost

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update src/aiidalab_qe/app/submission/global_settings/model.py

Co-authored-by: Jusong Yu <jusong.yeu@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixing warning message test.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jusong Yu <jusong.yeu@gmail.com>
---
 .../app/submission/global_settings/model.py      | 16 +++++-----------
 src/aiidalab_qe/common/widgets.py                | 14 ++++++++------
 tests/test_submit_qe_workchain.py                |  6 +-----
 3 files changed, 14 insertions(+), 22 deletions(-)

diff --git a/src/aiidalab_qe/app/submission/global_settings/model.py b/src/aiidalab_qe/app/submission/global_settings/model.py
index 8b2710aa8..5f7446894 100644
--- a/src/aiidalab_qe/app/submission/global_settings/model.py
+++ b/src/aiidalab_qe/app/submission/global_settings/model.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import os
-
 import traitlets as tl
 
 from aiida import orm
@@ -125,12 +123,8 @@ def check_resources(self):
         num_sites = len(self.input_structure.sites)
         volume = self.input_structure.get_cell_volume()
 
-        try:
-            localhost_cpus = len(os.sched_getaffinity(0))
-        except Exception:
-            # Fallback, in some OS os.sched_getaffinity(0) is not supported
-            # However, not so reliable in containers
-            localhost_cpus = os.cpu_count()
+        code = orm.load_node(pw_code_model.selected)
+        machine_cpus = code.computer.get_default_mpiprocs_per_machine()
 
         large_system = (
             num_sites > self._RUN_ON_LOCALHOST_NUM_SITES_WARN_THRESHOLD
@@ -173,11 +167,11 @@ def check_resources(self):
                 + suggestions["change_configuration"]
                 + "</ul>"
             )
-        if on_localhost and num_cpus / localhost_cpus > 0.8:
+        if on_localhost and num_cpus / machine_cpus > 0.8:
             # Warning-3: on localhost, more than half of the available cpus
             alert_message += (
                 "<span>&#9888;</span> Warning: the selected pw.x code will run locally, but "
-                f"the number of requested CPUs ({num_cpus}) is larger than the 80% of the available resources ({localhost_cpus}). "
+                f"the number of requested CPUs ({num_cpus}) is larger than the 80% of the available resources ({machine_cpus}). "
                 "Please be sure that your local "
                 "environment has enough free CPUs for the calculation. Consider the following: "
                 "<ul>"
@@ -188,7 +182,7 @@ def check_resources(self):
 
         self.submission_warning_messages = (
             ""
-            if (on_localhost and num_cpus / localhost_cpus) <= 0.8
+            if (on_localhost and num_cpus / machine_cpus) <= 0.8
             and (not large_system or estimated_CPUs <= num_cpus)
             else self._ALERT_MESSAGE.format(
                 alert_class="warning",
diff --git a/src/aiidalab_qe/common/widgets.py b/src/aiidalab_qe/common/widgets.py
index 89f8673a4..d358d14b4 100644
--- a/src/aiidalab_qe/common/widgets.py
+++ b/src/aiidalab_qe/common/widgets.py
@@ -707,19 +707,21 @@ def update_resources(self, change):
             self.set_resource_defaults(load_code(change["new"]).computer)
 
     def set_resource_defaults(self, computer=None):
-        import os
-
-        if computer is None or computer.hostname == "localhost":
+        if computer is None:
             self.num_nodes.disabled = True
             self.num_nodes.value = 1
-            self.num_cpus.max = os.cpu_count()
+            self.num_cpus.max = 1
             self.num_cpus.value = 1
             self.num_cpus.description = "CPUs"
         else:
             default_mpiprocs = computer.get_default_mpiprocs_per_machine()
-            self.num_nodes.disabled = False
+            self.num_nodes.disabled = (
+                True if computer.hostname == "localhost" else False
+            )
             self.num_cpus.max = default_mpiprocs
-            self.num_cpus.value = default_mpiprocs
+            self.num_cpus.value = (
+                1 if computer.hostname == "localhost" else default_mpiprocs
+            )
             self.num_cpus.description = "CPUs"
 
     @property
diff --git a/tests/test_submit_qe_workchain.py b/tests/test_submit_qe_workchain.py
index 90d06d58f..df85d4175 100644
--- a/tests/test_submit_qe_workchain.py
+++ b/tests/test_submit_qe_workchain.py
@@ -148,12 +148,8 @@ def test_warning_messages(
     global_model = submit_model.get_model("global")
 
     pw_code = global_model.get_model("quantumespresso.pw")
-    pw_code.num_cpus = 1
-    global_model.check_resources()
-    # no warning:
-    assert submit_model.submission_warning_messages == ""
 
-    # now we increase the resources, so we should have the Warning-3
+    # we increase the resources, so we should have the Warning-3
     pw_code.num_cpus = len(os.sched_getaffinity(0))
     global_model.check_resources()
     for suggestion in ["avoid_overloading", "go_remote"]: