Merge pull request #384 from ZacharyZcR/perf/auto-numa

plan NUMA interleave on multi-socket Linux

c/resource_plan.py

@@ -203,6 +203,22 @@ def physical_cpu_count():
     return os.cpu_count() or 1
 
 
+def cpu_socket_count():
+    """Return the number of physical CPU sockets visible to this process."""
+    if not sys.platform.startswith("linux"):
+        return 1
+    try:
+        result = subprocess.run(["lscpu", "-p=socket"], text=True,
+                                capture_output=True, check=True, timeout=5)
+        sockets = {int(line) for line in result.stdout.splitlines()
+                   if line and not line.startswith("#")}
+        if sockets:
+            return len(sockets)
+    except (OSError, ValueError, subprocess.SubprocessError):
+        pass
+    return 1
+
+
 POLICIES = {
     "quality": {"preserve_quantization": True, "preserve_router": True},
     "balanced": {"preserve_quantization": True, "preserve_router": True},
@@ -212,11 +228,12 @@ POLICIES = {
 
 def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
                available_memory=None, available_disk=None, gpus=None,
-               policy="quality", physical_cpus=None):
+               policy="quality", physical_cpus=None, cpu_sockets=None):
     if policy not in POLICIES:
         raise ValueError(f"unknown policy: {policy}")
     info = analyze_model(model)
     physical_cpus = physical_cpu_count() if physical_cpus is None else physical_cpus
+    cpu_sockets = cpu_socket_count() if cpu_sockets is None else cpu_sockets
     cfg = info["config"]
     available_memory = memory_available() if available_memory is None else available_memory
     if available_disk is None:
@@ -286,6 +303,7 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
                    "quality_preserving": policy != "experimental-fast"},
         "model": {key: value for key, value in info.items() if key != "config"},
         "cpu": {"physical_cores": max(1, int(physical_cpus)),
+                "sockets": max(1, int(cpu_sockets)),
                 "thread_policy": "physical-cores"},
         "tiers": {
             "disk": {"role": "cold-backing", "model_bytes": info["model_bytes"],
@@ -318,6 +336,10 @@ def environment_for_plan(plan, env=None, cuda_enabled=True):
         # ("Affinity not supported on this configuration"): non impostarle li'.
         result.setdefault("OMP_PROC_BIND", "spread")
         result.setdefault("OMP_PLACES", "cores")
+    if sys.platform.startswith("linux") and plan["cpu"].get("sockets", 1) > 1:
+        # Selectively interleave large expert/dense slabs across memory controllers.
+        # Unlike blanket numactl interleave, this leaves CUDA staging buffers local.
+        result.setdefault("COLI_NUMA", "1")
     if plan["policy"]["name"] == "balanced":
         result.setdefault("REPIN", "64")
     ram = plan["tiers"]["ram"]

c/tests/test_resource_plan.py

@@ -10,6 +10,7 @@ from resource_plan import (
     GB,
     analyze_model,
     build_plan,
+    cpu_socket_count,
     environment_for_plan,
     format_plan,
     memory_available,
@@ -66,15 +67,19 @@ class ResourcePlanTest(unittest.TestCase):
         # 0 slots/layer. The value must be a sane positive number of bytes.
         self.assertGreater(memory_available(), 0)
 
+    def test_cpu_socket_count_is_positive(self):
+        self.assertGreaterEqual(cpu_socket_count(), 1)
+
     def test_builds_bounded_three_tier_plan(self):
         gpus = [{"index": 0, "name": "test-gpu", "total_bytes": 12 * GB,
                  "free_bytes": 10 * GB}]
         plan = build_plan(self.model, ram_gb=16, context=32, vram_gb=20,
                           available_memory=32 * GB, available_disk=100 * GB, gpus=gpus,
-                          physical_cpus=24)
+                          physical_cpus=24, cpu_sockets=2)
         self.assertEqual(plan["version"], 2)
         self.assertEqual(plan["policy"]["name"], "quality")
         self.assertEqual(plan["cpu"]["physical_cores"], 24)
+        self.assertEqual(plan["cpu"]["sockets"], 2)
         self.assertTrue(plan["policy"]["preserve_quantization"])
         self.assertFalse(plan["tiers"]["vram"]["requires_host_backing"])
         self.assertEqual(plan["tiers"]["ram"]["budget_bytes"], 16 * GB)
@@ -105,7 +110,7 @@ class ResourcePlanTest(unittest.TestCase):
             {"index": 1, "name": "b", "total_bytes": 12 * GB, "free_bytes": 10 * GB},
         ]
         plan = build_plan(self.model, ram_gb=16, available_memory=32 * GB,
-                          available_disk=1, gpus=gpus)
+                          available_disk=1, gpus=gpus, cpu_sockets=2)
         env = environment_for_plan(plan, {"RAM_GB": "12", "PIN": "stats.txt",
                                                "COLI_GPUS": "1"})
         self.assertEqual(env["RAM_GB"], "12")
@@ -125,6 +130,16 @@ class ResourcePlanTest(unittest.TestCase):
                                                         "OMP_PROC_BIND": "close"})
         self.assertEqual(explicit_threads["OMP_NUM_THREADS"], "7")
         self.assertEqual(explicit_threads["OMP_PROC_BIND"], "close")
+
+        if sys.platform.startswith("linux"):
+            self.assertEqual(env["COLI_NUMA"], "1")
+            explicit_numa = environment_for_plan(plan, {"COLI_NUMA": "0"})
+            self.assertEqual(explicit_numa["COLI_NUMA"], "0")
+
+    def test_single_socket_plan_does_not_enable_numa(self):
+        plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
+                          gpus=[], physical_cpus=8, cpu_sockets=1)
+        self.assertNotIn("COLI_NUMA", environment_for_plan(plan))
     def test_cpu_binary_does_not_apply_gpu_tier(self):
         plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
                           gpus=[{"index": 0, "name": "a", "total_bytes": 8 * GB,