From 340c613ab56d550f2aafd10f2f2ee64c1dd78001 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 30 Nov 2025 17:00:33 +0800 Subject: [PATCH] Support numactl bind for CPU and memory before process starts (#14156) --- python/sglang/srt/entrypoints/engine.py | 5 +- python/sglang/srt/environ.py | 3 + python/sglang/srt/grpc/scheduler_launcher.py | 6 +- .../srt/managers/data_parallel_controller.py | 5 +- python/sglang/srt/managers/scheduler.py | 4 +- python/sglang/srt/utils/numa_utils.py | 56 +++++++++++++++++++ 6 files changed, 74 insertions(+), 5 deletions(-) create mode 100644 python/sglang/srt/utils/numa_utils.py diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 7ca884e41..4af385fa3 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -76,6 +76,7 @@ from sglang.srt.utils import ( kill_process_tree, launch_dummy_health_check_server, maybe_reindex_device_id, + numa_utils, prepare_model_and_tokenizer, set_prometheus_multiproc_dir, set_ulimit, @@ -846,7 +847,9 @@ def _launch_subprocesses( writer, ), ) - with memory_saver_adapter.configure_subprocess(): + with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess( + server_args, gpu_id + ): proc.start() scheduler_procs.append(proc) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e2facecae..2f9fcc9d1 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -338,6 +338,9 @@ class Envs: SGLANG_EXTERNAL_MM_MODEL_ARCH = EnvStr("") SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE = EnvStr("") + # Numa + SGLANG_NUMA_BIND_V2 = EnvBool(True) + # fmt: on diff --git a/python/sglang/srt/grpc/scheduler_launcher.py b/python/sglang/srt/grpc/scheduler_launcher.py index 77a62d8a6..6c8c5df66 100644 --- a/python/sglang/srt/grpc/scheduler_launcher.py +++ b/python/sglang/srt/grpc/scheduler_launcher.py @@ -15,7 +15,7 @@ from sglang.srt.managers.data_parallel_controller import ( ) from sglang.srt.managers.scheduler import run_scheduler_process from sglang.srt.server_args import PortArgs, ServerArgs -from sglang.srt.utils import configure_logger, prepare_model_and_tokenizer +from sglang.srt.utils import configure_logger, numa_utils, prepare_model_and_tokenizer from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter logger = logging.getLogger(__name__) @@ -135,7 +135,9 @@ def launch_scheduler_process_only( ), ) - with memory_saver_adapter.configure_subprocess(): + with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess( + server_args, gpu_id + ): proc.start() scheduler_procs.append(proc) diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index cb897a643..4ed72a8c0 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -49,6 +49,7 @@ from sglang.srt.tracing.trace import ( trace_slice_end, trace_slice_start, ) +from sglang.srt.utils import numa_utils from sglang.srt.utils.common import ( bind_port, configure_ipv6, @@ -440,7 +441,9 @@ class DataParallelController: writer, ), ) - with memory_saver_adapter.configure_subprocess(): + with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess( + server_args, gpu_id + ): proc.start() self.scheduler_procs.append(proc) scheduler_pipe_readers.append(reader) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 30f754581..38704541f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2631,7 +2631,9 @@ def run_scheduler_process( set_gpu_proc_affinity( server_args.pp_size, server_args.tp_size, server_args.nnodes, gpu_id ) - if (numa_node := server_args.numa_node) is not None: + if ( + numa_node := server_args.numa_node + ) is not None and not envs.SGLANG_NUMA_BIND_V2.get(): numa_bind_to_node(numa_node[gpu_id]) # Set up tracing diff --git a/python/sglang/srt/utils/numa_utils.py b/python/sglang/srt/utils/numa_utils.py new file mode 100644 index 000000000..a144d1002 --- /dev/null +++ b/python/sglang/srt/utils/numa_utils.py @@ -0,0 +1,56 @@ +import logging +import multiprocessing +import os +import random +import time +from contextlib import contextmanager +from pathlib import Path + +from sglang import ServerArgs +from sglang.srt.environ import envs + +logger = logging.getLogger(__name__) + + +@contextmanager +def configure_subprocess(server_args: ServerArgs, gpu_id: int): + if ( + numa_nodes := server_args.numa_node + ) is not None and envs.SGLANG_NUMA_BIND_V2.get(): + numa_node = numa_nodes[gpu_id] + numactl_args = f"--cpunodebind={numa_node} --membind={numa_node}" + executable, debug_str = _create_numactl_executable(numactl_args=numactl_args) + with _mp_set_executable(executable=executable, debug_str=debug_str): + yield + else: + 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.info(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.info(f"mp.set_executable revert to {old_executable}")