From 476308a6a28e4b71d8553c1a712286d74d00f046 Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 21:57:03 +0000 Subject: [PATCH 1/5] fix: align RDMA device parsing with NCCL --- checkpoint_engine/device_utils.py | 4 +++- tests/test_rdma_parser.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/checkpoint_engine/device_utils.py b/checkpoint_engine/device_utils.py index be5b81f..35fe0ca 100644 --- a/checkpoint_engine/device_utils.py +++ b/checkpoint_engine/device_utils.py @@ -140,6 +140,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 +165,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)] diff --git a/tests/test_rdma_parser.py b/tests/test_rdma_parser.py index e41b07f..a2900ea 100644 --- a/tests/test_rdma_parser.py +++ b/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", [ From 162829312ea61663dbf81ff2935babc79d665b10 Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 21:57:54 +0000 Subject: [PATCH 2/5] fix: keep parameter metadata state consistent --- checkpoint_engine/p2p_store.py | 2 +- checkpoint_engine/ps.py | 4 +++- tests/test_api.py | 15 ++++++++++++++ tests/test_p2p_store.py | 36 ++++++++++++++++++++++++++++++++++ 4 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 tests/test_p2p_store.py diff --git a/checkpoint_engine/p2p_store.py b/checkpoint_engine/p2p_store.py index 218e1b6..c247b92 100644 --- a/checkpoint_engine/p2p_store.py +++ b/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] diff --git a/checkpoint_engine/ps.py b/checkpoint_engine/ps.py index 1d8c5cf..21fdf23 100644 --- a/checkpoint_engine/ps.py +++ b/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" diff --git a/tests/test_api.py b/tests/test_api.py index 6d5f6c6..d629d4a 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -13,6 +13,7 @@ MemoryBufferMetas, ParameterMeta, ) +from checkpoint_engine.ps import ParameterServer _METAS_ADAPTER = TypeAdapter(dict[int, MemoryBufferMetaList]) @@ -137,3 +138,17 @@ 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()) diff --git a/tests/test_p2p_store.py b/tests/test_p2p_store.py new file mode 100644 index 0000000..6174b0c --- /dev/null +++ b/tests/test_p2p_store.py @@ -0,0 +1,36 @@ +import pytest +import torch + +from checkpoint_engine.p2p_store import P2PStore + + +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} From 8fd33b860aeef57364125fe4a431f203900bb35b Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 21:58:29 +0000 Subject: [PATCH 3/5] fix: close inference update HTTP client --- checkpoint_engine/api.py | 19 ++++++++------- tests/test_api.py | 52 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 60 insertions(+), 11 deletions(-) diff --git a/checkpoint_engine/api.py b/checkpoint_engine/api.py index 28aa812..d6fbac3 100644 --- a/checkpoint_engine/api.py +++ b/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() diff --git a/tests/test_api.py b/tests/test_api.py index d629d4a..bab3f36 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,13 +1,16 @@ """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, @@ -152,3 +155,48 @@ def test_load_metas_filters_empty_owners(fake_metas: dict[int, MemoryBufferMetaL 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, + ) From c75b190fd17776e1c281b54def75d2b14fd59641 Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 22:07:23 +0000 Subject: [PATCH 4/5] fix: preserve empty RDMA env selections --- checkpoint_engine/device_utils.py | 7 +++++-- tests/test_rdma_parser.py | 3 ++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/checkpoint_engine/device_utils.py b/checkpoint_engine/device_utils.py index 35fe0ca..6b0bc9b 100644 --- a/checkpoint_engine/device_utils.py +++ b/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: diff --git a/tests/test_rdma_parser.py b/tests/test_rdma_parser.py index a2900ea..1586704 100644 --- a/tests/test_rdma_parser.py +++ b/tests/test_rdma_parser.py @@ -143,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"]), ], ) From 43989bc9d7abb966504f50a1d1a6f145457865a1 Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 22:18:30 +0000 Subject: [PATCH 5/5] fix: make p2p registration cleanup idempotent --- checkpoint_engine/ps.py | 11 ++++++++--- tests/test_p2p_store.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/checkpoint_engine/ps.py b/checkpoint_engine/ps.py index 21fdf23..795ffde 100644 --- a/checkpoint_engine/ps.py +++ b/checkpoint_engine/ps.py @@ -718,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, diff --git a/tests/test_p2p_store.py b/tests/test_p2p_store.py index 6174b0c..034a210 100644 --- a/tests/test_p2p_store.py +++ b/tests/test_p2p_store.py @@ -1,7 +1,9 @@ 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: @@ -34,3 +36,36 @@ def batch_register_memory(self, addrs: list[int], caps: list[int]) -> int: 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