NUMA: probe numactl binding and fall back when --membind is rejected (#28401)
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
This commit is contained in:
chengcuiping
co-authored by
Mohammad Miadh Angkad
Mohammad Miadh Angkad
parent
7c9bb316cf
commit
e2b55bdbab
@@ -6,6 +6,7 @@ import multiprocessing
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
@@ -27,8 +28,34 @@ def configure_subprocess(server_args: ServerArgs, gpu_id: int):
|
||||
if envs.SGLANG_NUMA_BIND_V2.get():
|
||||
numa_node = get_numa_node_if_available(server_args, gpu_id)
|
||||
if numa_node is not None:
|
||||
# _numactl_cpu_mem_args returns None (warn/raise) on empty CPU intersection (#26983).
|
||||
numactl_args = _numactl_cpu_mem_args(numa_node, gpu_id)
|
||||
if numactl_args is not None:
|
||||
# Verify numactl can actually apply the binding before we exec it
|
||||
# in front of the interpreter; relax the memory policy if not.
|
||||
numactl_args, probe_err = _probe_numactl_args(numactl_args)
|
||||
if numactl_args is None:
|
||||
# numactl could not apply even a CPU-only binding (e.g.
|
||||
# set_mempolicy(2)/sched_setaffinity(2) blocked by seccomp,
|
||||
# which the read-only get_mempolicy(2) probe in
|
||||
# _can_set_mempolicy cannot detect). Reuse #26983's failure
|
||||
# semantics (warn-and-continue, or raise when
|
||||
# SGLANG_CRASH_ON_NUMA_BIND_FAILURE) with an explicit reason
|
||||
# carrying the captured stderr: the CPU intersection already
|
||||
# succeeded here, so the default "no CPU cores allowed"
|
||||
# message would mislead operators toward the wrong cause.
|
||||
probe_suffix = f": {probe_err}" if probe_err else ""
|
||||
_handle_numa_bind_failure(
|
||||
numa_node,
|
||||
reason=(
|
||||
f"numactl could not apply NUMA binding for node "
|
||||
f"{numa_node} (e.g. set_mempolicy/sched_setaffinity "
|
||||
f"blocked by seccomp, or cpuset rejects the policy)"
|
||||
f"{probe_suffix}; skipping NUMA binding for GPU {gpu_id}."
|
||||
),
|
||||
)
|
||||
yield
|
||||
return
|
||||
executable, debug_str = _create_numactl_executable(
|
||||
numactl_args=numactl_args
|
||||
)
|
||||
@@ -192,17 +219,120 @@ def _numactl_cpu_mem_args(node: int, gpu_id: int) -> Optional[str]:
|
||||
return f"--physcpubind={cpu_list} --membind={node}"
|
||||
|
||||
|
||||
def _handle_numa_bind_failure(
|
||||
node: int, allowed_cpus, gpu_id: Optional[int] = None
|
||||
) -> None:
|
||||
gpu_str = f" for GPU {gpu_id}" if gpu_id is not None else ""
|
||||
msg = (
|
||||
f"NUMA node {node} has no CPU cores allowed by the current affinity "
|
||||
f"{sorted(allowed_cpus)}, skipping NUMA binding{gpu_str}."
|
||||
def _strip_memory_args(numactl_args: str) -> str:
|
||||
"""Return ``numactl_args`` with the ``--membind`` segment removed, keeping
|
||||
only the CPU binding (``--cpunodebind`` / ``--physcpubind``)."""
|
||||
return " ".join(
|
||||
token for token in numactl_args.split() if not token.startswith("--membind")
|
||||
)
|
||||
logger.warning(msg)
|
||||
|
||||
|
||||
def _probe_numactl_args(numactl_args: str) -> tuple[Optional[str], str]:
|
||||
"""Dry-run ``numactl <args> true`` and fall back to a weaker binding when the
|
||||
kernel rejects the strongest one.
|
||||
|
||||
``configure_subprocess`` applies NUMA binding by exec-ing ``numactl`` in front
|
||||
of the Python interpreter (see ``_create_numactl_executable``), so a binding
|
||||
that ``numactl`` refuses kills the worker before Python starts, with no
|
||||
traceback. ``_can_set_mempolicy`` only probes ``get_mempolicy(2)`` (read),
|
||||
which does not catch ``set_mempolicy(2)`` being denied (e.g. by a seccomp
|
||||
profile) or a ``--membind`` that the cpuset rejects with ``EINVAL``.
|
||||
|
||||
To avoid that silent crash we probe the requested args and progressively relax
|
||||
the *memory* policy while keeping the CPU binding intact::
|
||||
|
||||
--membind=N -> --preferred=N -> drop the memory segment
|
||||
|
||||
Returns ``(args, last_stderr)``: ``args`` is the strongest binding that
|
||||
actually runs, or ``None`` if even CPU-only fails (or ``numactl`` is missing /
|
||||
errors out); ``last_stderr`` is the rejection reason numactl printed for the
|
||||
strongest binding that was rejected (empty on success), so the caller can
|
||||
surface it on the total-failure path.
|
||||
"""
|
||||
|
||||
def _probe(args: str):
|
||||
"""Run ``numactl <args> true``; return ``(succeeded, stderr_text)``."""
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
["numactl", *args.split(), "true"],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.PIPE,
|
||||
timeout=10,
|
||||
)
|
||||
stderr = proc.stderr.decode("utf-8", errors="replace").strip()
|
||||
if proc.returncode != 0:
|
||||
logger.debug(f"numactl probe for {args!r} rejected: {stderr!r}")
|
||||
return proc.returncode == 0, stderr
|
||||
except Exception as e:
|
||||
# Missing numactl, timeout, etc. Treat as "this binding does not work".
|
||||
logger.debug(f"numactl probe for {args!r} failed: {e}")
|
||||
return False, str(e)
|
||||
|
||||
def _suffix(err: str) -> str:
|
||||
return f": {err}" if err else ""
|
||||
|
||||
# 1. Strongest binding: exactly what was requested.
|
||||
ok, last_err = _probe(numactl_args)
|
||||
if ok:
|
||||
return numactl_args, ""
|
||||
|
||||
# 2. Relax a hard --membind=N to a soft --preferred=N. The memory segment here
|
||||
# is always a single node, which maps cleanly onto --preferred (single-node
|
||||
# only). MPOL_PREFERRED is a hint and can succeed where MPOL_BIND is denied.
|
||||
if "--membind=" in numactl_args:
|
||||
preferred_args = numactl_args.replace("--membind=", "--preferred=")
|
||||
ok, _ = _probe(preferred_args)
|
||||
if ok:
|
||||
logger.warning(
|
||||
f"numactl rejected hard memory binding ({numactl_args!r})"
|
||||
f"{_suffix(last_err)}; falling back to soft preferred policy "
|
||||
f"({preferred_args!r})."
|
||||
)
|
||||
return preferred_args, ""
|
||||
|
||||
# 3. Drop the memory segment entirely, keep only the CPU binding.
|
||||
cpu_only_args = _strip_memory_args(numactl_args)
|
||||
if cpu_only_args and cpu_only_args != numactl_args:
|
||||
ok, cpu_err = _probe(cpu_only_args)
|
||||
if ok:
|
||||
logger.warning(
|
||||
f"numactl rejected memory binding ({numactl_args!r})"
|
||||
f"{_suffix(last_err)}; falling back to CPU-only binding "
|
||||
f"({cpu_only_args!r})."
|
||||
)
|
||||
return cpu_only_args, ""
|
||||
last_err = cpu_err
|
||||
|
||||
# 4. Nothing worked.
|
||||
return None, last_err
|
||||
|
||||
|
||||
def _handle_numa_bind_failure(
|
||||
node: int,
|
||||
allowed_cpus=None,
|
||||
gpu_id: Optional[int] = None,
|
||||
*,
|
||||
reason: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Emit the NUMA-bind failure warning, or raise it when
|
||||
``SGLANG_CRASH_ON_NUMA_BIND_FAILURE`` is set.
|
||||
|
||||
Two call modes:
|
||||
* ``reason is None`` (default): the failure is an empty CPU intersection,
|
||||
so the message reports ``allowed_cpus`` (which must be provided).
|
||||
* ``reason`` provided: the failure is something else (e.g. numactl rejected
|
||||
the binding at runtime); the caller supplies the exact message and
|
||||
``allowed_cpus`` / ``gpu_id`` are not needed.
|
||||
"""
|
||||
if reason is None:
|
||||
gpu_str = f" for GPU {gpu_id}" if gpu_id is not None else ""
|
||||
reason = (
|
||||
f"NUMA node {node} has no CPU cores allowed by the current affinity "
|
||||
f"{sorted(allowed_cpus)}, skipping NUMA binding{gpu_str}."
|
||||
)
|
||||
logger.warning(reason)
|
||||
if envs.SGLANG_CRASH_ON_NUMA_BIND_FAILURE.get():
|
||||
raise RuntimeError(msg)
|
||||
raise RuntimeError(reason)
|
||||
|
||||
|
||||
def _can_set_mempolicy() -> bool:
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import ctypes
|
||||
import os
|
||||
import unittest
|
||||
from contextlib import ExitStack
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.srt.utils.numa_utils import (
|
||||
@@ -8,7 +9,10 @@ from sglang.srt.utils.numa_utils import (
|
||||
_is_numa_available,
|
||||
_node_cpus,
|
||||
_numactl_cpu_mem_args,
|
||||
_probe_numactl_args,
|
||||
_query_numa_node_for_gpu,
|
||||
_strip_memory_args,
|
||||
configure_subprocess,
|
||||
get_numa_node_if_available,
|
||||
numa_bind_to_node,
|
||||
)
|
||||
@@ -366,5 +370,201 @@ class TestNumaBindIntersection(unittest.TestCase):
|
||||
_handle_numa_bind_failure(0, {72, 73})
|
||||
|
||||
|
||||
def _run_result(returncode, stderr=b""):
|
||||
"""Build a fake subprocess.CompletedProcess-like object with a returncode
|
||||
and captured stderr (bytes, as subprocess.run(..., stderr=PIPE) returns)."""
|
||||
result = MagicMock()
|
||||
result.returncode = returncode
|
||||
result.stderr = stderr
|
||||
return result
|
||||
|
||||
|
||||
class TestProbeNumactlArgs(unittest.TestCase):
|
||||
"""Tests for _probe_numactl_args: dry-run numactl and relax the memory policy
|
||||
(--membind -> --preferred -> CPU-only) when the kernel rejects the binding.
|
||||
|
||||
subprocess.run is mocked and orchestrated by returncode; no real numactl or
|
||||
GPU is required. Returns ``(args, last_stderr)``."""
|
||||
|
||||
@patch("sglang.srt.utils.numa_utils.subprocess.run")
|
||||
def test_membind_probe_succeeds_returns_original(self, mock_run):
|
||||
# The requested binding works on the first probe.
|
||||
mock_run.side_effect = [_run_result(0)]
|
||||
args = "--cpunodebind=0 --membind=0"
|
||||
self.assertEqual(_probe_numactl_args(args), (args, ""))
|
||||
self.assertEqual(mock_run.call_count, 1)
|
||||
|
||||
@patch("sglang.srt.utils.numa_utils.subprocess.run")
|
||||
def test_membind_fails_preferred_succeeds(self, mock_run):
|
||||
# --membind rejected, --preferred accepted.
|
||||
mock_run.side_effect = [_run_result(1), _run_result(0)]
|
||||
with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm:
|
||||
result = _probe_numactl_args("--cpunodebind=0 --membind=0")
|
||||
self.assertEqual(result, ("--cpunodebind=0 --preferred=0", ""))
|
||||
self.assertTrue(any("preferred" in msg for msg in cm.output))
|
||||
# Second probe must have used the --preferred form.
|
||||
second_call_argv = mock_run.call_args_list[1].args[0]
|
||||
self.assertIn("--preferred=0", second_call_argv)
|
||||
self.assertNotIn("--membind=0", second_call_argv)
|
||||
|
||||
@patch("sglang.srt.utils.numa_utils.subprocess.run")
|
||||
def test_membind_and_preferred_fail_cpu_only_succeeds(self, mock_run):
|
||||
# --membind and --preferred rejected, CPU-only accepted.
|
||||
mock_run.side_effect = [_run_result(1), _run_result(1), _run_result(0)]
|
||||
with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm:
|
||||
result = _probe_numactl_args("--physcpubind=0,21,22 --membind=0")
|
||||
self.assertEqual(result, ("--physcpubind=0,21,22", ""))
|
||||
self.assertTrue(any("CPU-only" in msg for msg in cm.output))
|
||||
third_call_argv = mock_run.call_args_list[2].args[0]
|
||||
self.assertNotIn("--membind=0", third_call_argv)
|
||||
self.assertNotIn("--preferred=0", third_call_argv)
|
||||
|
||||
@patch("sglang.srt.utils.numa_utils.subprocess.run")
|
||||
def test_all_probes_fail_returns_none_with_last_stderr(self, mock_run):
|
||||
# Every binding, down to CPU-only, is rejected; the returned stderr is the
|
||||
# CPU-only (last / strongest-attempted) rejection reason.
|
||||
mock_run.side_effect = [
|
||||
_run_result(1, stderr=b"numactl: setting membind: Invalid argument"),
|
||||
_run_result(1, stderr=b"numactl: setting preferred: Invalid argument"),
|
||||
_run_result(1, stderr=b"numactl: cpunodebind: Operation not permitted"),
|
||||
]
|
||||
args, err = _probe_numactl_args("--cpunodebind=0 --membind=0")
|
||||
self.assertIsNone(args)
|
||||
self.assertIn("cpunodebind", err)
|
||||
self.assertEqual(mock_run.call_count, 3)
|
||||
|
||||
@patch("sglang.srt.utils.numa_utils.subprocess.run")
|
||||
def test_cpu_only_input_failure_returns_none(self, mock_run):
|
||||
# No --membind in the input: the requested args are already CPU-only, so
|
||||
# step 1 is the only probe and on failure we skip --preferred / strip.
|
||||
mock_run.side_effect = [
|
||||
_run_result(1, stderr=b"numactl: cpunodebind: Operation not permitted")
|
||||
]
|
||||
args, err = _probe_numactl_args("--cpunodebind=0")
|
||||
self.assertIsNone(args)
|
||||
self.assertIn("cpunodebind", err)
|
||||
self.assertEqual(mock_run.call_count, 1)
|
||||
|
||||
@patch("sglang.srt.utils.numa_utils.subprocess.run")
|
||||
def test_numactl_missing_returns_none(self, mock_run):
|
||||
# numactl not installed / raises: probe must not propagate, returns None.
|
||||
mock_run.side_effect = FileNotFoundError("numactl")
|
||||
args, _err = _probe_numactl_args("--cpunodebind=0 --membind=0")
|
||||
self.assertIsNone(args)
|
||||
|
||||
@patch("sglang.srt.utils.numa_utils.subprocess.run")
|
||||
def test_rejection_stderr_surfaces_in_fallback_warning(self, mock_run):
|
||||
# numactl prints the precise rejection reason to stderr (e.g.
|
||||
# "setting membind: Invalid argument"); the fallback warning must
|
||||
# surface it so operators can tell seccomp vs cpuset apart.
|
||||
mock_run.side_effect = [
|
||||
_run_result(1, stderr=b"numactl: setting membind: Invalid argument"),
|
||||
_run_result(0),
|
||||
]
|
||||
with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm:
|
||||
result = _probe_numactl_args("--cpunodebind=0 --membind=0")
|
||||
self.assertEqual(result, ("--cpunodebind=0 --preferred=0", ""))
|
||||
self.assertTrue(
|
||||
any("Invalid argument" in msg for msg in cm.output),
|
||||
f"expected numactl stderr in warning, got {cm.output}",
|
||||
)
|
||||
|
||||
|
||||
class TestStripMemoryArgs(unittest.TestCase):
|
||||
"""Direct tests for _strip_memory_args: drop --membind, keep CPU binding."""
|
||||
|
||||
def test_strips_membind_keeps_cpu(self):
|
||||
self.assertEqual(
|
||||
_strip_memory_args("--cpunodebind=0 --membind=0"),
|
||||
"--cpunodebind=0",
|
||||
)
|
||||
self.assertEqual(
|
||||
_strip_memory_args("--physcpubind=0,21,22 --membind=0"),
|
||||
"--physcpubind=0,21,22",
|
||||
)
|
||||
|
||||
def test_no_membind_returns_unchanged(self):
|
||||
self.assertEqual(_strip_memory_args("--cpunodebind=0"), "--cpunodebind=0")
|
||||
|
||||
|
||||
class TestConfigureSubprocessProbeFailure(unittest.TestCase):
|
||||
"""Tests the wiring in configure_subprocess when _probe_numactl_args gives up
|
||||
(returns None): the worker must start unbound (warn-and-yield) by default, or
|
||||
raise before yielding when SGLANG_CRASH_ON_NUMA_BIND_FAILURE=1.
|
||||
|
||||
get_numa_node_if_available / _numactl_cpu_mem_args / _probe_numactl_args are
|
||||
mocked to drive the probe-failure branch directly; _create_numactl_executable
|
||||
and _mp_set_executable are mocked to assert the failure path never installs a
|
||||
numactl executable. No real numactl or GPU is required."""
|
||||
|
||||
def _common_patches(self):
|
||||
return [
|
||||
patch(
|
||||
"sglang.srt.utils.numa_utils.get_numa_node_if_available",
|
||||
return_value=0,
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.utils.numa_utils._numactl_cpu_mem_args",
|
||||
return_value="--cpunodebind=0 --membind=0",
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.utils.numa_utils._probe_numactl_args",
|
||||
return_value=(
|
||||
None,
|
||||
"numactl: setting membind: Invalid argument",
|
||||
),
|
||||
),
|
||||
patch("sglang.srt.utils.numa_utils._create_numactl_executable"),
|
||||
patch("sglang.srt.utils.numa_utils._mp_set_executable"),
|
||||
]
|
||||
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"SGLANG_NUMA_BIND_V2": "1", "SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "0"},
|
||||
)
|
||||
def test_probe_none_warns_and_yields_unbound(self):
|
||||
with ExitStack() as stack:
|
||||
mocks = [stack.enter_context(p) for p in self._common_patches()]
|
||||
_mock_get, _mock_args, _mock_probe, mock_create, mock_mp = mocks
|
||||
server_args = MagicMock()
|
||||
with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm:
|
||||
with configure_subprocess(server_args, 0):
|
||||
pass # worker would start unbound here
|
||||
# The probe-failure path reuses #26983's failure helper (warn) and
|
||||
# must NOT install a numactl executable. The captured numactl stderr
|
||||
# is threaded into the warning so operators can see the rejection cause.
|
||||
self.assertTrue(
|
||||
any("could not apply NUMA binding" in msg for msg in cm.output),
|
||||
f"expected probe-failure warning, got {cm.output}",
|
||||
)
|
||||
self.assertTrue(
|
||||
any("Invalid argument" in msg for msg in cm.output),
|
||||
f"expected numactl stderr in warning, got {cm.output}",
|
||||
)
|
||||
mock_create.assert_not_called()
|
||||
mock_mp.assert_not_called()
|
||||
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"SGLANG_NUMA_BIND_V2": "1", "SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "1"},
|
||||
)
|
||||
def test_probe_none_raises_before_yield_when_crash_enabled(self):
|
||||
with ExitStack() as stack:
|
||||
mocks = [stack.enter_context(p) for p in self._common_patches()]
|
||||
_mock_get, _mock_args, _mock_probe, mock_create, mock_mp = mocks
|
||||
server_args = MagicMock()
|
||||
with self.assertRaises(RuntimeError) as cm:
|
||||
with configure_subprocess(server_args, 0):
|
||||
self.fail(
|
||||
"contextmanager must not yield when crash-on-failure is set"
|
||||
)
|
||||
# The RuntimeError carries the captured stderr so crash logs show the
|
||||
# rejection cause, not just the failure category.
|
||||
self.assertIn("could not apply NUMA binding", str(cm.exception))
|
||||
self.assertIn("Invalid argument", str(cm.exception))
|
||||
mock_create.assert_not_called()
|
||||
mock_mp.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user