find_avail_cores_cgroup.py

import os

def _clamp(x, min, max):
    if x < min:
        return min
    if x > max:
        return max
    return x

def _read_cpu_quota_limit_cgroup():
    quota = -1
    for cpudirname in ("cpu,cpuacct", "cpuacct,cpu", "cpu"):
        quota_file = "/sys/fs/cgroup/" + cpudirname + "/cpu.cfs_quota_us"
        if not os.file.exists(quota_file):
            continue
        with open(quota_file, "r") as fh:
            quota = int(fh.read())
        break
    if quota < 0:
        return None
    period = -1
    for cpudirname in ("cpu,cpuacct", "cpuacct,cpu", "cpu"):
        period_file = "/sys/fs/cgroup/" + cpudirname + "/cpu.cfs_period_us"
        if not os.file.exists(period_file):
            continue
        with open(period_file, "r") as fh:
            period = int(fh.read())
        break
    if period < 0:
        return None
    return float(quota)/float(period)


def available_cores(min_cores=1, max_cores=Inf):
    # Initial estimation: number of (logical) cpus:
    max_usable_cores = min(os.cpu_count(), max_cores)
    # CPU affinity may further restrict the number of cpus:
    max_usable_cores = min(len(os.sched_getaffinity(0)), max_usable_cores)
    # cgroup quota
    ac = max_usable_cores
    cgroup_quota = _read_cpu_quota_limit_cgroup()
    if cgroup_quota is not None:
        ac = _clamp(cgroup_quota, min_cores, ac)
    # cgroup cpu shares
    ## cgroup cpu shares not implemented
    ## Estimate CPU load: est_load
    ## Find shares for current cgroup: cur_shares
    ## Find total CPU shares (are those available?): tot_shares
    ## if est_load > os.cpu_count():
    ##     ac = min(max_usable_cores * (cur_shares / tot_shares), ac)
    ac = floor(_clamp(ac, min_cores, max_usable_cores))
    return ac