number of available cpus (in cpuset)

easybuild/framework/easyblock.py

@@ -63,7 +63,7 @@
 from easybuild.tools.modules import ROOT_ENV_VAR_NAME_PREFIX, VERSION_ENV_VAR_NAME_PREFIX, DEVEL_ENV_VAR_NAME_PREFIX
 from easybuild.tools.modules import get_software_root, modules_tool
 from easybuild.tools.toolchain import DUMMY_TOOLCHAIN_NAME
-from easybuild.tools.systemtools import get_core_count
+from easybuild.tools.systemtools import get_avail_core_count
 from easybuild.tools.utilities import remove_unwanted_chars
 from easybuild.tools.version import this_is_easybuild, VERBOSE_VERSION, VERSION
 
@@ -1040,7 +1040,7 @@ def set_parallelism(self, nr=None):
             except ValueError, err:
                 self.log.error("Parallelism %s not integer: %s" % (nr, err))
         else:
-            nr = get_core_count()
+            nr = get_avail_core_count()
             # check ulimit -u
             out, ec = run_cmd('ulimit -u')
             try:

easybuild/main.py

@@ -777,7 +777,7 @@ def get_build_stats(app, starttime):
                               ('host', os.uname()[1]),
                               ('platform' , platform.platform()),
                               ('cpu_model', systemtools.get_cpu_model()),
-                              ('core_count', systemtools.get_core_count()),
+                              ('core_count', systemtools.get_avail_core_count()),
                               ('timestamp', int(time.time())),
                               ('build_time', buildtime),
                               ('install_size', app.det_installsize()),

easybuild/tools/modules.py

@@ -530,6 +530,7 @@ def run_module(self, *args, **kwargs):
         # old versions of modulecmd.tcl spit out something like "exec '<file>'" for load commands,
         # which is not correct Python code (and it knows, as the comments in modulecmd.tcl indicate)
         # so, rewrite "exec '/tmp/modulescript_X'" to the correct "execfile('/tmp/modulescript_X')"
+        # this is required for the DEISA variant of modulecmd.tcl which is commonly used
         def tweak_stdout(txt):
             """Tweak stdout before it's exec'ed as Python code."""
             modulescript_regex = "^exec\s+[\"'](?P<modulescript>/tmp/modulescript_[0-9_]+)[\"']$"
@@ -552,7 +553,7 @@ def available(self, mod_name=None):
         """
         mods = super(EnvironmentModulesTcl, self).available(mod_name=mod_name)
         # strip off slash at beginning, if it's there
-        # under certain circumstances, modulecmd.tcl (DEISA variant) spits out available modules like this
+        # under certain circumstances, 'modulecmd.tcl avail' (DEISA variant) spits out available modules like this
         clean_mods = [mod.lstrip(os.path.sep) for mod in mods]
 
         return clean_mods

easybuild/tools/systemtools.py

@@ -32,6 +32,11 @@
 import platform
 import re
 from vsc import fancylogger
+try:
+    # this import fails with Python 2.4 because it requires the ctypes module (only in Python 2.5+)
+    from vsc.utils.affinity import sched_getaffinity
+except ImportError:
+    pass
 
 from easybuild.tools.filetools import read_file, run_cmd
 
@@ -53,50 +58,79 @@ class SystemToolsException(Exception):
     """raised when systemtools fails"""
 
 
-def get_core_count():
-    """Try to detect the number of virtual or physical CPUs on this system.
-
-    inspired by http://stackoverflow.com/questions/1006289/how-to-find-out-the-number-of-cpus-in-python/1006301#1006301
+def get_avail_core_count():
     """
-    # Python 2.6+
-    try:
-        from multiprocessing import cpu_count
-        return cpu_count()
-    except (ImportError, NotImplementedError):
-        pass
-
-    # POSIX
-    try:
-        cores = int(os.sysconf('SC_NPROCESSORS_ONLN'))
-        if cores > 0:
-            return cores
-    except (AttributeError, ValueError):
-        pass
+    Returns the number of available CPUs, according to cgroups and taskssets limits
+    """
+    # tiny inner function to help figure out number of available cores in a cpuset
+    def count_bits(n):
+        """Count the number of set bits for a given integer."""
+        bit_cnt = 0
+        while n > 0:
+            n &= n - 1
+            bit_cnt += 1
+        return bit_cnt
 
     os_type = get_os_type()
