sgl-project--sglang/python/sglang/srt/utils/numa_utils.py

import ctypes
import glob
import logging
import math
import multiprocessing
import os
import random
import shutil
import subprocess
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Optional

import torch

from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import is_cuda

_is_cuda = is_cuda()

logger = logging.getLogger(__name__)


@contextmanager
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
                )
                debug_str += (
                    f", logical_gpu_id={gpu_id}, "
                    f"physical_gpu_id={_get_nvml_device_index(gpu_id)}, "
                    f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}"
                )
                with _mp_set_executable(executable=executable, debug_str=debug_str):
                    yield
                    return
    yield


def _create_numactl_executable(numactl_args: str):
    old_executable = os.fsdecode(multiprocessing.spawn.get_executable())
    script = f'''#!/bin/sh
exec numactl {numactl_args} {old_executable} "$@"'''
    path = Path(
        f"/tmp/sglang_temp_file_{time.time()}_{random.randrange(0, 10000000)}.sh"
    )
    path.write_text(script)
    path.chmod(0o777)
    return str(path), f"{script=}"


@contextmanager
def _mp_set_executable(executable: str, debug_str: str):
    start_method = multiprocessing.get_start_method()
    assert start_method == "spawn", f"{start_method=}"

    old_executable = os.fsdecode(multiprocessing.spawn.get_executable())
    multiprocessing.spawn.set_executable(executable)
    logger.debug(f"mp.set_executable {old_executable} -> {executable} ({debug_str})")
    try:
        yield
    finally:
        assert (
            os.fsdecode(multiprocessing.spawn.get_executable()) == executable
        ), f"{multiprocessing.spawn.get_executable()=}"
        multiprocessing.spawn.set_executable(old_executable)
        logger.debug(f"mp.set_executable revert to {old_executable}")


def _get_nvml_device_index(device_id: int) -> int:
    # _get_nvml_device_index is an internal PyTorch helper, so fall back to
    # device_id directly if the helper is unavailable.
    get_nvml_device_index = getattr(torch.cuda, "_get_nvml_device_index", None)
    if get_nvml_device_index is None:
        logger.warning(
            "torch.cuda._get_nvml_device_index is unavailable; falling back to "
            f"device_id={device_id} as the NVML device index. This may select "
            "the wrong physical GPU when CUDA_VISIBLE_DEVICES reorders devices "
            f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')})."
        )
        return device_id
    return get_nvml_device_index(device_id)


def get_numa_node_if_available(server_args: ServerArgs, gpu_id: int) -> Optional[int]:
    """
    Returns the NUMA node for the given GPU id. If it is not set in the server_args, it will try to query the NUMA node for the GPU.
    If the NUMA node is not available, has already been configured externally, or the user lacks permission to set NUMA affinity, it will return None.

    Args:
        server_args: The server arguments.
        gpu_id: The GPU id.

    Returns:
        The NUMA node for the given GPU id or None if it is not available.
    """
    if server_args.numa_node is not None:
        return server_args.numa_node[gpu_id]
    if _is_numa_available():
        queried_numa_node = _query_numa_node_for_gpu(gpu_id)
        if len(queried_numa_node) == 0:
            return None
        if len(queried_numa_node) > 1:
            # get_numa_node_for_gpu could return multiple nodes, we use the first one for now.
            # I don't think there any hardware configs that would have more than one.
            logger.warning(
                f"Multiple NUMA nodes found for GPU {gpu_id}: {queried_numa_node}. Using the first one."
            )
        return queried_numa_node[0]
    return None


def get_libnuma():
    libnuma = None

    for libnuma_so in ["libnuma.so", "libnuma.so.1"]:
        try:
            libnuma = ctypes.CDLL(libnuma_so)
        except OSError as e:
            logger.debug(f"{e}")
            libnuma = None
        if libnuma is not None:
            break
    return libnuma


def numa_bind_to_node(node: int):
    libnuma = get_libnuma()

    if libnuma is None or libnuma.numa_available() < 0:
        logger.warning("numa not available on this system, skip bind action")
        return

    node_cpus = _node_cpus(node)
    if node_cpus:
        allowed_cpus = os.sched_getaffinity(0)
        target_cpus = node_cpus & allowed_cpus
        if not target_cpus:
            _handle_numa_bind_failure(node, allowed_cpus)
            return
        os.sched_setaffinity(0, target_cpus)
    else:
        libnuma.numa_run_on_node(ctypes.c_int(node))
    libnuma.numa_set_preferred(ctypes.c_int(node))


class _Bitmask(ctypes.Structure):
    _fields_ = [("size", ctypes.c_ulong), ("maskp", ctypes.POINTER(ctypes.c_ulong))]


