diff --git a/mooncake-store/include/transfer_task.h b/mooncake-store/include/transfer_task.h index df7c90e928..42961d05d4 100644 --- a/mooncake-store/include/transfer_task.h +++ b/mooncake-store/include/transfer_task.h @@ -410,8 +410,24 @@ class TransferSubmitter { const std::unordered_map>& batched_slices); + /** + * @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. + * + * Two endpoints identify the same process only when their ip:port (or + * full hostname) match exactly; same-host different-process pairs share + * an IP but not a port and must NOT be treated as locally addressable. + */ + static bool isSameProcessEndpoint(const std::string& handle_endpoint, + const std::string& local_endpoint); + private: TransferEngine& engine_; + // Cached at construction: the local transport endpoint never changes for + // the lifetime of the TransferSubmitter, so we avoid calling + // engine_.getLocalIpAndPort() (which allocates a string) on every transfer. + const std::string local_endpoint_; std::unique_ptr memcpy_pool_; std::unique_ptr fileread_pool_; bool memcpy_enabled_; diff --git a/mooncake-store/src/transfer_task.cpp b/mooncake-store/src/transfer_task.cpp index a5ab382c18..0941536b00 100644 --- a/mooncake-store/src/transfer_task.cpp +++ b/mooncake-store/src/transfer_task.cpp @@ -451,6 +451,7 @@ TransferSubmitter::TransferSubmitter(TransferEngine& engine, const std::string& local_hostname, TransferMetric* transfer_metric) : engine_(engine), + local_endpoint_(engine.getLocalIpAndPort()), memcpy_pool_(std::make_unique()), fileread_pool_(std::make_unique(backend)), local_hostname_(local_hostname), @@ -816,28 +817,67 @@ TransferStrategy TransferSubmitter::selectStrategy( return TransferStrategy::TRANSFER_ENGINE; } -bool TransferSubmitter::isLocalTransfer( - const AllocatedBuffer::Descriptor& handle) const { - if (handle.transport_endpoint_.empty()) return false; +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 ""; + } - // Metadata-service descriptors use the client hostname as the segment ID. - // If it matches this client's hostname, the buffer address is local. - if (!local_hostname_.empty() && - local_hostname_ == handle.transport_endpoint_) { - return true; + // 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 ""; + } + return endpoint.substr(1, closing_bracket - 1); + } + + // Handle IPv4 or hostname:port format. + size_t colon_pos = endpoint.rfind(':'); + if (colon_pos != std::string::npos) { + return endpoint.substr(0, colon_pos); } - // P2P descriptors use the transfer engine endpoint as the segment ID. - // If it matches this engine's endpoint, the buffer address is local. - std::string local_ep = engine_.getLocalIpAndPort(); - if (!local_ep.empty() && handle.transport_endpoint_ == local_ep) { + // No colon found, return the whole string (might be just IP or hostname). + return endpoint; +} +} // namespace + +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; + } + if (handle_endpoint == local_endpoint) { return true; } - // Without a local endpoint we cannot prove locality; disable memcpy. + 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; + } + return false; } +bool TransferSubmitter::isLocalTransfer( + const AllocatedBuffer::Descriptor& handle) const { + return isSameProcessEndpoint(handle.transport_endpoint_, local_hostname_) || + isSameProcessEndpoint(handle.transport_endpoint_, local_endpoint_); +} + bool TransferSubmitter::validateTransferParams( const AllocatedBuffer::Descriptor& handle, const std::vector& slices) const { diff --git a/mooncake-store/tests/transfer_task_test.cpp b/mooncake-store/tests/transfer_task_test.cpp index caa7c7ad7e..d9c2ee9ea9 100644 --- a/mooncake-store/tests/transfer_task_test.cpp +++ b/mooncake-store/tests/transfer_task_test.cpp @@ -144,6 +144,37 @@ TEST_F(TransferTaskTest, MemcpyWorkerPoolMultipleOperations) { } } +// Test the locality decision used by TransferSubmitter::isLocalTransfer. +// Same-host different-process pairs share an IP but have distinct ports; +// they must NOT be treated as locally addressable, otherwise memcpy in the +// caller process would dereference a virtual address belonging to a peer +// process and segfault. +TEST_F(TransferTaskTest, IsSameProcessEndpoint) { + // Empty inputs -> not same-process (cannot prove locality). + EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint("", "")); + EXPECT_FALSE( + TransferSubmitter::isSameProcessEndpoint("", "192.168.1.10:12345")); + EXPECT_FALSE( + TransferSubmitter::isSameProcessEndpoint("192.168.1.10:12345", "")); + + // Identical ip:port -> same process. + EXPECT_TRUE(TransferSubmitter::isSameProcessEndpoint("192.168.1.10:12345", + "192.168.1.10:12345")); + + // Same host, different port -> different process, NOT local. + // This is the regression case fixed by this change. + EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint( + "192.168.1.10:12345", "192.168.1.10:12346")); + + // Different hosts -> not local. + EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint( + "192.168.1.10:12345", "192.168.1.11:12345")); + + // Hostname endpoints (non-P2P metadata mode) compare as full strings. + EXPECT_TRUE(TransferSubmitter::isSameProcessEndpoint("host-a", "host-a")); + EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint("host-a", "host-b")); +} + // Test TransferStrategy enum and stream operator TEST_F(TransferTaskTest, TransferStrategyEnum) { // Test enum values diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh index 64adbe967a..2fc703b5bd 100755 --- a/scripts/run_tests.sh +++ b/scripts/run_tests.sh @@ -40,7 +40,11 @@ DUMMY_TEST_PID_2=$! wait $DUMMY_TEST_PID_1 $DUMMY_TEST_PID_2 kill $CLIENT_PID || true -pip install torch numpy safetensors packaging +pip install numpy safetensors packaging +# Keep the test torch aligned with the EP/PG variants packaged into the CI wheel. +pip install "${MOONCAKE_TEST_TORCH_SPEC:-torch==2.11.0+cu128}" \ + --index-url "${MOONCAKE_TEST_TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu128}" \ + --extra-index-url https://pypi.org/simple MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_put_get_tensor.py MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_safetensor_functions.py kill $MASTER_PID || true