Count the number of CPUs available to a process. In the case of containers (e.g. Docker) this will detect resources available to the container -- not the host machine.

container_cpu_count.py

"""
The MIT License

Copyright (c) 2020 Mark Douthwaite

Count the number of available CPUs to a process. In the case of containers (e.g. Docker) 
this will detect resources available to the container -- not the host machine. 

Note: in cases where _fractional_ CPU values are detected, this will _round up_ to the 
nearest integer value.
"""

import os
import math
import multiprocessing
from typing import Optional


def load(path: str) -> int:
    """Load a file containing a single integer and return it."""

    with open(path, "r") as file:
        return int(file.read().rstrip())


def cpu_count(
    quota_file: Optional[str] = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us",
    period_file: Optional[str] = "/sys/fs/cgroup/cpu/cpu.cfs_period_us",
    shares_file: Optional[str] = "/sys/fs/cgroup/cpu/cpu.shares",
) -> int:
    """
    Get the available CPU count of the current system.

    Note: this will find the CPU count of the system -- not the host machine.
    """

    if os.path.isfile(quota_file) and os.path.isfile(period_file):
        cfs_quota_us = load(quota_file)
        cfs_period_us = load(period_file)
        if cfs_quota_us > 0 and cfs_period_us > 0:
            return int(math.ceil(cfs_quota_us / cfs_period_us))

    if os.path.isfile(shares_file):
        cpu_shares = load(shares_file)
        return int(math.ceil(cpu_shares / 1024))

    return multiprocessing.cpu_count()