numa: bind within allowed CPUs when affinity is already constrained (#26983)

test/registered/utils/test_numa_utils.py

@@ -1,11 +1,16 @@
 import ctypes
+import os
 import unittest
 from unittest.mock import MagicMock, patch
 
 from sglang.srt.utils.numa_utils import (
+    _handle_numa_bind_failure,
     _is_numa_available,
+    _node_cpus,
+    _numactl_cpu_mem_args,
     _query_numa_node_for_gpu,
     get_numa_node_if_available,
+    numa_bind_to_node,
 )
 from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
 
@@ -26,63 +31,31 @@ class TestIsNumaAvailable(unittest.TestCase):
     def test_returns_false_when_no_numa_nodes(self, _mock_isdir):
         self.assertFalse(_is_numa_available())
 
-    @patch("sglang.srt.utils.numa_utils._is_cuda", True)
-    @patch("os.path.isdir", return_value=True)
-    @patch("sglang.srt.utils.numa_utils.psutil")
-    def test_returns_false_when_affinity_constrained(self, mock_psutil, _mock_isdir):
-        mock_process = MagicMock()
-        mock_process.cpu_affinity.return_value = [0, 1]
-        mock_psutil.Process.return_value = mock_process
-        mock_psutil.cpu_count.return_value = 128
-
-        self.assertFalse(_is_numa_available())
-
     @patch("sglang.srt.utils.numa_utils._can_set_mempolicy", return_value=True)
     @patch("sglang.srt.utils.numa_utils.shutil.which", return_value="/usr/bin/numactl")
     @patch("sglang.srt.utils.numa_utils._is_cuda", True)
     @patch("os.path.isdir", return_value=True)
-    @patch("sglang.srt.utils.numa_utils.psutil")
-    def test_returns_true_on_numa_system_with_full_affinity(
-        self, mock_psutil, _mock_isdir, _mock_which, _mock_mempolicy
+    def test_returns_true_on_numa_system(
+        self, _mock_isdir, _mock_which, _mock_mempolicy
     ):
-        all_cpus = list(range(128))
-        mock_process = MagicMock()
-        mock_process.cpu_affinity.return_value = all_cpus
-        mock_psutil.Process.return_value = mock_process
-        mock_psutil.cpu_count.return_value = 128
-
         self.assertTrue(_is_numa_available())
 
     @patch("sglang.srt.utils.numa_utils._can_set_mempolicy", return_value=False)
     @patch("sglang.srt.utils.numa_utils.shutil.which", return_value="/usr/bin/numactl")
     @patch("sglang.srt.utils.numa_utils._is_cuda", True)
     @patch("os.path.isdir", return_value=True)
-    @patch("sglang.srt.utils.numa_utils.psutil")
     def test_returns_false_when_mempolicy_not_permitted(
-        self, mock_psutil, _mock_isdir, _mock_which, _mock_mempolicy
+        self, _mock_isdir, _mock_which, _mock_mempolicy
     ):
-        all_cpus = list(range(128))
-        mock_process = MagicMock()
-        mock_process.cpu_affinity.return_value = all_cpus
-        mock_psutil.Process.return_value = mock_process
-        mock_psutil.cpu_count.return_value = 128
-
         self.assertFalse(_is_numa_available())
 
     @patch("sglang.srt.utils.numa_utils._can_set_mempolicy", return_value=True)
     @patch("sglang.srt.utils.numa_utils.shutil.which", return_value="/usr/bin/numactl")
     @patch("sglang.srt.utils.numa_utils._is_cuda", True)
     @patch("os.path.isdir", return_value=True)
-    @patch("sglang.srt.utils.numa_utils.psutil")
     def test_isdir_called_with_node1_path(
-        self, mock_psutil, mock_isdir, _mock_which, _mock_mempolicy
+        self, mock_isdir, _mock_which, _mock_mempolicy
     ):
-        all_cpus = list(range(8))
-        mock_process = MagicMock()
-        mock_process.cpu_affinity.return_value = all_cpus
-        mock_psutil.Process.return_value = mock_process
-        mock_psutil.cpu_count.return_value = 8
-
         _is_numa_available()
         mock_isdir.assert_called_with("/sys/devices/system/node/node1")
 
@@ -313,5 +286,85 @@ class TestB200NumaTopology(unittest.TestCase):
         )
 
 
+class TestNumaBindIntersection(unittest.TestCase):
+    """Tests for constraint-aware NUMA binding (node CPUs intersected with the
+    process's allowed CPUs)."""
+
+    @patch("sglang.srt.utils.numa_utils.get_libnuma", return_value=None)
+    def test_node_cpus_no_libnuma_returns_empty(self, _mock_lib):
+        self.assertEqual(_node_cpus(0), set())
+
+    @patch("os.sched_getaffinity", return_value=set(range(72)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    def test_numactl_args_unconstrained_uses_cpunodebind(self, _cpus, _aff):
+        self.assertEqual(_numactl_cpu_mem_args(0, 0), "--cpunodebind=0 --membind=0")
+
+    @patch("os.sched_getaffinity", return_value={0} | set(range(21, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    def test_numactl_args_constrained_uses_physcpubind(self, _cpus, _aff):
+        expected_cpus = ",".join(str(c) for c in [0] + list(range(21, 72)))
+        self.assertEqual(
+            _numactl_cpu_mem_args(0, 0),
+            f"--physcpubind={expected_cpus} --membind=0",
+        )
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "0"})
+    @patch("os.sched_getaffinity", return_value=set(range(72, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    def test_numactl_args_empty_intersection_returns_none(self, _cpus, _aff):
+        self.assertIsNone(_numactl_cpu_mem_args(0, 0))
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "1"})
+    @patch("os.sched_getaffinity", return_value=set(range(72, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    def test_numactl_args_empty_intersection_crashes_when_enabled(self, _cpus, _aff):
+        with self.assertRaises(RuntimeError):
+            _numactl_cpu_mem_args(0, 0)
+
+    @patch("os.sched_setaffinity")
+    @patch("os.sched_getaffinity", return_value={0} | set(range(21, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    @patch("sglang.srt.utils.numa_utils.get_libnuma")
+    def test_numa_bind_to_node_constrained_sets_intersection(
+        self, mock_libnuma, _cpus, _aff, mock_setaff
+    ):
+        lib = MagicMock()
+        lib.numa_available.return_value = 0
+        mock_libnuma.return_value = lib
+
+        numa_bind_to_node(0)
+
+        mock_setaff.assert_called_once_with(0, {0} | set(range(21, 72)))
+        lib.numa_set_preferred.assert_called_once()
+        lib.numa_run_on_node.assert_not_called()
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "0"})
+    @patch("os.sched_setaffinity")
+    @patch("os.sched_getaffinity", return_value=set(range(72, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    @patch("sglang.srt.utils.numa_utils.get_libnuma")
+    def test_numa_bind_to_node_empty_intersection_skips(
+        self, mock_libnuma, _cpus, _aff, mock_setaff
+    ):
+        lib = MagicMock()
+        lib.numa_available.return_value = 0
+        mock_libnuma.return_value = lib
+
+        numa_bind_to_node(0)
+
+        mock_setaff.assert_not_called()
+        lib.numa_set_preferred.assert_not_called()
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "1"})
+    def test_handle_failure_raises_when_enabled(self):
+        with self.assertRaises(RuntimeError):
+            _handle_numa_bind_failure(0, {72, 73})
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "0"})
+    def test_handle_failure_warns_when_disabled(self):
+        with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING"):
+            _handle_numa_bind_failure(0, {72, 73})
+
+
 if __name__ == "__main__":
     unittest.main()