Fix update-path crashes, stale P2P state, and NCCL RDMA parsing

checkpoint_engine/api.py

@@ -31,15 +31,16 @@ def request_inference_to_update(
         httpx.HTTPStatusError: If the response contains an HTTP error status.
         httpx.RequestError: If there was an issue while making the request.
     """
-    resp = httpx.Client(transport=httpx.HTTPTransport(uds=uds)).post(
-        url,
-        json={
-            "method": "update_weights_from_ipc",
-            "args": [socket_paths],
-            "timeout": timeout,
-        },
-        timeout=timeout,
-    )
+    with httpx.Client(transport=httpx.HTTPTransport(uds=uds)) as client:
+        resp = client.post(
+            url,
+            json={
+                "method": "update_weights_from_ipc",
+                "args": [socket_paths],
+                "timeout": timeout,
+            },
+            timeout=timeout,
+        )
     resp.raise_for_status()
 
 

checkpoint_engine/device_utils.py

@@ -77,8 +77,11 @@ def _get_rdma_devices() -> list[str]:
     if devices_str:
         return devices_str.split(",")
     # if PS_P2P_STORE_RDMA_DEVICES is not set, try to use NCCL_IB_HCA to get RDMA devices
-    hca = os.getenv("NCCL_IB_HCA", None)
-    return _parse_NCCL_IB_HCA(hca or "", _ibv_get_device_list()) or _ibv_get_device_list()
+    hca = os.getenv("NCCL_IB_HCA")
+    devices = _ibv_get_device_list()
+    if hca is not None:
+        return _parse_NCCL_IB_HCA(hca, devices)
+    return devices
 
 
 def _get_my_rdma_device(local_rank: int, gpu_count: int, devices: list[str]) -> str:
@@ -140,6 +143,8 @@ def _parse_NCCL_IB_HCA(value: str, available_devices: list[str]) -> list[str]:
         value = value.removeprefix("=")
 
     device_specs = [spec.strip() for spec in value.split(",") if spec.strip()]
+    if is_exclude and not device_specs:
+        return []
 
     result = _resolve_device_specs(device_specs, is_exact_match, available_devices)
     if is_exclude:
@@ -163,7 +168,7 @@ def _resolve_device_specs(
         # port = parts[1].strip() if len(parts) > 1 else None
         base_devices = (
             [device_name]
-            if device_name in available_devices
+            if is_exact_match and device_name in available_devices
             else []
             if is_exact_match
             else [dev for dev in available_devices if dev.startswith(device_name)]

checkpoint_engine/p2p_store.py

@@ -51,12 +51,12 @@ def addr(self) -> str:
     def register_named_tensors(self, named_tensors: dict[str, torch.Tensor]):
         buffer_addresses = [tensor.data_ptr() for tensor in named_tensors.values()]
         capacities = [tensor.nbytes for tensor in named_tensors.values()]
-        self.named_tensors.update(named_tensors)
         for i, name in enumerate(named_tensors.keys()):
             logger.info(
                 f"[rank{self.rank}] p2p store register tensor {name} with addr {hex(buffer_addresses[i])} and capacity {capacities[i]}"
             )
         assert self.engine.batch_register_memory(buffer_addresses, capacities) == 0
+        self.named_tensors.update(named_tensors)
 
     def unregister_named_tensors(self, names: list[str]) -> int:
         buffer_addresses = [self.named_tensors[name].data_ptr() for name in names]

checkpoint_engine/ps.py

@@ -262,7 +262,9 @@ def get_metas(self) -> dict[int, MemoryBufferMetaList]:
         return self._current_global_parameter_metas
 
     def load_metas(self, metas: dict[int, MemoryBufferMetaList]):
-        self._current_global_parameter_metas = metas
+        self._current_global_parameter_metas = {
+            rank: meta for rank, meta in metas.items() if meta.memory_buffer_metas_list
+        }
         self._remote_rdma_devices = defaultdict(set)
         for i, meta in self._current_global_parameter_metas.items():
             assert meta.rdma_device is not None, "meta.rdma_device should not be None"
@@ -716,9 +718,14 @@ def _unregister_parameters_from_p2p_store(self, checkpoint_name: str) -> int:
             if checkpoint_name != self._current_shared_memory_pool_user
             else self.shared_memory_pool_name
         )
-        return self._p2p_store.unregister_named_tensors(
-            [f"memory_pool_{unregister_name}_{idx}" for idx, _ in enumerate(pool)]
-        )
+        names = [
+            name
+            for idx, _ in enumerate(pool)
+            if (name := f"memory_pool_{unregister_name}_{idx}") in self._p2p_store.named_tensors
+        ]
+        if not names:
+            return 0
+        return self._p2p_store.unregister_named_tensors(names)
 
     def _update_per_bucket(
         self,

tests/test_api.py

@@ -1,18 +1,22 @@
 """CPU-only tests for the metas endpoints in api.py."""
 
-from unittest.mock import MagicMock
+from types import TracebackType
+from typing import Any
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 from fastapi.testclient import TestClient
 from pydantic import TypeAdapter
+from typing_extensions import Self
 
-from checkpoint_engine.api import _init_api
+from checkpoint_engine.api import _init_api, request_inference_to_update
 from checkpoint_engine.data_types import (
     MemoryBufferMetaList,
     MemoryBufferMetas,
     ParameterMeta,
 )
+from checkpoint_engine.ps import ParameterServer
 
 
 _METAS_ADAPTER = TypeAdapter(dict[int, MemoryBufferMetaList])
@@ -137,3 +141,62 @@ def test_round_trip_get_then_load(
     )
     assert load_resp.status_code == 200
     ps_mock.load_metas.assert_called_once_with(fake_metas)
+
+
+def test_load_metas_filters_empty_owners(fake_metas: dict[int, MemoryBufferMetaList]) -> None:
+    ps = ParameterServer.__new__(ParameterServer)
+    empty_meta = MemoryBufferMetaList(
+        p2p_store_addr="192.168.1.2:12345",
+        rdma_device="mlx5_2",
+        memory_buffer_metas_list=[],
+    )
+
+    ps.load_metas({**fake_metas, 2: empty_meta})
+
+    assert ps.get_metas() == fake_metas
+    assert all(2 not in ranks for ranks in ps._remote_rdma_devices.values())
+
+
+def test_request_inference_to_update_closes_httpx_client() -> None:
+    class FakeResponse:
+        def raise_for_status(self) -> None:
+            pass
+
+    class FakeClient:
+        closed = False
+        payload = None
+
+        def __init__(self, *, transport: Any):
+            self.transport = transport
+
+        def __enter__(self) -> Self:
+            return self
+
+        def __exit__(
+            self,
+            exc_type: type[BaseException] | None,
+            exc: BaseException | None,
+            tb: TracebackType | None,
+        ) -> None:
+            type(self).closed = True
+
+        def post(self, url: str, *, json: dict[str, Any], timeout: float) -> FakeResponse:
+            type(self).payload = (url, json, timeout)
+            return FakeResponse()
+
+    with (
+        patch("checkpoint_engine.api.httpx.HTTPTransport", return_value="transport"),
+        patch("checkpoint_engine.api.httpx.Client", FakeClient),
+    ):
+        request_inference_to_update("http://example/update", {"GPU-0": "ipc://x"}, timeout=1.5)
+
+    assert FakeClient.closed is True
+    assert FakeClient.payload == (
+        "http://example/update",
+        {
+            "method": "update_weights_from_ipc",
+            "args": [{"GPU-0": "ipc://x"}],
+            "timeout": 1.5,
+        },
+        1.5,
+    )

tests/test_p2p_store.py

@@ -0,0 +1,71 @@
+import pytest
+import torch
+
+from checkpoint_engine.data_types import MemoryBuffer
+from checkpoint_engine.p2p_store import P2PStore
+from checkpoint_engine.ps import ParameterServer
+
+
+def test_register_named_tensors_does_not_mutate_state_on_engine_failure() -> None:
+    class FakeEngine:
+        def batch_register_memory(self, addrs: list[int], caps: list[int]) -> int:
+            return -1
+
+    store = P2PStore.__new__(P2PStore)
+    store.rank = 0
+    store.engine = FakeEngine()
+    store.named_tensors = {}
+
+    with pytest.raises(AssertionError):
+        store.register_named_tensors({"w": torch.zeros(1)})
+
+    assert store.named_tensors == {}
+
+
+def test_register_named_tensors_records_state_after_engine_success() -> None:
+    class FakeEngine:
+        def batch_register_memory(self, addrs: list[int], caps: list[int]) -> int:
+            return 0
+
+    store = P2PStore.__new__(P2PStore)
+    store.rank = 0
+    store.engine = FakeEngine()
+    store.named_tensors = {}
+    tensor = torch.zeros(1)
+
+    store.register_named_tensors({"w": tensor})
+
+    assert store.named_tensors == {"w": tensor}
+
+
+def test_register_checkpoint_cleans_memory_pool_after_p2p_registration_failure(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class FakeDeviceManager:
+        device_type = "npu"
+
+    class FakeP2PStore:
+        def __init__(self) -> None:
+            self.named_tensors: dict[str, torch.Tensor] = {}
+
+        def register_named_tensors(self, named_tensors: dict[str, torch.Tensor]) -> None:
+            raise AssertionError("engine register failed")
+
+        def unregister_named_tensors(self, names: list[str]) -> int:
+            raise AssertionError("should not unregister names that were never registered")
+
+    buffer = MemoryBuffer(buffer=torch.zeros(1), size=4, metas=[])
+
+    ps = ParameterServer.__new__(ParameterServer)
+    ps.device_manager = FakeDeviceManager()
+    ps._rank = 0
+    ps._p2p_store = FakeP2PStore()
+    ps._memory_pool = {ps.shared_memory_pool_name: []}
+    ps._current_shared_memory_pool_user = ""
+
+    monkeypatch.setattr("checkpoint_engine.ps._register_checkpoint", lambda **_: [buffer])
+
+    with pytest.raises(AssertionError, match="engine register failed"):
+        ps.register_checkpoint("ckpt", named_tensors={"w": torch.zeros(1)})
+
+    assert "ckpt" not in ps._memory_pool

tests/test_rdma_parser.py

@@ -57,8 +57,8 @@ def test_get_rdma_devices_no_env_vars(mock_available_devices: list[str]):
         pytest.param("", ["mlx5_0", "mlx5_1", "mlx4_0", "mlx4_1"], id="empty string"),
         pytest.param("   \t\n  ", ["mlx5_0", "mlx5_1", "mlx4_0", "mlx4_1"], id="whitespace"),
         pytest.param("None", [], id="None string"),
-        pytest.param("^", ["mlx5_0", "mlx5_1", "mlx4_0", "mlx4_1"], id="caret"),
-        pytest.param("^=", ["mlx5_0", "mlx5_1", "mlx4_0", "mlx4_1"], id="caret-equals"),
+        pytest.param("^", [], id="caret"),
+        pytest.param("^=", [], id="caret-equals"),
         pytest.param("=^", [], id="equals-caret"),
         pytest.param("^^", ["mlx5_0", "mlx5_1", "mlx4_0", "mlx4_1"], id="double-caret"),
         pytest.param("=", [], id="equals"),
@@ -103,6 +103,12 @@ def test_parse_various_patterns(
     assert result == expected
 
 
+def test_parse_prefix_when_prefix_is_also_device_name() -> None:
+    """Bare NCCL_IB_HCA tokens are prefixes unless the value starts with '='."""
+    result = _parse_NCCL_IB_HCA("mlx5", ["mlx5", "mlx5_0", "mlx5_1"])
+    assert result == ["mlx5", "mlx5_0", "mlx5_1"]
+
+
 @pytest.mark.parametrize(
     "input_value,expected_result,expected_warning",
     [
@@ -137,7 +143,8 @@ def test_parse_exact_match_with_nonexistent_device(
         ("NCCL_IB_HCA", "mlx5", ["mlx5_0", "mlx5_1"]),
         ("NCCL_IB_HCA", "mlx5_0,mlx5_1", ["mlx5_0", "mlx5_1"]),
         ("NCCL_IB_HCA", "^mlx5_0", ["mlx5_1", "mlx4_0", "mlx4_1"]),
-        ("NCCL_IB_HCA", "mlx6", ["mlx5_0", "mlx5_1", "mlx4_0", "mlx4_1"]),
+        ("NCCL_IB_HCA", "mlx6", []),
+        ("NCCL_IB_HCA", "^=", []),
         ("NCCL_IB_HCA", "", ["mlx5_0", "mlx5_1", "mlx4_0", "mlx4_1"]),
     ],
 )