def _node_cpus(node: int) -> set:
    libnuma = get_libnuma()
    if libnuma is None or libnuma.numa_available() < 0:
        return set()
    libnuma.numa_allocate_cpumask.restype = ctypes.POINTER(_Bitmask)
    libnuma.numa_node_to_cpus.argtypes = [ctypes.c_int, ctypes.POINTER(_Bitmask)]
    libnuma.numa_node_to_cpus.restype = ctypes.c_int
    libnuma.numa_bitmask_isbitset.argtypes = [ctypes.POINTER(_Bitmask), ctypes.c_uint]
    libnuma.numa_bitmask_isbitset.restype = ctypes.c_int
    libnuma.numa_bitmask_free.argtypes = [ctypes.POINTER(_Bitmask)]
    mask = libnuma.numa_allocate_cpumask()
    try:
        if libnuma.numa_node_to_cpus(node, mask) != 0:
            return set()
        return {
            i
            for i in range(mask.contents.size)
            if libnuma.numa_bitmask_isbitset(mask, i)
        }
    finally:
        libnuma.numa_bitmask_free(mask)


def _numactl_cpu_mem_args(node: int, gpu_id: int) -> Optional[str]:
    node_cpus = _node_cpus(node)
    if not node_cpus:
        return f"--cpunodebind={node} --membind={node}"
    allowed_cpus = os.sched_getaffinity(0)
    target_cpus = node_cpus & allowed_cpus
    if not target_cpus:
        _handle_numa_bind_failure(node, allowed_cpus, gpu_id)
        return None
    if target_cpus == node_cpus:
        return f"--cpunodebind={node} --membind={node}"
    cpu_list = ",".join(str(c) for c in sorted(target_cpus))
    return f"--physcpubind={cpu_list} --membind={node}"


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")
    )


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(reason)


def _can_set_mempolicy() -> bool:
    """Check if the process has permission to use NUMA memory policy syscalls."""
    try:
        libnuma = get_libnuma()
        if libnuma is None or libnuma.numa_available() < 0:
            return False
        mode = ctypes.c_int()
        ret = libnuma.get_mempolicy(
            ctypes.byref(mode), None, ctypes.c_ulong(0), None, ctypes.c_ulong(0)
        )
        return ret == 0
    except Exception:
        return False


def _is_numa_available() -> bool:
    """
    Check if NUMA is available and not already configured externally.
    """
    if not _is_cuda:
        return False

    # Check if this is a numa system.
    if not os.path.isdir("/sys/devices/system/node/node1"):
        return False

    if not shutil.which("numactl") and envs.SGLANG_NUMA_BIND_V2.get():
        logger.debug(
            "numactl command not found, skipping NUMA node configuration for GPU. Install numactl (e.g., apt-get install numactl) to enable automatic NUMA binding."
        )
        return False

    if not _can_set_mempolicy():
        logger.warning(
            "User lacks permission to set NUMA affinity, skipping NUMA node configuration for GPU. If using docker, try adding --cap-add SYS_NICE to your docker run command."
        )
        return False

    return True


def _query_numa_node_for_gpu(device_id: int):
    """
    Get the NUMA node affinity list for a GPU device.

    Args:
        device_id: CUDA logical device index (post-CUDA_VISIBLE_DEVICES).
    Returns:
        List of NUMA node IDs that have affinity with the device.
    """
    try:
        import pynvml
    except ModuleNotFoundError:
        logger.warning("pynvml not installed, skipping NUMA node configuration for GPU")
        return []

    try:
        pynvml.nvmlInit()

        # device_id is a CUDA logical index. Convert it to the corresponding
        # NVML index so reordered CUDA_VISIBLE_DEVICES maps to the right GPU.
        # _get_nvml_device_index takes CUDA_VISIBLE_DEVICES into account.
        nvml_device_id = _get_nvml_device_index(device_id)
        handle = pynvml.nvmlDeviceGetHandleByIndex(nvml_device_id)
        numa_node_count = len(glob.glob("/sys/devices/system/node/node[0-9]*"))

        c_ulong_bits = ctypes.sizeof(ctypes.c_ulong) * 8
        node_set_size = max(1, math.ceil(numa_node_count / c_ulong_bits))
        node_set = pynvml.nvmlDeviceGetMemoryAffinity(
            handle,
            node_set_size,
            pynvml.NVML_AFFINITY_SCOPE_NODE,
        )

        # Decode the bitmask into a list of NUMA node IDs
        numa_nodes = []
        for node_id in range(numa_node_count):
            mask_array_index = node_id // c_ulong_bits
            mask_bit_index = node_id % c_ulong_bits
            if node_set[mask_array_index] & (1 << mask_bit_index):
                numa_nodes.append(node_id)
        return numa_nodes
    except pynvml.NVMLError as e:
        logger.warning(
            f"NVML error querying memory affinity for GPU {device_id}: {e}, skipping NUMA node configuration for GPU"
        )
        return []
    finally:
        try:
            pynvml.nvmlShutdown()
        except Exception:
            pass  # Ignore shutdown errors