From 58bab43ab3c8dfcdc05d3148aebcdb435327d4c3 Mon Sep 17 00:00:00 2001
From: determine123 <determine@sjtu.edu.cn>
Date: Wed, 21 Oct 2026 14:46:43 +0800
Subject: [PATCH] fix(topology): release allocated NUMA bitmasks

---
 deep_cluster/topology/bind.py                | 32 +++++++----
 tests/deep_cluster/test_numa_mask_cleanup.py | 58 ++++++++++++++++++++
 2 files changed, 78 insertions(+), 12 deletions(-)
 create mode 100644 tests/deep_cluster/test_numa_mask_cleanup.py

diff --git a/deep_cluster/topology/bind.py b/deep_cluster/topology/bind.py
index 7d180fd..89ecfa8 100644
--- a/deep_cluster/topology/bind.py
+++ b/deep_cluster/topology/bind.py
@@ -30,6 +30,8 @@ def _libnuma() -> ctypes.CDLL:
         raise OSError("libnuma.so.1 is required for NUMA binding") from error
     library.numa_allocate_cpumask.restype = ctypes.POINTER(_Bitmask)
     library.numa_allocate_nodemask.restype = ctypes.POINTER(_Bitmask)
+    library.numa_bitmask_free.argtypes = [ctypes.POINTER(_Bitmask)]
+    library.numa_bitmask_free.restype = None
     library.numa_node_to_cpus.argtypes = [ctypes.c_int, ctypes.POINTER(_Bitmask)]
     library.numa_sched_setaffinity.argtypes = [ctypes.c_int, ctypes.POINTER(_Bitmask)]
     library.numa_bitmask_setbit.argtypes = [ctypes.POINTER(_Bitmask), ctypes.c_uint]
@@ -79,18 +81,24 @@ def _perror(message: str) -> None:
 def bind_numa(node: int) -> int:
     library = _libnuma()
     cpumask = library.numa_allocate_cpumask()
-    if library.numa_node_to_cpus(node, cpumask):
-        _perror("get cpus in numa node failed")
-        return -1
-    if library.numa_sched_setaffinity(os.getpid(), cpumask):
-        _perror("failed to set affinity")
-        return -1
-    nodemask = library.numa_allocate_nodemask()
-    library.numa_bitmask_setbit(nodemask, node)
-    if library.set_mempolicy(MPOL_BIND, nodemask.contents.maskp, nodemask.contents.size + 1):
-        _perror("failed to set memory policy")
-        return -1
-    return 0
+    try:
+        if library.numa_node_to_cpus(node, cpumask):
+            _perror("get cpus in numa node failed")
+            return -1
+        if library.numa_sched_setaffinity(os.getpid(), cpumask):
+            _perror("failed to set affinity")
+            return -1
+        nodemask = library.numa_allocate_nodemask()
+        try:
+            library.numa_bitmask_setbit(nodemask, node)
+            if library.set_mempolicy(MPOL_BIND, nodemask.contents.maskp, nodemask.contents.size + 1):
+                _perror("failed to set memory policy")
+                return -1
+            return 0
+        finally:
+            library.numa_bitmask_free(nodemask)
+    finally:
+        library.numa_bitmask_free(cpumask)
 
 
 def set_cpu_affinity(cpus: Iterable[int]) -> None:
diff --git a/tests/deep_cluster/test_numa_mask_cleanup.py b/tests/deep_cluster/test_numa_mask_cleanup.py
new file mode 100644
index 0000000..27bd2c5
--- /dev/null
+++ b/tests/deep_cluster/test_numa_mask_cleanup.py
@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import ctypes
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from deep_cluster.topology import bind
+
+pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
+
+
+@pytest.fixture
+def numa(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace:
+    word = ctypes.c_ulong(0)
+    cpu = ctypes.pointer(bind._Bitmask(1, ctypes.pointer(word)))
+    node = ctypes.pointer(bind._Bitmask(1, ctypes.pointer(word)))
+    library = SimpleNamespace(
+        numa_allocate_cpumask=Mock(return_value=cpu),
+        numa_allocate_nodemask=Mock(return_value=node),
+        numa_node_to_cpus=Mock(return_value=0),
+        numa_sched_setaffinity=Mock(return_value=0),
+        numa_bitmask_setbit=Mock(),
+        set_mempolicy=Mock(return_value=0),
+        numa_bitmask_free=Mock(),
+    )
+    monkeypatch.setattr(bind, "_libnuma", lambda: library)
+    monkeypatch.setattr(bind, "_perror", lambda message: None)
+    return library
+
+
+@pytest.mark.parametrize("failure", [None, "numa_node_to_cpus", "numa_sched_setaffinity", "set_mempolicy"])
+def test_masks_are_freed_on_success_and_native_errors(numa: SimpleNamespace, failure: str | None) -> None:
+    if failure is not None:
+        getattr(numa, failure).return_value = -1
+    assert bind.bind_numa(0) == (0 if failure is None else -1)
+    allocated = [numa.numa_allocate_cpumask.return_value]
+    if numa.numa_allocate_nodemask.called:
+        allocated.append(numa.numa_allocate_nodemask.return_value)
+    freed = [call.args[0] for call in numa.numa_bitmask_free.call_args_list]
+    assert len(freed) == len(allocated)
+    assert {id(mask) for mask in freed} == {id(mask) for mask in allocated}
+
+
+@pytest.mark.parametrize("failure", ["numa_node_to_cpus", "numa_bitmask_setbit", "set_mempolicy"])
+def test_masks_are_freed_when_calls_raise(numa: SimpleNamespace, failure: str) -> None:
+    getattr(numa, failure).side_effect = RuntimeError("binding interrupted")
+    with pytest.raises(RuntimeError, match="binding interrupted"):
+        bind.bind_numa(0)
+    expected = 2 if numa.numa_allocate_nodemask.called else 1
+    assert numa.numa_bitmask_free.call_count == expected
+
+
+def test_repeated_binding_releases_each_allocation(numa: SimpleNamespace) -> None:
+    for _ in range(3):
+        assert bind.bind_numa(0) == 0
+    assert numa.numa_bitmask_free.call_count == 6