Support numactl bind for CPU and memory before process starts (#14156)

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)

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
 
 

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)

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)

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

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