ContainerHelper.h#

Fully qualified name: omni/extras/ContainerHelper.h

File members: omni/extras/ContainerHelper.h

// SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: LicenseRef-NvidiaProprietary
//
// NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
// property and proprietary rights in and to this material, related
// documentation and any modifications thereto. Any use, reproduction,
// disclosure or distribution of this material and related documentation
// without an express license agreement from NVIDIA CORPORATION or
// its affiliates is strictly prohibited.
#pragma once

#include "../../carb/AlignSize.h"
#include "../../carb/MinMax.h"
#include "../../carb/Platform.h"

#if CARB_PLATFORM_LINUX
#    include <fcntl.h>
#    include <unistd.h>

#    include <cstdio>
#    include <cstdlib>
#    include <cstring>
#    include <string>
#    include <utility>
#    include <vector>
#endif

namespace omni
{
namespace extras
{

#if CARB_PLATFORM_LINUX || defined(DOXYGEN_BUILD)

#    ifndef DOXYGEN_SHOULD_SKIP_THIS
namespace detail
{

inline bool readLineFromFile(const char* file, char* buffer, size_t len) noexcept
{
    auto fd = ::open(file, O_RDONLY);
    if (fd == -1)
    {
        return false;
    }

    auto size = CARB_RETRY_EINTR(::read(fd, buffer, len - 1));
    ::close(fd);

    if (size <= 0)
    {
        return false;
    }

    buffer[size] = '\0';
    return true;
}

inline int32_t readIntFromFile(const char* file) noexcept
{
    char buffer[64];

    if (!readLineFromFile(file, buffer, CARB_COUNTOF(buffer)))
    {
        return -1;
    }

    return std::atoi(buffer);
}

struct CpuQuota
{
    int32_t quota = -1;

