sage-num-threads.py

#!/usr/bin/env python
#
# Determine the number of threads to be used by Sage.
#
# Outputs three space-separated numbers:
# 1) The number of threads to use for Sage, based on MAKE, MAKEFLAGS
#    and SAGE_NUM_THREADS
# 2) The number of threads to use when parallel execution is explicitly
#    asked for (e.g. sage -tp)
# 3) The number of CPU cores in the system, as determined by
#    multiprocessing.cpu_count()
#
# AUTHOR: Jeroen Demeyer (2011-12-08): Trac ticket #12016
#

import os
import multiprocessing
import re
import math

def number_of_cores():
    """
    Try to determine the number of CPU cores in this system.
    If successful return that number. Otherwise return 1.
    """
    # If the environment variable SAGE_NUM_CORES exists, use that value.
    # This is useful for testing purposes.
    try:
        n = int(os.environ["SAGE_NUM_CORES"])
        if n > 0:
            return n
    except (ValueError, KeyError):
        pass

    try:
        n = multiprocessing.cpu_count()
        if n > 0:
            return n
    except NotImplementedError:
        pass

    try:  # Solaris fix
        from subprocess import Popen, PIPE
        p = Popen(['sysctl','-n','hw.ncpu'], stdin=PIPE, stdout=PIPE, stderr=PIPE, close_fds=True)
        n = int(p.stdout.read().strip())
        if n > 0:
            return n
    except (ValueError, OSError):
        pass

    return 1

def parse_jobs_from_MAKE(MAKE, unlimited=999999):
    """
    Parse an environment variable like :envvar:`MAKE` for the number of
    jobs specified. This looks at arguments ``-j``, ``--jobs``, ``-l``,
    ``--load-average``.

    INPUT:

    - ``MAKE`` -- The value of :envvar:`MAKE` or :envvar:`MAKEFLAGS`.

    - ``unlimited`` -- The value to return when ``MAKE`` contains ``-j``
      without argument and no ``-l`` option.  Normally this is interpreted
      as "unlimited".

    OUTPUT:

    The number of jobs specified by that variable.  Raise ``KeyError``
    if no number of jobs is specified in ``MAKE``.
    """
    # First, find value of -j
    # Since this is doing a greedy match on the left and non-greedy on the right,
    # we find the last -j or --jobs
    (j,num) = re.subn(r'^(.* )?(-j *|--jobs(=(?=[0-9])| +))([0-9]*)( .*?)?$', r'\4', MAKE, count=1)
    if num < 1:
        # No replacement done, i.e. no -j option found
        raise KeyError, "No number of jobs specified"
    elif j == "":
        # j is empty: unlimited number of jobs! :-)
        j = unlimited
    else:
        j = int(j)
        if j <= 0:
            raise ValueError, "Non-positive value specified for -j"

    # Next, find the value of -l
    # If it is specified, use this as an upper bound on j
    (l,num) = re.subn(r'^(.* )?(-l *|--(load-average|max-load)(=(?=[0-9])| +))([0-9.]*)( .*?)?$', r'\5', MAKE, count=1)
    if num < 1:
        # No replacement done, i.e. no -l option found
        pass
    elif l == "":
        # No load limit specified
        pass
    else:
        l = int(math.ceil(float(l)))
        # A load limit will never prevent starting at least one job
        if l <= 1:
            l = 1
        j = min(j,l)

    return j

def num_threads():
    """
    Determine the number of threads from the environment variables
    (in decreasing priority) :envvar:`SAGE_NUM_THREADS`, :envvar:`MAKE`,
    :envvar:`MAKEFLAGS` and :envvar:`MFLAGS`.

    If :envvar:`SAGE_NUM_THREADS` is 0 and neither :envvar:`MAKE` nor
    :envvar:`MAKEFLAGS` specifies a number of jobs, the use a default
    of ``min(8, number_of_cores)``.

    OUTPUT:

    a tuple (num_threads, num_threads_parallel, num_cores)
    """
    num_cores = number_of_cores()

    num_threads = None
    # Handle MFLAGS only for backwards compatibility
    try:
        num_threads = parse_jobs_from_MAKE(os.environ["MFLAGS"], unlimited=2)
    except (ValueError, KeyError):
        pass

    try:
        # Prepend hyphen to MAKEFLAGS if it does not start with one
        MAKEFLAGS=os.environ["MAKEFLAGS"]
        if MAKEFLAGS[0] != '-':
            MAKEFLAGS = '-' + MAKEFLAGS
        # In MAKEFLAGS, "-j" does not mean unlimited.  It probably
        # means an inherited number of jobs, let us use 2 for safety.
        num_threads = parse_jobs_from_MAKE(MAKEFLAGS, unlimited=2)
    except (ValueError, KeyError, IndexError):
        pass

    try:
        num_threads = parse_jobs_from_MAKE(os.environ["MAKE"])
    except (ValueError, KeyError):
        pass

    # Number of threads to use when parallel execution is explicitly
    # asked for
    num_threads_parallel = num_threads
    if num_threads_parallel is None:
        num_threads_parallel = max(min(8, num_cores), 2)

    try:
        sage_num_threads = int(os.environ["SAGE_NUM_THREADS"])
        if sage_num_threads == 0:
            # If SAGE_NUM_THREADS is 0, use the default only
            # if none of the above variables specified anything.
            if num_threads is None:
                num_threads = min(8, num_cores)
        elif sage_num_threads > 0:
            # SAGE_NUM_THREADS overrides everything
            num_threads = sage_num_threads
            num_threads_parallel = sage_num_threads
    except (ValueError, KeyError):
        pass

    # If we still don't know, use 1 thread
    if num_threads is None:
        num_threads = 1

    # Finally, use SAGE_NUM_THREADS_PARALLEL if set.  A user isn't
    # likely to set this, but it ensures that sage-env is idempotent
    # if called more than once.
    try:
        sage_num_threads = int(os.environ["SAGE_NUM_THREADS_PARALLEL"])
        if sage_num_threads > 0:
            num_threads_parallel = sage_num_threads
    except (ValueError, KeyError):
        pass

    return (num_threads, num_threads_parallel, num_cores)

print "%i %i %i"%num_threads()