-
     if os_type == LINUX:
+        try:
+            # the preferred approach is via sched_getaffinity (yields a long, so cast it down to int)
+            num_cores = int(sum(sched_getaffinity().cpus))
+            return num_cores
+        except NameError:
+            pass
+
+        # in case sched_getaffinity isn't available, fall back to relying on /proc/cpuinfo
+
+        # determine total number of cores via /proc/cpuinfo
         try:
             txt = read_file('/proc/cpuinfo', log_error=False)
             # sometimes this is uppercase
-            res = txt.lower().count('processor\t:')
-            if res > 0:
-                return res
+            max_num_cores = txt.lower().count('processor\t:')
         except IOError, err:
-            raise SystemToolsException("An error occured while determining core count: %s" % err)
+            raise SystemToolsException("An error occured while determining total core count: %s" % err)
+
+        # determine cpuset we're in (if any)
+        mypid = os.getpid()
+        try:
+            f = open("/proc/%s/status" % mypid, 'r')
+            txt = f.read()
+            f.close()
+            cpuset = re.search("^Cpus_allowed:\s*([0-9,a-f]+)", txt, re.M|re.I)
+        except IOError:
+            cpuset = None
+
+        if cpuset is not None:
+            # use cpuset mask to determine actual number of available cores
+            mask_as_int = long(cpuset.group(1).replace(',', ''), 16)
+            num_cores_in_cpuset = count_bits((2**max_num_cores - 1) & mask_as_int)
+            _log.info("In cpuset with %s CPUs" % num_cores_in_cpuset)
+            return num_cores_in_cpuset
+        else:
+            _log.debug("No list of allowed CPUs found, not in a cpuset.")
+            return max_num_cores
     else:
         # BSD
         try:
             out, _ = run_cmd('sysctl -n hw.ncpu')
-            cores = int(out)
-            if cores > 0:
-                return cores
+            num_cores = int(out)
+            if num_cores > 0:
+                return num_cores
         except ValueError:
             pass
 
     raise SystemToolsException('Can not determine number of cores on this system')
 
 
+def get_core_count():
+    """
+    Try to detect the number of virtual or physical CPUs on this system
+    (DEPRECATED, use get_avail_core_count instead)
+    """
+    _log.deprecated("get_core_count() is deprecated, use get_avail_core_count() instead", '2.0')
+    return get_avail_core_count()
+
+
 def get_cpu_vendor():
     """Try to detect the cpu identifier
 

test/framework/modules.py

@@ -30,13 +30,15 @@
 @author: Stijn De Weirdt (Ghent University)
 """
 
+import copy
 import os
 import re
 import tempfile
 import shutil
 
 import easybuild.tools.options as eboptions
 from easybuild.tools import config
+from easybuild.tools.environment import modify_env
 from easybuild.tools.modules import modules_tool
 from unittest import TestCase, TestLoader, main
 
@@ -50,6 +52,9 @@ class ModulesTest(TestCase):
 
     def setUp(self):
         """set up everything for a unit test."""
+        # keep track of original environment, so we can restore it
+        self.orig_environ = copy.deepcopy(os.environ)
+
         # initialize configuration so config.get_modules_tool function works
         eb_go = eboptions.parse_options()
         config.init(eb_go.options, eb_go.get_options_by_section('config'))
@@ -185,6 +190,7 @@ def tearDown(self):
         os.environ['MODULEPATH'] = os.pathsep.join(self.orig_modulepaths)
         # reinitialize a modules tool, to trigger 'module use' on module paths
         modules_tool()
