-
Notifications
You must be signed in to change notification settings - Fork 143
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[TransferEngine] feature: registerLocalMemory support the "*" locatio…
…n. (#86) * [TransferEngine] feature: registerLocalMemory support the "*" location. 1. try best to recognize the cpu numa node for now, 2. use all nic when failed to get the numa node, 3. will support cuda memory in the feature. Signed-off-by: doujiang24 <doujiang24@gmail.com> * fix test when there is only one numa node Signed-off-by: doujiang24 <doujiang24@gmail.com> --------- Signed-off-by: doujiang24 <doujiang24@gmail.com>
- Loading branch information
1 parent
19aafbe
commit 86a8649
Showing
11 changed files
with
342 additions
and
67 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| // Copyright 2024 KVCache.AI | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| #ifndef MEMORY_LOCATION_H | ||
| #define MEMORY_LOCATION_H | ||
|
|
||
| #include <glog/logging.h> | ||
|
|
||
| #include <memory> | ||
|
|
||
| #include "common.h" | ||
|
|
||
| const int pagesize = 4096; | ||
|
|
||
| namespace mooncake { | ||
| struct MemoryLocationEntry { | ||
| uint64_t start; | ||
| size_t len; | ||
| std::string location; | ||
| }; | ||
|
|
||
| // Get CPU numa node id | ||
| // TODO: support getting cuda device id from unified address. | ||
| const std::vector<MemoryLocationEntry> getMemoryLocation(void *start, | ||
| size_t len); | ||
|
|
||
| } // namespace mooncake | ||
|
|
||
| #endif // MEMORY_LOCATION_H |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,67 @@ | ||
| // Copyright 2024 KVCache.AI | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| #include "memory_location.h" | ||
|
|
||
| namespace mooncake { | ||
|
|
||
| uintptr_t alignPage(uintptr_t address) { return address & ~(pagesize - 1); } | ||
|
|
||
| std::string genCpuNodeName(int node) { | ||
| if (node >= 0) return "cpu:" + std::to_string(node); | ||
|
|
||
| // use "*" when failed to get the numa node. | ||
| return "*"; | ||
| } | ||
|
|
||
| const std::vector<MemoryLocationEntry> getMemoryLocation(void *start, | ||
| size_t len) { | ||
| std::vector<MemoryLocationEntry> entries; | ||
|
|
||
| // start and end address may not be page aligned. | ||
| uintptr_t aligned_start = alignPage((uintptr_t)start); | ||
| int n = (uintptr_t(start) - aligned_start + len + pagesize - 1) / pagesize; | ||
| void **pages = (void **)malloc(sizeof(void *) * n); | ||
| int *status = (int *)malloc(sizeof(int) * n); | ||
|
|
||
| for (int i = 0; i < n; i++) { | ||
| pages[i] = (void *)((char *)aligned_start + i * pagesize); | ||
| } | ||
|
|
||
| int rc = numa_move_pages(0, n, pages, nullptr, status, 0); | ||
| if (rc != 0) { | ||
| PLOG(WARNING) << "Failed to get NUMA node, addr: " << start | ||
| << ", len: " << len; | ||
| entries.push_back({(uint64_t)start, len, "*"}); | ||
| return entries; | ||
| } | ||
|
|
||
| int node = status[0]; | ||
| uint64_t start_addr = (uint64_t)start; | ||
| uint64_t new_start_addr; | ||
| for (int i = 1; i < n; i++) { | ||
| if (status[i] != node) { | ||
| new_start_addr = alignPage((uint64_t)start) + i * pagesize; | ||
| entries.push_back({start_addr, size_t(new_start_addr - start_addr), | ||
| genCpuNodeName(node)}); | ||
| start_addr = new_start_addr; | ||
| node = status[i]; | ||
| } | ||
| } | ||
| entries.push_back( | ||
| {start_addr, (uint64_t)start + len - start_addr, genCpuNodeName(node)}); | ||
| return entries; | ||
| } | ||
|
|
||
| } // namespace mooncake |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
123 changes: 123 additions & 0 deletions
123
mooncake-transfer-engine/tests/memory_location_test.cpp
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,123 @@ | ||
| #include "memory_location.h" | ||
|
|
||
| #include <glog/logging.h> | ||
| #include <gtest/gtest.h> | ||
| #include <numa.h> | ||
| #include <numaif.h> | ||
| #include <sys/mman.h> | ||
|
|
||
| TEST(MemoryLocationTest, MallocSimpleNode0) { | ||
| int size = 4096 * 10; | ||
| void *addr = numa_alloc_onnode(size, 0); | ||
| ASSERT_NE(addr, nullptr); | ||
|
|
||
| auto entries = mooncake::getMemoryLocation(addr, size); | ||
| ASSERT_EQ(entries.size(), 1); | ||
|
|
||
| // check the memory location, no node before page fault | ||
| EXPECT_EQ(entries[0].start, (uint64_t)addr); | ||
| EXPECT_EQ(entries[0].location, "*"); | ||
| EXPECT_EQ(entries[0].len, size); | ||
|
|
||
| // trigger page fault | ||
| memset(addr, 0, size); | ||
|
|
||
| entries = mooncake::getMemoryLocation(addr, size); | ||
| ASSERT_EQ(entries.size(), 1); | ||
|
|
||
| // check the memory location, node 0 after page fault | ||
| EXPECT_EQ(entries[0].start, (uint64_t)addr); | ||
| EXPECT_EQ(entries[0].location, "cpu:0"); | ||
| EXPECT_EQ(entries[0].len, size); | ||
|
|
||
| numa_free(addr, size); | ||
| } | ||
|
|
||
| TEST(MemoryLocationTest, MallocSimpleNodeLargest) { | ||
| int node = numa_max_node(); | ||
| LOG(INFO) << "node: " << node; | ||
|
|
||
| std::string location = "cpu:" + std::to_string(node); | ||
|
|
||
| int size = 4096 * 10; | ||
| void *addr = numa_alloc_onnode(size, node); | ||
| ASSERT_NE(addr, nullptr); | ||
|
|
||
| // trigger page fault | ||
| memset(addr, 0, size); | ||
|
|
||
| auto entries = mooncake::getMemoryLocation(addr, size); | ||
| ASSERT_EQ(entries.size(), 1); | ||
|
|
||
| // check the memory location | ||
| EXPECT_EQ(entries[0].start, (uint64_t)addr); | ||
| EXPECT_EQ(entries[0].location, location); | ||
| EXPECT_EQ(entries[0].len, size); | ||
|
|
||
| numa_free(addr, size); | ||
| } | ||
|
|
||
| TEST(MemoryLocationTest, MallocMultipleNodes) { | ||
| int nodea = 0; | ||
| int nodeb = numa_max_node(); | ||
| LOG(INFO) << "node a: " << nodea << " node b: " << nodeb; | ||
|
|
||
| std::string locationa = "cpu:" + std::to_string(nodea); | ||
| std::string locationb = "cpu:" + std::to_string(nodeb); | ||
|
|
||
| int size = 4096 * 10; | ||
| void *addr = numa_alloc_onnode(size, nodea); | ||
| ASSERT_NE(addr, nullptr); | ||
| ASSERT_EQ((uint64_t)addr % 4096, 0); // page aligned | ||
|
|
||
| // trigger page fault | ||
| memset(addr, 0, size); | ||
|
|
||
| int rc; | ||
|
|
||
| // move first two pages & last one page to nodeb | ||
| void *pages[3] = {addr, (void *)((uint64_t)addr + 4096), | ||
| (void *)((uint64_t)addr + 4096 * 9)}; | ||
| int nodes[3] = {nodeb, nodeb, nodeb}; | ||
| int status[3]; | ||
| rc = numa_move_pages(0, 3, pages, nodes, status, MPOL_MF_MOVE); | ||
| if (rc != 0) { | ||
| PLOG(ERROR) << "numa_move_pages failed, rc: " << rc; | ||
| } | ||
| ASSERT_EQ(rc, 0); | ||
|
|
||
| // not page aligned | ||
| void *start = (void *)((uint64_t)addr + 1024 * 2); | ||
|
|
||
| auto entries = mooncake::getMemoryLocation(start, size - 1024 * 4); | ||
|
|
||
| if (nodea == nodeb) { | ||
| // only one numa node | ||
| ASSERT_EQ(entries.size(), 1); | ||
|
|
||
| // check the first memory location | ||
| EXPECT_EQ(entries[0].start, (uint64_t)start); | ||
| EXPECT_EQ(entries[0].location, locationa); | ||
| EXPECT_EQ(entries[0].len, size - 1024 * 4); | ||
|
|
||
| } else { | ||
| ASSERT_EQ(entries.size(), 3); | ||
|
|
||
| // check the first memory location | ||
| EXPECT_EQ(entries[0].start, (uint64_t)start); | ||
| EXPECT_EQ(entries[0].location, locationb); | ||
| EXPECT_EQ(entries[0].len, 4096 * 2 - 1024 * 2); | ||
|
|
||
| // check the second memory location | ||
| EXPECT_EQ(entries[1].start, (uint64_t)addr + 4096 * 2); | ||
| EXPECT_EQ(entries[1].location, locationa); | ||
| EXPECT_EQ(entries[1].len, 4096 * 7); | ||
|
|
||
| // check the third memory location | ||
| EXPECT_EQ(entries[2].start, (uint64_t)addr + 4096 * 9); | ||
| EXPECT_EQ(entries[2].location, locationb); | ||
| EXPECT_EQ(entries[2].len, 4096 - 1024 * 2); | ||
| } | ||
|
|
||
| numa_free(addr, size); | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.