    int32_t period = 100'000;
};

inline CpuQuota readCpuQuotaFromFile(const char* file) noexcept
{
    char buffer[64];
    int quota;
    int period;
    char* endp = nullptr;

    if (!readLineFromFile(file, buffer, CARB_COUNTOF(buffer)))
    {
        return { -1, -1 };
    }

    // attempt to parse two values from the line.  The first value may be 'max' or a number
    // between 0 and 100000.  The second number is often 100000 but could be different depending
    // on scheduler settings.  We expect this file to always be properly formatted.

    // set to use the maximum quota for all available CPU cores => fail.
    if (strncmp(buffer, "max ", 4) == 0)
    {
        return { -1, std::atoi(&buffer[4]) };
    }

    quota = (int)strtol(buffer, &endp, 10);
    period = std::atoi(endp);

    return { quota, period };
}

inline int32_t getCpuSetCoreCount(const char* buffer, std::vector<int32_t>* cores = nullptr) noexcept
{
    std::string str = buffer;
    size_t pos = 0;
    size_t next;
    int32_t coreStart;
    int32_t coreEnd;
    int32_t total = 0;

    while (1)
    {
        std::string start;
        std::string end;

        next = str.find_first_of(",-", pos);

        // last item in the list (must be a single core) or a single core => parse and add it.
        if (next == std::string::npos || str[next] == ',')
        {
            size_t limit = (next == std::string::npos) ? str.length() : next;

            start = str.substr(pos, limit - pos);
            coreStart = std::atoi(start.c_str());
            total++;

            if (cores != nullptr)
            {
                cores->push_back(coreStart);
            }

            // done
            if (next == std::string::npos)
            {
                break;
            }
        }

        // found a core range => parse the start and end and add them.
        else if (str[next] == '-')
        {
            size_t posEnd = next + 1;

            start = str.substr(pos, next - pos);
            coreStart = std::atoi(start.c_str());

            next = str.find_first_of(",", posEnd);
            end = str.substr(posEnd, next - posEnd);
            coreEnd = std::atoi(end.c_str());

            // add the total core count.  Note that the ranges are always inclusive so we always
            // need to add 1 to the calculated range.
            total += coreEnd - coreStart + 1;

            // add all the cores to the list (inclusive on both ends of the range).
            if (cores != nullptr)
            {
                for (int32_t i = coreStart; i <= coreEnd; i++)
                {
                    cores->push_back(i);
                }
            }

            // nothing left in the source string => done.
            if (next == std::string::npos)
            {
                break;
            }
        }

        pos = next + 1;
    }

    return total;
}

inline int readCpuSetCoreCountFromFile(const char* file, char* buffer = nullptr, size_t len = 0) noexcept
{
    char localBuffer[1024];

    if (buffer == nullptr || len == 0)
    {
        buffer = localBuffer;
        len = CARB_COUNTOF(localBuffer);
    }

    // read the first line from the file.  This will contain the CPU core list.  This list is
    // comma separated where each value may either be a single core index or a core index range.
    // For example, "0-1,3" would indicate that cores 0, 1, and 3 are available and "0-3,7,9"
    // would indicate that cores 0, 1, 2, 3, 7, and 9 are available.
    if (!readLineFromFile(file, buffer, len))
    {
        return -1;
    }

    if (std::strspn(buffer, " \t\r\n") == std::strlen(buffer))
    {
        return -1;
    }

    // return the total number of cores included in the CPU set.
    return getCpuSetCoreCount(buffer);
}

inline bool controllerListContains(const char* controllers, const char* controller) noexcept
{
    if (!controllers || !controller || *controllers == '\0' || *controller == '\0')
        return false;

    const char* p = std::strstr(controllers, controller);
    if (!p)
        return false;

    if (!(p == controllers || *(p - 1) == ','))
        return false;

    const size_t len = std::strlen(controller);
    return p[len] == '\0' || p[len] == ',';
}

inline bool getCurrentCgroupPath(const char* controller,
                                 std::string& path,
                                 const char* procCgroupFile = "/proc/self/cgroup") noexcept
{
    FILE* fp = std::fopen(procCgroupFile, "r");
    if (!fp)
        return false;

    bool found = false;
    char line[4096];

    while (std::fgets(line, CARB_COUNTOF(line), fp) != nullptr)
    {
        char* firstColon = std::strchr(line, ':');
        if (!firstColon)
            continue;

        char* secondColon = std::strchr(firstColon + 1, ':');
        if (!secondColon)
            continue;

        *firstColon = '\0';
        *secondColon = '\0';

        const char* controllers = firstColon + 1;
        char* cgroupPath = secondColon + 1;
        if (char* newline = std::strpbrk(cgroupPath, "\r\n"); newline)
            *newline = '\0';

        const bool isUnifiedCgroup = secondColon == firstColon + 1;
        const bool needsUnifiedCgroup = !controller || controller[0] == '\0';

        if ((needsUnifiedCgroup && isUnifiedCgroup) ||
            (!needsUnifiedCgroup && !isUnifiedCgroup && controllerListContains(controllers, controller)))
        {
            path = cgroupPath;
            found = true;
            break;
        }
    }

    std::fclose(fp);
    return found;
}

inline std::string makeCgroupFilePath(const char* cgroupRoot, const std::string& cgroupPath, const char* file) noexcept
{
    std::string result = cgroupRoot;

    if (!cgroupPath.empty() && cgroupPath != "/")
    {
        if (cgroupPath.front() != '/')
            result.push_back('/');

        result += cgroupPath;
    }

    if (result.empty() || result.back() != '/')
        result.push_back('/');

    result += file;
    return result;
}

inline bool cpuQuotaIsValid(const CpuQuota& quota) noexcept
{
    return quota.quota > 0 && quota.period > 0;
}

inline bool cpuQuotaIsLower(const CpuQuota& lhs, const CpuQuota& rhs) noexcept
{
    return int64_t(lhs.quota) * int64_t(rhs.period) < int64_t(rhs.quota) * int64_t(lhs.period);
}

inline bool readCgroupV2CpuQuotaFromPath(const std::string& cgroupPath,
                                         CpuQuota& quota,
                                         const char* cgroupRoot = "/sys/fs/cgroup") noexcept
{
    bool found = false;
    std::string current = cgroupPath.empty() ? "/" : cgroupPath;

    for (;;)
    {
        const std::string cpuMaxPath = makeCgroupFilePath(cgroupRoot, current, "cpu.max");
        CpuQuota currentQuota = readCpuQuotaFromFile(cpuMaxPath.c_str());
        if (cpuQuotaIsValid(currentQuota) && (!found || cpuQuotaIsLower(currentQuota, quota)))
        {
            quota = currentQuota;
            found = true;
        }

        if (current == "/")
            break;

        const size_t slash = current.find_last_of('/');
        if (slash == std::string::npos || slash == 0)
            current = "/";
        else
            current.resize(slash);
    }

    return found;
}

inline bool readCurrentCgroupCpuQuota(CpuQuota& quota) noexcept
{
    std::string cgroupPath;

    if (getCurrentCgroupPath("", cgroupPath) && readCgroupV2CpuQuotaFromPath(cgroupPath, quota))
        return true;

    if (getCurrentCgroupPath("cpu", cgroupPath))
    {
        const std::string quotaPath = makeCgroupFilePath("/sys/fs/cgroup/cpu", cgroupPath, "cpu.cfs_quota_us");
        const std::string periodPath = makeCgroupFilePath("/sys/fs/cgroup/cpu", cgroupPath, "cpu.cfs_period_us");
        quota.quota = readIntFromFile(quotaPath.c_str());
        quota.period = readIntFromFile(periodPath.c_str());

        if (cpuQuotaIsValid(quota))
            return true;
    }

    return false;
}

inline int readCgroupV2CpuSetCoreCountFromPath(const std::string& cgroupPath,
                                               char* buffer,
                                               size_t len,
                                               const char* cgroupRoot = "/sys/fs/cgroup") noexcept
{
    std::string current = cgroupPath.empty() ? "/" : cgroupPath;

    for (;;)
    {
        std::string path = makeCgroupFilePath(cgroupRoot, current, "cpuset.cpus.effective");
        int coreCount = readCpuSetCoreCountFromFile(path.c_str(), buffer, len);
        if (coreCount >= 0)
            return coreCount;

        path = makeCgroupFilePath(cgroupRoot, current, "cpuset.cpus");
        coreCount = readCpuSetCoreCountFromFile(path.c_str(), buffer, len);
        if (coreCount >= 0)
            return coreCount;

        if (current == "/")
            break;

        const size_t slash = current.find_last_of('/');
        if (slash == std::string::npos || slash == 0)
            current = "/";
        else
            current.resize(slash);
    }

    return -1;
}

inline int readCpuSetCoreCountFromCurrentCgroup(char* buffer = nullptr, size_t len = 0) noexcept
{
    std::string cgroupPath;

    if (getCurrentCgroupPath("", cgroupPath))
    {
        int coreCount = readCgroupV2CpuSetCoreCountFromPath(cgroupPath, buffer, len);
        if (coreCount >= 0)
            return coreCount;
    }

    if (getCurrentCgroupPath("cpuset", cgroupPath))
    {
        std::string path = makeCgroupFilePath("/sys/fs/cgroup/cpuset", cgroupPath, "cpuset.effective_cpus");
        int coreCount = readCpuSetCoreCountFromFile(path.c_str(), buffer, len);
        if (coreCount >= 0)
            return coreCount;

        path = makeCgroupFilePath("/sys/fs/cgroup/cpuset", cgroupPath, "cpuset.cpus");
        coreCount = readCpuSetCoreCountFromFile(path.c_str(), buffer, len);
        if (coreCount >= 0)
            return coreCount;
    }

    return -1;
}

inline bool isRunningInContainer() noexcept
{
    FILE* fp;

    // first (and easiest) check is to check whether the `/.dockerenv` file exists.  This file
    // is not necessarily always present though.  Also check for other marker files that
    // various runtimes create when launching a container.
    if (access("/.dockerenv", F_OK) == 0 || access("/run/.containerenv", F_OK) == 0 ||
        access("/run/systemd/container", F_OK) == 0)
    {
        return true;
    }

    // a more reliable but more expensive check is to verify the control group of `init`.  If
    // running under docker, all of the entries will have a path that starts with `/docker` or
    // `/lxc` instead of just `/`.
    // Kubernetes seems to use `:/kubepods.slice`.
    fp = fopen("/proc/1/cgroup", "r");

    if (fp != nullptr)
    {
        char line[256];

        while (fgets(line, CARB_COUNTOF(line), fp) != nullptr)
        {
            if (feof(fp) || ferror(fp))
                break;

            if (strstr(line, ":/docker") != nullptr || strstr(line, ":/lxc") != nullptr ||
                strstr(line, ":/kubepods") != nullptr || strstr(line, ":/containerd") != nullptr ||
                strstr(line, ":/crio") != nullptr || strstr(line, ":/libpod") != nullptr)
            {
                return true;
            }
        }

        fclose(fp);
    }

    return false;
}

inline bool getCgroupCpuQuota(CpuQuota& quota)
{
    if (readCurrentCgroupCpuQuota(quota))
    {
        return true;
    }

    // attempt to read the CPU quota from cgroup v2 first.
    quota = readCpuQuotaFromFile("/sys/fs/cgroup/cpu.max");

    if (cpuQuotaIsValid(quota))
    {
        return true;
    }

    // attempt to read the CPU quota from cgroup v1 next.
    quota.quota = detail::readIntFromFile("/sys/fs/cgroup/cpu/cpu.cfs_quota_us");
    quota.period = detail::readIntFromFile("/sys/fs/cgroup/cpu/cpu.cfs_period_us");

    if (cpuQuotaIsValid(quota))
    {
        return true;
    }

    // CPU quota not available.
    return false;
}

inline int32_t getCgroupCpuQuota() noexcept
{
    CpuQuota quota;

    // attempt to read the CPU quota from cgroup v2 and cgroup v1.  This can affect how much
    // of the assigned cores scheduled time the processes in the container are effectively allowed
    // to make use of.
    if (getCgroupCpuQuota(quota))
    {
        return ::carb_max(1, (quota.quota + (quota.period / 2)) / quota.period);
    }

    return -1;
}

inline int getDockerCpuLimit() noexcept
{
    // See:
    // https://docs.docker.com/config/containers/resource_constraints/#cpu
    // https://engineering.squarespace.com/blog/2017/understanding-linux-container-scheduling
    // https://docs.kernel.org/admin-guide/cgroup-v1/cpusets.html
    // https://docs.kernel.org/admin-guide/cgroup-v2.html

    // ****** read the CPU core limit if defined ******
    // first attempt to read from the cgroup v2 and v1 CPU set lists to get the total core count.
    // We'll try the 'effective' CPU set list first.  If they doesn't exist or is empty, we'll fall
    // back to the main CPU set lists.  Note that the v1 and v2 cgroups will never simultaneously
    // exist in the container and that Docker will always respect the host system's cgroup version
    // regardless of the container's base image version.
    int32_t coreCount;

    // try to read the process cgroup CPU set.
    coreCount = detail::readCpuSetCoreCountFromCurrentCgroup();

    // fall back to the cgroup v2 effective CPU set.
    if (coreCount < 0)
        coreCount = detail::readCpuSetCoreCountFromFile("/sys/fs/cgroup/cpuset.cpus.effective");

    // fall back to the cgroup v2 main CPU set.
    if (coreCount < 0)
        coreCount = detail::readCpuSetCoreCountFromFile("/sys/fs/cgroup/cpuset.cpus");

    // fall back to the cgroup v1 effective CPU set.
    if (coreCount < 0)
        coreCount = detail::readCpuSetCoreCountFromFile("/sys/fs/cgroup/cpuset/cpuset.effective_cpus");

    // fall back to the cgroup v1 main CPU set.
    if (coreCount < 0)
        coreCount = detail::readCpuSetCoreCountFromFile("/sys/fs/cgroup/cpuset/cpuset.cpus");

    // ****** read the CPU usage quota if defined ******
    // next attempt to read the CPU quota from cgroup v2 and cgroup v1.  This can affect how much
    // of the assigned cores scheduled time the processes in the container are effectively allowed
    // to make use of.
    int32_t coreQuota = getCgroupCpuQuota();

    // ****** calculate the effective CPU core limit ******
    // a CPU set has been assigned and has limited the core count for the container => calculate
    //   the effective core count from the CPU set and the CPU quota.
    if (coreCount > 0)
    {
        // a CPU quota has been defined as well => caclulate the effective core count from both
        //   values.  Note that the CPU usage quota can oversubscribe the assigned core count.
        //   In that case we'd want to return the minimum of the CPU set core count and the
        //   effective count from the CPU quota.
        if (coreQuota > 0)
            return ::carb_min(coreQuota, coreCount);

        // no CPU usage quota defined => just return the CPU set core count.
        return coreCount;
    }

    // no CPU set core count defined => just return the effective core count from the CPU quota.
    else if (coreQuota > 0)
    {
        return coreQuota;
    }

    // no CPU usage limit.  This will effectively limit it to the host's bare metal CPU core
    // count.
    return -1;
}

} // namespace detail
#    endif

inline int getDockerCpuLimit() noexcept
{
    static int s_coreCount = detail::getDockerCpuLimit();
    return s_coreCount;
}

inline bool isRunningInContainer()
{
    static bool s_inContainer = detail::isRunningInContainer();
    return s_inContainer;
}

#else

inline int getDockerCpuLimit() noexcept
{
    return -1;
}

inline bool isRunningInContainer() noexcept
{
    return false;
}

#endif

} // namespace extras
} // namespace omni