fix: respect CPU affinity and cgroup CPU quota for default thread counts

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- fix: respect CPU affinity and cgroup CPU quota for default thread counts by @manunicholasjacob in #2385
+
 ## [0.3.36]
 
 - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098

llama_cpp/_utils.py

@@ -1,7 +1,9 @@
 import os
 import sys
+import math
+import multiprocessing
 
-from typing import Any, Dict
+from typing import Any, Dict, List, Optional
 
 # Avoid "LookupError: unknown encoding: ascii" when open() called in a destructor
 outnull_file = open(os.devnull, "w")
@@ -97,3 +99,85 @@ class Singleton(object, metaclass=MetaSingleton):
 
     def __init__(self):
         super(Singleton, self).__init__()
+
+
+def _read_cgroup_cpu_limit(cgroup_dir: str, v2: bool) -> Optional[float]:
+    """Return the CPU limit (in CPUs) set on one cgroup directory, or None."""
+    try:
+        if v2:
+            with open(os.path.join(cgroup_dir, "cpu.max")) as f:
+                quota, period = f.read().split()[:2]
+            if quota == "max":
+                return None
+        else:
+            with open(os.path.join(cgroup_dir, "cpu.cfs_quota_us")) as f:
+                quota = f.read().strip()
+            with open(os.path.join(cgroup_dir, "cpu.cfs_period_us")) as f:
+                period = f.read().strip()
+        if int(quota) <= 0 or int(period) <= 0:
+            return None
+        return int(quota) / int(period)
+    except (OSError, ValueError):
+        return None
+
+
+def _cgroup_cpu_limit(
+    proc_cgroup: str = "/proc/self/cgroup",
+    cgroup_root: str = "/sys/fs/cgroup",
+) -> Optional[float]:
+    """Return the tightest CPU quota (in CPUs) on this process's cgroup, or None.
+
+    Reads cgroup v2 `cpu.max` or cgroup v1 `cpu.cfs_quota_us`/`cpu.cfs_period_us`
+    for the process's cgroup and its ancestors (e.g. `docker run --cpus` or a
+    Kubernetes CPU limit). Returns None when no quota is set or on non-Linux.
+    """
+    try:
+        with open(proc_cgroup) as f:
+            lines = f.read().splitlines()
+    except OSError:
+        return None
+    limits: List[float] = []
+    for line in lines:
+        parts = line.split(":", 2)
+        if len(parts) != 3:
+            continue
+        controllers, path = parts[1], parts[2]
+        if controllers == "":
+            mounts, v2 = [cgroup_root], True
+        elif "cpu" in controllers.split(","):
+            names = ("cpu", "cpu,cpuacct", "cpuacct,cpu")
+            mounts, v2 = [os.path.join(cgroup_root, n) for n in names], False
+        else:
+            continue
+        for mount in mounts:
+            # Walk from the process cgroup up to the mount root. The mount root
+            # is also the process cgroup inside containers without a cgroup
+            # namespace, where `path` does not exist under the mount.
+            rel = path.strip("/")
+            while True:
+                limit = _read_cgroup_cpu_limit(os.path.join(mount, rel), v2)
+                if limit is not None:
+                    limits.append(limit)
+                if not rel:
+                    break
+                rel = os.path.dirname(rel)
+    return min(limits) if limits else None
+
+
+def default_cpu_count() -> int:
+    """Number of CPUs this process can actually use.
+
+    `multiprocessing.cpu_count()` reports every CPU on the host, even when the
+    process is pinned to fewer (`taskset`, `docker run --cpuset-cpus`) or
+    limited by a cgroup CPU quota (`docker run --cpus`, Kubernetes limits).
+    This uses the CPU affinity mask where available and caps it at the
+    cgroup quota rounded up.
+    """
+    try:
+        count = len(os.sched_getaffinity(0))  # type: ignore[attr-defined]
+    except (AttributeError, OSError):
+        count = multiprocessing.cpu_count()
+    limit = _cgroup_cpu_limit()
+    if limit is not None:
+        count = min(count, math.ceil(limit))
+    return max(count, 1)