+        modify_env(os.environ, self.orig_environ)
 
 def suite():
     """ returns all the testcases in this module """

test/framework/options.py

@@ -28,6 +28,7 @@
 @author: Kenneth Hoste (Ghent University)
 """
 
+import copy
 import os
 import re
 import shutil
@@ -41,6 +42,7 @@
 from easybuild.framework.easyconfig import BUILD, CUSTOM, DEPENDENCIES, EXTENSIONS, FILEMANAGEMENT, LICENSE
 from easybuild.framework.easyconfig import MANDATORY, MODULES, OTHER, TOOLCHAIN
 from easybuild.tools import config
+from easybuild.tools.environment import modify_env
 from easybuild.tools.filetools import read_file, write_file
 from easybuild.tools.modules import modules_tool
 from easybuild.tools.options import EasyBuildOptions
@@ -237,6 +239,9 @@ def test_force(self):
     def test_skip(self):
         """Test skipping installation of module (--skip, -k)."""
 
+        # keep track of original environment to restore after purging *all* loaded modules
+        orig_environ = copy.deepcopy(os.environ)
+
         # use temporary paths for build/install paths, make sure sources can be found
         buildpath = tempfile.mkdtemp()
         installpath = tempfile.mkdtemp()
@@ -311,6 +316,7 @@ def test_skip(self):
             os.environ.pop('MODULEPATH')
         # reinitialize modules tool with original $MODULEPATH, to avoid problems with future tests
         modules_tool()
+        modify_env(os.environ, orig_environ)
 
         # cleanup
         shutil.rmtree(buildpath)

test/framework/systemtools.py

@@ -31,7 +31,8 @@
 from unittest import TestCase, TestLoader, main
 
 from easybuild.tools.systemtools import AMD, ARM, DARWIN, INTEL, LINUX, UNKNOWN
-from easybuild.tools.systemtools import get_core_count, get_cpu_model, get_cpu_speed, get_cpu_vendor
+from easybuild.tools.systemtools import get_avail_core_count, get_core_count
+from easybuild.tools.systemtools import get_cpu_model, get_cpu_speed, get_cpu_vendor
 from easybuild.tools.systemtools import get_os_type, get_shared_lib_ext, get_platform_name, get_os_name, get_os_version
 
 
@@ -40,9 +41,9 @@ class SystemToolsTest(TestCase):
 
     def test_core_count(self):
         """Test getting core count."""
-        core_count = get_core_count()
-        self.assertTrue(isinstance(core_count, int))
-        self.assertTrue(core_count > 0)
+        for core_count in [get_avail_core_count(), get_core_count()]:
+            self.assertTrue(isinstance(core_count, int), "core_count has type int: %s, %s" % (core_count, type(core_count)))
+            self.assertTrue(core_count > 0, "core_count %d > 0" % core_count)
 
     def test_cpu_model(self):
         """Test getting CPU model."""

vsc/README.md

@@ -1,3 +1,3 @@
 Code from https://github.com/hpcugent/vsc-base
 
-based on 19c7ebcdc3f329b20450b5f703bd846c5b222ca5
+based on 2f42ee01237ec7a7c1ef0247c2dbd64a5dae086c

vsc/__init__.py

@@ -25,15 +25,15 @@
 # along with vsc-base. If not, see <http://www.gnu.org/licenses/>.
 ##
 """
-@author: Jens Timmerman (Ghent University)
-
 Initialize vsc package.
-the vsc namespace is used in different folders allong the system
+The vsc namespace is used in different folders allong the system
 so explicitly declare this is also the vsc namespace
+
+@author: Jens Timmerman (Ghent University)
 """
-# not in EB
-# import pkg_resources
-# pkg_resources.declare_namespace(__name__)
+# avoid that EasyBuild uses vsc package from somewhere else, e.g. a system vsc-base installation
+#import pkg_resources
+#pkg_resources.declare_namespace(__name__)
 
 # here for backwards compatibility
 from vsc.utils import fancylogger