diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..fcfd800f5 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/llama_cpp/_utils.py b/llama_cpp/_utils.py index 08ede6420..f6a2e2606 100644 --- a/llama_cpp/_utils.py +++ b/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) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 14e2f8500..774afb975 100644 --- a/llama_cpp/llama.py +++ b/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 diff --git a/llama_cpp/server/settings.py b/llama_cpp/server/settings.py index 78dd7cdeb..f6a6fa84f 100644 --- a/llama_cpp/server/settings.py +++ b/llama_cpp/server/settings.py @@ -1,7 +1,5 @@ from __future__ import annotations -import multiprocessing - from typing import Any, Optional, List, Literal, Union, Dict, cast from typing_extensions import Self @@ -9,6 +7,7 @@ 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 diff --git a/tests/test_llama.py b/tests/test_llama.py index c1c16e30c..2f9592834 100644 --- a/tests/test_llama.py +++ b/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) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 000000000..135f76b6d --- /dev/null +++ b/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