llama_cpp/llama.py

@@ -11,7 +11,6 @@
 import fnmatch
 import warnings
 import contextlib
-import multiprocessing
 
 from typing import (
     Any,
@@ -49,7 +48,7 @@
 
 import llama_cpp._internals as internals
 from ._logger import set_verbose
-from ._utils import suppress_stdout_stderr
+from ._utils import suppress_stdout_stderr, default_cpu_count
 
 
 class Llama:
@@ -309,8 +308,8 @@ def __init__(
             self.model_params.kv_overrides = self._kv_overrides_array
 
         self.n_batch = min(n_ctx, n_batch)  # ???
-        self.n_threads = n_threads or max(multiprocessing.cpu_count() // 2, 1)
-        self.n_threads_batch = n_threads_batch or multiprocessing.cpu_count()
+        self.n_threads = n_threads or max(default_cpu_count() // 2, 1)
+        self.n_threads_batch = n_threads_batch or default_cpu_count()
 
         # Used by the sampler
         self._seed = seed or llama_cpp.LLAMA_DEFAULT_SEED

llama_cpp/server/settings.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
-import multiprocessing
-
 from typing import Any, Optional, List, Literal, Union, Dict, cast
 from typing_extensions import Self
 
 from pydantic import Field, model_validator
 from pydantic_settings import BaseSettings
 
 import llama_cpp
+from llama_cpp._utils import default_cpu_count
 
 # Disable warning for model and model_alias settings
 BaseSettings.model_config["protected_namespaces"] = ()
@@ -74,12 +73,12 @@ class ModelSettings(BaseSettings):
         default=512, ge=1, description="The physical batch size used by llama.cpp"
     )
     n_threads: int = Field(
-        default=max(multiprocessing.cpu_count() // 2, 1),
+        default=max(default_cpu_count() // 2, 1),
         ge=1,
         description="The number of threads to use. Use -1 for max cpu threads",
     )
     n_threads_batch: int = Field(
-        default=max(multiprocessing.cpu_count(), 1),
+        default=default_cpu_count(),
         ge=0,
         description="The number of threads to use when batch processing. Use -1 for max cpu threads",
     )
@@ -193,8 +192,8 @@ class ModelSettings(BaseSettings):
         mode="before"
     )  # pre=True to ensure this runs before any other validation
     def set_dynamic_defaults(self) -> Self:
-        # If n_threads or n_threads_batch is -1, set it to multiprocessing.cpu_count()
-        cpu_count = multiprocessing.cpu_count()
+        # If n_threads or n_threads_batch is -1, set it to the usable CPU count
+        cpu_count = default_cpu_count()
         values = cast(Dict[str, int], self)
         if values.get("n_threads", 0) == -1:
             values["n_threads"] = cpu_count

tests/test_llama.py

@@ -20,6 +20,14 @@ def test_llama_cpp_version():
     assert llama_cpp.__version__
 
 
+def test_llama_default_threads_use_usable_cpu_count(monkeypatch):
+    monkeypatch.setattr(llama_cpp.llama, "default_cpu_count", lambda: 6)
+    llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True, verbose=False)
+
+    assert llama.n_threads == 3
+    assert llama.n_threads_batch == 6
+
+
 def test_llama_cpp_tokenization():
     llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True, verbose=False)
 

tests/test_utils.py

@@ -0,0 +1,89 @@
+import os
+
+import pytest
+
+import llama_cpp._utils as utils
+
+
+def _write(path, text):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(text)
+
+
+def _proc_cgroup(tmp_path, text):
+    proc = tmp_path / "proc_self_cgroup"
+    proc.write_text(text)
+    return str(proc)
+
+
+def test_cgroup_v2_quota(tmp_path):
+    root = tmp_path / "cgroup"
+    _write(root / "cpu.max", "200000 100000\n")
+    proc = _proc_cgroup(tmp_path, "0::/\n")
+    assert utils._cgroup_cpu_limit(proc, str(root)) == 2.0
+
+
+def test_cgroup_v2_no_quota(tmp_path):
+    root = tmp_path / "cgroup"
+    _write(root / "cpu.max", "max 100000\n")
+    proc = _proc_cgroup(tmp_path, "0::/\n")
+    assert utils._cgroup_cpu_limit(proc, str(root)) is None
+
+
+def test_cgroup_v2_uses_tightest_ancestor(tmp_path):
+    root = tmp_path / "cgroup"
+    _write(root / "kubepods.slice" / "cpu.max", "150000 100000\n")
+    _write(root / "kubepods.slice" / "pod" / "cpu.max", "max 100000\n")
+    _write(root / "other.slice" / "cpu.max", "50000 100000\n")
+    proc = _proc_cgroup(tmp_path, "0::/kubepods.slice/pod\n")
+    assert utils._cgroup_cpu_limit(proc, str(root)) == 1.5
+
+
+def test_cgroup_v1_quota(tmp_path):
+    root = tmp_path / "cgroup"
+    # Without a cgroup namespace the container's own cgroup is the mount root
+    # and the host path from /proc/self/cgroup does not exist under it.
+    _write(root / "cpu,cpuacct" / "cpu.cfs_quota_us", "50000\n")
+    _write(root / "cpu,cpuacct" / "cpu.cfs_period_us", "100000\n")
+    proc = _proc_cgroup(tmp_path, "12:memory:/docker/abc\n4:cpu,cpuacct:/docker/abc\n")
+    assert utils._cgroup_cpu_limit(proc, str(root)) == 0.5
+
+
+def test_cgroup_v1_no_quota(tmp_path):
+    root = tmp_path / "cgroup"
+    _write(root / "cpu" / "cpu.cfs_quota_us", "-1\n")
+    _write(root / "cpu" / "cpu.cfs_period_us", "100000\n")
+    proc = _proc_cgroup(tmp_path, "4:cpu,cpuacct:/\n")
+    assert utils._cgroup_cpu_limit(proc, str(root)) is None
+
+
+def test_cgroup_missing_files(tmp_path):
+    root = tmp_path / "cgroup"
+    root.mkdir()
+    assert utils._cgroup_cpu_limit(str(tmp_path / "missing"), str(root)) is None
+    proc = _proc_cgroup(tmp_path, "0::/\n")
+    assert utils._cgroup_cpu_limit(proc, str(root)) is None
+
+
+@pytest.mark.parametrize(
+    "limit, expected",
+    [(None, 20), (2.0, 2), (1.5, 2), (0.5, 1), (64.0, 20)],
+)
+def test_default_cpu_count_caps_affinity_at_quota(monkeypatch, limit, expected):
+    monkeypatch.setattr(os, "sched_getaffinity", lambda pid: set(range(20)), False)
+    monkeypatch.setattr(utils, "_cgroup_cpu_limit", lambda: limit)
+    assert utils.default_cpu_count() == expected
+
+
+def test_default_cpu_count_uses_affinity(monkeypatch):
+    monkeypatch.setattr(os, "sched_getaffinity", lambda pid: {0, 1}, False)
+    monkeypatch.setattr(utils.multiprocessing, "cpu_count", lambda: 20)
+    monkeypatch.setattr(utils, "_cgroup_cpu_limit", lambda: None)
+    assert utils.default_cpu_count() == 2
+
+
+def test_default_cpu_count_without_affinity(monkeypatch):
+    monkeypatch.delattr(os, "sched_getaffinity", raising=False)
+    monkeypatch.setattr(utils.multiprocessing, "cpu_count", lambda: 8)
+    monkeypatch.setattr(utils, "_cgroup_cpu_limit", lambda: None)
+    assert utils.default_cpu_count() == 8