[Store] fix: prevent cross-process memcpy segfault when MC_STORE_MEMCPY auto-enables - #2001
Conversation
isLocalTransfer compared only the IP of handle.transport_endpoint_ to the local endpoint, so two processes on the same host (same IP, different ports) were treated as LOCAL_MEMCPY-eligible. The memcpy worker then dereferenced handle.buffer_address_, which is a virtual address only valid in the owning process, and segfaulted inside __memcpy_avx512_unaligned_erms. This was latent before kvcache-ai#1936 (MC_STORE_MEMCPY defaulted to off). The TCP-only auto-enable exposed it on multi-process workloads such as the TorchSpec inference/trainer pipeline. Compare the full transport endpoint instead, matching the check already used by Client::IsReplicaOnLocalMemory. Cross-process same-host transfers now correctly fall through to TRANSFER_ENGINE; same-process transfers still take the memcpy fast path. Fixes the crash reported with MC_STORE_MEMCPY auto-enabled on TCP-only hosts.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the isLocalTransfer logic in mooncake-store/src/transfer_task.cpp to require an exact match of the transport endpoint instead of just the IP address. This ensures that local memory copies are only attempted within the same process, avoiding potential segmentation faults from accessing virtual addresses in different process spaces. A performance optimization was suggested to use a constant reference for the local endpoint to avoid unnecessary string allocations in the transfer hot path.
| if (handle.transport_endpoint_.empty()) { | ||
| return false; | ||
| } | ||
| std::string local_ep = engine_.getLocalIpAndPort(); |
There was a problem hiding this comment.
Calling engine_.getLocalIpAndPort() on every transfer request can be expensive, especially if it involves string allocations or internal lookups. Since the local endpoint is typically constant for the lifetime of the TransferSubmitter, consider caching this value in a member variable during construction to improve performance in the hot path. For now, using a const std::string& can at least avoid an extra copy if the engine returns a reference.
| std::string local_ep = engine_.getLocalIpAndPort(); | |
| const std::string& local_ep = engine_.getLocalIpAndPort(); |
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
…lity check - Cache engine.getLocalIpAndPort() in TransferSubmitter::local_endpoint_ at construction, removing the per-transfer string allocation in the hot path (Gemini review feedback). - Extract the endpoint comparison into a static TransferSubmitter::isSameProcessEndpoint helper so the locality decision is testable without spinning up a real TransferEngine. - Add transfer_task_test.cpp coverage for the early-return branches and the regression case (same host, different port -> not local).
| namespace { | ||
| // Helper function to extract IP address from endpoint string (ip:port format) | ||
| // Supports both IPv4 (ip:port) and IPv6 ([ipv6]:port) formats | ||
| std::string extractIpAddress(const std::string& endpoint) { | ||
| if (endpoint.empty()) { | ||
| return ""; | ||
| } | ||
|
|
||
| // Handle IPv6 format: [ipv6]:port | ||
| if (endpoint[0] == '[') { | ||
| size_t closing_bracket = endpoint.find(']'); | ||
| if (closing_bracket == std::string::npos) { | ||
| LOG(WARNING) << "Invalid IPv6 endpoint format: " << endpoint; | ||
| // Return empty to disable local memcpy optimization | ||
| return ""; | ||
| } | ||
| return endpoint.substr(1, closing_bracket - 1); // Extract IPv6 address | ||
| } | ||
|
|
||
| // Handle IPv4 or hostname:port format | ||
| // Find the last colon (to handle IPv6 addresses without brackets) | ||
| size_t colon_pos = endpoint.rfind(':'); | ||
| if (colon_pos != std::string::npos) { | ||
| return endpoint.substr(0, colon_pos); | ||
| bool TransferSubmitter::isSameProcessEndpoint( | ||
| const std::string& handle_endpoint, const std::string& local_endpoint) { | ||
| // Local memcpy requires that handle.buffer_address_ is a virtual address | ||
| // valid in THIS process. Same host is not enough — two processes on the | ||
| // same host share an IP but have distinct virtual address spaces, so a | ||
| // memcpy on a peer process's address would segfault. Require the full | ||
| // transport endpoint to match, which uniquely identifies the owning | ||
| // process. | ||
| if (handle_endpoint.empty() || local_endpoint.empty()) { | ||
| return false; | ||
| } | ||
|
|
There was a problem hiding this comment.
Please keep this function. I think it is useful for other scenarios.
…oss-process-segfault # Conflicts: # mooncake-store/src/transfer_task.cpp
|
@Yeuvoir Sorry for the late review. Please fix the conflicts. |
…oss-process-segfault # Conflicts: # mooncake-store/src/transfer_task.cpp
There was a problem hiding this comment.
Pull request overview
Fixes a crash in mooncake-store where the memcpy fast-path could be incorrectly selected for same-host, different-process transfers (leading to dereferencing a virtual address that isn’t valid in the current process). The change tightens the locality decision to require an exact endpoint match (ip:port / full hostname), and adds a focused unit test for the comparison helper.
Changes:
- Cache the local transfer-engine endpoint at
TransferSubmitterconstruction time and use exact endpoint equality to decide memcpy eligibility. - Introduce
TransferSubmitter::isSameProcessEndpoint()and supporting endpoint parsing/logging to avoid treating same-host/different-process as local. - Add unit coverage for endpoint equality semantics (empty, same ip:port, same IP different port, hostname cases).
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| mooncake-store/src/transfer_task.cpp | Updates locality logic to require exact endpoint match and adds helper + logging for same-host/different-endpoint cases. |
| mooncake-store/include/transfer_task.h | Exposes isSameProcessEndpoint for unit testing and caches local_endpoint_ in the submitter. |
| mooncake-store/tests/transfer_task_test.cpp | Adds unit test covering the endpoint-comparison behavior used for memcpy locality decisions. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| const std::string handle_ip = extractIpAddress(handle_endpoint); | ||
| const std::string local_ip = extractIpAddress(local_endpoint); | ||
| if (!handle_ip.empty() && handle_ip == local_ip) { | ||
| VLOG(2) << "Disabling local memcpy for same-host endpoints with " | ||
| "different process endpoints: handle=" | ||
| << handle_endpoint << ", local=" << local_endpoint; |
| * @brief Pure comparison helper: returns true iff both endpoints are | ||
| * non-empty and identical. Exposed for unit testing of the locality | ||
| * decision without instantiating a full TransferEngine. |
Head branch was pushed to by a user without write access
…PY auto-enables (kvcache-ai#2001) * [Store] fix: require same-process endpoint for LOCAL_MEMCPY strategy isLocalTransfer compared only the IP of handle.transport_endpoint_ to the local endpoint, so two processes on the same host (same IP, different ports) were treated as LOCAL_MEMCPY-eligible. The memcpy worker then dereferenced handle.buffer_address_, which is a virtual address only valid in the owning process, and segfaulted inside __memcpy_avx512_unaligned_erms. This was latent before kvcache-ai#1936 (MC_STORE_MEMCPY defaulted to off). The TCP-only auto-enable exposed it on multi-process workloads such as the TorchSpec inference/trainer pipeline. Compare the full transport endpoint instead, matching the check already used by Client::IsReplicaOnLocalMemory. Cross-process same-host transfers now correctly fall through to TRANSFER_ENGINE; same-process transfers still take the memcpy fast path. Fixes the crash reported with MC_STORE_MEMCPY auto-enabled on TCP-only hosts. --------- Co-authored-by: Teng Ma <teng-ma@linux.alibaba.com>
Summary
MemcpyWorkerPool::workerThreadreported whenMC_STORE_MEMCPYauto-enables on TCP-only hosts ([Store] auto-enable MC_STORE_MEMCPY in TCP-only environments #1936 regression).isLocalTransferwas comparing only the IP, so peer-process buffers on the same host were treated asLOCAL_MEMCPY-eligible and the worker dereferenced a virtual address valid only in the owning process.Client::IsReplicaOnLocalMemory); cross-process same-host transfers now correctly fall through toTRANSFER_ENGINE.Test plan
put_from/get_intofrom separate processes (e.g., TorchSpec Qwen 8B pipeline) and confirm it no longer segfaults withMC_STORE_MEMCPYunset.ctestformooncake-storepasses.