forked from mooncake-track/Mooncake
Compare commits
397 Commits
P2P-Moonca
...
main
| Author | SHA1 | Date |
|---|---|---|
|
|
24e29df083 | |
|
|
db7bce3056 | |
|
|
a14e0b600a | |
|
|
7442626169 | |
|
|
e1d6d6f6f4 | |
|
|
255e287bc1 | |
|
|
cfea2cb0f5 | |
|
|
8a61d5b47c | |
|
|
741cf0adff | |
|
|
1e9fa36703 | |
|
|
5078873532 | |
|
|
38c3975138 | |
|
|
f9dd50c543 | |
|
|
0f22234d0b | |
|
|
c58d1f90b9 | |
|
|
ac53c874ba | |
|
|
952da65651 | |
|
|
cd67a36da0 | |
|
|
1fbe35c2fe | |
|
|
3e7c78de9b | |
|
|
62651b325b | |
|
|
32329c8356 | |
|
|
7d443489c2 | |
|
|
28464f3aee | |
|
|
e878cb2312 | |
|
|
a6cbc1a417 | |
|
|
9ce5d25292 | |
|
|
0a9c9937c7 | |
|
|
dc965d3121 | |
|
|
8d3beecb28 | |
|
|
15d99a3002 | |
|
|
020942121d | |
|
|
4a8684bc3e | |
|
|
c44a052944 | |
|
|
d85f9bd82f | |
|
|
c2573405bb | |
|
|
041ddb1794 | |
|
|
aadcc7daf1 | |
|
|
cfc5cf4c2a | |
|
|
3a69fa4b4d | |
|
|
01900be50b | |
|
|
a175e9aba8 | |
|
|
738b375338 | |
|
|
c0d07af568 | |
|
|
1a30e4e110 | |
|
|
37df986f64 | |
|
|
d3ace8fb20 | |
|
|
510cd4ee24 | |
|
|
c3d428b902 | |
|
|
9d35047fcb | |
|
|
a2207f464b | |
|
|
245cfb6e3a | |
|
|
e8e8e05d55 | |
|
|
4ae8c7e8bd | |
|
|
83ac8fc620 | |
|
|
4592c86c71 | |
|
|
b0c472dd72 | |
|
|
7e09bd1edb | |
|
|
9933470d39 | |
|
|
f6f7e51398 | |
|
|
ff7ceed623 | |
|
|
223405db96 | |
|
|
17d118ad03 | |
|
|
a7518f382d | |
|
|
30da37a554 | |
|
|
875b6c652e | |
|
|
6393672343 | |
|
|
6da25727f8 | |
|
|
8a6bd5f035 | |
|
|
c16d113107 | |
|
|
e9f288713a | |
|
|
da67fad95f | |
|
|
483c9970f7 | |
|
|
18c59612f7 | |
|
|
7d0978b9ca | |
|
|
5ff565dcab | |
|
|
a85800b83f | |
|
|
2e3b3c5769 | |
|
|
e426e350e4 | |
|
|
fdee379b21 | |
|
|
cd1ef9b536 | |
|
|
ba5f3ac2db | |
|
|
f6a7c6b798 | |
|
|
f0a1dc8717 | |
|
|
81100d93c9 | |
|
|
ccf68da040 | |
|
|
830d84127a | |
|
|
ba862909f5 | |
|
|
a56eb6a119 | |
|
|
c78d4e9512 | |
|
|
9351a54014 | |
|
|
6105f9cda8 | |
|
|
d32bbc4048 | |
|
|
aebd5eff4a | |
|
|
b9a593c5e8 | |
|
|
cf835bed51 | |
|
|
4b3d44f39f | |
|
|
41d40dabd7 | |
|
|
7113245ee4 | |
|
|
77ae1034a6 | |
|
|
8f39661a31 | |
|
|
4d163b026d | |
|
|
e1220bb309 | |
|
|
af733e1209 | |
|
|
418bb774ad | |
|
|
50bc7cf976 | |
|
|
5821a19547 | |
|
|
fad70b90d9 | |
|
|
5002ae6ac1 | |
|
|
49ec245620 | |
|
|
1818be8594 | |
|
|
4c19802458 | |
|
|
ece427ed2a | |
|
|
86fef8cbeb | |
|
|
6f0fbbfd1f | |
|
|
ef34131c7d | |
|
|
11fd29abd2 | |
|
|
688849bedb | |
|
|
bc9947873e | |
|
|
ea2824a474 | |
|
|
98d60c59b7 | |
|
|
79266ffb4d | |
|
|
cc2ecf6b1a | |
|
|
75438df15f | |
|
|
692ccdabc8 | |
|
|
c0f190d6f2 | |
|
|
8c647a7b85 | |
|
|
fdb7040787 | |
|
|
d2bff9b9a8 | |
|
|
d7f82d81e4 | |
|
|
d004bc0a23 | |
|
|
79e0ae1140 | |
|
|
ad1bdcf74b | |
|
|
59c9f90151 | |
|
|
8833f4ec3c | |
|
|
c424824eea | |
|
|
f40407eb59 | |
|
|
41b4fbaff0 | |
|
|
30e9fbba2d | |
|
|
84686a6841 | |
|
|
7136c45787 | |
|
|
e1605e7dd2 | |
|
|
47b016d00b | |
|
|
335de42cd2 | |
|
|
99562454a7 | |
|
|
22ecbd794e | |
|
|
a811c288ea | |
|
|
f55a59f83e | |
|
|
3a36980661 | |
|
|
fd73547e26 | |
|
|
71886f2792 | |
|
|
bd18da6cdd | |
|
|
3f249c5a72 | |
|
|
5ad7076415 | |
|
|
bd4a183ad6 | |
|
|
1aeeffcf92 | |
|
|
0c310237d7 | |
|
|
8e40563711 | |
|
|
911fd2d3bb | |
|
|
9bd773e6fe | |
|
|
3680022876 | |
|
|
13e9845a21 | |
|
|
4e608cd80c | |
|
|
b20fe888be | |
|
|
ad8c3c2355 | |
|
|
68c1da7c62 | |
|
|
e7f6bad7fa | |
|
|
e43dd7dec1 | |
|
|
8887c354bf | |
|
|
97987eca42 | |
|
|
ff6274ff16 | |
|
|
3194a7bd59 | |
|
|
d864d122c5 | |
|
|
9968d05452 | |
|
|
1bc3c651bc | |
|
|
5a47ce013a | |
|
|
93493562e8 | |
|
|
ea59594071 | |
|
|
82d21747f6 | |
|
|
a8fde4c767 | |
|
|
4c4f1e31d7 | |
|
|
34dacaec18 | |
|
|
834c416097 | |
|
|
7a96121aa7 | |
|
|
64627804f1 | |
|
|
59ae11fa0e | |
|
|
2456b0c32b | |
|
|
ea8f1b902b | |
|
|
99e9ecee6d | |
|
|
41b1969e47 | |
|
|
6ca7375136 | |
|
|
0a220f67e5 | |
|
|
3878b98768 | |
|
|
a402dc715f | |
|
|
33cb4efa7a | |
|
|
f6c2c78982 | |
|
|
cbc32b9c20 | |
|
|
4dee57eedd | |
|
|
5fe8e21e7c | |
|
|
9e4e9490b2 | |
|
|
752ff3d5f8 | |
|
|
2a9e64dd74 | |
|
|
b5eddeb721 | |
|
|
94013a9d36 | |
|
|
7af8470f07 | |
|
|
8f054b6e5f | |
|
|
ba8fd0985b | |
|
|
b843433b04 | |
|
|
7386f21f96 | |
|
|
a05cfc2fae | |
|
|
09a91318ac | |
|
|
c0a1e6570d | |
|
|
58b933ca4c | |
|
|
0bc07265d1 | |
|
|
1a7df2624f | |
|
|
492284ee93 | |
|
|
40cebe30b3 | |
|
|
b381320156 | |
|
|
ac732acf6d | |
|
|
1c9da7167f | |
|
|
1abb593ad8 | |
|
|
8ed67617e8 | |
|
|
4ff4dc1e27 | |
|
|
53e750b2b5 | |
|
|
99da058d61 | |
|
|
30e48cfd93 | |
|
|
845ccc43dc | |
|
|
3897e1b53e | |
|
|
95e4c489f2 | |
|
|
b847d5c41f | |
|
|
6e40f95935 | |
|
|
140e9c2210 | |
|
|
b1fbcd182d | |
|
|
330863e389 | |
|
|
8d5adc21bf | |
|
|
3e67c9e443 | |
|
|
1eec5a5b6b | |
|
|
dc432620cf | |
|
|
45305495a3 | |
|
|
393ad48967 | |
|
|
38d2fbb18b | |
|
|
f21b3791ac | |
|
|
b84aca829a | |
|
|
3cdbbdc6d0 | |
|
|
feb39a91bc | |
|
|
1f61d40c7d | |
|
|
8c60e6c4cb | |
|
|
3218467e77 | |
|
|
fd2da4fae5 | |
|
|
4136d2b73b | |
|
|
ab4af4f2da | |
|
|
83384e9c0c | |
|
|
c724318217 | |
|
|
e87fb51e3f | |
|
|
bf46a3702c | |
|
|
f236abc4b7 | |
|
|
4779cb0e7d | |
|
|
a00f757424 | |
|
|
24f7dd4d9c | |
|
|
4c620377a4 | |
|
|
62d28faddb | |
|
|
5c5f18c817 | |
|
|
e5d9e7a503 | |
|
|
adfa9374ab | |
|
|
01d1c6a8b3 | |
|
|
17b9df8fb8 | |
|
|
61536b95b6 | |
|
|
680ae5c8df | |
|
|
c6ca713ff4 | |
|
|
e9f523cb8c | |
|
|
fd84d4784a | |
|
|
2e601a1bb1 | |
|
|
020a906295 | |
|
|
ff59e990bf | |
|
|
b4a9b284c1 | |
|
|
8f438716ec | |
|
|
5b9243280b | |
|
|
9b242111aa | |
|
|
53f4f9ccb4 | |
|
|
589684e9a2 | |
|
|
5cedb8feb2 | |
|
|
9bf48112d5 | |
|
|
4f76f12b3c | |
|
|
45bfd19e11 | |
|
|
17efcb1719 | |
|
|
23181faa51 | |
|
|
20494ee644 | |
|
|
500e8e71b6 | |
|
|
47633c4661 | |
|
|
67d0afc4a9 | |
|
|
c8c7b0a025 | |
|
|
56241b8482 | |
|
|
17a1102c17 | |
|
|
6406edae02 | |
|
|
8383b93e5e | |
|
|
46c451b005 | |
|
|
8284f3d893 | |
|
|
5920954661 | |
|
|
0b8e1d63c3 | |
|
|
63778abf88 | |
|
|
2c25a6bd3b | |
|
|
47501380fe | |
|
|
a233cd2ce3 | |
|
|
b5e610a2eb | |
|
|
2937342e50 | |
|
|
c36bf8bf33 | |
|
|
b880bca594 | |
|
|
c6a151d5eb | |
|
|
2a37b91528 | |
|
|
6f6c95e480 | |
|
|
24c632f84a | |
|
|
1e32df7054 | |
|
|
619b03a954 | |
|
|
39d7baa94e | |
|
|
b68e7fcd98 | |
|
|
c9d64c91c4 | |
|
|
45a0381329 | |
|
|
7c6508a1b2 | |
|
|
3bc66839b0 | |
|
|
b064a662e4 | |
|
|
27045681f7 | |
|
|
9523563c22 | |
|
|
0267e354ec | |
|
|
65019f94ab | |
|
|
96b82a012b | |
|
|
d50eeaca69 | |
|
|
e43dce6f72 | |
|
|
05247e0c40 | |
|
|
23380988c8 | |
|
|
8510cc8b45 | |
|
|
f19033c03e | |
|
|
3df2dfc4a7 | |
|
|
914ed1d414 | |
|
|
af3c2a8908 | |
|
|
731f0ea2c7 | |
|
|
971c99ddca | |
|
|
fc2f688661 | |
|
|
15930ac8ad | |
|
|
b98b7e6655 | |
|
|
c921c02087 | |
|
|
30ca0868d5 | |
|
|
e5c844f6b5 | |
|
|
9afaf5dee4 | |
|
|
b6743066f6 | |
|
|
e89652985f | |
|
|
257eb718f3 | |
|
|
3cf3d5a34d | |
|
|
617c50b304 | |
|
|
5b3fd4deff | |
|
|
03a55034f0 | |
|
|
ba2bbd7f6c | |
|
|
abb429f935 | |
|
|
4e3c1555d3 | |
|
|
4fce690446 | |
|
|
086f1e6821 | |
|
|
2b802a62a2 | |
|
|
fdf08f108b | |
|
|
640bb207a6 | |
|
|
d9a915b65a | |
|
|
e1007e7bcd | |
|
|
5738f80752 | |
|
|
6fa22a23dc | |
|
|
0280d8337b | |
|
|
b8dbc2654d | |
|
|
f357e3ded8 | |
|
|
4591728264 | |
|
|
c755916ebc | |
|
|
d088b40501 | |
|
|
ea17a5a723 | |
|
|
2749bc3225 | |
|
|
f98cf979cb | |
|
|
f8ad6d2f1a | |
|
|
69663d117b | |
|
|
5dec16615d | |
|
|
060f2bda53 | |
|
|
1719043de4 | |
|
|
fe3eee8f13 | |
|
|
e0ddd91629 | |
|
|
e9a9241f7e | |
|
|
f06a2ab8e5 | |
|
|
9217f29117 | |
|
|
47d928c02b | |
|
|
70488aae3f | |
|
|
0b149e6c13 | |
|
|
7d196747a2 | |
|
|
dc40512726 | |
|
|
a6827f2bd4 | |
|
|
9f15ba4c58 | |
|
|
823064aa76 | |
|
|
ed2fa9dd22 | |
|
|
0807ca99bb | |
|
|
ddc9bb9d87 | |
|
|
f7f65aa140 | |
|
|
fdfd74818b | |
|
|
b159da3de5 | |
|
|
0815e450cd | |
|
|
2719179bcb |
|
|
@ -0,0 +1,512 @@
|
|||
<!-- Let's perfect this skill together. -->
|
||||
|
||||
# Mooncake Python API Skill
|
||||
|
||||
Use this skill to help users work with Mooncake Python APIs for distributed storage and high-performance data transfer.
|
||||
|
||||
## When to Use This Skill
|
||||
|
||||
Use this skill when users ask about:
|
||||
- Using Mooncake Store for distributed KV cache storage
|
||||
- Using Transfer Engine for RDMA/TCP data transfers
|
||||
- Setting up Mooncake services (master, metadata server)
|
||||
- Working with PyTorch tensors in Mooncake Store
|
||||
- Zero-copy operations and buffer management
|
||||
- Batch operations and replication configuration
|
||||
- Mooncake EP (Expert Parallelism) and Mooncake Backend
|
||||
- Troubleshooting Mooncake Python API issues
|
||||
|
||||
## Core Components
|
||||
|
||||
### 1. Mooncake Store (Distributed KV Cache)
|
||||
|
||||
**Import:**
|
||||
```python
|
||||
from mooncake.store import MooncakeDistributedStore, ReplicateConfig
|
||||
```
|
||||
|
||||
**Basic Setup:**
|
||||
```python
|
||||
store = MooncakeDistributedStore()
|
||||
store.setup(
|
||||
"localhost", # local_hostname
|
||||
"http://localhost:8080/metadata", # metadata_server
|
||||
512*1024*1024, # global_segment_size (512MB)
|
||||
128*1024*1024, # local_buffer_size (128MB)
|
||||
"tcp", # protocol ("tcp" or "rdma")
|
||||
"", # rdma_devices (empty for auto-select)
|
||||
"localhost:50051" # master_server_address
|
||||
)
|
||||
```
|
||||
|
||||
**Common Operations:**
|
||||
```python
|
||||
# Put/Get
|
||||
store.put("key", b"value")
|
||||
data = store.get("key")
|
||||
|
||||
# Batch operations
|
||||
store.put_batch(["key1", "key2"], [b"val1", b"val2"])
|
||||
values = store.get_batch(["key1", "key2"])
|
||||
|
||||
# Check existence
|
||||
exists = store.is_exist("key") # Returns 1 (exists), 0 (not exists), -1 (error)
|
||||
|
||||
# Remove
|
||||
store.remove("key")
|
||||
store.remove_by_regex("^prefix_.*")
|
||||
store.remove_all()
|
||||
|
||||
# Cleanup
|
||||
store.close()
|
||||
```
|
||||
|
||||
**Zero-Copy Operations (Advanced):**
|
||||
```python
|
||||
import numpy as np
|
||||
|
||||
# Create and register buffer
|
||||
buffer = np.zeros(100*1024*1024, dtype=np.uint8)
|
||||
buffer_ptr = buffer.ctypes.data
|
||||
store.register_buffer(buffer_ptr, buffer.nbytes)
|
||||
|
||||
# Zero-copy put
|
||||
store.put_from("key", buffer_ptr, buffer.nbytes)
|
||||
|
||||
# Zero-copy get
|
||||
recv_buffer = np.empty(100*1024*1024, dtype=np.uint8)
|
||||
recv_ptr = recv_buffer.ctypes.data
|
||||
store.register_buffer(recv_ptr, recv_buffer.nbytes)
|
||||
bytes_read = store.get_into("key", recv_ptr, recv_buffer.nbytes)
|
||||
|
||||
# Cleanup
|
||||
store.unregister_buffer(buffer_ptr)
|
||||
store.unregister_buffer(recv_ptr)
|
||||
```
|
||||
|
||||
**PyTorch Tensor Operations:**
|
||||
```python
|
||||
import torch
|
||||
|
||||
# Simple tensor operations
|
||||
tensor = torch.randn(100, 100)
|
||||
store.put_tensor("my_tensor", tensor)
|
||||
retrieved = store.get_tensor("my_tensor")
|
||||
|
||||
# Batch tensor operations
|
||||
tensors = [torch.randn(100, 100) for _ in range(3)]
|
||||
store.batch_put_tensor(["t1", "t2", "t3"], tensors)
|
||||
retrieved_tensors = store.batch_get_tensor(["t1", "t2", "t3"])
|
||||
|
||||
# Tensor Parallelism (TP) support
|
||||
store.put_tensor_with_tp("model_weights", tensor, tp_rank=0, tp_size=4, split_dim=0)
|
||||
shard = store.get_tensor_with_tp("model_weights", tp_rank=0, tp_size=4)
|
||||
```
|
||||
|
||||
**Replication Configuration:**
|
||||
```python
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = 3 # Number of replicas
|
||||
config.with_soft_pin = True # Keep in memory longer
|
||||
config.preferred_segment = "host:port" # Preferred location
|
||||
|
||||
store.put("key", b"value", config)
|
||||
```
|
||||
|
||||
### 2. Transfer Engine (High-Performance Data Transfer)
|
||||
|
||||
**Import:**
|
||||
```python
|
||||
from mooncake.engine import TransferEngine, TransferOpcode, TransferNotify
|
||||
```
|
||||
|
||||
**Basic Setup:**
|
||||
```python
|
||||
engine = TransferEngine()
|
||||
engine.initialize(
|
||||
"127.0.0.1:12345", # local_hostname
|
||||
"127.0.0.1:2379", # metadata_server (or "etcd://...")
|
||||
"tcp", # protocol ("tcp" or "rdma")
|
||||
"" # device_name (empty for all devices)
|
||||
)
|
||||
```
|
||||
|
||||
**Buffer Management:**
|
||||
```python
|
||||
# Allocate managed buffer
|
||||
buffer_size = 1024 * 1024 # 1MB
|
||||
buffer_addr = engine.allocate_managed_buffer(buffer_size)
|
||||
|
||||
# Write/read bytes
|
||||
data = b"Hello, Transfer Engine!"
|
||||
engine.write_bytes_to_buffer(buffer_addr, data, len(data))
|
||||
read_data = engine.read_bytes_from_buffer(buffer_addr, len(data))
|
||||
|
||||
# Free buffer
|
||||
engine.free_managed_buffer(buffer_addr, buffer_size)
|
||||
```
|
||||
|
||||
**Data Transfer Operations:**
|
||||
```python
|
||||
# Synchronous write
|
||||
result = engine.transfer_sync_write(
|
||||
"target_host:port", # target_hostname
|
||||
local_buffer_addr, # buffer
|
||||
remote_buffer_addr, # peer_buffer_address
|
||||
data_length # length
|
||||
)
|
||||
|
||||
# Synchronous read
|
||||
result = engine.transfer_sync_read(
|
||||
"target_host:port",
|
||||
local_buffer_addr,
|
||||
remote_buffer_addr,
|
||||
data_length
|
||||
)
|
||||
|
||||
# Asynchronous write
|
||||
batch_id = engine.transfer_submit_write(
|
||||
"target_host:port",
|
||||
local_buffer_addr,
|
||||
remote_buffer_addr,
|
||||
data_length
|
||||
)
|
||||
|
||||
# Check status
|
||||
status = engine.transfer_check_status(batch_id)
|
||||
# Returns: 1 (completed), 0 (in progress), -1 (failed), -2 (timeout)
|
||||
```
|
||||
|
||||
**Batch Transfer Operations:**
|
||||
```python
|
||||
# Batch synchronous write
|
||||
local_addrs = [addr1, addr2, addr3]
|
||||
remote_addrs = [remote1, remote2, remote3]
|
||||
lengths = [len1, len2, len3]
|
||||
|
||||
result = engine.batch_transfer_sync_write(
|
||||
"target_host:port",
|
||||
local_addrs,
|
||||
remote_addrs,
|
||||
lengths
|
||||
)
|
||||
|
||||
# Batch asynchronous operations
|
||||
batch_id = engine.batch_transfer_async_write(
|
||||
"target_host:port",
|
||||
local_addrs,
|
||||
remote_addrs,
|
||||
lengths
|
||||
)
|
||||
|
||||
# Wait for completion
|
||||
result = engine.get_batch_transfer_status([batch_id])
|
||||
```
|
||||
|
||||
**Memory Registration (for RDMA):**
|
||||
```python
|
||||
import numpy as np
|
||||
|
||||
buffer = np.ones(1024*1024, dtype=np.uint8)
|
||||
buffer_ptr = buffer.ctypes.data
|
||||
buffer_size = buffer.nbytes
|
||||
|
||||
# Register memory
|
||||
engine.register_memory(buffer_ptr, buffer_size)
|
||||
|
||||
# Use buffer for transfers...
|
||||
|
||||
# Unregister when done
|
||||
engine.unregister_memory(buffer_ptr)
|
||||
```
|
||||
|
||||
### 3. Mooncake EP & Backend (Expert Parallelism)
|
||||
|
||||
**Mooncake Backend (Fault-Tolerant Collectives):**
|
||||
```python
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from mooncake import pg
|
||||
|
||||
# Initialize with fault tolerance
|
||||
active_ranks = torch.ones((world_size,), dtype=torch.int32, device="cuda")
|
||||
dist.init_process_group(
|
||||
backend="mooncake",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
pg_options=pg.MooncakeBackendOptions(active_ranks),
|
||||
)
|
||||
|
||||
# Use standard PyTorch distributed APIs
|
||||
dist.all_gather(...)
|
||||
dist.all_reduce(...)
|
||||
|
||||
# Check for failures
|
||||
assert active_ranks.all() # Verify no ranks are broken
|
||||
```
|
||||
|
||||
**Mooncake EP (Expert Parallelism):**
|
||||
```python
|
||||
from mooncake.mooncake_ep_buffer import Buffer
|
||||
import torch.distributed as dist
|
||||
|
||||
# Calculate buffer size
|
||||
num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint(
|
||||
num_max_dispatch_tokens_per_rank=1024,
|
||||
hidden=4096,
|
||||
num_ranks=8,
|
||||
num_experts=64
|
||||
)
|
||||
|
||||
# Create buffer (must be Mooncake Backend process group)
|
||||
buffer = Buffer(group=dist.group.WORLD, num_ep_buffer_bytes=num_ep_buffer_bytes)
|
||||
|
||||
# Dispatch/combine operations
|
||||
active_ranks = torch.ones((num_ranks,), dtype=torch.int32, device="cuda")
|
||||
buffer.dispatch(..., active_ranks=active_ranks, timeout_us=1000000)
|
||||
buffer.combine(..., active_ranks=active_ranks, timeout_us=1000000)
|
||||
```
|
||||
|
||||
## Starting Services
|
||||
|
||||
### Start Master Service (with HTTP metadata server)
|
||||
```bash
|
||||
mooncake_master \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_host=0.0.0.0 \
|
||||
--http_metadata_server_port=8080 \
|
||||
--default_kv_lease_ttl=5000
|
||||
```
|
||||
|
||||
### Using External etcd (Production)
|
||||
```bash
|
||||
# Start etcd
|
||||
etcd --listen-client-urls http://0.0.0.0:2379 \
|
||||
--advertise-client-urls http://0.0.0.0:2379
|
||||
|
||||
# Start master
|
||||
mooncake_master --default_kv_lease_ttl=5000
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
### Transfer Engine
|
||||
- `MC_METADATA_SERVER`: Metadata server URL
|
||||
- `MC_FORCE_TCP`: Force TCP transport (set to "true")
|
||||
- `MC_LOG_LEVEL`: Logging level (0=INFO, 1=WARNING, 2=ERROR)
|
||||
- `MC_MS_AUTO_DISC`: Enable RDMA device auto-discovery (set to "1")
|
||||
- `MC_MS_FILTERS`: Filter RDMA devices (e.g., "mlx5_0,mlx5_2")
|
||||
- `MC_TRANSFER_TIMEOUT`: Transfer timeout in seconds (default: 30)
|
||||
|
||||
### Mooncake Store
|
||||
- `MC_STORE_CLUSTER_ID`: Cluster identifier (default: "mooncake")
|
||||
- `MC_STORE_USE_HUGEPAGE`: Enable hugepage support
|
||||
- `MC_STORE_MEMCPY`: Enable local memcpy optimization (set to "1")
|
||||
- `MC_STORE_CLIENT_METRIC`: Enable client metrics (enabled by default)
|
||||
- `MC_YLT_LOG_LEVEL`: Log level (trace/debug/info/warn/error/critical)
|
||||
|
||||
## Common Patterns
|
||||
|
||||
### Pattern 1: Simple KV Store
|
||||
```python
|
||||
from mooncake.store import MooncakeDistributedStore
|
||||
|
||||
store = MooncakeDistributedStore()
|
||||
store.setup("localhost", "http://localhost:8080/metadata",
|
||||
512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051")
|
||||
|
||||
# Store and retrieve
|
||||
store.put("config", b'{"model": "llama-7b"}')
|
||||
config = store.get("config")
|
||||
|
||||
store.close()
|
||||
```
|
||||
|
||||
### Pattern 2: High-Performance Tensor Storage
|
||||
```python
|
||||
import torch
|
||||
from mooncake.store import MooncakeDistributedStore, ReplicateConfig
|
||||
|
||||
store = MooncakeDistributedStore()
|
||||
store.setup("localhost", "http://localhost:8080/metadata",
|
||||
512*1024*1024, 128*1024*1024, "rdma", "mlx5_0", "localhost:50051")
|
||||
|
||||
# Configure replication
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = 2
|
||||
config.with_soft_pin = True
|
||||
|
||||
# Store tensor with replication
|
||||
tensor = torch.randn(1000, 1000)
|
||||
store.put_tensor("weights", tensor, config)
|
||||
|
||||
# Retrieve
|
||||
retrieved = store.get_tensor("weights")
|
||||
|
||||
store.close()
|
||||
```
|
||||
|
||||
### Pattern 3: Zero-Copy Batch Operations
|
||||
```python
|
||||
import numpy as np
|
||||
from mooncake.store import MooncakeDistributedStore
|
||||
|
||||
store = MooncakeDistributedStore()
|
||||
store.setup("localhost", "http://localhost:8080/metadata",
|
||||
512*1024*1024, 16*1024*1024, "rdma", "", "localhost:50051")
|
||||
|
||||
# Prepare buffers
|
||||
num_buffers = 10
|
||||
buffers = [np.random.randn(1024*1024).astype(np.float32) for _ in range(num_buffers)]
|
||||
buffer_ptrs = [buf.ctypes.data for buf in buffers]
|
||||
sizes = [buf.nbytes for buf in buffers]
|
||||
|
||||
# Register all buffers
|
||||
for ptr, size in zip(buffer_ptrs, sizes):
|
||||
store.register_buffer(ptr, size)
|
||||
|
||||
# Batch put
|
||||
keys = [f"tensor_{i}" for i in range(num_buffers)]
|
||||
results = store.batch_put_from(keys, buffer_ptrs, sizes)
|
||||
|
||||
# Batch get
|
||||
recv_buffers = [np.empty(1024*1024, dtype=np.float32) for _ in range(num_buffers)]
|
||||
recv_ptrs = [buf.ctypes.data for buf in recv_buffers]
|
||||
|
||||
for ptr, size in zip(recv_ptrs, sizes):
|
||||
store.register_buffer(ptr, size)
|
||||
|
||||
results = store.batch_get_into(keys, recv_ptrs, sizes)
|
||||
|
||||
# Cleanup
|
||||
for ptr in buffer_ptrs + recv_ptrs:
|
||||
store.unregister_buffer(ptr)
|
||||
|
||||
store.close()
|
||||
```
|
||||
|
||||
### Pattern 4: Transfer Engine Direct Transfer
|
||||
```python
|
||||
from mooncake.engine import TransferEngine
|
||||
import numpy as np
|
||||
|
||||
# Setup engines on both nodes
|
||||
engine = TransferEngine()
|
||||
engine.initialize("127.0.0.1:12345", "127.0.0.1:2379", "tcp", "")
|
||||
|
||||
# Allocate and register buffer
|
||||
buffer = np.ones(1024*1024, dtype=np.uint8)
|
||||
buffer_ptr = buffer.ctypes.data
|
||||
engine.register_memory(buffer_ptr, buffer.nbytes)
|
||||
|
||||
# Get remote buffer address (from peer via metadata exchange)
|
||||
remote_addr = engine.get_first_buffer_address("target_host:port")
|
||||
|
||||
# Transfer data
|
||||
data = b"Hello from Transfer Engine!"
|
||||
engine.write_bytes_to_buffer(buffer_ptr, data, len(data))
|
||||
result = engine.transfer_sync_write("target_host:port", buffer_ptr, remote_addr, len(data))
|
||||
|
||||
# Cleanup
|
||||
engine.unregister_memory(buffer_ptr)
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
|
||||
All methods return status codes:
|
||||
- `0`: Success
|
||||
- Negative values: Error codes
|
||||
|
||||
Common checks:
|
||||
```python
|
||||
# Store operations
|
||||
result = store.put("key", b"value")
|
||||
if result != 0:
|
||||
print(f"Put failed with error code: {result}")
|
||||
|
||||
# Existence check
|
||||
exists = store.is_exist("key")
|
||||
if exists == 1:
|
||||
print("Key exists")
|
||||
elif exists == 0:
|
||||
print("Key not found")
|
||||
else:
|
||||
print("Error checking existence")
|
||||
|
||||
# Transfer operations
|
||||
result = engine.transfer_sync_write(...)
|
||||
if result == 0:
|
||||
print("Transfer successful")
|
||||
else:
|
||||
print(f"Transfer failed with code: {result}")
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Connection Issues
|
||||
```python
|
||||
# Force TCP for testing without RDMA
|
||||
import os
|
||||
os.environ["MC_FORCE_TCP"] = "true"
|
||||
|
||||
# Enable verbose logging
|
||||
os.environ["MC_LOG_LEVEL"] = "0"
|
||||
os.environ["MC_YLT_LOG_LEVEL"] = "debug"
|
||||
```
|
||||
|
||||
### Memory Issues
|
||||
```python
|
||||
# Check buffer registration before zero-copy ops
|
||||
result = store.register_buffer(buffer_ptr, size)
|
||||
if result != 0:
|
||||
raise RuntimeError(f"Failed to register buffer: {result}")
|
||||
```
|
||||
|
||||
### Service Connectivity
|
||||
```bash
|
||||
# Check master is running
|
||||
curl http://localhost:50051
|
||||
|
||||
# Check metadata server
|
||||
curl http://localhost:8080/metadata
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
|
||||
1. **Always close stores**: Call `store.close()` when done
|
||||
2. **Register buffers for zero-copy**: Required for RDMA operations
|
||||
3. **Use batch operations**: Better throughput for multiple operations
|
||||
4. **Configure replication**: Use `ReplicateConfig` for important data
|
||||
5. **Use soft pinning**: For frequently accessed objects
|
||||
6. **Choose protocol wisely**: TCP for dev/test, RDMA for production
|
||||
7. **Monitor leases**: Objects have TTL, renew if needed
|
||||
8. **Handle errors**: Check return codes and handle failures
|
||||
|
||||
## Quick Reference
|
||||
|
||||
### Mooncake Store Methods
|
||||
- `setup()`: Initialize store
|
||||
- `put()`, `get()`: Basic operations
|
||||
- `put_batch()`, `get_batch()`: Batch operations
|
||||
- `put_from()`, `get_into()`: Zero-copy operations
|
||||
- `put_tensor()`, `get_tensor()`: PyTorch tensors
|
||||
- `register_buffer()`, `unregister_buffer()`: Buffer management
|
||||
- `is_exist()`, `remove()`: Metadata operations
|
||||
- `close()`: Cleanup
|
||||
|
||||
### Transfer Engine Methods
|
||||
- `initialize()`: Setup engine
|
||||
- `allocate_managed_buffer()`, `free_managed_buffer()`: Buffer allocation
|
||||
- `transfer_sync_write()`, `transfer_sync_read()`: Synchronous transfers
|
||||
- `transfer_submit_write()`, `transfer_check_status()`: Async transfers
|
||||
- `batch_transfer_sync_write()`, `batch_transfer_sync_read()`: Batch sync
|
||||
- `batch_transfer_async_write()`, `get_batch_transfer_status()`: Batch async
|
||||
- `register_memory()`, `unregister_memory()`: Memory registration
|
||||
- `write_bytes_to_buffer()`, `read_bytes_from_buffer()`: Buffer I/O
|
||||
|
||||
## Documentation Links
|
||||
|
||||
- Full API Reference: https://kvcache-ai.github.io/Mooncake/
|
||||
- Mooncake Store: docs/source/python-api-reference/mooncake-store.md
|
||||
- Transfer Engine: docs/source/python-api-reference/transfer-engine.md
|
||||
- EP Backend: docs/source/python-api-reference/ep-backend.md
|
||||
|
|
@ -0,0 +1,266 @@
|
|||
---
|
||||
name: mooncake-ci-local
|
||||
description: Run Mooncake CI test suite locally — maps GitHub Actions CI steps to local commands. Use this skill whenever the user wants to run tests locally, reproduce a CI failure, check if their changes break tests, or run any subset of the CI test suite (C++ unit tests via ctest, Python integration tests, code format checks, or the full test pipeline). Trigger on phrases like "run tests", "run CI locally", "reproduce CI failure", "check my changes", "test before PR", "run ctest", "run python tests", "run all tests".
|
||||
---
|
||||
|
||||
# Mooncake CI Local Test Runner
|
||||
|
||||
You help users run the Mooncake CI test suite locally. The CI has three test layers. Map what the user wants to the right layer, check prerequisites, and run the tests.
|
||||
|
||||
## CI Test Layers
|
||||
|
||||
### Layer 1 — C++ Unit Tests (ctest)
|
||||
**CI equivalent:** `build` job in `ci.yml` — "Test (in build env) with coverage"
|
||||
|
||||
**Prerequisite services:**
|
||||
```bash
|
||||
# 1. etcd (port 2379)
|
||||
etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 &
|
||||
sleep 2
|
||||
etcdctl --endpoints=http://127.0.0.1:2379 endpoint health # verify
|
||||
|
||||
# 2. HTTP metadata server (port 8080)
|
||||
cd mooncake-transfer-engine/example/http-metadata-server-python
|
||||
pip install aiohttp
|
||||
python ./bootstrap_server.py &
|
||||
cd -
|
||||
```
|
||||
|
||||
**Run:**
|
||||
```bash
|
||||
cd build
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
|
||||
```
|
||||
|
||||
**Run specific test:**
|
||||
```bash
|
||||
cd build
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -R <test_name_pattern> --output-on-failure
|
||||
# List all available tests: ctest -N
|
||||
```
|
||||
|
||||
### Layer 2 — Python Integration Tests
|
||||
**CI equivalent:** `test-wheel-ubuntu` job — `run_tests.sh`
|
||||
|
||||
**Prerequisite:** Mooncake wheel must be installed (either via `pip install` or via `make install` after build).
|
||||
|
||||
**Check install:**
|
||||
```bash
|
||||
python -c "import mooncake; print('OK')"
|
||||
which mooncake_master # must NOT be /usr/local/bin (must be from Python package)
|
||||
```
|
||||
|
||||
**Run full suite:**
|
||||
```bash
|
||||
# Start metadata server first
|
||||
mooncake_http_metadata_server --port 8080 &
|
||||
sleep 1
|
||||
|
||||
cd mooncake-wheel/tests
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 MC_FORCE_TCP=true \
|
||||
bash ../../scripts/run_tests.sh
|
||||
```
|
||||
|
||||
**Individual Python tests** (all require metadata server + mooncake_master on port 50051):
|
||||
```bash
|
||||
# Setup shared services
|
||||
mooncake_http_metadata_server --port 8080 &
|
||||
mooncake_master --default_kv_lease_ttl=500 &
|
||||
sleep 2
|
||||
|
||||
cd mooncake-wheel/tests
|
||||
export MC_METADATA_SERVER=http://127.0.0.1:8080/metadata
|
||||
export DEFAULT_KV_LEASE_TTL=500
|
||||
export MC_FORCE_TCP=true
|
||||
|
||||
# Pick any test:
|
||||
python test_distributed_object_store.py
|
||||
python test_replicated_distributed_object_store.py
|
||||
python test_put_get_tensor.py # requires torch + numpy
|
||||
python test_safetensor_functions.py # requires safetensors
|
||||
python test_dummy_client.py
|
||||
python test_cli.py
|
||||
python test_distributed_object_store_cxl.py # requires CXL build
|
||||
```
|
||||
|
||||
**Transfer engine tests specifically:**
|
||||
```bash
|
||||
cd mooncake-wheel/tests
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true python transfer_engine_target.py &
|
||||
TARGET_PID=$!
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true python transfer_engine_initiator_test.py
|
||||
kill $TARGET_PID
|
||||
```
|
||||
|
||||
**Scripts-based tests** (from `test-wheel-ubuntu` job):
|
||||
```bash
|
||||
# Tensor API perf test
|
||||
export MOONCAKE_MASTER="127.0.0.1:50051"
|
||||
export MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata"
|
||||
export MOONCAKE_PROTOCOL="tcp"
|
||||
export LOCAL_HOSTNAME="127.0.0.1"
|
||||
python scripts/test_tensor_api.py -n 1
|
||||
python scripts/test_async_store.py
|
||||
python scripts/test_copy_move_api.py
|
||||
```
|
||||
|
||||
### Layer 3 — Static Checks (no services needed)
|
||||
**CI equivalent:** `clang-format` and `spell-check` jobs
|
||||
|
||||
**Code format (changed files vs main):**
|
||||
```bash
|
||||
./scripts/code_format.sh --check --base origin/main
|
||||
# Auto-fix:
|
||||
./scripts/code_format.sh --base origin/main
|
||||
```
|
||||
|
||||
**Spell check:**
|
||||
```bash
|
||||
# Requires typos tool: cargo install typos-cli
|
||||
typos
|
||||
```
|
||||
|
||||
**Pre-commit (runs all hooks):**
|
||||
```bash
|
||||
pip install pre-commit
|
||||
pre-commit run --all-files
|
||||
# Or just on staged files:
|
||||
pre-commit run
|
||||
```
|
||||
|
||||
## Build Configurations (from CI)
|
||||
|
||||
If the user needs to build first, here are the CI-equivalent cmake flags:
|
||||
|
||||
**Standard build with coverage (mirrors `build` job):**
|
||||
```bash
|
||||
mkdir build && cd build
|
||||
cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DCMAKE_BUILD_TYPE=Debug
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
```
|
||||
|
||||
**All features ON (mirrors `build-flags` job):**
|
||||
```bash
|
||||
mkdir build && cd build
|
||||
cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
```
|
||||
|
||||
**Transfer engine only:**
|
||||
```bash
|
||||
cd mooncake-transfer-engine
|
||||
mkdir build && cd build
|
||||
cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON
|
||||
cmake --build .
|
||||
```
|
||||
|
||||
## Workflow: Diagnosing and Running Tests
|
||||
|
||||
### Step 1 — Understand what the user wants
|
||||
|
||||
Ask (or infer from context):
|
||||
- All tests, or a specific subset?
|
||||
- Did a specific CI job fail? Which one?
|
||||
- Is the build already done, or do they need to build first?
|
||||
|
||||
### Step 2 — Check and Fix Prerequisites
|
||||
|
||||
**One-command setup** — this script checks all prerequisites and auto-fixes issues:
|
||||
|
||||
```bash
|
||||
bash .claude/skills/mooncake-ci-local/scripts/check-prerequisites.sh
|
||||
```
|
||||
|
||||
**What it checks:**
|
||||
1. ✓ Build directory exists
|
||||
2. ✓ mooncake package installed (auto-installs via cmake --install if missing)
|
||||
3. ✓ ctest available
|
||||
4. ✓ Restarts all services (etcd, metadata server) in clean state
|
||||
5. ✓ Verifies all services are healthy
|
||||
|
||||
**If you need to build first:**
|
||||
```bash
|
||||
mkdir build && cd build
|
||||
cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CXL=ON -DSTORE_USE_ETCD=ON -DCMAKE_BUILD_TYPE=Debug
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
```
|
||||
|
||||
**If script fails:**
|
||||
- Build issues: See "Build Configurations" section below
|
||||
- mooncake install fails: Try `pip install mooncake-wheel/dist/*.whl` manually
|
||||
- etcd install fails: Download from https://github.com/etcd-io/etcd/releases
|
||||
|
||||
### Step 3 — Run and report
|
||||
|
||||
Run the relevant test layer. On failure:
|
||||
1. Show the exact error message
|
||||
2. Check if it's a service/env issue (most common) vs a real test failure
|
||||
3. Suggest the fix (see common issues below)
|
||||
|
||||
## Common Local Test Issues
|
||||
|
||||
**"mooncake_master found in /usr/local/bin" error in run_tests.sh:**
|
||||
The test expects mooncake_master to come from the Python package, not a system install.
|
||||
```bash
|
||||
# Remove the system-installed binary:
|
||||
sudo rm /usr/local/bin/mooncake_master
|
||||
# Or use the wheel-installed one:
|
||||
pip install mooncake-wheel/dist/*.whl
|
||||
```
|
||||
|
||||
**etcd port conflict:**
|
||||
```bash
|
||||
pkill etcd && sleep 1
|
||||
etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 &
|
||||
```
|
||||
|
||||
**Metadata server port conflict:**
|
||||
```bash
|
||||
pkill -f bootstrap_server.py
|
||||
pkill -f mooncake_http_metadata_server
|
||||
```
|
||||
|
||||
**Tests hang (master not responding):**
|
||||
```bash
|
||||
pkill mooncake_master
|
||||
sleep 2
|
||||
mooncake_master --default_kv_lease_ttl=500 &
|
||||
sleep 1
|
||||
```
|
||||
|
||||
**torch/numpy not installed for tensor tests:**
|
||||
```bash
|
||||
pip install torch numpy safetensors packaging
|
||||
```
|
||||
|
||||
**ctest shows no tests found:**
|
||||
```bash
|
||||
# Rebuild with unit tests enabled:
|
||||
cd build
|
||||
cmake .. -DBUILD_UNIT_TESTS=ON
|
||||
cmake --build .
|
||||
```
|
||||
|
||||
## Quick One-Liners
|
||||
|
||||
```bash
|
||||
# Run ALL C++ tests (after building with etcd + metadata server running):
|
||||
# Note: full suite takes 5-15 minutes depending on hardware
|
||||
cd build && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
|
||||
|
||||
# Run only fast tests (skip slow integration tests):
|
||||
cd build && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure --exclude-regex "etcd|ha_test|redis"
|
||||
|
||||
# Run ALL Python tests:
|
||||
mooncake_http_metadata_server --port 8080 & sleep 1 && cd mooncake-wheel/tests && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true bash ../../scripts/run_tests.sh
|
||||
|
||||
# Check code format (changed files only):
|
||||
./scripts/code_format.sh --check --base origin/main
|
||||
|
||||
# Full pre-commit check:
|
||||
pre-commit run --all-files
|
||||
```
|
||||
|
|
@ -0,0 +1,101 @@
|
|||
#!/bin/bash
|
||||
# Mooncake CI Local Test Prerequisites Check
|
||||
# Usage: bash check-prerequisites.sh
|
||||
# This script checks and auto-fixes all prerequisites for running Mooncake CI tests locally.
|
||||
|
||||
set -e
|
||||
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
NC='\033[0m'
|
||||
|
||||
echo "🔍 Checking Mooncake CI test prerequisites..."
|
||||
|
||||
# 1. Check build directory
|
||||
if [ ! -f build/CMakeCache.txt ]; then
|
||||
echo -e "${RED}✗ Build directory not found or not built${NC}"
|
||||
echo " → Run: mkdir build && cd build && cmake .. && cmake --build ."
|
||||
exit 1
|
||||
fi
|
||||
echo -e "${GREEN}✓ Build exists${NC}"
|
||||
|
||||
# 2. Check mooncake installation
|
||||
if ! python -c "import mooncake" 2>/dev/null; then
|
||||
echo -e "${RED}✗ mooncake package not installed${NC}"
|
||||
echo " → Fixing: Installing mooncake package..."
|
||||
cd build && sudo cmake --install . && cd - >/dev/null
|
||||
if ! python -c "import mooncake" 2>/dev/null; then
|
||||
echo " → Alternative: pip install mooncake-wheel/dist/*.whl"
|
||||
exit 1
|
||||
fi
|
||||
echo -e "${GREEN}✓ mooncake package installed${NC}"
|
||||
else
|
||||
echo -e "${GREEN}✓ mooncake package already installed${NC}"
|
||||
fi
|
||||
|
||||
# 3. Check ctest availability
|
||||
if ! command -v ctest &> /dev/null; then
|
||||
echo -e "${RED}✗ ctest not found${NC}"
|
||||
exit 1
|
||||
fi
|
||||
echo -e "${GREEN}✓ ctest available${NC}"
|
||||
|
||||
# 4. Kill and restart services (safest approach for local testing)
|
||||
echo -e "\n${YELLOW}Cleaning up and restarting services...${NC}"
|
||||
pkill -f "^etcd" || true
|
||||
pkill -f bootstrap_server.py || true
|
||||
pkill -f mooncake_http_metadata_server || true
|
||||
sleep 1
|
||||
|
||||
# 5. Start etcd
|
||||
if ! command -v etcd &> /dev/null; then
|
||||
echo -e "${YELLOW}⚠ etcd not found, installing...${NC}"
|
||||
ETCD_VER=v3.6.1
|
||||
OS=$(uname -s | tr '[:upper:]' '[:lower:]')
|
||||
ARCH=$(uname -m)
|
||||
[ "$ARCH" = "x86_64" ] && ARCH="amd64"
|
||||
DOWNLOAD_URL="https://github.com/etcd-io/etcd/releases/download/${ETCD_VER}/etcd-${ETCD_VER}-${OS}-${ARCH}.tar.gz"
|
||||
echo " Downloading from: $DOWNLOAD_URL"
|
||||
cd /tmp
|
||||
wget -q "$DOWNLOAD_URL" && tar xzf "etcd-${ETCD_VER}-${OS}-${ARCH}.tar.gz" && \
|
||||
sudo mv "etcd-${ETCD_VER}-${OS}-${ARCH}"/etcd* /usr/local/bin/
|
||||
cd - >/dev/null
|
||||
echo -e "${GREEN}✓ etcd installed${NC}"
|
||||
fi
|
||||
|
||||
etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 >/dev/null 2>&1 &
|
||||
ETCD_PID=$!
|
||||
sleep 2
|
||||
if ! etcdctl --endpoints=http://127.0.0.1:2379 endpoint health &>/dev/null; then
|
||||
echo -e "${RED}✗ etcd failed to start${NC}"
|
||||
kill $ETCD_PID 2>/dev/null || true
|
||||
exit 1
|
||||
fi
|
||||
echo -e "${GREEN}✓ etcd running (PID: $ETCD_PID)${NC}"
|
||||
|
||||
# 6. Start HTTP metadata server
|
||||
if [ -f "mooncake-transfer-engine/example/http-metadata-server-python/bootstrap_server.py" ]; then
|
||||
cd mooncake-transfer-engine/example/http-metadata-server-python
|
||||
pip install -q aiohttp 2>/dev/null || true
|
||||
python ./bootstrap_server.py >/dev/null 2>&1 &
|
||||
METADATA_PID=$!
|
||||
cd - >/dev/null
|
||||
sleep 1
|
||||
if curl -s http://127.0.0.1:8080/metadata > /dev/null 2>&1; then
|
||||
echo -e "${GREEN}✓ HTTP Metadata server running (PID: $METADATA_PID)${NC}"
|
||||
else
|
||||
echo -e "${RED}✗ HTTP Metadata server failed to start${NC}"
|
||||
kill $METADATA_PID $ETCD_PID 2>/dev/null || true
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo -e "${YELLOW}⚠ Metadata server script not found, skipping${NC}"
|
||||
fi
|
||||
|
||||
echo -e "\n${GREEN}✅ All prerequisites ready!${NC}"
|
||||
echo "Service PIDs: etcd=$ETCD_PID"
|
||||
[ -n "$METADATA_PID" ] && echo "Metadata server PID: $METADATA_PID"
|
||||
echo -e "\n${YELLOW}To kill services:${NC}"
|
||||
echo " pkill -f '^etcd'"
|
||||
echo " pkill -f bootstrap_server"
|
||||
|
|
@ -0,0 +1,366 @@
|
|||
---
|
||||
name: mooncake-troubleshoot
|
||||
description: Automatically diagnose Mooncake deployment and runtime issues. Checks services (mooncake_master, metadata server), RDMA devices, environment variables, connectivity, memory limits, and analyzes logs for common error patterns. Use when Mooncake deployment fails, services won't start, connections fail, or you encounter runtime errors like "Error from etcd client", "No matched device found", "Failed to register memory", "NO_AVAILABLE_HANDLE", or any RDMA/networking issues. Also use when user asks to troubleshoot, debug, diagnose, or fix Mooncake problems.
|
||||
---
|
||||
|
||||
# Mooncake Deployment Troubleshooting
|
||||
|
||||
You are a Mooncake deployment troubleshooting specialist. Your job is to systematically diagnose issues and provide actionable solutions based on the comprehensive troubleshooting knowledge from Mooncake documentation.
|
||||
|
||||
## Diagnostic Strategy
|
||||
|
||||
Run checks systematically, reporting findings as you go. Start with simple checks (services, connectivity) before diving into complex issues (RDMA, memory registration).
|
||||
|
||||
### 1. Service Status Check
|
||||
|
||||
Check if critical services are running:
|
||||
|
||||
```bash
|
||||
# Check mooncake_master
|
||||
ps aux | grep mooncake_master | grep -v grep
|
||||
|
||||
# Check port usage
|
||||
netstat -tuln | grep -E '(50051|8080|2379|9003)'
|
||||
|
||||
# If using etcd
|
||||
ps aux | grep etcd | grep -v grep
|
||||
```
|
||||
|
||||
**Common issues:**
|
||||
- `bind address already in use` → Port conflict, use different port with `--rpc_port`
|
||||
- Master not running → Check startup logs for errors
|
||||
|
||||
### 2. Metadata Server Connectivity
|
||||
|
||||
The metadata server is critical for node discovery and coordination.
|
||||
|
||||
```bash
|
||||
# Test etcd connectivity
|
||||
curl -s http://127.0.0.1:2379/version
|
||||
|
||||
# Or test custom metadata server
|
||||
curl -s $MC_METADATA_SERVER
|
||||
|
||||
# Check for proxy interference
|
||||
echo "http_proxy: $http_proxy"
|
||||
echo "https_proxy: $https_proxy"
|
||||
```
|
||||
|
||||
**Common issues:**
|
||||
- `Error from etcd client` → Metadata server unreachable
|
||||
- **Fix:** Ensure etcd is bound to `0.0.0.0` not `127.0.0.1`:
|
||||
```bash
|
||||
etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<your_ip>:2379
|
||||
```
|
||||
- **Fix:** Disable HTTP proxy:
|
||||
```bash
|
||||
unset http_proxy https_proxy
|
||||
```
|
||||
|
||||
### 3. Environment Variables Check
|
||||
|
||||
Verify critical environment variables are set correctly:
|
||||
|
||||
```bash
|
||||
# Display all MC_* variables
|
||||
env | grep ^MC_
|
||||
|
||||
# Key variables to check:
|
||||
echo "MC_METADATA_SERVER: $MC_METADATA_SERVER"
|
||||
echo "MC_FORCE_TCP: $MC_FORCE_TCP"
|
||||
echo "MC_LOG_LEVEL: $MC_LOG_LEVEL"
|
||||
echo "MC_YLT_LOG_LEVEL: $MC_YLT_LOG_LEVEL"
|
||||
echo "MC_MS_AUTO_DISC: $MC_MS_AUTO_DISC"
|
||||
echo "MC_MS_FILTERS: $MC_MS_FILTERS"
|
||||
echo "MC_GID_INDEX: $MC_GID_INDEX"
|
||||
echo "MC_MTU: $MC_MTU"
|
||||
echo "MC_IB_PORT: $MC_IB_PORT"
|
||||
echo "MC_ENABLE_DEST_DEVICE_AFFINITY: $MC_ENABLE_DEST_DEVICE_AFFINITY"
|
||||
```
|
||||
|
||||
**Key variables:**
|
||||
- `MC_METADATA_SERVER` - Metadata server URL (required)
|
||||
- `MC_FORCE_TCP=true` - Force TCP for testing without RDMA
|
||||
- `MC_LOG_LEVEL=0` - Enable verbose logging (0=INFO, 1=WARNING, 2=ERROR)
|
||||
- `MC_YLT_LOG_LEVEL=debug` - yalantinglibs log level
|
||||
- `MC_MS_AUTO_DISC=1` - Enable topology auto-discovery (default)
|
||||
- `MC_MS_FILTERS` - Filter specific RDMA devices (e.g., "mlx5_1,mlx5_2")
|
||||
- `MC_GID_INDEX` - RDMA GID index (set if GID is all zeros)
|
||||
- `MC_MTU` - RDMA MTU size
|
||||
- `MC_ENABLE_DEST_DEVICE_AFFINITY=1` - Reduce QP creation (fix "Failed to create QP")
|
||||
|
||||
### 4. RDMA Device Check
|
||||
|
||||
Only run if RDMA is being used (skip if `MC_FORCE_TCP=true`):
|
||||
|
||||
```bash
|
||||
# List RDMA devices
|
||||
ibv_devices
|
||||
|
||||
# Check device details and status
|
||||
ibv_devinfo
|
||||
|
||||
# Check for ACTIVE ports
|
||||
ibv_devinfo | grep -A 10 "state:"
|
||||
|
||||
# Check GID addresses (should NOT be all zeros)
|
||||
ibv_devinfo | grep -A 20 "GID"
|
||||
|
||||
# Check peer memory modules
|
||||
lsmod | grep peer_mem
|
||||
lsmod | grep nvidia_peer_mem
|
||||
|
||||
# Check QP count (if "Failed to create QP" error)
|
||||
rdma resource show qp
|
||||
```
|
||||
|
||||
**Common issues:**
|
||||
- `No matched device found` → RDMA device name in config doesn't exist
|
||||
- **Fix:** Use `ibv_devices` to get correct device names
|
||||
- `Device XXX port not active` → RDMA port not in ACTIVE state
|
||||
- **Fix:** Check cable connections, verify with `ibv_devinfo | grep state`
|
||||
- **Fix:** Try different port with `MC_IB_PORT` environment variable
|
||||
- GID all zeros → Wrong GID index
|
||||
- **Fix:** Set `MC_GID_INDEX=1` (or 2, 3 depending on network)
|
||||
- `Failed to create QP: Cannot allocate memory` → Too many QPs created
|
||||
- **Fix:** Set `MC_ENABLE_DEST_DEVICE_AFFINITY=1`
|
||||
|
||||
### 5. Memory and Resource Limits
|
||||
|
||||
Check system limits that affect RDMA memory registration:
|
||||
|
||||
```bash
|
||||
# Check ulimits
|
||||
ulimit -a
|
||||
|
||||
# Focus on max locked memory
|
||||
ulimit -l
|
||||
|
||||
# Check RDMA device memory limits
|
||||
ibv_devinfo -v | grep max_mr_size
|
||||
|
||||
# Check dmesg for memory errors
|
||||
dmesg -T | tail -50 | grep -i "out of mr size"
|
||||
```
|
||||
|
||||
**Common issues:**
|
||||
- `Failed to register memory: Input/output error` → Memory registration limit exceeded
|
||||
- **Diagnostic:** Check `max_mr_size` with `ibv_devinfo -v`
|
||||
- **Fix:** Reduce memory allocation or split into smaller chunks
|
||||
- Cannot allocate memory → ulimit restriction
|
||||
- **Fix:** Set unlimited locked memory:
|
||||
```bash
|
||||
ulimit -l unlimited
|
||||
```
|
||||
- **Permanent fix:** Add to `/etc/security/limits.conf`:
|
||||
```
|
||||
* soft memlock unlimited
|
||||
* hard memlock unlimited
|
||||
```
|
||||
|
||||
### 6. Network Connectivity
|
||||
|
||||
Test connectivity between nodes:
|
||||
|
||||
```bash
|
||||
# Test basic RDMA connectivity
|
||||
ib_write_bw -d <device_name> -R
|
||||
|
||||
# On peer node:
|
||||
ib_write_bw -d <device_name> -R <server_ip>
|
||||
|
||||
# Test GPU Direct RDMA (if CUDA enabled)
|
||||
ib_write_bw -d <device_name> -R -x gdr
|
||||
|
||||
# On peer node:
|
||||
ib_write_bw -d <device_name> -R -x gdr <server_ip>
|
||||
|
||||
# Test DNS resolution
|
||||
nslookup <connectable_name>
|
||||
ping <connectable_name>
|
||||
```
|
||||
|
||||
**Common issues:**
|
||||
- `connection refused` → Incorrect `connectable_name` or `rpc_port`
|
||||
- **Fix:** Ensure `connectable_name` is NOT loopback (127.0.0.1/localhost)
|
||||
- **Fix:** Use actual LAN/WAN IP or valid hostname
|
||||
- `Failed to exchange handshake` → RDMA connection setup failure
|
||||
- **Fix:** Verify MTU matches: set `MC_MTU` environment variable
|
||||
- **Fix:** Verify GID is valid (not all zeros)
|
||||
- **Fix:** Test with `ib_send_bw` between nodes first
|
||||
|
||||
### 7. Log Analysis
|
||||
|
||||
Search logs for common error patterns and their meanings:
|
||||
|
||||
**Metadata/Connectivity Errors:**
|
||||
- `Error from etcd client` → Cannot connect to metadata server
|
||||
- `ERR_METADATA` → Metadata server communication failed
|
||||
- `ERR_DNS` → Invalid `local_server_name` (not valid DNS/IP)
|
||||
|
||||
**RDMA Errors:**
|
||||
- `No matched device found` → RDMA device name doesn't exist
|
||||
- `Device XXX port not active` → RDMA port not in ACTIVE state
|
||||
- `Failed to exchange handshake description` → RDMA handshake failed
|
||||
- `Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num` → MTU/GID mismatch
|
||||
- `Failed to register memory` → Memory registration limit exceeded
|
||||
- `Failed to create QP` → Too many QPs, enable `MC_ENABLE_DEST_DEVICE_AFFINITY=1`
|
||||
- `Worker: Process failed for slice` → Network instability
|
||||
- `work request flushed error` → Cascading error (find first error)
|
||||
|
||||
**Store Errors:**
|
||||
- `NO_AVAILABLE_HANDLE` (-200) → Memory pool exhausted
|
||||
- **Fix:** Increase `global_segment_size` in setup
|
||||
- **Fix:** Check eviction is working (look for eviction logs)
|
||||
- `LEASE_EXPIRED` (-707) → Lease expired during transfer
|
||||
- **Fix:** Increase `default_kv_lease_ttl` in master startup
|
||||
- `OBJECT_NOT_FOUND` (-704) → Object doesn't exist
|
||||
- `SEGMENT_NOT_FOUND` (-101) → No available segments
|
||||
- `Failed to get description of XXX` → Segment name mismatch
|
||||
- **Fix:** Ensure segment name matches `local_hostname` from peer
|
||||
|
||||
**Port/Service Errors:**
|
||||
- `bind address already in use` → Port conflict
|
||||
- **Fix:** Use different port: `--rpc_port=50052`
|
||||
|
||||
### 8. Configuration Validation
|
||||
|
||||
Verify configuration is correct:
|
||||
|
||||
```bash
|
||||
# Check connectable_name is not loopback
|
||||
hostname -I
|
||||
|
||||
# Verify master startup flags
|
||||
ps aux | grep mooncake_master
|
||||
|
||||
# Check if using correct protocol
|
||||
env | grep MC_FORCE_TCP
|
||||
```
|
||||
|
||||
**Critical checks:**
|
||||
- `connectable_name` must be non-loopback IP or valid hostname
|
||||
- MTU and GID configurations must match network environment
|
||||
- RDMA device names must exist on the machine
|
||||
- Ports must not be in use by other services
|
||||
|
||||
## Error Code Quick Reference
|
||||
|
||||
### Transfer Engine Error Codes
|
||||
|
||||
| Code | Name | Meaning | Fix |
|
||||
|------|------|---------|-----|
|
||||
| 0 | Success | Normal execution | - |
|
||||
| -12 | ERR_ADDRESS_NOT_REGISTERED | Memory not registered | Register memory before use |
|
||||
| -14 | ERR_DEVICE_NOT_FOUND | RDMA device not found | Check device name with `ibv_devices` |
|
||||
| -16 | ERR_DNS | Invalid local_server_name | Use valid IP/hostname |
|
||||
| -19 | ERR_REJECT_HANDSHAKE | Peer rejected handshake | Check peer logs for reason |
|
||||
| -20 | ERR_METADATA | Metadata server unreachable | Check etcd/HTTP server |
|
||||
|
||||
### Store Error Codes
|
||||
|
||||
| Code | Name | Meaning | Fix |
|
||||
|------|------|---------|-----|
|
||||
| 0 | Success | Operation successful | - |
|
||||
| -200 | NO_AVAILABLE_HANDLE | Memory pool exhausted | Increase segment size |
|
||||
| -707 | LEASE_EXPIRED | Lease expired | Increase lease TTL |
|
||||
| -704 | OBJECT_NOT_FOUND | Object doesn't exist | Check object key |
|
||||
| -101 | SEGMENT_NOT_FOUND | No available segments | Check segment registration |
|
||||
| -900 | RPC_FAIL | RPC failed | Check network/master |
|
||||
| -1000 | ETCD_OPERATION_ERROR | etcd operation failed | Check etcd status |
|
||||
|
||||
## Output Format
|
||||
|
||||
Provide a structured diagnostic report:
|
||||
|
||||
```
|
||||
🔍 MOONCAKE DEPLOYMENT DIAGNOSTICS
|
||||
==================================
|
||||
|
||||
✅ PASSED CHECKS:
|
||||
- Service status: mooncake_master running on port 50051
|
||||
- Metadata server: etcd accessible at http://127.0.0.1:2379
|
||||
- Environment: MC_METADATA_SERVER set correctly
|
||||
- [other passing checks]
|
||||
|
||||
❌ FAILED CHECKS:
|
||||
- RDMA device: mlx5_0 port not ACTIVE (state: PORT_DOWN)
|
||||
- Memory limits: max locked memory is 64KB (too low)
|
||||
- [other failures with specific error messages]
|
||||
|
||||
⚠️ WARNINGS:
|
||||
- GID index not set, may cause connection issues
|
||||
- HTTP proxy variables set, may interfere with metadata server
|
||||
- [other potential issues]
|
||||
|
||||
🔧 RECOMMENDED FIXES:
|
||||
|
||||
1. Fix RDMA port status:
|
||||
- Check physical cable connections
|
||||
- Verify driver configuration
|
||||
- Command: ibv_devinfo | grep -A 10 "state:"
|
||||
|
||||
2. Increase memory limits:
|
||||
ulimit -l unlimited
|
||||
# Or permanently in /etc/security/limits.conf:
|
||||
* soft memlock unlimited
|
||||
* hard memlock unlimited
|
||||
|
||||
3. Set GID index:
|
||||
export MC_GID_INDEX=1
|
||||
|
||||
4. Disable HTTP proxy:
|
||||
unset http_proxy https_proxy
|
||||
|
||||
📋 SUMMARY:
|
||||
[Brief 2-3 sentence conclusion about deployment health and next steps]
|
||||
```
|
||||
|
||||
## Troubleshooting Workflow
|
||||
|
||||
1. **Start simple**: Check services and basic connectivity first
|
||||
2. **Read logs carefully**: First error is usually root cause (subsequent errors cascade)
|
||||
3. **Test incrementally**: Use `MC_FORCE_TCP=true` to isolate RDMA issues
|
||||
4. **Verify basics**: Check connectable_name, ports, env vars before deep diving
|
||||
5. **Use diagnostic tools**: ibv_devices, ibv_devinfo, ib_write_bw, curl
|
||||
6. **Reference documentation**: Check error codes and troubleshooting guide
|
||||
|
||||
## Quick Fix Commands
|
||||
|
||||
**Start metadata server properly:**
|
||||
```bash
|
||||
etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<your_ip>:2379
|
||||
```
|
||||
|
||||
**Enable verbose logging:**
|
||||
```bash
|
||||
export MC_LOG_LEVEL=0
|
||||
export MC_YLT_LOG_LEVEL=debug
|
||||
```
|
||||
|
||||
**Force TCP mode for testing:**
|
||||
```bash
|
||||
export MC_FORCE_TCP=true
|
||||
```
|
||||
|
||||
**Fix memory limits:**
|
||||
```bash
|
||||
ulimit -l unlimited
|
||||
```
|
||||
|
||||
**Fix too many QPs:**
|
||||
```bash
|
||||
export MC_ENABLE_DEST_DEVICE_AFFINITY=1
|
||||
```
|
||||
|
||||
**Fix GID issues:**
|
||||
```bash
|
||||
export MC_GID_INDEX=1 # or 2, 3 depending on network
|
||||
```
|
||||
|
||||
**Use different port:**
|
||||
```bash
|
||||
mooncake_master --rpc_port=50052
|
||||
```
|
||||
|
||||
Now execute the diagnostic checks systematically and provide the structured report.
|
||||
|
|
@ -23,11 +23,20 @@ RUN apt-get install -y libibverbs-dev \
|
|||
libhiredis-dev \
|
||||
libyaml-cpp-dev \
|
||||
libjemalloc-dev \
|
||||
libzstd-dev \
|
||||
libmsgpack-dev \
|
||||
libgflags-dev \
|
||||
pkg-config \
|
||||
patchelf
|
||||
|
||||
RUN wget https://go.dev/dl/go1.22.12.linux-amd64.tar.gz \
|
||||
&& tar -C /usr/local -xzf go1.22.12.linux-amd64.tar.gz
|
||||
RUN GO_VERSION="1.23.8" && \
|
||||
ARCH=$(uname -m) && \
|
||||
if [ "$ARCH" = "aarch64" ]; then GOARCH="arm64"; \
|
||||
elif [ "$ARCH" = "x86_64" ]; then GOARCH="amd64"; \
|
||||
else echo "Unsupported architecture: $ARCH" && exit 1; fi && \
|
||||
wget https://go.dev/dl/go${GO_VERSION}.linux-${GOARCH}.tar.gz \
|
||||
&& tar -C /usr/local -xzf go${GO_VERSION}.linux-${GOARCH}.tar.gz \
|
||||
&& rm go${GO_VERSION}.linux-${GOARCH}.tar.gz
|
||||
|
||||
RUN git clone https://github.com/alibaba/yalantinglibs.git \
|
||||
&& cd yalantinglibs \
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
**/.dockerignore
|
||||
**/*.Dockerfile
|
||||
|
|
@ -4,13 +4,20 @@
|
|||
# Transfer engine: @alogfans renfeng.chn@outlook.com
|
||||
# Store: @ykwd yangke@approaching.ai
|
||||
# EP: @UNIDY2002 UNIDY2002@outlook.com
|
||||
# PG: @UNIDY2002 UNIDY2002@outlook.com
|
||||
|
||||
.github @stmatengss @ykwd @Ann-1024
|
||||
.github @stmatengss @ykwd @Ann-1024 @luketong777
|
||||
/docs @ShangmingCai @stmatengss @ykwd
|
||||
/mooncake-ep @UNIDY2002 @ympcMark
|
||||
/mooncake-integration/ep @UNIDY2002 @ympcMark
|
||||
/mooncake-integration/transfer_engine @ShangmingCai @alogfans
|
||||
/mooncake-integration/store @ykwd @stmatengss
|
||||
/mooncake-store @ykwd @stmatengss @XucSh
|
||||
/mooncake-ep @UNIDY2002 @ympcMark @yuechen-sys
|
||||
/mooncake-integration/transfer_engine @ShangmingCai @alogfans
|
||||
/mooncake-integration/store @ykwd @stmatengss
|
||||
/mooncake-pg @UNIDY2002 @ympcMark @yuechen-sys
|
||||
/mooncake-store @ykwd @stmatengss @XucSh @YiXR
|
||||
/mooncake-store/*/ha/ @Libotry @YiXR @00fish0
|
||||
/mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q
|
||||
/mooncake-transfer-engine/*/transport/hip_transport/ @alogfans @amd-arozanov
|
||||
/mooncake-transfer-engine/*/transport/ascend_transport/ @alogfans @ascend-direct-dev
|
||||
/mooncake-transfer-engine/*/transport/efa_transport/ @alogfans @whn09
|
||||
/mooncake-wheel @ShangmingCai @stmatengss
|
||||
/scripts/tone_tests @luketong777
|
||||
/scripts/ascend/ @ascend-direct-dev @VNightMare @MingYang119
|
||||
|
|
|
|||
|
|
@ -11,3 +11,14 @@ Transfer Engine:
|
|||
- changed-files:
|
||||
- any-glob-to-any-file: 'mooncake-transfer-engine/**/*'
|
||||
|
||||
PyTorch Backend:
|
||||
- changed-files:
|
||||
- any-glob-to-any-file: 'mooncake-pg/**/*'
|
||||
|
||||
Mooncake EP:
|
||||
- changed-files:
|
||||
- any-glob-to-any-file: 'mooncake-ep/**/*'
|
||||
|
||||
Installation:
|
||||
- changed-files:
|
||||
- any-glob-to-any-file: 'mooncake-wheel/**/*'
|
||||
|
|
|
|||
|
|
@ -2,21 +2,28 @@
|
|||
|
||||
<!-- A clear and concise description of the changes. Link to any relevant issues. -->
|
||||
|
||||
## Module
|
||||
|
||||
- [ ] Transfer Engine (`mooncake-transfer-engine`)
|
||||
- [ ] Mooncake Store (`mooncake-store`)
|
||||
- [ ] Mooncake EP (`mooncake-ep`)
|
||||
- [ ] Integration (`mooncake-integration`)
|
||||
- [ ] P2P Store (`mooncake-p2p-store`)
|
||||
- [ ] Python Wheel (`mooncake-wheel`)
|
||||
- [ ] PyTorch Backend (`mooncake-pg`)
|
||||
- [ ] Mooncake RL (`mooncake-rl`)
|
||||
- [ ] CI/CD
|
||||
- [ ] Docs
|
||||
- [ ] Other
|
||||
|
||||
## Type of Change
|
||||
|
||||
* Types
|
||||
- [ ] Bug fix
|
||||
- [ ] New feature
|
||||
- [ ] Transfer Engine
|
||||
- [ ] Mooncake Store
|
||||
- [ ] Mooncake EP
|
||||
- [ ] Integration
|
||||
- [ ] P2P Store
|
||||
- [ ] Python Wheel
|
||||
- [ ] Breaking change
|
||||
- [ ] CI/CD
|
||||
- [ ] Documentation update
|
||||
- [ ] Other
|
||||
- [ ] Bug fix
|
||||
- [ ] New feature
|
||||
- [ ] Refactor
|
||||
- [ ] Breaking change
|
||||
- [ ] Documentation update
|
||||
- [ ] Other
|
||||
|
||||
## How Has This Been Tested?
|
||||
|
||||
|
|
@ -25,5 +32,6 @@
|
|||
## Checklist
|
||||
|
||||
- [ ] I have performed a self-review of my own code.
|
||||
- [ ] I have formatted my own code using `./scripts/code_format.sh` before submitting.
|
||||
- [ ] I have updated the documentation.
|
||||
- [ ] I have added tests to prove my changes are effective.
|
||||
|
|
|
|||
|
|
@ -6,23 +6,37 @@ on:
|
|||
pull_request:
|
||||
branches: [ "main" ]
|
||||
types: [opened, synchronize, reopened, labeled]
|
||||
workflow_dispatch: {}
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
build:
|
||||
needs: [spell-check, clang-format, check-paths]
|
||||
if: >-
|
||||
github.event_name == 'push' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci')
|
||||
(needs.check-paths.outputs.should-run-downstream == 'true' ||
|
||||
github.event_name == 'workflow_dispatch') &&
|
||||
(github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
|
||||
runs-on: ubuntu-22.04
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.10', '3.12']
|
||||
env:
|
||||
CI: "true"
|
||||
SCCACHE_GHA_ENABLED: "true"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
|
|
@ -53,6 +67,19 @@ jobs:
|
|||
method: 'network'
|
||||
sub-packages: '["nvcc"]'
|
||||
|
||||
- name: Install coverage tools and build utilities
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y lcov gcovr ninja-build
|
||||
|
||||
- name: Set up coverage compilation flags
|
||||
run: |
|
||||
echo "Setting up coverage compilation flags..."
|
||||
echo "CXXFLAGS=--coverage" >> $GITHUB_ENV
|
||||
echo "CFLAGS=--coverage" >> $GITHUB_ENV
|
||||
echo "LDFLAGS=--coverage" >> $GITHUB_ENV
|
||||
shell: bash
|
||||
|
||||
- name: Run sccache-cache
|
||||
uses: mozilla-actions/sccache-action@v0.0.9
|
||||
|
||||
|
|
@ -67,20 +94,20 @@ jobs:
|
|||
shell: bash
|
||||
run: ${SCCACHE_PATH} --show-stats
|
||||
|
||||
- name: Configure project
|
||||
- name: Configure project with coverage support
|
||||
run: |
|
||||
sudo apt update -y
|
||||
sudo bash -x dependencies.sh -y
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON
|
||||
cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_UB=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
|
||||
shell: bash
|
||||
|
||||
- name: Build project
|
||||
run: |
|
||||
cd build
|
||||
make -j
|
||||
sudo make install
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
shell: bash
|
||||
|
||||
- name: Build nvlink_allocator.so
|
||||
|
|
@ -98,35 +125,131 @@ jobs:
|
|||
python ./bootstrap_server.py &
|
||||
shell: bash
|
||||
|
||||
- name: Test (in build env)
|
||||
- name: Run Go store binding integration tests
|
||||
run: |
|
||||
$GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \
|
||||
--eviction_high_watermark_ratio=0.95 \
|
||||
--cluster_id=ci_go_test_cluster \
|
||||
--port 50051 &
|
||||
MASTER_PID=$!
|
||||
sleep 3
|
||||
cd mooncake-store/go
|
||||
export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-common:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
|
||||
export CGO_ENABLED=1
|
||||
export CGO_CFLAGS="-I$GITHUB_WORKSPACE/mooncake-store/include -I$GITHUB_WORKSPACE/mooncake-transfer-engine/include"
|
||||
export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov"
|
||||
# Link cudart if CUDA is available (needed for D2H staging in mooncake_store)
|
||||
if [ -d /usr/local/cuda/lib64 ]; then export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart"; fi
|
||||
ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata go test -v ./tests/...
|
||||
kill $MASTER_PID 2>/dev/null || true
|
||||
shell: bash
|
||||
|
||||
- name: Test (in build env) with coverage
|
||||
run: |
|
||||
cd build
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
|
||||
ldconfig -v || echo "always continue"
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 make test -j ARGS="-V"
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
|
||||
shell: bash
|
||||
|
||||
- name: Drain HTTP E2E test
|
||||
if: matrix.python-version == '3.12'
|
||||
run: |
|
||||
cd build
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
|
||||
# Keep the sanitizer gate on the C++ integration test. The Python
|
||||
# drain script is manual/nightly only because pybind + ASan teardown in
|
||||
# a Python host process is not stable.
|
||||
DEFAULT_KV_LEASE_TTL=500 ./mooncake-store/tests/task_integration_test --gtest_filter='TaskExecutorIntegrationTest.DrainJobCompleteFlow'
|
||||
shell: bash
|
||||
|
||||
- name: Generate coverage report
|
||||
id: coverage
|
||||
run: |
|
||||
cd build
|
||||
echo "=== Starting coverage report generation ==="
|
||||
echo "Current directory: $(pwd)"
|
||||
|
||||
echo "=== Looking for .gcda files ==="
|
||||
find . -name "*.gcda" 2>/dev/null | head -10 || echo "No .gcda files found"
|
||||
|
||||
echo "=== Running lcov ==="
|
||||
lcov --capture --directory . --output-file coverage.info 2>&1 || {
|
||||
echo "WARNING: lcov failed to capture coverage data"
|
||||
echo "Creating minimal lcov-compliant coverage file to allow CI to continue"
|
||||
echo "TN:dummy" > coverage.filtered.info
|
||||
echo "SF:/dev/null" >> coverage.filtered.info
|
||||
echo "DA:0,0" >> coverage.filtered.info
|
||||
echo "end_of_record" >> coverage.filtered.info
|
||||
echo "coverage_failed=true" >> $GITHUB_OUTPUT
|
||||
exit 0 # Exit successfully, do not block CI
|
||||
}
|
||||
|
||||
echo "=== Processing coverage data ==="
|
||||
lcov --remove coverage.info '/usr/*' '*/test/*' '*/third_party/*' --output-file coverage.filtered.info 2>&1 || true
|
||||
|
||||
echo "=== Generating HTML report ==="
|
||||
genhtml coverage.filtered.info --output-directory coverage_report 2>&1 || echo "genhtml failed, continuing..."
|
||||
|
||||
echo "=== Coverage summary ==="
|
||||
lcov --list coverage.filtered.info 2>&1 || echo "lcov list failed"
|
||||
|
||||
echo "=== Coverage report generation completed ==="
|
||||
shell: bash
|
||||
|
||||
- name: Upload coverage to Codecov
|
||||
uses: codecov/codecov-action@v4
|
||||
with:
|
||||
files: build/coverage.filtered.info
|
||||
flags: unittests
|
||||
name: code-coverage-report
|
||||
token: ${{ secrets.CODECOV_TOKEN }}
|
||||
fail_ci_if_error: false
|
||||
continue-on-error: true
|
||||
|
||||
- name: Check coverage status
|
||||
if: always()
|
||||
run: |
|
||||
if [ "${{ steps.coverage.outputs.coverage_failed }}" = "true" ]; then
|
||||
echo "⚠️ Coverage collection failed but CI continued"
|
||||
echo "::warning::Code coverage collection failed. Please check the build logs."
|
||||
else
|
||||
echo "✅ Coverage collected successfully"
|
||||
fi
|
||||
|
||||
- name: Generate Python version tag
|
||||
id: generate_tag_build
|
||||
run: |
|
||||
echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
|
||||
shell: bash
|
||||
|
||||
# In CI, build_wheel.sh removes build/ to free disk (CI=true); set FREE_BUILD_DIR=1 locally to enable.
|
||||
- name: Build Python wheel
|
||||
run: |
|
||||
# Build wheel with specific Python version
|
||||
PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_build.outputs.python_version_tag }} ./scripts/build_wheel.sh
|
||||
shell: bash
|
||||
|
||||
- name: Upload wheel for ZMQ test job
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: wheel-build-py${{ steps.generate_tag_build.outputs.python_version_tag }}
|
||||
path: mooncake-wheel/dist-py${{ steps.generate_tag_build.outputs.python_version_tag }}/*.whl
|
||||
|
||||
build-musa:
|
||||
needs: [spell-check, clang-format, check-paths]
|
||||
if: >-
|
||||
github.event_name == 'push' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci')
|
||||
(needs.check-paths.outputs.should-run-downstream == 'true' ||
|
||||
github.event_name == 'workflow_dispatch') &&
|
||||
(github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
|
||||
runs-on: ubuntu-22.04
|
||||
container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Mark repository as safe
|
||||
run: git config --global --add safe.directory $GITHUB_WORKSPACE
|
||||
|
|
@ -135,22 +258,29 @@ jobs:
|
|||
- name: Configure project
|
||||
run: |
|
||||
apt update -y
|
||||
apt install -y ninja-build
|
||||
bash -x dependencies.sh -y
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. -DUSE_MUSA=ON -DUSE_MNNVL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DUSE_CXL=ON -DUSE_TCP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF
|
||||
cmake -G Ninja .. -DUSE_MUSA=ON -DUSE_MNNVL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DUSE_CXL=ON -DUSE_TCP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF
|
||||
shell: bash
|
||||
|
||||
- name: Build project
|
||||
run: |
|
||||
cd build
|
||||
source ~/.bashrc
|
||||
make -j
|
||||
make install
|
||||
cmake --build .
|
||||
cmake --install .
|
||||
shell: bash
|
||||
|
||||
test-wheel-ubuntu:
|
||||
needs: build-flags
|
||||
needs: [spell-check, clang-format, build-flags]
|
||||
if: >-
|
||||
needs.build-flags.result == 'success' &&
|
||||
(github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
|
||||
strategy:
|
||||
matrix:
|
||||
ubuntu-version: [ubuntu-22.04, ubuntu-24.04]
|
||||
|
|
@ -158,6 +288,8 @@ jobs:
|
|||
runs-on: ${{ matrix.ubuntu-version }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
|
|
@ -213,6 +345,14 @@ jobs:
|
|||
|
||||
- name: Run tests with ssd
|
||||
run: |
|
||||
# Reserve port 50052 (mooncake_client RPC port) so the kernel never
|
||||
# auto-allocates it as ephemeral source port for other outbound
|
||||
# connections in the test suite. Without this, a random Python test
|
||||
# connection can pick src_port=50052, leave a TIME_WAIT on
|
||||
# <eth0_ip>:50052 for 60s, and block mooncake_client's bind to
|
||||
# 0.0.0.0:50052 even with SO_REUSEADDR (Linux only relaxes
|
||||
# TIME_WAIT+bind conflict for same-IP or loopback).
|
||||
sudo sysctl -w net.ipv4.ip_local_reserved_ports=50052
|
||||
source test_env/bin/activate
|
||||
MC_STORE_MEMCPY=false TEST_SSD_OFFLOAD_IN_EVICT=true ./scripts/run_tests.sh
|
||||
rm -rf /tmp/mooncake_test_ssd
|
||||
|
|
@ -241,6 +381,40 @@ jobs:
|
|||
python scripts/test_tensor_api.py -n 1
|
||||
shell: bash
|
||||
|
||||
- name: Run Python Async API Test (CI check)
|
||||
env:
|
||||
MOONCAKE_MASTER: "127.0.0.1:50051"
|
||||
MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata"
|
||||
MOONCAKE_PROTOCOL: "tcp"
|
||||
LOCAL_HOSTNAME: "127.0.0.1"
|
||||
run: |
|
||||
source test_env/bin/activate
|
||||
python scripts/test_async_store.py
|
||||
shell: bash
|
||||
|
||||
- name: Test Mooncake Copy/Move API
|
||||
env:
|
||||
MOONCAKE_MASTER: "127.0.0.1:50051"
|
||||
MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata"
|
||||
MOONCAKE_PROTOCOL: "tcp"
|
||||
LOCAL_HOSTNAME: "127.0.0.1"
|
||||
run: |
|
||||
source test_env/bin/activate
|
||||
python scripts/test_copy_move_api.py
|
||||
shell: bash
|
||||
|
||||
- name: Run Python Drain HTTP E2E Test (CI check)
|
||||
env:
|
||||
MOONCAKE_MASTER: "127.0.0.1:50051"
|
||||
MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata"
|
||||
MOONCAKE_PROTOCOL: "tcp"
|
||||
LOCAL_HOSTNAME: "127.0.0.1"
|
||||
run: |
|
||||
source test_env/bin/activate
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
|
||||
python scripts/test_drain_http_api.py --timeout-sec 90
|
||||
shell: bash
|
||||
|
||||
- name: Run RPC Communicator Bandwidth Test
|
||||
run: |
|
||||
source test_env/bin/activate
|
||||
|
|
@ -251,90 +425,44 @@ jobs:
|
|||
kill $SERVER_PID 2>/dev/null || true
|
||||
wait $SERVER_PID 2>/dev/null || true
|
||||
|
||||
- name: Test Mooncake EP Backend (CPU Only)
|
||||
- name: Test Mooncake PyTorch Backend (CPU Only)
|
||||
env:
|
||||
MC_FORCE_TCP: "true"
|
||||
run: |
|
||||
source test_env/bin/activate
|
||||
python -m unittest mooncake-wheel.tests.test_mooncake_backend_cpu
|
||||
# Disable these tests in CI as they fail occasionally.
|
||||
# python -m unittest mooncake-wheel.tests.test_mooncake_backend_elastic
|
||||
# python -m unittest mooncake-wheel.tests.test_mooncake_backend_p2p_cpu
|
||||
shell: bash
|
||||
|
||||
test-sglang-integration:
|
||||
needs: build-flags
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
tone_user_name: ${{ secrets.TONE_USER_NAME }}
|
||||
steps:
|
||||
- name: trigger T-one test
|
||||
if: ${{ env.tone_user_name != '' }}
|
||||
run: |
|
||||
curl -L -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" https://api.github.com/repos/${{ github.repository }}/actions/runs/${{ github.run_id }}/artifacts > artifact.json
|
||||
cat artifact.json
|
||||
artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | .id" artifact.json)
|
||||
signature="${{ secrets.TONE_USER_NAME }}|${{ secrets.TONE_USER_TOKEN }}|$(python3 -c "import time;print(time.time())")"
|
||||
signature="$(python3 -c "import base64;print(base64.b64encode(\"$signature\".encode('utf-8')).decode('utf-8'))")"
|
||||
curl -s -H 'Content-Type: application/json' -X POST -d "{\"workspace\":\"mooncake_test\",\"project\":\"mooncake-ci\",\"template\":\"mooncake-ci-test\",\"name\":\"mooncake-ci-${{ github.sha }}\",\"username\":\"${{ secrets.TONE_USER_NAME }}\",\"env_ifs\":\" \",\"env_info\":\"ARTIFACT_ID=${artifact_id} GIT_REPO=${{ github.repository }}\",\"signature\":\"$signature\"}" https://tone.openanolis.cn/api/job/create/ > job.json
|
||||
if [ "$(jq .code job.json)" == 200 ]; then
|
||||
echo "job created"
|
||||
else
|
||||
echo "job create failed"
|
||||
exit 1
|
||||
fi
|
||||
job_id=$(jq .data.id job.json)
|
||||
echo "check job status here and remember to cancel it before restart the job !"
|
||||
echo "job_url: https://tone.openanolis.cn/ws/gclfnh19/test_result/${job_id}?tab=4"
|
||||
echo "job_id=${job_id}" >> $GITHUB_ENV
|
||||
shell: bash
|
||||
|
||||
- name: qurey job results
|
||||
if: ${{ env.tone_user_name != '' }}
|
||||
run: |
|
||||
time=0
|
||||
while true; do
|
||||
if [ $time -gt 720 ]; then
|
||||
echo "timeout"
|
||||
exit 1
|
||||
fi
|
||||
signature="${{ secrets.TONE_USER_NAME }}|${{ secrets.TONE_USER_TOKEN }}|$(python3 -c "import time;print(time.time())")"
|
||||
signature="$(python3 -c "import base64;print(base64.b64encode(\"$signature\".encode('utf-8')).decode('utf-8'))")"
|
||||
curl -s -H 'Content-Type: application/json' -X POST -d "{\"username\":\"${{ secrets.TONE_USER_NAME }}\", \"signature\":\"$signature\", \"job_id\": \"${job_id}\"}" https://tone.openanolis.cn/api/job/query/ > job_status.json
|
||||
if ! [ "$(jq .code job_status.json)" == 200 ]; then
|
||||
echo "job query failed"
|
||||
exit 1
|
||||
fi
|
||||
job_status=$(jq .data.job_second_state job_status.json)
|
||||
if [[ $job_status =~ "pass" ]]; then
|
||||
echo "job successful !"
|
||||
exit 0
|
||||
elif [[ $job_status =~ "fail" ]] ; then
|
||||
echo "job failed or stopped !"
|
||||
exit 1
|
||||
fi
|
||||
time=$(( time + 1))
|
||||
sleep 10
|
||||
done
|
||||
shell: bash
|
||||
- name: Test Safetensor Functions
|
||||
run: |
|
||||
source test_env/bin/activate
|
||||
pip install safetensors
|
||||
python -m unittest mooncake-wheel.tests.test_safetensor_functions
|
||||
shell: bash
|
||||
|
||||
build-flags:
|
||||
needs: [spell-check, clang-format, check-paths]
|
||||
if: >-
|
||||
github.event_name == 'push' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci')
|
||||
(needs.check-paths.outputs.should-run-downstream == 'true' ||
|
||||
github.event_name == 'workflow_dispatch') &&
|
||||
(github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
|
||||
runs-on: ubuntu-22.04
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.10', '3.12']
|
||||
env:
|
||||
CI: "true"
|
||||
BUILD_WITH_EP: "1"
|
||||
EP_TORCH_VERSIONS: "2.8.0;2.9.0;2.9.1"
|
||||
TORCH_CUDA_ARCH_LIST: "8.0;9.0"
|
||||
SCCACHE_GHA_ENABLED: "true"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
|
|
@ -375,10 +503,14 @@ jobs:
|
|||
- name: Install dependencies
|
||||
run: |
|
||||
sudo apt update -y
|
||||
sudo apt install -y ninja-build
|
||||
sudo bash -x dependencies.sh -y
|
||||
df -h
|
||||
shell: bash
|
||||
|
||||
- name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: Build transfer engine only
|
||||
run: |
|
||||
cd mooncake-transfer-engine
|
||||
|
|
@ -386,9 +518,9 @@ jobs:
|
|||
cd build
|
||||
export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
|
||||
cmake .. -DUSE_ETCD=OFF -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
|
||||
make -j4
|
||||
sudo make install
|
||||
cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
df -h
|
||||
shell: bash
|
||||
|
||||
|
|
@ -396,7 +528,7 @@ jobs:
|
|||
run: |
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. -DUSE_ETCD=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_EP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
|
||||
cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
|
||||
shell: bash
|
||||
# TODO: lack USE_NVMEOF,USE_MNNVL
|
||||
|
||||
|
|
@ -405,32 +537,39 @@ jobs:
|
|||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
|
||||
cd build
|
||||
make -j4
|
||||
sudo make install
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
df -h
|
||||
shell: bash
|
||||
|
||||
- name: Configure project with unit tests and examples
|
||||
run: |
|
||||
cd build
|
||||
cmake .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON
|
||||
cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DWITH_STORE_RUST=ON -DENABLE_SCCACHE=ON
|
||||
shell: bash
|
||||
# TODO: lack WITH_RUST_EXAMPLE
|
||||
|
||||
- name: Build project with unit tests and examples
|
||||
run: |
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
|
||||
cd build
|
||||
make -j4
|
||||
sudo make install
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
shell: bash
|
||||
|
||||
- name: Check Mooncake Store Rust bindings and example
|
||||
run: |
|
||||
cd mooncake-store/rust
|
||||
MOONCAKE_STORE_LIB_DIR=$GITHUB_WORKSPACE/build/mooncake-store/src \
|
||||
MOONCAKE_STORE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-store/include \
|
||||
cargo check --example basic_usage --tests
|
||||
shell: bash
|
||||
|
||||
- name: Configure project
|
||||
run: |
|
||||
cd build
|
||||
rm -r */tests
|
||||
cmake .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON
|
||||
cmake -G Ninja .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0"
|
||||
shell: bash
|
||||
|
||||
- name: Build project
|
||||
|
|
@ -438,8 +577,8 @@ jobs:
|
|||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
|
||||
cd build
|
||||
make -j4
|
||||
sudo make install
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
shell: bash
|
||||
|
||||
- name: Build nvlink_allocator.so
|
||||
|
|
@ -458,9 +597,9 @@ jobs:
|
|||
echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
|
||||
shell: bash
|
||||
|
||||
# In CI, build_wheel.sh removes build/ to free disk (CI=true); set FREE_BUILD_DIR=1 locally to enable.
|
||||
- name: Build Python wheel
|
||||
run: |
|
||||
# Build wheel with specific Python version
|
||||
PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }} ./scripts/build_wheel.sh
|
||||
shell: bash
|
||||
|
||||
|
|
@ -472,30 +611,43 @@ jobs:
|
|||
|
||||
build-docker:
|
||||
name: Build Docker Image
|
||||
needs: [spell-check, clang-format, check-paths]
|
||||
if: >-
|
||||
github.event_name == 'push' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci')
|
||||
(needs.check-paths.outputs.should-run-downstream == 'true' ||
|
||||
github.event_name == 'workflow_dispatch') &&
|
||||
(github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
|
||||
runs-on: ubuntu-22.04
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v2
|
||||
|
||||
- name: Build Docker image
|
||||
run: docker build -t mooncake-app .
|
||||
run: |
|
||||
docker build -f docker/mooncake.Dockerfile \
|
||||
--build-arg PYTHON_VERSION=3.10 \
|
||||
--build-arg EP_TORCH_VERSIONS="2.9.1" \
|
||||
-t mooncake:from-source .
|
||||
|
||||
spell-check:
|
||||
name: Spell Check with Typos
|
||||
if: >-
|
||||
github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci')
|
||||
runs-on: ubuntu-22.04
|
||||
steps:
|
||||
- name: Checkout Actions Repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
- name: Spell Check Repo
|
||||
uses: crate-ci/typos@v1.30.2
|
||||
|
||||
|
|
@ -503,12 +655,16 @@ jobs:
|
|||
name: Check code format
|
||||
if: >-
|
||||
github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci')
|
||||
runs-on: ubuntu-22.04
|
||||
steps:
|
||||
- name: Checkout Actions Repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0 # Need full history for branch comparison
|
||||
persist-credentials: false
|
||||
|
||||
- name: Install clang-format 20
|
||||
run: |
|
||||
|
|
@ -517,17 +673,118 @@ jobs:
|
|||
sudo ./llvm.sh 20
|
||||
sudo apt-get install -y clang-format-20
|
||||
|
||||
- name: run clang-format-20
|
||||
- name: Check code format
|
||||
run: |
|
||||
# the old clang-format-14 which is the default version in ubuntu 22.04,
|
||||
# is inconsistent with clang-format-20.
|
||||
ls -lh /usr/bin/clang-format*
|
||||
clang-format --version
|
||||
clang-format-20 --version
|
||||
# skip cachelib_memory_allocator
|
||||
find . -type f \( -name "*.h" -o -name "*.cpp" \) | grep -v cachelib_memory_allocator | xargs clang-format-20 -style=file -i
|
||||
if ! git diff --exit-code; then
|
||||
echo "Please follow the .clang-format code style, try clang-format -i FILENAME"
|
||||
# Check script exists and is executable
|
||||
if [[ ! -x ./scripts/code_format.sh ]]; then
|
||||
echo "Error: code_format.sh not found or not executable"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Determine base ref for comparison
|
||||
if [ "${{ github.event_name }}" == "pull_request" ]; then
|
||||
# For PRs: compare against the target branch
|
||||
BASE_REF="origin/${{ github.base_ref }}"
|
||||
else
|
||||
# For push events: use github.event.before to handle multi-commit pushes
|
||||
BEFORE_SHA="${{ github.event.before }}"
|
||||
if [ "${BEFORE_SHA}" == "0000000000000000000000000000000000000000" ]; then
|
||||
# New branch push, compare against default branch
|
||||
BASE_REF="origin/${{ github.event.repository.default_branch }}"
|
||||
else
|
||||
# Normal push (single or multiple commits)
|
||||
BASE_REF="${BEFORE_SHA}"
|
||||
fi
|
||||
fi
|
||||
echo "Comparing against: ${BASE_REF}"
|
||||
./scripts/code_format.sh --check --base "${BASE_REF}"
|
||||
shell: bash
|
||||
|
||||
|
||||
check-paths:
|
||||
if: >-
|
||||
github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci')
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
should-run-downstream: ${{ steps.dispatch-override.outputs.src || steps.filter.outputs.src }}
|
||||
steps:
|
||||
# workflow_dispatch has no PR/push diff context — skip paths-filter and default to true
|
||||
- name: Default to true for workflow_dispatch
|
||||
id: dispatch-override
|
||||
if: github.event_name == 'workflow_dispatch'
|
||||
run: echo "src=true" >> $GITHUB_OUTPUT
|
||||
- uses: actions/checkout@v4
|
||||
if: github.event_name != 'workflow_dispatch'
|
||||
with:
|
||||
fetch-depth: 2
|
||||
persist-credentials: false
|
||||
- uses: dorny/paths-filter@v3
|
||||
if: github.event_name != 'workflow_dispatch'
|
||||
id: filter
|
||||
with:
|
||||
filters: |
|
||||
src:
|
||||
- 'mooncake-*/**'
|
||||
- 'extern/**'
|
||||
- 'CMakeLists.txt'
|
||||
- 'dependencies.sh'
|
||||
- 'scripts/**'
|
||||
- '.github/workflows/**'
|
||||
|
||||
build-wheel-cu13:
|
||||
needs: [spell-check, clang-format, check-paths]
|
||||
if: >-
|
||||
(needs.check-paths.outputs.should-run-downstream == 'true' ||
|
||||
github.event_name == 'workflow_dispatch') &&
|
||||
(github.event_name == 'push' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.action == 'opened' ||
|
||||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
|
||||
uses: ./.github/workflows/ci_cu13.yml
|
||||
secrets: inherit
|
||||
|
||||
ascend-test:
|
||||
needs: [build, check-paths]
|
||||
if: needs.check-paths.outputs.should-run-downstream == 'true'
|
||||
uses: ./.github/workflows/ci_ascend.yml
|
||||
secrets: inherit
|
||||
|
||||
integration-test:
|
||||
needs: [build, check-paths]
|
||||
if: needs.check-paths.outputs.should-run-downstream == 'true'
|
||||
uses: ./.github/workflows/integration-test.yml
|
||||
secrets: inherit
|
||||
|
||||
ci-gate:
|
||||
name: CI Gate
|
||||
if: always()
|
||||
needs:
|
||||
- spell-check
|
||||
- clang-format
|
||||
- build
|
||||
- build-musa
|
||||
- build-flags
|
||||
- build-docker
|
||||
- test-wheel-ubuntu
|
||||
- build-wheel-cu13
|
||||
- ascend-test
|
||||
- integration-test
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check required job results
|
||||
run: |
|
||||
failing=$(echo "$NEEDS_JSON" | jq -r '
|
||||
to_entries[] |
|
||||
select(.value.result != "success" and .value.result != "skipped") |
|
||||
"\(.key): \(.value.result)"')
|
||||
if [ -n "$failing" ]; then
|
||||
echo "::error::The following jobs failed or were cancelled:"
|
||||
echo "$failing"
|
||||
exit 1
|
||||
fi
|
||||
echo "All checks passed or were acceptably skipped."
|
||||
env:
|
||||
NEEDS_JSON: ${{ toJSON(needs) }}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,364 @@
|
|||
name: 'CI Test on ASCEND Platform'
|
||||
|
||||
on:
|
||||
workflow_call:
|
||||
inputs:
|
||||
checkout_ref:
|
||||
description: 'Git ref to checkout (PR head SHA for pull_request_target)'
|
||||
required: false
|
||||
type: string
|
||||
|
||||
jobs:
|
||||
build-and-test:
|
||||
if: github.repository == 'kvcache-ai/Mooncake'
|
||||
runs-on: self-hosted
|
||||
|
||||
container:
|
||||
image: localhost:5000/mooncake-hixl-ci:v5
|
||||
options: --privileged --user 0:0 --device /dev/davinci0 --device /dev/davinci1 --device /dev/davinci2 --device /dev/davinci3
|
||||
--device /dev/davinci4 --device /dev/davinci5 --device /dev/davinci6 --device /dev/davinci7
|
||||
--device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc --ulimit nproc=65535:65535
|
||||
env:
|
||||
GITHUB_ACTIONS: "true"
|
||||
LD_PRELOAD: "/usr/lib64/libjemalloc.so.2:"
|
||||
volumes:
|
||||
- /usr/local/dcmi:/usr/local/dcmi
|
||||
- /usr/local/Ascend/driver/:/usr/local/Ascend/driver/
|
||||
- /etc/ascend_install.info:/etc/ascend_install.info
|
||||
- /etc/hccn.conf:/etc/hccn.conf
|
||||
|
||||
steps:
|
||||
- name: Configure GitHub fetch defaults
|
||||
shell: bash
|
||||
run: |
|
||||
git config --global protocol.version 2
|
||||
git config --global http.version HTTP/1.1
|
||||
git config --global http.lowSpeedLimit 1024
|
||||
git config --global http.lowSpeedTime 30
|
||||
|
||||
- name: Checkout code
|
||||
id: checkout_code
|
||||
continue-on-error: true
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.checkout_ref || github.sha }}
|
||||
fetch-depth: 1
|
||||
persist-credentials: false
|
||||
|
||||
- name: Retry checkout via GitHub mirror
|
||||
if: steps.checkout_code.outcome == 'failure'
|
||||
shell: bash
|
||||
env:
|
||||
ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
|
||||
CHECKOUT_REF: ${{ inputs.checkout_ref || github.sha }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
if [ -z "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
|
||||
echo "Checkout from GitHub failed and ASCEND_GITHUB_MIRROR_URLS is not set"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
normalize_base() {
|
||||
local base="$1"
|
||||
base="${base#${base%%[![:space:]]*}}"
|
||||
base="${base%${base##*[![:space:]]}}"
|
||||
[ -n "$base" ] || return 1
|
||||
[ "$base" != "https://github.com/" ] && base="${base%/}/"
|
||||
printf '%s\n' "$base"
|
||||
}
|
||||
|
||||
candidates=()
|
||||
while IFS= read -r raw; do
|
||||
base="$(normalize_base "$raw" || true)"
|
||||
[ -n "$base" ] || continue
|
||||
[ "$base" = "https://github.com/" ] && continue
|
||||
candidates+=("$base")
|
||||
done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
|
||||
|
||||
if [ ${#candidates[@]} -eq 0 ]; then
|
||||
echo "Checkout from GitHub failed and no valid mirror candidates were configured"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
workdir="${GITHUB_WORKSPACE}"
|
||||
git config --global --add safe.directory "$workdir"
|
||||
|
||||
for base in "${candidates[@]}"; do
|
||||
mirror_url="${base}https://github.com/${GITHUB_REPOSITORY}.git"
|
||||
echo "Retrying checkout with ${mirror_url}"
|
||||
|
||||
find "$workdir" -mindepth 1 -maxdepth 1 -exec rm -rf {} +
|
||||
git init "$workdir"
|
||||
git -C "$workdir" remote add origin "$mirror_url"
|
||||
|
||||
if git -C "$workdir" fetch --depth=1 origin "$CHECKOUT_REF" && \
|
||||
git -C "$workdir" checkout --force --detach FETCH_HEAD; then
|
||||
echo "Mirror checkout succeeded via ${base}"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Mirror checkout failed via ${base}"
|
||||
rm -rf "$workdir/.git"
|
||||
done
|
||||
|
||||
echo "Direct GitHub checkout failed and all mirror retries failed"
|
||||
exit 1
|
||||
|
||||
- name: Configure CMake
|
||||
shell: bash
|
||||
env:
|
||||
ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
|
||||
run: |
|
||||
source /usr/local/Ascend/cann-9.0.0/set_env.sh
|
||||
pwd
|
||||
|
||||
submodule_updated=false
|
||||
if git submodule update --init --recursive; then
|
||||
submodule_updated=true
|
||||
elif [ -n "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
|
||||
normalize_base() {
|
||||
local base="$1"
|
||||
base="${base#${base%%[![:space:]]*}}"
|
||||
base="${base%${base##*[![:space:]]}}"
|
||||
[ -n "$base" ] || return 1
|
||||
[ "$base" != "https://github.com/" ] && base="${base%/}/"
|
||||
printf '%s\n' "$base"
|
||||
}
|
||||
|
||||
while IFS= read -r raw; do
|
||||
base="$(normalize_base "$raw" || true)"
|
||||
[ -n "$base" ] || continue
|
||||
[ "$base" = "https://github.com/" ] && continue
|
||||
|
||||
echo "Retrying submodule update with ${base}"
|
||||
if git -c url."${base}https://github.com/".insteadOf=https://github.com/ \
|
||||
submodule update --init --recursive; then
|
||||
submodule_updated=true
|
||||
break
|
||||
fi
|
||||
done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
|
||||
fi
|
||||
|
||||
if [ "$submodule_updated" != true ]; then
|
||||
if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then
|
||||
echo "git submodule update failed (mirrors also exhausted), trying to cp pybind11..."
|
||||
if [ -d "../pybind11" ]; then
|
||||
cp -r ../pybind11 extern/
|
||||
else
|
||||
echo "Error: ../pybind11 does not exist. Cannot copy pybind11."
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "Detected that extern/pybind11 already exists, continuing execution...."
|
||||
fi
|
||||
fi
|
||||
|
||||
bash scripts/ascend/dependencies_ascend_installation.sh
|
||||
echo "Configuring CMake..."
|
||||
rm -rf build
|
||||
mkdir -p build
|
||||
cd build
|
||||
|
||||
cmake .. \
|
||||
-DUSE_ASCEND_DIRECT=ON \
|
||||
-DBUILD_EXAMPLES=OFF \
|
||||
-DBUILD_UNIT_TESTS=OFF
|
||||
|
||||
- name: Build
|
||||
shell: bash
|
||||
run: |
|
||||
source /usr/local/Ascend/cann-9.0.0/set_env.sh
|
||||
echo "Building..."
|
||||
cd build
|
||||
cmake --build . -j$(nproc)
|
||||
cmake --install .
|
||||
echo "Mooncake installed successfully."
|
||||
|
||||
- name: Run Hixl Mooncake Store Test
|
||||
shell: bash
|
||||
run: |
|
||||
source /usr/local/Ascend/cann-9.0.0/set_env.sh
|
||||
set -e
|
||||
export ASCEND_PROCESS_LOG_PATH=/tmp/hixl-test-log/
|
||||
export ASCEND_GLOBAL_LOG_LEVEL=3
|
||||
echo "=== Cloning Hixl repository ==="
|
||||
cd ..
|
||||
rm -rf hixl
|
||||
git clone https://gitcode.com/cann/hixl.git
|
||||
cd hixl/examples/third_parties/mooncake_store/python/
|
||||
|
||||
export LD_LIBRARY_PATH=/usr/local/lib:${LD_LIBRARY_PATH}
|
||||
echo "=== Starting Mooncake Master ==="
|
||||
|
||||
# Find mooncake_master binary
|
||||
MOONCAKE_MASTER=$(find /usr/local/bin /usr/bin -name "mooncake_master" -type f 2>/dev/null | head -1)
|
||||
if [ -z "$MOONCAKE_MASTER" ]; then
|
||||
# Try finding in build directory
|
||||
MOONCAKE_MASTER=$(find $GITHUB_WORKSPACE/build -name "mooncake_master" -type f 2>/dev/null | head -1)
|
||||
fi
|
||||
|
||||
if [ -z "$MOONCAKE_MASTER" ]; then
|
||||
echo "Error: mooncake_master binary not found"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Found mooncake_master at: $MOONCAKE_MASTER"
|
||||
|
||||
# Start Mooncake master in background
|
||||
$MOONCAKE_MASTER \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_host=0.0.0.0 \
|
||||
--http_metadata_server_port=8080 \
|
||||
> /tmp/mooncake_master.log 2>&1 &
|
||||
MASTER_PID=$!
|
||||
echo "Mooncake Master started with PID: $MASTER_PID"
|
||||
|
||||
# Wait for master to be ready
|
||||
echo "Waiting for Mooncake Master to initialize..."
|
||||
sleep 5
|
||||
|
||||
# Check if master is running
|
||||
if ! kill -0 $MASTER_PID 2>/dev/null; then
|
||||
echo "Error: Mooncake Master failed to start"
|
||||
cat /tmp/mooncake_master.log
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Mooncake Master is running"
|
||||
echo "=== Running Hixl Mooncake Store Tests ==="
|
||||
# List of test cases to run
|
||||
TEST_CASES=(
|
||||
"batch_put_get_sample.py"
|
||||
"batch_put_get_multi_buffers_sample.py"
|
||||
)
|
||||
|
||||
# List of test scenarios (HCCL_INTRA_ROCE_ENABLE settings)
|
||||
TEST_SCENARIOS=(
|
||||
"HCCL_INTRA_ROCE_ENABLE=1"
|
||||
"HCCL_INTRA_ROCE_ENABLE_UNSET"
|
||||
)
|
||||
|
||||
# Track test results
|
||||
FAILED_TESTS=()
|
||||
PASSED_TESTS=()
|
||||
|
||||
# Run each test scenario
|
||||
for scenario in "${TEST_SCENARIOS[@]}"; do
|
||||
echo ""
|
||||
echo "========================================="
|
||||
echo "Running scenario: $scenario"
|
||||
echo "========================================="
|
||||
|
||||
# Configure environment variables for the current scenario
|
||||
if [ "$scenario" = "HCCL_INTRA_ROCE_ENABLE=1" ]; then
|
||||
export HCCL_INTRA_ROCE_ENABLE=1
|
||||
unset ASCEND_BUFFER_POOL
|
||||
echo "HCCL_INTRA_ROCE_ENABLE is set to 1, ASCEND_BUFFER_POOL is unset"
|
||||
else
|
||||
unset HCCL_INTRA_ROCE_ENABLE
|
||||
export ASCEND_BUFFER_POOL=4:8
|
||||
echo "HCCL_INTRA_ROCE_ENABLE is not set, ASCEND_BUFFER_POOL is set to 4:8"
|
||||
fi
|
||||
|
||||
# Run each test case in the current scenario
|
||||
for test_case in "${TEST_CASES[@]}"; do
|
||||
echo ""
|
||||
echo "-----------------------------------------"
|
||||
echo "Test: $test_case"
|
||||
echo "-----------------------------------------"
|
||||
|
||||
if [ ! -f "$test_case" ]; then
|
||||
echo "Warning: Test file $test_case not found, skipping..."
|
||||
continue
|
||||
fi
|
||||
|
||||
# Run the test with 2 devices in distributed mode
|
||||
# Run rank 0 on device 0
|
||||
python3 $test_case \
|
||||
--device_id=0 \
|
||||
--rank=0 \
|
||||
--world_size=2 \
|
||||
--distributed \
|
||||
2>&1 | tee "/tmp/hixl_test_${scenario//=/}_${test_case%.py}_rank0.log" &
|
||||
PID0=$!
|
||||
|
||||
# Run rank 1 on device 1
|
||||
python3 $test_case \
|
||||
--device_id=2 \
|
||||
--rank=1 \
|
||||
--world_size=2 \
|
||||
--distributed \
|
||||
2>&1 | tee "/tmp/hixl_test_${scenario//=/}_${test_case%.py}_rank1.log" &
|
||||
PID1=$!
|
||||
|
||||
# Wait for both processes to complete
|
||||
wait $PID0
|
||||
TEST_RESULT0=$?
|
||||
wait $PID1
|
||||
TEST_RESULT1=$?
|
||||
|
||||
# Check test results
|
||||
if [ $TEST_RESULT0 -eq 0 ] && [ $TEST_RESULT1 -eq 0 ]; then
|
||||
echo "✓ $test_case PASSED (scenario: $scenario)"
|
||||
PASSED_TESTS+=("$scenario:$test_case")
|
||||
else
|
||||
echo "✗ $test_case FAILED (scenario: $scenario)"
|
||||
if [ $TEST_RESULT0 -ne 0 ]; then
|
||||
echo " Rank 0 failed with code: $TEST_RESULT0"
|
||||
fi
|
||||
if [ $TEST_RESULT1 -ne 0 ]; then
|
||||
echo " Rank 1 failed with code: $TEST_RESULT1"
|
||||
fi
|
||||
FAILED_TESTS+=("$scenario:$test_case")
|
||||
fi
|
||||
done
|
||||
done
|
||||
|
||||
echo ""
|
||||
echo "========================================="
|
||||
echo "Test Summary"
|
||||
echo "========================================="
|
||||
echo "Passed tests: ${#PASSED_TESTS[@]}"
|
||||
for test in "${PASSED_TESTS[@]}"; do
|
||||
echo " ✓ $test"
|
||||
done
|
||||
|
||||
echo ""
|
||||
echo "Failed tests: ${#FAILED_TESTS[@]}"
|
||||
for test in "${FAILED_TESTS[@]}"; do
|
||||
echo " ✗ $test"
|
||||
done
|
||||
|
||||
# Cleanup: Stop Mooncake Master
|
||||
echo ""
|
||||
echo "Stopping Mooncake Master..."
|
||||
kill $MASTER_PID 2>/dev/null || true
|
||||
wait $MASTER_PID 2>/dev/null || true
|
||||
|
||||
# Exit with error if any tests failed
|
||||
if [ ${#FAILED_TESTS[@]} -gt 0 ]; then
|
||||
echo ""
|
||||
echo "Some tests failed!"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "All Hixl Mooncake Store tests completed successfully!"
|
||||
|
||||
|
||||
- name: Test Summary
|
||||
if: always()
|
||||
shell: bash
|
||||
run: |
|
||||
echo "CI Test completed"
|
||||
|
||||
- name: Upload Test Logs
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: test-logs-${{ github.run_number }}
|
||||
path: |
|
||||
/tmp/hixl-test-log/*
|
||||
retention-days: 30
|
||||
if-no-files-found: warn
|
||||
|
|
@ -0,0 +1,124 @@
|
|||
name: 'Build Wheel (CUDA 13)'
|
||||
|
||||
on:
|
||||
workflow_call: {}
|
||||
|
||||
jobs:
|
||||
build-wheel-cu13:
|
||||
runs-on: ubuntu-22.04
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.10', '3.12']
|
||||
env:
|
||||
BUILD_WITH_EP: "1"
|
||||
CU13_BUILD: "1"
|
||||
TORCH_CUDA_ARCH_LIST: "8.0;9.0"
|
||||
SCCACHE_GHA_ENABLED: "true"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
- name: Free up disk space
|
||||
run: |
|
||||
sudo rm -rf /usr/share/dotnet
|
||||
sudo rm -rf /opt/ghc
|
||||
sudo rm -rf /opt/hostedtoolcache/CodeQL
|
||||
sudo rm -rf /usr/local/lib/android
|
||||
df -h
|
||||
|
||||
- name: Install CUDA Toolkit
|
||||
uses: Jimver/cuda-toolkit@v0.2.29
|
||||
with:
|
||||
cuda: '13.0.2'
|
||||
linux-local-args: '["--toolkit"]'
|
||||
method: 'network'
|
||||
sub-packages: '["nvcc", "nvrtc-dev"]'
|
||||
non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
|
||||
|
||||
- name: Run sccache-cache
|
||||
uses: mozilla-actions/sccache-action@v0.0.9
|
||||
|
||||
- name: Configure sccache
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
|
||||
core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
|
||||
|
||||
- name: Run sccache stat for check
|
||||
shell: bash
|
||||
run: ${SCCACHE_PATH} --show-stats
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
sudo apt update -y
|
||||
sudo apt install -y ninja-build
|
||||
sudo bash -x dependencies.sh -y
|
||||
df -h
|
||||
shell: bash
|
||||
|
||||
- name: Configure project
|
||||
run: |
|
||||
mkdir build
|
||||
cd build
|
||||
cmake -G Ninja .. \
|
||||
-DUSE_ETCD=ON \
|
||||
-DUSE_REDIS=ON \
|
||||
-DUSE_HTTP=ON \
|
||||
-DWITH_STORE=ON \
|
||||
-DWITH_P2P_STORE=ON \
|
||||
-DWITH_EP=ON \
|
||||
-DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" \
|
||||
-DWITH_METRICS=ON \
|
||||
-DBUILD_UNIT_TESTS=OFF \
|
||||
-DBUILD_EXAMPLES=ON \
|
||||
-DENABLE_SCCACHE=ON \
|
||||
-DBUILD_BENCHMARK=ON \
|
||||
-DUSE_CUDA=ON \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
|
||||
shell: bash
|
||||
|
||||
- name: Build project
|
||||
run: |
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
|
||||
cd build
|
||||
cmake --build .
|
||||
sudo cmake --install .
|
||||
df -h
|
||||
shell: bash
|
||||
|
||||
- name: Build nvlink_allocator.so
|
||||
run: |
|
||||
mkdir -p build/mooncake-transfer-engine/nvlink-allocator
|
||||
cd mooncake-transfer-engine/nvlink-allocator
|
||||
export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
|
||||
bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
|
||||
shell: bash
|
||||
|
||||
- name: Generate Python version tag
|
||||
id: generate_tag
|
||||
run: |
|
||||
echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
|
||||
shell: bash
|
||||
|
||||
- name: Build Python wheel
|
||||
run: |
|
||||
PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag.outputs.python_version_tag }} ./scripts/build_wheel.sh
|
||||
shell: bash
|
||||
|
||||
- name: Upload Python wheel artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: mooncake-wheel-cu130-ubuntu-py${{ steps.generate_tag.outputs.python_version_tag }}
|
||||
path: mooncake-wheel/dist-py${{ steps.generate_tag.outputs.python_version_tag }}/*.whl
|
||||
|
|
@ -4,6 +4,10 @@ on:
|
|||
# Runs on pushes targeting the default branch
|
||||
push:
|
||||
branches: ["main"]
|
||||
paths:
|
||||
- 'docs/**'
|
||||
- 'requirements_docs.txt'
|
||||
- '.github/workflows/deploy.yml'
|
||||
# Allows you to run this workflow manually from the Actions tab
|
||||
workflow_dispatch:
|
||||
|
||||
|
|
@ -31,7 +35,9 @@ jobs:
|
|||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Setup Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,94 @@
|
|||
name: E2E CI
|
||||
|
||||
on:
|
||||
pull_request_target:
|
||||
branches: ["main"]
|
||||
types: [labeled]
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
pr_number:
|
||||
description: 'PR number that triggered this'
|
||||
required: false
|
||||
type: string
|
||||
pr_sha:
|
||||
description: 'PR head SHA to checkout'
|
||||
required: false
|
||||
type: string
|
||||
triggered_by:
|
||||
description: 'User who triggered this'
|
||||
required: false
|
||||
type: string
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: e2e-ci-${{ github.event.pull_request.number || inputs.pr_number || github.sha }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
ascend-test:
|
||||
if: >
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.label.name == 'run-e2e-ci'
|
||||
uses: ./.github/workflows/ci_ascend.yml
|
||||
with:
|
||||
checkout_ref: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
|
||||
secrets: inherit
|
||||
|
||||
integration-test:
|
||||
if: >
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event.label.name == 'run-e2e-ci'
|
||||
uses: ./.github/workflows/integration-test.yml
|
||||
with:
|
||||
pr_sha: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
|
||||
pr_number: ${{ inputs.pr_number || github.event.pull_request.number }}
|
||||
secrets: inherit
|
||||
|
||||
e2e-gate:
|
||||
name: E2E Gate
|
||||
if: >
|
||||
always() &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
github.event.label.name == 'run-e2e-ci')
|
||||
needs:
|
||||
- ascend-test
|
||||
- integration-test
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check E2E results
|
||||
run: |
|
||||
echo "PR: #${{ inputs.pr_number || github.event.pull_request.number }}"
|
||||
echo "SHA: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}"
|
||||
failing=$(echo "$NEEDS_JSON" | jq -r '
|
||||
to_entries[] |
|
||||
select(.value.result != "success" and .value.result != "skipped") |
|
||||
"\(.key): \(.value.result)"')
|
||||
if [ -n "$failing" ]; then
|
||||
echo "::error::The following E2E jobs failed:"
|
||||
echo "$failing"
|
||||
exit 1
|
||||
fi
|
||||
echo "All E2E checks passed."
|
||||
env:
|
||||
NEEDS_JSON: ${{ toJSON(needs) }}
|
||||
|
||||
cleanup-label:
|
||||
name: Cleanup E2E Label
|
||||
if: >
|
||||
always() &&
|
||||
github.event_name != 'workflow_dispatch' &&
|
||||
github.event.label.name == 'run-e2e-ci'
|
||||
needs:
|
||||
- e2e-gate
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Remove run-e2e-ci label
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
gh pr edit ${{ github.event.pull_request.number }} \
|
||||
--repo ${{ github.repository }} \
|
||||
--remove-label "run-e2e-ci" 2>/dev/null || true
|
||||
|
|
@ -0,0 +1,129 @@
|
|||
name: 'Integration test (Linux)'
|
||||
|
||||
on:
|
||||
workflow_call:
|
||||
inputs:
|
||||
pr_sha:
|
||||
description: 'PR head SHA (passed from parent workflow for workflow_dispatch)'
|
||||
required: false
|
||||
type: string
|
||||
pr_number:
|
||||
description: 'PR number (passed from parent workflow for workflow_dispatch)'
|
||||
required: false
|
||||
type: string
|
||||
|
||||
jobs:
|
||||
test-sglang-integration:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
tone_user_name: ${{ secrets.TONE_USER_NAME }}
|
||||
steps:
|
||||
- name: trigger T-one test
|
||||
if: ${{ env.tone_user_name != '' }}
|
||||
run: |
|
||||
# Priority: explicit inputs > PR event context > push SHA
|
||||
SHA="${{ inputs.pr_sha || github.event.pull_request.head.sha || github.sha }}"
|
||||
PR_ID="${{ inputs.pr_number || github.event.pull_request.number }}"
|
||||
|
||||
if [ "${{ github.event_name }}" = "push" ]; then
|
||||
SHA="${{ github.sha }}"
|
||||
PR_ID=""
|
||||
fi
|
||||
echo "PR_ID=${PR_ID}"
|
||||
max_attempts=120
|
||||
attempt=1
|
||||
while [ $attempt -le $max_attempts ]; do
|
||||
echo "Attempt $attempt: Fetching artifact..."
|
||||
if curl -L -fs -o artifact.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" https://api.github.com/repos/${{ github.repository }}/actions/artifacts?per_page=100; then
|
||||
artifact_id=""
|
||||
if jq empty artifact.json >/dev/null 2>&1; then
|
||||
artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | select(.name | contains(\"mooncake\") ) | select(.name | contains(\"cu130\") | not) | select(.workflow_run.head_sha == \"$SHA\" ) | .id" artifact.json | head -n 1)
|
||||
else
|
||||
echo "Failed to download artifact list. Retrying..."
|
||||
fi
|
||||
if [ -n "$artifact_id" ]; then
|
||||
echo "Successfully fetched expected artifact id $artifact_id"
|
||||
break
|
||||
else
|
||||
echo "Failed to fetch expected artifact. Retrying..."
|
||||
if [ $attempt -lt $max_attempts ]; then
|
||||
sleep $((attempt * 60 < 600 ? attempt * 60 : 600))
|
||||
fi
|
||||
fi
|
||||
else
|
||||
echo "Failed to fetch artifacts. Retrying..."
|
||||
if [ $attempt -lt $max_attempts ]; then
|
||||
sleep $((attempt * 60))
|
||||
fi
|
||||
fi
|
||||
attempt=$((attempt + 1))
|
||||
done
|
||||
if [ $attempt -gt $max_attempts ]; then
|
||||
echo "Failed to fetch artifacts after $max_attempts attempts"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
ENV_INFO="ARTIFACT_ID=${artifact_id} GIT_REPO=${{ github.repository }}"
|
||||
if [ -n "$PR_ID" ]; then
|
||||
ENV_INFO="${ENV_INFO} PR_ID=${PR_ID}"
|
||||
fi
|
||||
signature="${{ secrets.TONE_USER_NAME }}|${{ secrets.TONE_USER_TOKEN }}|$(python3 -c "import time;print(time.time())")"
|
||||
signature="$(python3 -c "import base64;print(base64.b64encode(\"$signature\".encode('utf-8')).decode('utf-8'))")"
|
||||
curl -s -H 'Content-Type: application/json' -X POST -d "{\"workspace\":\"mooncake_test\",\"project\":\"mooncake-ci\",\"template\":\"mooncake-ci-test\",\"name\":\"mooncake-ci-${SHA}\",\"username\":\"${{ secrets.TONE_USER_NAME }}\",\"env_ifs\":\" \",\"env_info\":\"${ENV_INFO}\",\"signature\":\"$signature\"}" https://tone.openanolis.cn/api/job/create/ > job.json
|
||||
if [ "$(jq .code job.json)" == 200 ]; then
|
||||
echo "job created"
|
||||
else
|
||||
echo "job create failed"
|
||||
exit 1
|
||||
fi
|
||||
job_id=$(jq .data.id job.json)
|
||||
echo "check job status here and remember to cancel it before restart the job !"
|
||||
echo "job_url: https://tone.openanolis.cn/ws/gclfnh19/test_result/${job_id}?tab=4"
|
||||
echo "job_id=${job_id}" >> $GITHUB_ENV
|
||||
shell: bash
|
||||
|
||||
- name: qurey job results
|
||||
if: ${{ env.tone_user_name != '' }}
|
||||
run: |
|
||||
total_time=0
|
||||
max_total_time=2880
|
||||
time=0
|
||||
max_time=240
|
||||
|
||||
while true; do
|
||||
if [ $total_time -gt $max_total_time ]; then
|
||||
echo "Total timeout reached (24 hours)"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ $time -gt $max_time ]; then
|
||||
echo "Current running job timeout reached (720 attempts)"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
signature="${{ secrets.TONE_USER_NAME }}|${{ secrets.TONE_USER_TOKEN }}|$(python3 -c "import time;print(time.time())")"
|
||||
signature="$(python3 -c "import base64;print(base64.b64encode(\"$signature\".encode('utf-8')).decode('utf-8'))")"
|
||||
curl -s -H 'Content-Type: application/json' -X POST -d "{\"username\":\"${{ secrets.TONE_USER_NAME }}\", \"signature\":\"$signature\", \"job_id\": \"${job_id}\"}" https://tone.openanolis.cn/api/job/query/ > job_status.json
|
||||
if ! [ "$(jq .code job_status.json)" == 200 ]; then
|
||||
echo "job query failed"
|
||||
exit 1
|
||||
fi
|
||||
job_state=$(jq .data.job_state job_status.json)
|
||||
job_status=$(jq .data.job_second_state job_status.json)
|
||||
|
||||
if [[ $job_status =~ "pass" ]]; then
|
||||
echo "job successful !"
|
||||
exit 0
|
||||
elif [[ $job_status =~ "fail" ]] ; then
|
||||
echo "job failed or stopped !"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ $job_state =~ "running" ]]; then
|
||||
time=$(( time + 1))
|
||||
fi
|
||||
|
||||
total_time=$(( total_time + 1 ))
|
||||
sleep 30
|
||||
done
|
||||
shell: bash
|
||||
|
|
@ -0,0 +1,145 @@
|
|||
name: Release CUDA 13
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'v*'
|
||||
|
||||
env:
|
||||
SCCACHE_GHA_ENABLED: "true"
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-22.04
|
||||
permissions:
|
||||
contents: write
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.10', '3.11', '3.12', '3.13']
|
||||
env:
|
||||
BUILD_WITH_EP: "1"
|
||||
CU13_BUILD: "1"
|
||||
TORCH_CUDA_ARCH_LIST: "8.0;9.0"
|
||||
steps:
|
||||
- name: Checkout source
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
- name: Free up disk space
|
||||
run: |
|
||||
sudo rm -rf /usr/share/dotnet
|
||||
sudo rm -rf /opt/ghc
|
||||
sudo rm -rf /opt/hostedtoolcache/CodeQL
|
||||
sudo rm -rf /usr/local/lib/android
|
||||
df -h
|
||||
|
||||
- name: Install CUDA Toolkit 13
|
||||
uses: Jimver/cuda-toolkit@v0.2.29
|
||||
with:
|
||||
cuda: '13.0.2'
|
||||
linux-local-args: '["--toolkit"]'
|
||||
method: 'network'
|
||||
sub-packages: '["nvcc", "nvrtc-dev"]'
|
||||
non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
|
||||
|
||||
- name: Run sccache-cache
|
||||
uses: mozilla-actions/sccache-action@v0.0.9
|
||||
|
||||
- name: Configure sccache
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
|
||||
core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
|
||||
|
||||
- name: Run sccache stat for check
|
||||
shell: bash
|
||||
run: ${SCCACHE_PATH} --show-stats
|
||||
|
||||
- name: Configure project
|
||||
run: |
|
||||
sudo apt update -y
|
||||
sudo bash -x dependencies.sh -y
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
|
||||
shell: bash
|
||||
|
||||
- name: Build project
|
||||
run: |
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
|
||||
cd build
|
||||
make -j
|
||||
sudo make install
|
||||
shell: bash
|
||||
|
||||
- name: Build nvlink_allocator.so
|
||||
run: |
|
||||
export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
|
||||
mkdir -p build/mooncake-transfer-engine/nvlink-allocator
|
||||
cd mooncake-transfer-engine/nvlink-allocator
|
||||
bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
|
||||
shell: bash
|
||||
|
||||
- name: Generate Python version tag
|
||||
id: generate_tag_release
|
||||
run: |
|
||||
echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
|
||||
shell: bash
|
||||
|
||||
- name: Build Python wheel
|
||||
run: |
|
||||
# Set LD_LIBRARY_PATH for wheel building
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
|
||||
PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
|
||||
env:
|
||||
VERSION: ${{ env.VERSION }}
|
||||
|
||||
- name: Upload Python wheel artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: mooncake-wheel-cuda13-py${{ steps.generate_tag_release.outputs.python_version_tag }}
|
||||
path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
|
||||
|
||||
publish-release:
|
||||
needs: build
|
||||
runs-on: ubuntu-22.04
|
||||
permissions:
|
||||
contents: write
|
||||
id-token: write
|
||||
steps:
|
||||
- name: Checkout source
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Download all wheel artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
path: mooncake-wheel/dist-all
|
||||
pattern: mooncake-wheel-cuda13-py*
|
||||
|
||||
- name: Prepare wheels for release
|
||||
run: |
|
||||
# Move all wheels to a single directory
|
||||
mkdir -p mooncake-wheel/dist-release
|
||||
find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \;
|
||||
ls -la mooncake-wheel/dist-release/
|
||||
# List all collected wheels
|
||||
echo "Collected wheels for release:"
|
||||
ls -la mooncake-wheel/dist-release/
|
||||
|
||||
- name: Upload wheels to GitHub Release
|
||||
uses: softprops/action-gh-release@v1
|
||||
with:
|
||||
files: mooncake-wheel/dist-release/*.whl
|
||||
|
||||
- name: Publish package to PyPI
|
||||
if: github.repository == 'kvcache-ai/Mooncake'
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: mooncake-wheel/dist-release/
|
||||
password: ${{ secrets.PYPI_CU13_API_TOKEN }}
|
||||
|
|
@ -14,7 +14,7 @@ jobs:
|
|||
contents: write
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
|
||||
python-version: ['3.10', '3.11', '3.12', '3.13']
|
||||
env:
|
||||
BUILD_WITH_EP: "0"
|
||||
NON_CUDA_BUILD: "1"
|
||||
|
|
|
|||
|
|
@ -14,10 +14,9 @@ jobs:
|
|||
contents: write
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
|
||||
python-version: ['3.10', '3.11', '3.12', '3.13']
|
||||
env:
|
||||
BUILD_WITH_EP: "1"
|
||||
EP_TORCH_VERSIONS: "2.8.0;2.9.0;2.9.1"
|
||||
TORCH_CUDA_ARCH_LIST: "8.0;9.0"
|
||||
steps:
|
||||
- name: Checkout source
|
||||
|
|
@ -33,6 +32,8 @@ jobs:
|
|||
sudo rm -rf /usr/share/dotnet
|
||||
sudo rm -rf /opt/ghc
|
||||
sudo rm -rf /opt/hostedtoolcache/CodeQL
|
||||
sudo rm -rf /usr/local/lib/android
|
||||
df -h
|
||||
|
||||
- name: Install CUDA Toolkit
|
||||
uses: Jimver/cuda-toolkit@v0.2.24
|
||||
|
|
@ -63,7 +64,7 @@ jobs:
|
|||
sudo bash -x dependencies.sh -y
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
|
||||
cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
|
||||
shell: bash
|
||||
|
||||
- name: Build project
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ build_ofed4
|
|||
old
|
||||
local_test
|
||||
go.sum
|
||||
!mooncake-common/etcd/go.sum
|
||||
*.so
|
||||
bin
|
||||
mod
|
||||
|
|
@ -193,6 +194,7 @@ cmake-build
|
|||
libetcd_wrapper.h
|
||||
|
||||
mooncake-wheel/mooncake/allocator.py
|
||||
mooncake-wheel/mooncake/allocator_ascend_npu.py
|
||||
mooncake-wheel/mooncake/mooncake_master
|
||||
mooncake-wheel/mooncake/transfer_engine_bench
|
||||
|
||||
|
|
@ -201,3 +203,6 @@ CLAUDE.md
|
|||
|
||||
# CodeQL
|
||||
_codeql_detected_source_root
|
||||
|
||||
# CodeBuddy Memory
|
||||
.codebuddy/
|
||||
|
|
|
|||
|
|
@ -2,3 +2,7 @@
|
|||
path = extern/pybind11
|
||||
url = https://github.com/pybind/pybind11.git
|
||||
branch = stable
|
||||
[submodule "extern/yalantinglibs"]
|
||||
path = extern/yalantinglibs
|
||||
url = https://github.com/alibaba/yalantinglibs.git
|
||||
branch = v0.5.7
|
||||
|
|
|
|||
|
|
@ -23,6 +23,16 @@ repos:
|
|||
- id: check-added-large-files
|
||||
args: ['--maxkb=1024']
|
||||
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: mooncake-code-format
|
||||
name: Run Mooncake code format script
|
||||
entry: ./scripts/code_format.sh
|
||||
language: system
|
||||
pass_filenames: false
|
||||
always_run: true
|
||||
require_serial: true
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.6.9
|
||||
hooks:
|
||||
|
|
@ -37,10 +47,10 @@ repos:
|
|||
hooks:
|
||||
- id: codespell
|
||||
exclude: '^(extern/|FAST25-release/)'
|
||||
args: ['--ignore-words-list=te,mooncake,KVCache']
|
||||
args: ['--ignore-words-list=te,mooncake,KVCache,cann']
|
||||
|
||||
- repo: https://github.com/pre-commit/mirrors-clang-format
|
||||
rev: v19.1.0
|
||||
rev: v20.1.8
|
||||
hooks:
|
||||
- id: clang-format
|
||||
files: '\.(c|cc|cpp|cxx|h|hpp)$'
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
[default]
|
||||
extend-ignore-words = ["CANN", "ASO", "fre"]
|
||||
extend-ignore-words = ["CANN", "ASO", "fre", "wqs"]
|
||||
|
||||
[default.extend-words]
|
||||
CANN = "CANN"
|
||||
ASO = "ASO"
|
||||
fre = "fre"
|
||||
wqs = "wqs"
|
||||
|
||||
[files]
|
||||
extend-exclude = [
|
||||
|
|
|
|||
114
CMakeLists.txt
114
CMakeLists.txt
|
|
@ -14,12 +14,14 @@ endif()
|
|||
|
||||
option(WITH_TE "build mooncake transfer engine and sample code" ON)
|
||||
option(WITH_STORE "build mooncake store library and sample code" ON)
|
||||
option(WITH_STORE_GO "build Go bindings for mooncake store" OFF)
|
||||
option(WITH_P2P_STORE "build p2p store library and sample code" OFF)
|
||||
option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF)
|
||||
option(WITH_STORE_RUST "build the Rust bindings for the Mooncake Store" ON)
|
||||
option(WITH_EP "build mooncake with expert parallelism support" OFF)
|
||||
|
||||
include(${CMAKE_CURRENT_SOURCE_DIR}/mooncake-common/SetupPython.cmake)
|
||||
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11)
|
||||
set(PYTHON_EXECUTABLE "python3")
|
||||
execute_process(
|
||||
COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print(sys.path[-1])"
|
||||
OUTPUT_VARIABLE PYTHON_SYS_PATH
|
||||
|
|
@ -39,9 +41,26 @@ option(STORE_USE_ETCD "build mooncake store with etcd" OFF)
|
|||
if (STORE_USE_ETCD)
|
||||
add_compile_definitions(STORE_USE_ETCD)
|
||||
endif()
|
||||
option(STORE_USE_REDIS "build mooncake store with redis" OFF)
|
||||
if (STORE_USE_REDIS)
|
||||
add_compile_definitions(STORE_USE_REDIS)
|
||||
endif()
|
||||
option(STORE_USE_K8S_LEASE "build mooncake store with K8s Lease leader election" OFF)
|
||||
if (STORE_USE_K8S_LEASE)
|
||||
if (STORE_USE_ETCD)
|
||||
message(FATAL_ERROR "STORE_USE_K8S_LEASE and STORE_USE_ETCD cannot be enabled together because both build Go c-shared HA backends.")
|
||||
endif()
|
||||
if (USE_ETCD AND NOT USE_ETCD_LEGACY)
|
||||
message(FATAL_ERROR "STORE_USE_K8S_LEASE cannot be enabled with non-legacy USE_ETCD because both build Go c-shared libraries in the same process.")
|
||||
endif()
|
||||
add_compile_definitions(STORE_USE_K8S_LEASE)
|
||||
endif()
|
||||
|
||||
option(STORE_USE_JEMALLOC "Use jemalloc in mooncake store master" OFF)
|
||||
|
||||
# Define ASIO macros before building targets that include ASIO headers.
|
||||
add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK)
|
||||
|
||||
add_subdirectory(mooncake-common)
|
||||
include_directories(mooncake-common/etcd)
|
||||
include_directories(mooncake-common/include)
|
||||
|
|
@ -57,12 +76,103 @@ if (WITH_STORE)
|
|||
include_directories(mooncake-store/include)
|
||||
endif()
|
||||
|
||||
if (WITH_STORE_RUST)
|
||||
if (NOT WITH_STORE)
|
||||
message(FATAL_ERROR "WITH_STORE_RUST=ON requires WITH_STORE=ON")
|
||||
endif()
|
||||
message(STATUS "Mooncake Store Rust bindings will be built")
|
||||
add_subdirectory(mooncake-store/rust)
|
||||
endif()
|
||||
|
||||
option(EP_USE_IDE "Enable intelligent indexing for IDEs" OFF)
|
||||
if (WITH_EP)
|
||||
message(WARNING "Option `WITH_EP` is deprecated. Mooncake EP now builds with setuptools. Please set environment variable BUILD_WITH_EP=1 to enable.")
|
||||
if (EP_USE_IDE)
|
||||
message(WARNING "EP_USE_IDE enabled. DO NOT USE IN PRODUCTION!")
|
||||
add_subdirectory(mooncake-ep)
|
||||
include_directories(mooncake-ep/include)
|
||||
add_subdirectory(mooncake-pg)
|
||||
include_directories(mooncake-pg/include)
|
||||
else ()
|
||||
message(STATUS "WITH_EP enabled: building Mooncake EP and PG Python extensions")
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
message(STATUS "Detected CUDA version: ${CUDAToolkit_VERSION}")
|
||||
|
||||
# EP_TORCH_VERSIONS: semicolon-separated list of PyTorch versions to build for.
|
||||
# Can be set via -DEP_TORCH_VERSIONS="2.9.1;2.8.0" or the EP_TORCH_VERSIONS env var.
|
||||
# Empty means build with the currently-installed torch.
|
||||
if(NOT EP_TORCH_VERSIONS)
|
||||
set(EP_TORCH_VERSIONS "$ENV{EP_TORCH_VERSIONS}")
|
||||
endif()
|
||||
set(EP_TORCH_VERSIONS "${EP_TORCH_VERSIONS}" CACHE STRING
|
||||
"PyTorch versions for EP/PG extensions, semicolon-separated (empty = use currently-installed torch)")
|
||||
|
||||
# TORCH_CUDA_ARCH_LIST forwarded to the torch CUDA extension build.
|
||||
if(NOT TORCH_CUDA_ARCH_LIST)
|
||||
set(TORCH_CUDA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}")
|
||||
endif()
|
||||
if(NOT TORCH_CUDA_ARCH_LIST)
|
||||
set(TORCH_CUDA_ARCH_LIST "8.0;9.0")
|
||||
endif()
|
||||
set(TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}" CACHE STRING
|
||||
"CUDA arch list for EP/PG extension builds (e.g. \"8.0;9.0\")")
|
||||
|
||||
# Staging directory: EP/PG .so files are placed here during make and later
|
||||
# injected into the wheel AFTER auditwheel, so patchelf never touches the
|
||||
# CUDA fatbins (which would cause cudaErrorInvalidKernelImage at runtime).
|
||||
set(EP_PG_STAGING_DIR "${CMAKE_BINARY_DIR}/ep_pg_staging")
|
||||
|
||||
# Convert semicolon-separated lists to pipe-separated strings so they survive
|
||||
# CMake's COMMAND list-splitting (semicolons are CMake list separators).
|
||||
string(REPLACE ";" "|" _ep_torch_versions_pipe "${EP_TORCH_VERSIONS}")
|
||||
string(REPLACE ";" "|" _torch_cuda_arch_list_pipe "${TORCH_CUDA_ARCH_LIST}")
|
||||
|
||||
add_custom_target(mooncake_ep_ext ALL
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory "${EP_PG_STAGING_DIR}"
|
||||
COMMAND ${CMAKE_COMMAND}
|
||||
"-DSOURCE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep"
|
||||
"-DEP_CUDA_MAJOR=${CUDAToolkit_VERSION_MAJOR}"
|
||||
"-DEP_CUDA_MINOR=${CUDAToolkit_VERSION_MINOR}"
|
||||
"-DEP_TORCH_VERSIONS=${_ep_torch_versions_pipe}"
|
||||
"-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
|
||||
"-DSTAGING_DIR=${EP_PG_STAGING_DIR}"
|
||||
"-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
|
||||
-P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep/BuildEpExt.cmake"
|
||||
COMMENT "Building Mooncake EP Python extension(s)"
|
||||
DEPENDS engine
|
||||
VERBATIM
|
||||
)
|
||||
|
||||
add_custom_target(mooncake_pg_ext ALL
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory "${EP_PG_STAGING_DIR}"
|
||||
COMMAND ${CMAKE_COMMAND}
|
||||
"-DSOURCE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg"
|
||||
"-DEP_CUDA_MAJOR=${CUDAToolkit_VERSION_MAJOR}"
|
||||
"-DEP_CUDA_MINOR=${CUDAToolkit_VERSION_MINOR}"
|
||||
"-DEP_TORCH_VERSIONS=${_ep_torch_versions_pipe}"
|
||||
"-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
|
||||
"-DSTAGING_DIR=${EP_PG_STAGING_DIR}"
|
||||
"-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
|
||||
-P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg/BuildPgExt.cmake"
|
||||
COMMENT "Building Mooncake PG Python extension(s)"
|
||||
DEPENDS engine mooncake_ep_ext
|
||||
VERBATIM
|
||||
)
|
||||
endif ()
|
||||
endif()
|
||||
|
||||
add_subdirectory(mooncake-integration)
|
||||
|
||||
if (WITH_STORE_GO AND WITH_STORE)
|
||||
add_custom_target(build_store_go DEPENDS mooncake_store transfer_engine)
|
||||
add_custom_command(
|
||||
TARGET build_store_go
|
||||
COMMAND bash build.sh ${CMAKE_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR} ${USE_ETCD} ${USE_REDIS} ${USE_HTTP} ${USE_ETCD_LEGACY}
|
||||
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mooncake-store/go
|
||||
)
|
||||
set_property(TARGET build_store_go PROPERTY EXCLUDE_FROM_ALL FALSE)
|
||||
message(STATUS "Mooncake Store Go bindings will be built")
|
||||
endif()
|
||||
|
||||
if (WITH_P2P_STORE)
|
||||
add_subdirectory(mooncake-p2p-store)
|
||||
message(STATUS "P2P Store will be built")
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ Mooncake uses [pre-commit](https://pre-commit.com/) to enforce consistent format
|
|||
| Type | Tool | Purpose |
|
||||
|------|------|---------|
|
||||
| Generic | trailing-whitespace / end-of-file-fixer | Basic hygiene |
|
||||
| Project | `./scripts/code_format.sh` | Enforce Mooncake C/C++ formatting script before commit |
|
||||
| Python | ruff / ruff-format | Lint + format (includes import sorting) |
|
||||
| Spelling | codespell | Catch common typos (ignores domain-specific words) |
|
||||
| C/C++ | clang-format | Apply style from the repository's `.clang-format` |
|
||||
|
|
@ -53,6 +54,8 @@ pip install -r requirements-dev.txt
|
|||
pre-commit install
|
||||
```
|
||||
|
||||
After installation, every commit will run `./scripts/code_format.sh` automatically. If it rewrites files, re-stage the changes and commit again.
|
||||
|
||||
#### Usage
|
||||
Run on all files (first run will install hook environments):
|
||||
```bash
|
||||
|
|
|
|||
25
Dockerfile
25
Dockerfile
|
|
@ -1,25 +0,0 @@
|
|||
# Base Image from Alibaba Cloud AC2
|
||||
FROM ac2-registry.cn-hangzhou.cr.aliyuncs.com/ac2/pytorch-ubuntu:2.3.0-cuda12.1.1-ubuntu22.04
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY . .
|
||||
|
||||
ENV GOROOT=/usr/local/go
|
||||
ENV PATH=$GOROOT/bin:$PATH
|
||||
|
||||
RUN apt update \
|
||||
&& apt install -y unzip wget cmake git sudo \
|
||||
&& pip install pybind11
|
||||
|
||||
# Execute installation in the container
|
||||
RUN bash dependencies.sh \
|
||||
&& apt autoremove -y \
|
||||
&& apt clean -y \
|
||||
&& rm -rf /tmp/* /var/tmp/* \
|
||||
&& find /var/cache/apt/archives /var/lib/apt/lists -not -name lock -type f -delete \
|
||||
&& find /var/cache -type f -delete
|
||||
|
||||
RUN rm -rf build || true \
|
||||
&& mkdir build && cd build \
|
||||
&& cmake .. -DSTORE_USE_ETCD=ON && make -j$(nproc) && make install
|
||||
|
|
@ -12,5 +12,6 @@ Current list of codeowners on this project:
|
|||
| ---------------------------- | ---------------------------- | ---------------------------- | ---------------------------- |
|
||||
| <img src="image/partners/approaching_ai_logo.png" width="120"/> | <img src="image/partners/ant_group_logo.png" width="120"/> | <img src="image/partners/huawei_logo.png" width="120"/> | <img src="image/partners/nvidia_logo.png" width="120"/> |
|
||||
| <img src="image/partners/moore_thread_logo.jpg" width="120"/> | <img src="image/partners/tencent_logo.png" width="120"/> | <img src="image/partners/volcengine_logo.png" width="120"/> | <img src="image/partners/amd_logo.png" width="120"/> |
|
||||
| <img src="image/partners/ieitsystems_logo.png" width="120"/> | | | |
|
||||
|
||||
Want to include your company logo? Just open a Pull Request!
|
||||
148
README.md
148
README.md
|
|
@ -8,13 +8,15 @@
|
|||
| <a href="FAST25-release/traces" target="_blank"><strong>Traces</strong></a>
|
||||
| <a href="https://arxiv.org/abs/2407.00079" target="_blank"><strong>Technical Report</strong></a>
|
||||
| <a href="https://kvcache-ai.github.io/Mooncake/" target="_blank"><strong>Blog</strong></a>
|
||||
| <a href="https://join.slack.com/t/mooncake-project/shared_invite/zt-3ig4fjai8-KH1zIm3x8Vm8WqyH0i_JaA" target="_blank"><strong>Slack</strong></a>
|
||||
| <a href="https://join.slack.com/t/mooncake-project/shared_invite/zt-3qx4x35ea-zSSTqTHItHJs9SCoXLOSPA" target="_blank"><strong>Slack</strong></a>
|
||||
<br />
|
||||
<br />
|
||||
|
||||
[](https://kvcache-ai.github.io/Mooncake/)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine-cuda13)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine)
|
||||
[](https://deepwiki.com/kvcache-ai/Mooncake)
|
||||
[](https://github.com/kvcache-ai/Mooncake/graphs/commit-activity)
|
||||
|
|
@ -25,14 +27,26 @@
|
|||
|
||||
Mooncake is the serving platform for <a href="https://kimi.ai/"><img src="image/kimi.png" alt="icon" style="height: 16px; vertical-align: middle;"> Kimi</a>, a leading LLM service provided by <a href="https://www.moonshot.cn/"><img src="image/moonshot.jpg" alt="icon" style="height: 16px; vertical-align: middle;"> Moonshot AI</a>.
|
||||
Now both the Transfer Engine and Mooncake Store are open-sourced!
|
||||
This repository also hosts its technical report and the open sourced traces.
|
||||
This repository also hosts its technical report and the open-sourced traces.
|
||||
|
||||
<h2 id="updates">🔄 Updates</h2>
|
||||
|
||||
- **Mar 19, 2026**: [TorchSpec: Speculative Decoding Training at Scale](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale) is [open sourced](https://github.com/torchspec-project/TorchSpec), using Mooncake to decouple inference and training via efficient hidden states management.
|
||||
- **Mar 5, 2026**: [LightX2V](https://github.com/ModelTC/LightX2V/pull/893) now supports disaggregated deployment based on Mooncake, enabling encoder/transformer service decoupling with Mooncake Transfer Engine for high-performance cross-device and cross-machine data transfer.
|
||||
- **Feb 25, 2026**: [SGLang](https://github.com/sgl-project/sglang) merged [Encoder Global Cache Manager](https://github.com/sgl-project/sglang/pull/16137), introducing a Mooncake-powered global multimodal embedding cache that enables cross-instance sharing of ViT embeddings to avoid redundant GPU computation.
|
||||
- **Feb 24, 2026**: [vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/disaggregated_inference/) introduces disaggregated inference connectors with support for both `MooncakeStoreConnector` and `MooncakeTransferEngineConnector` for multi-node omni-modality pipelines.
|
||||
- **Feb 12, 2026**: [Mooncake Joins PyTorch Ecosystem](https://pytorch.org/blog/mooncake-joins-pytorch-ecosystem/) We are thrilled to announce that Mooncake has officially joined the PyTorch Ecosystem!
|
||||
- **Jan 28, 2026**: [FlexKV](https://github.com/taco-project/FlexKV), a distributed KV store and cache system from Tencent and NVIDIA in collaboration with the community, now supports [distributed KVCache reuse](https://github.com/taco-project/FlexKV/blob/main/docs/dist_reuse/README_en.md) with the Mooncake Transfer Engine.
|
||||
- **Dec 27, 2025**: Collaboration with [ROLL](https://github.com/alibaba/ROLL)! Check out the paper [here](https://arxiv.org/abs/2512.22560).
|
||||
- **Dec 23, 2025**: SGLang introduces [Encode-Prefill-Decode (EPD) Disaggregation](https://lmsys.org/blog/2026-01-12-epd/) with Mooncake as a transfer backend. This integration allows decoupling compute-intensive multimodal encoders (e.g., Vision Transformers) from language model nodes, utilizing Mooncake's RDMA engine for zero-copy transfer of large multimodal embeddings.
|
||||
- **Dec 19, 2025**: Mooncake Transfer Engine has been [integrated into TensorRT LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils) for KVCache transfer in PD-disaggregated inference.
|
||||
- **Dec 19, 2025**: Mooncake Transfer Engine has been directly integrated into vLLM v1 as a [KV Connector](https://docs.vllm.ai/en/latest/features/mooncake_connector_usage/) in PD-disaggregated setups.
|
||||
- **Nov 07, 2025**: [RBG + SGLang HiCache + Mooncake](https://github.com/sgl-project/rbg/blob/main/keps/74-mooncake-integration/README.md), a role-based out-of-the-box solution for cloud native deployment, which is elastic, scalable, and high-performance.
|
||||
- **Sept 18, 2025**: Mooncake Store empowers vLLM Ascend by serving as [the distributed KV cache pool backend](https://docs.vllm.ai/projects/ascend/zh-cn/main/user_guide/feature_guide/kv_pool.html).
|
||||
- **Sept 10, 2025**: SGLang officially supports Mooncake Store as a [hierarchical KV caching storage backend](https://lmsys.org/blog/2025-09-10-sglang-hicache/). The integration extends RadixAttention with multi-tier KV cache storage across device, host, and remote storage layers.
|
||||
- **Sept 10, 2025**: The official & high-performance version of Mooncake P2P Store is open-sourced as [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine/). It has been successfully applied in K1.5 and K2 production training, updating Kimi-K2 model (1T parameters) across thousands of GPUs in ~20s.
|
||||
- **Aug 23, 2025**: [xLLM](https://github.com/jd-opensource/xllm) high-performance inference engine builds hybrid KV cache management based on Mooncake, supporting global KV cache management with intelligent offloading and prefetching.
|
||||
- **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://docs.vllm.ai/projects/ascend/en/latest/tutorials/multi_node_pd_disaggregation_mooncake.html) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs.
|
||||
- **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/feature_guide/disaggregated_prefill.html) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs.
|
||||
- **Jul 20, 2025**: Mooncake powers [the deployment of Kimi K2](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/) on 128 H200 GPUs with PD disaggregation and large-scale expert parallelism, achieving 224k tokens/sec prefill throughput and 288k tokens/sec decode throughput.
|
||||
- **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](https://kvcache-ai.github.io/Mooncake/getting_started/examples/lmdeploy-integration-v0.9.html) for LMDeploy.
|
||||
- **May 9, 2025**: NIXL officially supports Mooncake Transfer Engine as [a backend plugin](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/mooncake/README.md).
|
||||
|
|
@ -40,54 +54,57 @@ This repository also hosts its technical report and the open sourced traces.
|
|||
- **May 5, 2025**: Supported by Mooncake Team, SGLang release <a href="https://lmsys.org/blog/2025-05-05-large-scale-ep/" target="_blank">guidance</a> to deploy DeepSeek with PD Disaggregation on 96 H100 GPUs.
|
||||
- **Apr 22, 2025**: LMCache officially supports Mooncake Store as a <a href="https://blog.lmcache.ai/2025-04-22-tencent/" target="_blank">remote connector</a>.
|
||||
- **Apr 10, 2025**: SGLang officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer.
|
||||
- **Mar 7, 2025**: We open sourced the Mooncake Store, a distributed KVCache based on Transfer Engine. vLLM's xPyD disaggregated prefilling & decoding based on Mooncake Store will be released soon.
|
||||
- **Mar 7, 2025**: We open-sourced the Mooncake Store, a distributed KVCache based on Transfer Engine. vLLM's xPyD disaggregated prefilling & decoding based on Mooncake Store will be released soon.
|
||||
- **Feb 25, 2025**: Mooncake receives the **Best Paper Award** at **FAST 2025**!
|
||||
- **Feb 21, 2025**: The updated <a href="FAST25-release/traces" target="_blank">traces</a> used in our FAST'25 paper have been released.
|
||||
- **Dec 16, 2024**: vLLM officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer.
|
||||
- **Nov 28, 2024**: We open sourced the Transfer Engine, the central component of Mooncake. We also provide two demonstrations of Transfer Engine: a P2P Store and vLLM integration.
|
||||
- **July 9, 2024**: We open sourced the trace as a <a href="https://github.com/kvcache-ai/Mooncake/blob/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl" target="_blank">jsonl file</a>.
|
||||
- **Nov 28, 2024**: We open-sourced the Transfer Engine, the central component of Mooncake. We also provide two demonstrations of Transfer Engine: a P2P Store and vLLM integration.
|
||||
- **July 9, 2024**: We open-sourced the trace as a <a href="https://github.com/kvcache-ai/Mooncake/blob/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl" target="_blank">JSONL file</a>.
|
||||
- **June 27, 2024**: We present a series of Chinese blogs with more discussions on <a href="https://zhuanlan.zhihu.com/p/705754254">zhihu 1</a>, <a href="https://zhuanlan.zhihu.com/p/705910725">2</a>, <a href="https://zhuanlan.zhihu.com/p/706204757">3</a>, <a href="https://zhuanlan.zhihu.com/p/707997501">4</a>, <a href="https://zhuanlan.zhihu.com/p/9461861451">5</a>, <a href="https://zhuanlan.zhihu.com/p/1939988652114580803">6</a>, <a href="https://zhuanlan.zhihu.com/p/1959366095443064318">7</a>.
|
||||
- **June 26, 2024**: Initial technical report release.
|
||||
|
||||
|
||||
<h2 id="overview">🎉 Overview</h2>
|
||||
|
||||
Mooncake features a KVCache-centric disaggregated architecture that separates the prefill and decoding clusters. It also leverages the underutilized CPU, DRAM, and SSD resources of the GPU cluster to implement a disaggregated cache of KVCache.
|
||||
Mooncake features a KVCache-centric disaggregated architecture that separates the prefill and decoding clusters. It also leverages the underutilized CPU, DRAM, and SSD resources of the GPU cluster to implement a disaggregated KVCache pool.
|
||||
|
||||

|
||||
|
||||
The core of Mooncake is its KVCache-centric scheduler, which balances maximizing overall effective throughput while meeting latency-related Service Level Objectives (SLOs) requirements. Unlike traditional studies that assume all requests will be processed, Mooncake faces challenges due to highly overloaded scenarios. To mitigate these, we developed a prediction-based early rejection policy. Experiments show that Mooncake excels in long-context scenarios. Compared to the baseline method, Mooncake can achieve up to a 525% increase in throughput in certain simulated scenarios while adhering to SLOs. Under real workloads, Mooncake’s innovative architecture enables <a href="https://kimi.ai/">Kimi</a> to handle 75% more requests.
|
||||
The core of Mooncake is its KVCache-centric scheduler, which balances maximizing overall effective throughput while meeting latency-related Service Level Objectives (SLOs). Unlike traditional studies that assume all requests will be processed, Mooncake faces challenges in highly overloaded scenarios. To mitigate these, we developed a prediction-based early rejection policy. Experiments show that Mooncake excels in long-context scenarios. Compared to the baseline method, Mooncake can achieve up to a 525% increase in throughput in certain simulated scenarios while adhering to SLOs. Under real workloads, Mooncake’s innovative architecture enables <a href="https://kimi.ai/">Kimi</a> to handle 75% more requests.
|
||||
|
||||
<h2 id="components">🧩 Components</h2>
|
||||
|
||||
<!--  -->
|
||||
<img src=image/components.png width=74% />
|
||||
|
||||
**Mooncake Core Component: Transfer Engine (TE)**
|
||||
**Mooncake Core Component: Transfer Engine (TE)**
|
||||
The core of Mooncake is the Transfer Engine (TE), which provides a unified interface for batched data transfer across various storage devices and network links. Supporting multiple protocols including TCP, RDMA, CXL/shared-memory, and NVMe over Fabric (NVMe-of), TE is designed to enable fast and reliable data transfer for AI workloads. Compared to Gloo (used by Distributed PyTorch) and traditional TCP, TE achieves significantly lower I/O latency, making it a superior solution for efficient data transmission.
|
||||
|
||||
**P2P Store and Mooncake Store**
|
||||
**P2P Store and Mooncake Store**
|
||||
Both P2P Store and Mooncake Store are built on the Transfer Engine and provide key/value caching for different scenarios. P2P Store focuses on sharing temporary objects (e.g., checkpoint files) across nodes in a cluster, preventing bandwidth saturation on a single machine. Mooncake Store, on the other hand, supports distributed pooled KVCache, specifically designed for XpYd disaggregation to enhance resource utilization and system performance.
|
||||
|
||||
**Mooncake Integration with Leading LLM Inference Systems**
|
||||
**Mooncake Integration with Leading LLM Inference Systems**
|
||||
Mooncake has been seamlessly integrated with several popular large language model (LLM) inference systems. Through collaboration with the vLLM and SGLang teams, Mooncake now officially supports prefill-decode disaggregation. By leveraging the high-efficiency communication capabilities of RDMA devices, Mooncake significantly improves inference efficiency in prefill-decode disaggregation scenarios, providing robust technical support for large-scale distributed inference tasks.
|
||||
In addition, Mooncake has been successfully integrated with SGLang's Hierarchical KV Caching, vLLM's prefill serving, and LMCache, augmenting KV cache management capabilities across large-scale inference scenarios.
|
||||
|
||||
**Elastic Expert Parallelism Support**
|
||||
Mooncake adds elasticity and fault tolerance support for MoE model inference, enabling inference systems to remain responsive and recoverable in the event of GPU failures or changes in resource configuration. This functionality includes automatic faulty rank detection and can incorporate with the EPLB module to dynamically route tokens to healthy ranks during inference.
|
||||
**Elastic Expert Parallelism Support**
|
||||
Mooncake adds elasticity and fault tolerance support for MoE model inference, enabling inference systems to remain responsive and recoverable in the event of GPU failures or changes in resource configuration. This functionality includes automatic faulty rank detection and can work with the EPLB module to dynamically route tokens to healthy ranks during inference.
|
||||
|
||||
**Tensor-Centric Ecosystem**
|
||||
Mooncake establishes a full-stack, Tensor-oriented AI infrastructure where Tensors serve as the fundamental data carrier. The ecosystem spans from the Transfer Engine, which accelerates Tensor data movement across heterogeneous storage (DRAM/VRAM/NVMe), to the P2P Store and Mooncake Store for distributed management of Tensor objects (e.g., Checkpoints and KVCache), up to the Mooncake Backend enabling Tensor-based elastic distributed computing. This architecture is designed to maximize Tensor processing efficiency for large-scale model inference and training.
|
||||
|
||||
<h2 id="show-cases">🔥 Show Cases</h2>
|
||||
|
||||
### Use Transfer Engine Standalone ([Guide](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html))
|
||||
|
||||
Transfer Engine is a high-performance data transfer framework. Transfer Engine provides a unified interface to transfer data from DRAM, VRAM or NVMe, while the technical details related to hardware are hidden. Transfer Engine supports TCP, RDMA (InfiniBand/RoCEv2/eRDMA/NVIDIA GPUDirect) and NVMe over Fabric (NVMe-of) protocols.
|
||||
Transfer Engine is a high-performance data transfer framework. Transfer Engine provides a unified interface to transfer data from DRAM, VRAM or NVMe, while the technical details related to hardware are hidden. Transfer Engine supports multiple communication protocols including TCP, RDMA (InfiniBand/RoCEv2/eRDMA/NVIDIA GPUDirect), NVMe over Fabric (NVMe-of), NVLink, HIP, CXL, and Ascend. When built with the corresponding runtime, Transfer Engine can also detect and route accelerator memory on CUDA, MUSA, HIP, and Cambricon MLU devices. For a complete list of supported protocols and configuration guide, see the [Supported Protocols Documentation](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html).
|
||||
|
||||
#### Highlights
|
||||
- **Efficient use of multiple RDMA NIC devices.** Transfer Engine supports the use of multiple RDMA NIC devices to achieve the *aggregation of transfer bandwidth*.
|
||||
|
||||
- **Topology aware path selection.** Transfer Engine can *select optimal devices* based on the location (NUMA affinity, etc.) of both source and destination.
|
||||
|
||||
- **More robust on temporary network error.** Once transmission fails, Transfer Engine will try to use alternative paths for data delivery automatically.
|
||||
- **More robust against temporary network errors.** Once transmission fails, Transfer Engine will try to use alternative paths for data delivery automatically.
|
||||
|
||||
#### Performance
|
||||
With 40 GB of data (equivalent to the size of the KVCache generated by 128k tokens in the LLaMA3-70B model), Mooncake Transfer Engine delivers up to **87 GB/s** and **190 GB/s** of bandwidth in 4×200 Gbps and 8×400 Gbps RoCE networks respectively, which are about **2.4x and 4.6x faster** than the TCP protocol.
|
||||
|
|
@ -96,7 +113,7 @@ With 40 GB of data (equivalent to the size of the KVCache generated by 128k toke
|
|||
<img src=image/transfer-engine-performance.png width=75% />
|
||||
|
||||
### P2P Store ([Guide](https://kvcache-ai.github.io/Mooncake/design/p2p-store.html))
|
||||
P2P Store is built on the Transfer Engine and supports sharing temporary objects between peer nodes in a cluster. P2P Store is ideal for scenarios like checkpoint transfer, where data needs to be rapidly and efficiently shared across a cluster.
|
||||
P2P Store is built on the Transfer Engine and supports sharing temporary objects between peer nodes in a cluster. P2P Store is ideal for scenarios like checkpoint transfer, where data needs to be rapidly and efficiently shared across a cluster.
|
||||
**P2P Store has been used in the checkpoint transfer service of Moonshot AI.**
|
||||
|
||||
#### Highlights
|
||||
|
|
@ -132,7 +149,7 @@ SGLang officially supports Mooncake Store as a [HiCache storage backend](https:/
|
|||
### vLLM Integration ([Guide v0.2](https://kvcache-ai.github.io/Mooncake/getting_started/examples/vllm-integration/vllm-integration-v0.2.html))
|
||||
To optimize LLM inference, the vLLM community is working on supporting [disaggregated prefilling (PR 10502)](https://github.com/vllm-project/vllm/pull/10502). This feature allows separating the **prefill** phase from the **decode** phase in different processes. The vLLM uses `nccl` and `gloo` as the transport layer by default, but currently it cannot efficiently decouple both phases in different machines.
|
||||
|
||||
We have implemented vLLM integration, which uses Transfer Engine as the network layer instead of `nccl` and `gloo`, to support **inter-node KVCache transfer** [(PR 10884)](https://github.com/vllm-project/vllm/pull/10884). Transfer Engine provides simpler interfaces and more efficient use of RDMA devices.
|
||||
We have implemented vLLM integration, which uses Transfer Engine as the network layer instead of `nccl` and `gloo`, to support **inter-node KVCache transfer** [(PR 10884)](https://github.com/vllm-project/vllm/pull/10884). Transfer Engine provides simpler interfaces and more efficient use of RDMA devices.
|
||||
|
||||
We will soon release the new vLLM integration based on Mooncake Store, which supports xPyD prefill/decode disaggregation.
|
||||
|
||||
|
|
@ -149,7 +166,7 @@ In the future, we will further improve TTFT through GPUDirect RDMA and zero-copy
|
|||
|
||||
- Click [here](https://kvcache-ai.github.io/Mooncake/performance/vllm-benchmark-results-v0.2.html) to access detailed benchmark results.
|
||||
|
||||
**More advanced features will coming soon, so stay tuned!**
|
||||
**More advanced features are coming soon, so stay tuned!**
|
||||
|
||||
<h2 id="quick-start">🚀 Quick Start</h2>
|
||||
|
||||
|
|
@ -157,18 +174,25 @@ In the future, we will further improve TTFT through GPUDirect RDMA and zero-copy
|
|||
|
||||
Mooncake is designed and optimized for high-speed RDMA networks. Though Mooncake supports TCP-only data transfer, we **strongly** recommend users to evaluate the functionality and performance of Mooncake with RDMA network support.
|
||||
|
||||
The following needs to be installed before running any component of Mooncake:
|
||||
- RDMA Driver & SDK, such as Mellanox OFED.
|
||||
The following need to be installed before running any component of Mooncake:
|
||||
- RDMA Driver & SDK, such as Mellanox OFED.
|
||||
- Python 3.10, virtual environment is recommended.
|
||||
- CUDA 12.1 and above, including NVIDIA GPUDirect Storage Support, if the package is build with `-DUSE_CUDA` (disabled by default). *You may install them from [here](https://developer.nvidia.com/cuda-downloads)*.
|
||||
- CUDA 12.1 and above, including NVIDIA GPUDirect Storage Support, if the package is built with `-DUSE_CUDA` (disabled by default). *You may install them from [here](https://developer.nvidia.com/cuda-downloads)*.
|
||||
- Cambricon Neuware, if the package is built with `-DUSE_MLU`. By default Mooncake looks for Neuware under `NEUWARE_HOME` or `/usr/local/neuware`.
|
||||
|
||||
### Use Python package
|
||||
The most simple way to use Mooncake Transfer Engine is using `pip`:
|
||||
The simplest way to use Mooncake Transfer Engine is using `pip`:
|
||||
|
||||
**For CUDA-enabled systems:**
|
||||
|
||||
- CUDA < 13.0
|
||||
```bash
|
||||
pip install mooncake-transfer-engine
|
||||
```
|
||||
- CUDA >= 13.0
|
||||
```bash
|
||||
pip install mooncake-transfer-engine-cuda13
|
||||
```
|
||||
|
||||
**For non-CUDA systems:**
|
||||
```bash
|
||||
|
|
@ -178,21 +202,41 @@ pip install mooncake-transfer-engine-non-cuda
|
|||
> [!IMPORTANT]
|
||||
> - The CUDA version (`mooncake-transfer-engine`) includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+.
|
||||
> - The non-CUDA version (`mooncake-transfer-engine-non-cuda`) is for environments without CUDA dependencies.
|
||||
> - MLU support is currently available through source builds with `-DUSE_MLU=ON`; there is no dedicated prebuilt MLU wheel yet.
|
||||
> - If users encounter problems such as missing `lib*.so`, they should uninstall the package they installed and build the binaries manually.
|
||||
|
||||
### Use Docker image
|
||||
Mooncake supports Docker-based deployment, see [Build Guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html) in detail.
|
||||
|
||||
To produce an image that compiles Mooncake from source, builds the wheel via `scripts/build_wheel.sh`, and installs that wheel inside the container, use `build-wheel.dockerfile`:
|
||||
|
||||
```bash
|
||||
docker build -f docker/mooncake.Dockerfile \
|
||||
--build-arg PYTHON_VERSION=3.10 \
|
||||
--build-arg EP_TORCH_VERSIONS="2.9.1" \
|
||||
-t mooncake:from-source .
|
||||
```
|
||||
|
||||
The resulting image already has a virtual environment at `/opt/venv` with the freshly built wheel installed. Launch it with GPU/RDMA access as needed, for example:
|
||||
|
||||
```bash
|
||||
docker run --gpus all --network host -it mooncake:from-source /bin/bash
|
||||
```
|
||||
|
||||
> [!NOTE]
|
||||
> Make sure you build the image from the repository root so that Git metadata and submodules are available inside the build context.
|
||||
|
||||
### Build and use binaries
|
||||
The following are additional dependencies for building Mooncake:
|
||||
- Build essentials, including gcc, g++ (9.4+) and cmake (3.16+).
|
||||
- Go 1.20+, if you want to build with `-DWITH_P2P_STORE`, `-DUSE_ETCD` (enabled by default to use etcd as metadata servers), or `-DSTORE_USE_ETCD` (use etcd for the failover of the store master).
|
||||
- CUDA 12.1 and above, including NVIDIA GPUDirect Storage Support, if the package is built with `-DUSE_CUDA`. *This is NOT included in the `dependencies.sh` script. You may install them from [here](https://developer.nvidia.com/cuda-downloads)*.
|
||||
- Cambricon Neuware, if you want to build with `-DUSE_MLU`. *This is NOT included in the `dependencies.sh` script.* Mooncake resolves it from `NEUWARE_HOME` or `/usr/local/neuware` by default, and also supports overriding `MLU_INCLUDE_DIR` / `MLU_LIB_DIR` during CMake configure.
|
||||
- [Optional] Rust Toolchain, if you want to build with `-DWITH_RUST_EXAMPLE`. *This is NOT included in the `dependencies.sh` script.*
|
||||
- [Optional] `hiredis`, if you want to build with `-DUSE_REDIS` to use Redis instead of etcd as metadata servers.
|
||||
- [Optional] `curl`, if you want to build with `-DUSE_HTTP` to use HTTP instead of etcd as metadata servers.
|
||||
|
||||
The building and installation steps are the following:
|
||||
The build and installation steps are as follows:
|
||||
1. Retrieve source code from GitHub repo
|
||||
```bash
|
||||
git clone https://github.com/kvcache-ai/Mooncake.git
|
||||
|
|
@ -213,6 +257,14 @@ The building and installation steps are the following:
|
|||
sudo make install # optional, make it ready to be used by vLLM/SGLang
|
||||
```
|
||||
|
||||
For Cambricon MLU builds, configure CMake with `-DUSE_MLU=ON`. For example:
|
||||
```bash
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=/usr/local/neuware
|
||||
make -j
|
||||
```
|
||||
|
||||
|
||||
<h2 id="milestones"> 🛣️ Incoming Milestones</h2>
|
||||
|
||||
|
|
@ -244,31 +296,37 @@ The above presents two samples from our trace dataset. The trace includes the ti
|
|||
Please kindly cite our paper if you find the paper or the traces are useful:
|
||||
|
||||
```bibtex
|
||||
@article{qin2024mooncake,
|
||||
title={Mooncake: A kvcache-centric disaggregated architecture for llm serving},
|
||||
author={Qin, Ruoyu and Li, Zheming and He, Weiran and Cui, Jialei and Tang, Heyi and Ren, Feng and Ma, Teng and Cai, Shangming and Zhang, Yineng and Zhang, Mingxing and others},
|
||||
journal={ACM Transactions on Storage},
|
||||
year={2024},
|
||||
publisher={ACM New York, NY}
|
||||
@article{qin2025mooncake_tos,
|
||||
author = {Qin Ruoyu and Li Zheming and He Weiran and Cui Jialei and Tang Heyi and Ren Feng and Ma Teng and Cai Shangming and Zhang Yineng and Zhang Mingxing and Wu Yongwei and Zheng Weimin and Xu Xinran},
|
||||
title = {Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving},
|
||||
year = {2025},
|
||||
publisher = {Association for Computing Machinery},
|
||||
address = {New York, NY, USA},
|
||||
issn = {1553-3077},
|
||||
url = {https://doi.org/10.1145/3773772},
|
||||
doi = {10.1145/3773772},
|
||||
journal = {ACM Trans. Storage},
|
||||
month = {nov},
|
||||
keywords = {Machine learning system, LLM serving, KVCache},
|
||||
}
|
||||
|
||||
@article{qin2024mooncake,
|
||||
title = {Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving},
|
||||
author = {Ruoyu Qin, Zheming Li, Weiran He, Mingxing Zhang, Yongwei Wu, Weimin Zheng, and Xinran Xu},
|
||||
year = {2024},
|
||||
url = {https://arxiv.org/abs/2407.00079}
|
||||
@inproceedings{qin2025mooncake,
|
||||
author = {Ruoyu Qin and Zheming Li and Weiran He and Jialei Cui and Feng Ren and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu},
|
||||
title = {Mooncake: Trading More Storage for Less Computation {\textemdash} A {KVCache-centric} Architecture for Serving {LLM} Chatbot},
|
||||
booktitle = {23rd USENIX Conference on File and Storage Technologies (FAST 25)},
|
||||
year = {2025},
|
||||
isbn = {978-1-939133-45-8},
|
||||
address = {Santa Clara, CA},
|
||||
pages = {155--170},
|
||||
url = {https://www.usenix.org/conference/fast25/presentation/qin},
|
||||
publisher = {USENIX Association},
|
||||
month = {feb},
|
||||
}
|
||||
|
||||
@inproceedings {qin2025mooncake,
|
||||
author = {Ruoyu Qin and Zheming Li and Weiran He and Jialei Cui and Feng Ren and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu},
|
||||
title = {Mooncake: Trading More Storage for Less Computation {\textemdash} A {KVCache-centric} Architecture for Serving {LLM} Chatbot},
|
||||
booktitle = {23rd USENIX Conference on File and Storage Technologies (FAST 25)},
|
||||
year = {2025},
|
||||
isbn = {978-1-939133-45-8},
|
||||
address = {Santa Clara, CA},
|
||||
pages = {155--170},
|
||||
url = {https://www.usenix.org/conference/fast25/presentation/qin},
|
||||
publisher = {USENIX Association},
|
||||
month = feb
|
||||
@article{qin2024mooncake_arxiv,
|
||||
title = {Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving},
|
||||
author = {Ruoyu Qin and Zheming Li and Weiran He and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu},
|
||||
year = {2024},
|
||||
url = {https://arxiv.org/abs/2407.00079},
|
||||
}
|
||||
```
|
||||
|
|
|
|||
|
|
@ -0,0 +1,956 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""
|
||||
Mooncake KVCache Storage Benchmark Tool
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
import os
|
||||
import statistics
|
||||
import random
|
||||
import errno
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
# ============================================================================
|
||||
# Constants
|
||||
# ============================================================================
|
||||
|
||||
BLOCK_SIZE_TOKENS = 512 # Number of tokens per block
|
||||
DEFAULT_BYTES_PER_TOKEN = 2048 # 7B model FP16 (2KB per token)
|
||||
BLOCK_SIZE_BYTES = BLOCK_SIZE_TOKENS * DEFAULT_BYTES_PER_TOKEN # 1MB per block
|
||||
MIN_LATENCY_MS = 0.001 # Minimum latency in milliseconds (1 microsecond)
|
||||
|
||||
# Model KVCache sizes (bytes per token, based on LMCache calculator)
|
||||
# Source: https://lmcache.ai/kv_cache_calculator.html
|
||||
MODEL_BYTES_PER_TOKEN = {
|
||||
"llama-3.1-405b": 327680,
|
||||
"qwen3-32b": 81920,
|
||||
"deepseek-v3": 1748992,
|
||||
"glm-4.6": 157013,
|
||||
"default": DEFAULT_BYTES_PER_TOKEN,
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Data Structures
|
||||
# ============================================================================
|
||||
|
||||
@dataclass
|
||||
class KVCacheRequest:
|
||||
"""KVCache request
|
||||
|
||||
Attributes:
|
||||
timestamp: Request timestamp in milliseconds
|
||||
hash_ids: List of block IDs (each ID corresponds to a 512-token block)
|
||||
input_length: Input token count
|
||||
output_length: Output token count
|
||||
"""
|
||||
timestamp: float
|
||||
hash_ids: List[int]
|
||||
input_length: int
|
||||
output_length: int
|
||||
|
||||
# ============================================================================
|
||||
# Storage Layer: Offset Allocator
|
||||
# ============================================================================
|
||||
|
||||
class OffsetAllocatorStorage:
|
||||
"""High-performance block storage based on Offset Allocator
|
||||
|
||||
Architecture:
|
||||
-----------
|
||||
1. Single large file stores all blocks (avoids file explosion)
|
||||
2. Uses offset to manage file space (similar to Mooncake's OffsetAllocator)
|
||||
3. hash_id -> offset mapping stored in memory (fast lookup)
|
||||
|
||||
Block Organization:
|
||||
-----------
|
||||
Each block corresponds to 512 tokens, fixed size 1MB:
|
||||
- hash_id[0] -> block_0 (tokens [0...511]) -> offset 0
|
||||
- hash_id[1] -> block_1 (tokens [512...1023]) -> offset 1
|
||||
- hash_id[i] -> block_i (tokens [i*512...(i+1)*512-1]) -> offset i
|
||||
|
||||
Performance Advantages:
|
||||
-----------
|
||||
- Only one file, no file explosion
|
||||
- Offset reuse, reduces memory allocation
|
||||
- pread/pwrite, thread-safe, no seek needed
|
||||
- Keep fd open, reduces open/close overhead
|
||||
- Metadata in memory, O(1) lookup
|
||||
|
||||
Attributes:
|
||||
storage_dir: Storage directory path
|
||||
block_size_bytes: Block size in bytes
|
||||
max_blocks: Maximum number of blocks
|
||||
hash_id_to_offset: hash_id -> offset mapping
|
||||
free_offsets: List of reusable offsets
|
||||
next_offset: Next allocatable offset
|
||||
"""
|
||||
|
||||
def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
|
||||
max_blocks: int = 100000, block_size_tokens: int = 512,
|
||||
fsync_mode: str = 'batch', fsync_batch_size: int = 100):
|
||||
"""Initialize Offset Allocator storage
|
||||
|
||||
Args:
|
||||
storage_dir: Storage directory path
|
||||
bytes_per_token: Bytes per token
|
||||
max_blocks: Maximum number of blocks (determines file size)
|
||||
block_size_tokens: Number of tokens per block
|
||||
fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
|
||||
fsync_batch_size: Number of writes between fsync in batch mode
|
||||
"""
|
||||
self.storage_dir = Path(storage_dir)
|
||||
self.bytes_per_token = bytes_per_token
|
||||
self.block_size_tokens = block_size_tokens
|
||||
self.block_size_bytes = self.block_size_tokens * self.bytes_per_token
|
||||
self.max_blocks = max_blocks
|
||||
|
||||
# Fsync configuration
|
||||
self.fsync_mode = fsync_mode
|
||||
self.fsync_batch_size = fsync_batch_size
|
||||
self.pending_sync_count = 0
|
||||
|
||||
# Create storage directory
|
||||
self.storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Single large file
|
||||
self.storage_file = self.storage_dir / "kvcache_storage.bin"
|
||||
self.file_size = self.max_blocks * self.block_size_bytes
|
||||
|
||||
# Initialize storage file
|
||||
if not self.storage_file.exists():
|
||||
self._init_storage_file()
|
||||
|
||||
# hash_id -> offset mapping (metadata, in memory)
|
||||
self.hash_id_to_offset: Dict[int, int] = {}
|
||||
|
||||
# Offset allocator (free list)
|
||||
self.free_offsets: List[int] = []
|
||||
self.next_offset = 0
|
||||
|
||||
# File descriptor (keep open, avoid repeated open/close)
|
||||
self.fd = None
|
||||
|
||||
# Pre-allocated data buffer with pattern to avoid SSD compression artifacts
|
||||
# Using a repeating pattern that looks like realistic data (not all zeros)
|
||||
# Pattern: 64-byte repeated sequence mixed with some variation
|
||||
pattern = bytes([(i & 0xFF) for i in range(256)]) # 0-255 byte pattern
|
||||
pattern_repeats = (self.block_size_bytes // len(pattern)) + 1
|
||||
self._data_buffer = (pattern * pattern_repeats)[:self.block_size_bytes]
|
||||
|
||||
# Statistics
|
||||
self.stats = {
|
||||
'read_count': 0,
|
||||
'write_count': 0,
|
||||
'read_bytes': 0,
|
||||
'write_bytes': 0,
|
||||
'read_latencies_ms': [],
|
||||
'write_latencies_ms': [],
|
||||
'sync_count': 0, # Number of fsync operations performed
|
||||
}
|
||||
|
||||
# ========================================================================
|
||||
# Internal Methods
|
||||
# ========================================================================
|
||||
|
||||
def _init_storage_file(self):
|
||||
"""Initialize storage file (pre-allocate space)
|
||||
|
||||
Create sparse file to avoid actual disk space usage until data is written
|
||||
"""
|
||||
with open(self.storage_file, 'wb') as f:
|
||||
f.seek(self.file_size - 1)
|
||||
f.write(b'\0')
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
|
||||
def _get_fd(self):
|
||||
"""Get file descriptor (lazy open)
|
||||
|
||||
Returns:
|
||||
int: File descriptor
|
||||
"""
|
||||
if self.fd is None:
|
||||
# Use O_RDWR | O_CREAT, no O_DIRECT (Python compatibility)
|
||||
self.fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT)
|
||||
return self.fd
|
||||
|
||||
def _allocate_offset(self) -> int:
|
||||
"""Allocate a new offset
|
||||
|
||||
Prioritize reusing freed offsets, otherwise allocate new offset
|
||||
|
||||
Returns:
|
||||
int: Allocated offset
|
||||
"""
|
||||
if self.free_offsets:
|
||||
return self.free_offsets.pop()
|
||||
offset = self.next_offset
|
||||
self.next_offset += 1
|
||||
return offset
|
||||
|
||||
def _free_offset(self, offset: int):
|
||||
"""Free offset for reuse
|
||||
|
||||
Args:
|
||||
offset: Offset to free
|
||||
"""
|
||||
self.free_offsets.append(offset)
|
||||
|
||||
# ========================================================================
|
||||
# Public Interface
|
||||
# ========================================================================
|
||||
|
||||
def block_exists(self, hash_id: int) -> bool:
|
||||
"""Check if block exists
|
||||
|
||||
Args:
|
||||
hash_id: Unique block identifier
|
||||
|
||||
Returns:
|
||||
bool: Whether block exists
|
||||
"""
|
||||
return hash_id in self.hash_id_to_offset
|
||||
|
||||
def read_block(self, hash_id: int) -> float:
|
||||
"""Read block using pread
|
||||
|
||||
Args:
|
||||
hash_id: Unique block identifier
|
||||
|
||||
Returns:
|
||||
float: Read latency in milliseconds, or 0 if block doesn't exist
|
||||
"""
|
||||
if hash_id not in self.hash_id_to_offset:
|
||||
return 0.0 # Block doesn't exist, no latency to measure
|
||||
|
||||
offset = self.hash_id_to_offset[hash_id]
|
||||
file_offset = offset * self.block_size_bytes
|
||||
|
||||
start = time.perf_counter()
|
||||
|
||||
try:
|
||||
fd = self._get_fd()
|
||||
data = os.pread(fd, self.block_size_bytes, file_offset)
|
||||
latency_ms = (time.perf_counter() - start) * 1000.0
|
||||
|
||||
self.stats['read_count'] += 1
|
||||
self.stats['read_bytes'] += len(data)
|
||||
self.stats['read_latencies_ms'].append(latency_ms)
|
||||
return latency_ms
|
||||
except OSError as e:
|
||||
print(f"Error reading block {hash_id} at offset {file_offset}: {e}")
|
||||
return 0.0 # Error case, don't pollute stats
|
||||
|
||||
def write_block(self, hash_id: int) -> float:
|
||||
"""Write block using pwrite
|
||||
|
||||
Args:
|
||||
hash_id: Unique block identifier
|
||||
|
||||
Returns:
|
||||
float: Write latency in milliseconds
|
||||
"""
|
||||
# Allocate offset
|
||||
offset = self._allocate_offset()
|
||||
file_offset = offset * self.block_size_bytes
|
||||
|
||||
# Use pre-allocated buffer (much faster than os.urandom)
|
||||
data = self._data_buffer
|
||||
|
||||
start = time.perf_counter()
|
||||
|
||||
try:
|
||||
fd = self._get_fd()
|
||||
written = os.pwrite(fd, data, file_offset)
|
||||
|
||||
write_done = time.perf_counter()
|
||||
|
||||
# Conditional fsync based on mode
|
||||
if self.fsync_mode == 'always':
|
||||
# Include fsync in latency measurement
|
||||
os.fsync(fd)
|
||||
self.stats['sync_count'] += 1
|
||||
self.pending_sync_count = 0
|
||||
latency_ms = (time.perf_counter() - start) * 1000.0
|
||||
# Evict from page cache AFTER fsync to ensure reads measure actual SSD performance
|
||||
os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
|
||||
elif self.fsync_mode == 'batch':
|
||||
# For batch mode, only measure write time (fsync is deferred)
|
||||
self.pending_sync_count += 1
|
||||
if self.pending_sync_count >= self.fsync_batch_size:
|
||||
os.fsync(fd)
|
||||
self.stats['sync_count'] += 1
|
||||
self.pending_sync_count = 0
|
||||
latency_ms = (write_done - start) * 1000.0 # Only write time
|
||||
# Evict from page cache after each write
|
||||
os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
|
||||
elif self.fsync_mode == 'none':
|
||||
latency_ms = (write_done - start) * 1000.0
|
||||
# Evict from page cache even when not syncing
|
||||
os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
|
||||
else: # 'end' mode
|
||||
latency_ms = (write_done - start) * 1000.0
|
||||
# Evict from page cache (fsync will happen at the end)
|
||||
os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
|
||||
|
||||
# Update mapping
|
||||
self.hash_id_to_offset[hash_id] = offset
|
||||
|
||||
self.stats['write_count'] += 1
|
||||
self.stats['write_bytes'] += written
|
||||
self.stats['write_latencies_ms'].append(latency_ms)
|
||||
return latency_ms
|
||||
except OSError as e:
|
||||
if e.errno == errno.ENOSPC:
|
||||
print(f"Error: Disk full when writing block {hash_id} at offset {file_offset}")
|
||||
else:
|
||||
print(f"Error writing block {hash_id} at offset {file_offset}: {e}")
|
||||
return 0.0 # Error case, don't pollute stats
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry"""
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit - ensures cleanup"""
|
||||
# Perform final fsync before closing for 'end' and 'batch' modes
|
||||
self._finalize_sync()
|
||||
self.close(force_sync=False) # Already synced above
|
||||
return False
|
||||
|
||||
def _finalize_sync(self):
|
||||
"""Perform final fsync before closing (for 'end' mode and pending batch writes)"""
|
||||
if self.fd is not None:
|
||||
if self.fsync_mode == 'end':
|
||||
try:
|
||||
os.fsync(self.fd)
|
||||
self.stats['sync_count'] += 1
|
||||
except OSError:
|
||||
pass
|
||||
elif self.fsync_mode == 'batch' and self.pending_sync_count > 0:
|
||||
# Flush remaining pending writes
|
||||
try:
|
||||
os.fsync(self.fd)
|
||||
self.stats['sync_count'] += 1
|
||||
self.pending_sync_count = 0
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def close(self, force_sync: bool = True):
|
||||
"""Close file
|
||||
|
||||
Args:
|
||||
force_sync: Whether to force fsync before closing
|
||||
"""
|
||||
# For backward compatibility with non-context-manager usage
|
||||
if force_sync:
|
||||
self._finalize_sync()
|
||||
|
||||
if self.fd is not None:
|
||||
os.close(self.fd)
|
||||
self.fd = None
|
||||
|
||||
def get_stats(self) -> Dict:
|
||||
"""Get statistics
|
||||
|
||||
Returns:
|
||||
Dict: Dictionary containing read/write statistics
|
||||
"""
|
||||
def calc_stats(latencies):
|
||||
"""Calculate latency statistics"""
|
||||
if not latencies:
|
||||
return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
|
||||
return {
|
||||
'avg_ms': statistics.mean(latencies),
|
||||
**calc_percentiles(latencies),
|
||||
}
|
||||
|
||||
return {
|
||||
'read': {
|
||||
'count': self.stats['read_count'],
|
||||
'mb': self.stats['read_bytes'] / 1024 / 1024,
|
||||
**calc_stats(self.stats['read_latencies_ms'])
|
||||
},
|
||||
'write': {
|
||||
'count': self.stats['write_count'],
|
||||
'mb': self.stats['write_bytes'] / 1024 / 1024,
|
||||
**calc_stats(self.stats['write_latencies_ms'])
|
||||
},
|
||||
'sync_count': self.stats['sync_count'],
|
||||
'total_blocks': len(self.hash_id_to_offset),
|
||||
'free_blocks': len(self.free_offsets),
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Benchmark Layer
|
||||
# ============================================================================
|
||||
|
||||
class StorageBenchmark:
|
||||
"""KVCache storage benchmark
|
||||
|
||||
Based on Mooncake OffsetAllocator + vLLM PagedAttention implementation:
|
||||
|
||||
Example:
|
||||
-----
|
||||
Request A: [1, 2, 4]
|
||||
-> hash_id 1 -> not exist, write block_1 (offset=0, 1MB)
|
||||
-> hash_id 2 -> not exist, write block_2 (offset=1, 1MB)
|
||||
-> hash_id 4 -> not exist, write block_4 (offset=2, 1MB)
|
||||
|
||||
Request B: [1, 2, 4, 6]
|
||||
-> hash_id 1 -> exists, read block_1 (offset=0) ✓ prefix reuse
|
||||
-> hash_id 2 -> exists, read block_2 (offset=1) ✓ prefix reuse
|
||||
-> hash_id 4 -> exists, read block_4 (offset=2) ✓ prefix reuse
|
||||
-> hash_id 6 -> not exist, write block_6 (offset=3, 1MB)
|
||||
|
||||
Performance Advantages:
|
||||
---------
|
||||
- Single file operation, no file explosion
|
||||
- Offset reuse, reduces memory allocation
|
||||
- pread/pwrite, thread-safe
|
||||
"""
|
||||
|
||||
def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
|
||||
max_blocks: int = 100000, block_size_tokens: int = 512,
|
||||
fsync_mode: str = 'batch', fsync_batch_size: int = 100):
|
||||
"""Initialize benchmark
|
||||
|
||||
Args:
|
||||
storage_dir: Storage directory
|
||||
bytes_per_token: Bytes per token
|
||||
max_blocks: Maximum number of blocks
|
||||
block_size_tokens: Number of tokens per block
|
||||
fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
|
||||
fsync_batch_size: Number of writes between fsync in batch mode
|
||||
"""
|
||||
self.storage = OffsetAllocatorStorage(
|
||||
storage_dir, bytes_per_token, max_blocks,
|
||||
block_size_tokens, fsync_mode, fsync_batch_size
|
||||
)
|
||||
self.bytes_per_token = bytes_per_token
|
||||
self.block_size_tokens = block_size_tokens
|
||||
|
||||
# Statistics
|
||||
self.stats = {
|
||||
'total_requests': 0,
|
||||
'total_blocks': 0,
|
||||
'read_blocks': 0,
|
||||
'write_blocks': 0,
|
||||
'prefix_hit_blocks': 0, # Number of prefix hit blocks
|
||||
'request_latencies_ms': [],
|
||||
}
|
||||
|
||||
def process_request(self, req: KVCacheRequest) -> float:
|
||||
"""Process a KVCache request
|
||||
|
||||
Based on vLLM's prefix caching mechanism:
|
||||
- Each hash_id corresponds to an independent block
|
||||
- Prefix reuse achieved through hash_id matching
|
||||
|
||||
Args:
|
||||
req: KVCache request
|
||||
|
||||
Returns:
|
||||
float: Request latency in milliseconds
|
||||
"""
|
||||
self.stats['total_requests'] += 1
|
||||
self.stats['total_blocks'] += len(req.hash_ids)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
total_latency = 0.0
|
||||
|
||||
# Process each hash_id (in order)
|
||||
for hash_id in req.hash_ids:
|
||||
if self.storage.block_exists(hash_id):
|
||||
# Block exists, read (reuse cached block)
|
||||
total_latency += self.storage.read_block(hash_id)
|
||||
self.stats['read_blocks'] += 1
|
||||
self.stats['prefix_hit_blocks'] += 1 # Count all cache hits as prefix reuse
|
||||
else:
|
||||
# Block doesn't exist, write (new block)
|
||||
total_latency += self.storage.write_block(hash_id)
|
||||
self.stats['write_blocks'] += 1
|
||||
|
||||
latency_ms = total_latency if total_latency > 0 else MIN_LATENCY_MS
|
||||
self.stats['request_latencies_ms'].append(latency_ms)
|
||||
|
||||
return latency_ms
|
||||
|
||||
def get_stats(self) -> Dict:
|
||||
"""Get statistics
|
||||
|
||||
Returns:
|
||||
Dict: Statistics dictionary
|
||||
"""
|
||||
storage_stats = self.storage.get_stats()
|
||||
|
||||
request_latencies = self.stats['request_latencies_ms']
|
||||
|
||||
if request_latencies:
|
||||
latency_stats = {
|
||||
'avg_ms': statistics.mean(request_latencies),
|
||||
**calc_percentiles(request_latencies),
|
||||
}
|
||||
else:
|
||||
latency_stats = {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
|
||||
|
||||
total_blocks = self.stats['total_blocks']
|
||||
read_blocks = self.stats['read_blocks']
|
||||
write_blocks = self.stats['write_blocks']
|
||||
|
||||
return {
|
||||
'total_requests': self.stats['total_requests'],
|
||||
'total_blocks': total_blocks,
|
||||
'read_blocks': read_blocks,
|
||||
'write_blocks': write_blocks,
|
||||
'prefix_hit_blocks': self.stats['prefix_hit_blocks'],
|
||||
'block_hit_rate': read_blocks / total_blocks if total_blocks > 0 else 0,
|
||||
'write_ratio': write_blocks / total_blocks if total_blocks > 0 else 0,
|
||||
'tokens_per_block': self.block_size_tokens, # Configurable block size in tokens
|
||||
'latency': latency_stats,
|
||||
'storage': storage_stats,
|
||||
}
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry"""
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit - ensures cleanup"""
|
||||
self.close()
|
||||
return False
|
||||
|
||||
def close(self, force_sync: bool = True):
|
||||
"""Close storage
|
||||
|
||||
Args:
|
||||
force_sync: Whether to force final sync before closing
|
||||
"""
|
||||
self.storage.close(force_sync=force_sync)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Utility Functions
|
||||
# ============================================================================
|
||||
|
||||
def calc_percentiles(data: List[float]) -> Dict[str, float]:
|
||||
"""Calculate latency percentiles
|
||||
|
||||
Uses linear interpolation for accurate percentile calculation.
|
||||
This is more accurate than statistics.quantiles() for small datasets.
|
||||
|
||||
Args:
|
||||
data: List of latency values in milliseconds
|
||||
|
||||
Returns:
|
||||
Dict containing p50, p95, p99 percentiles
|
||||
"""
|
||||
if not data:
|
||||
return {'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
|
||||
|
||||
# Sort data for percentile calculation
|
||||
sorted_data = sorted(data)
|
||||
n = len(sorted_data)
|
||||
|
||||
def get_percentile(p: float) -> float:
|
||||
"""Get percentile using linear interpolation
|
||||
|
||||
Args:
|
||||
p: Percentile (0-100)
|
||||
|
||||
Returns:
|
||||
Value at percentile
|
||||
"""
|
||||
index = (n - 1) * p / 100
|
||||
lower = int(index)
|
||||
upper = min(lower + 1, n - 1)
|
||||
|
||||
if lower == upper:
|
||||
return sorted_data[lower]
|
||||
|
||||
# Linear interpolation
|
||||
weight = index - lower
|
||||
return sorted_data[lower] * (1 - weight) + sorted_data[upper] * weight
|
||||
|
||||
return {
|
||||
'p50_ms': get_percentile(50),
|
||||
'p95_ms': get_percentile(95),
|
||||
'p99_ms': get_percentile(99),
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Trace Loader
|
||||
# ============================================================================
|
||||
|
||||
class TraceLoader:
|
||||
"""Load KVCache trace"""
|
||||
|
||||
def __init__(self, trace_path: str):
|
||||
"""Initialize trace loader
|
||||
|
||||
Args:
|
||||
trace_path: Trace file path
|
||||
"""
|
||||
self.trace_path = trace_path
|
||||
self.requests = []
|
||||
self._load_trace()
|
||||
|
||||
def _load_trace(self):
|
||||
"""Load trace file with error handling"""
|
||||
line_num = 0
|
||||
try:
|
||||
with open(self.trace_path, 'r') as f:
|
||||
for line in f:
|
||||
line_num += 1
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
req = json.loads(line)
|
||||
# Validate required fields
|
||||
if not all(k in req for k in ['timestamp', 'hash_ids', 'input_length', 'output_length']):
|
||||
print(f"Warning: Line {line_num} missing required fields, skipping")
|
||||
continue
|
||||
if not isinstance(req['hash_ids'], list):
|
||||
print(f"Warning: Line {line_num} has invalid hash_ids (not a list), skipping")
|
||||
continue
|
||||
self.requests.append(KVCacheRequest(
|
||||
timestamp=float(req['timestamp']),
|
||||
hash_ids=req['hash_ids'],
|
||||
input_length=int(req['input_length']),
|
||||
output_length=int(req['output_length'])
|
||||
))
|
||||
except (json.JSONDecodeError, ValueError, KeyError) as e:
|
||||
print(f"Warning: Line {line_num} has invalid format: {e}, skipping")
|
||||
continue
|
||||
except FileNotFoundError:
|
||||
raise FileNotFoundError(f"Trace file not found: {self.trace_path}")
|
||||
except OSError as e:
|
||||
raise OSError(f"Error reading trace file {self.trace_path}: {e}")
|
||||
|
||||
def get_requests(self) -> List[KVCacheRequest]:
|
||||
"""Get request list
|
||||
|
||||
Returns:
|
||||
List[KVCacheRequest]: Request list
|
||||
"""
|
||||
return self.requests
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Benchmark Runner
|
||||
# ============================================================================
|
||||
|
||||
def run_benchmark(trace_path: str, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
|
||||
max_requests: Optional[int] = None, max_blocks: int = 100000,
|
||||
replay_timestamps: bool = False, time_scale: float = 1.0,
|
||||
block_size_tokens: int = 512,
|
||||
fsync_mode: str = 'batch', fsync_batch_size: int = 100) -> Dict:
|
||||
"""Run benchmark
|
||||
|
||||
Args:
|
||||
trace_path: Trace file path
|
||||
storage_dir: Storage directory
|
||||
bytes_per_token: Bytes per token
|
||||
max_requests: Maximum number of requests (None = all)
|
||||
max_blocks: Maximum number of blocks
|
||||
replay_timestamps: Whether to replay timestamps from trace (simulate realistic timing)
|
||||
time_scale: Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed)
|
||||
block_size_tokens: Number of tokens per block
|
||||
fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
|
||||
fsync_batch_size: Number of writes between fsync in batch mode
|
||||
|
||||
Returns:
|
||||
Dict: Benchmark results
|
||||
"""
|
||||
block_size_bytes = block_size_tokens * bytes_per_token
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"Running: {Path(trace_path).name}")
|
||||
print(f"Architecture: Offset Allocator (Mooncake style)")
|
||||
print(f"Block size: {block_size_tokens} tokens/block ({block_size_bytes:,} bytes)")
|
||||
print(f"Storage: Single large file with offset-based block management")
|
||||
print(f"Bytes per token: {bytes_per_token}")
|
||||
print(f"Max blocks: {max_blocks}")
|
||||
print(f"Fsync mode: {fsync_mode}" + (f" (batch_size={fsync_batch_size})" if fsync_mode == 'batch' else ''))
|
||||
print(f"Timestamp replay: {'Enabled' if replay_timestamps else 'Disabled'}")
|
||||
if replay_timestamps:
|
||||
scale_desc = 'real-time' if time_scale == 1.0 else f'{1/time_scale:.1f}x speed' if time_scale < 1.0 else f'{time_scale}x slower'
|
||||
print(f"Time scale: {time_scale}x ({scale_desc})")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Load trace
|
||||
loader = TraceLoader(trace_path)
|
||||
requests = loader.get_requests()
|
||||
|
||||
if max_requests:
|
||||
requests = requests[:max_requests]
|
||||
|
||||
print(f"Loaded {len(requests)} requests")
|
||||
|
||||
# Show timestamp range
|
||||
if replay_timestamps and requests:
|
||||
timestamps = [req.timestamp for req in requests]
|
||||
time_span_ms = max(timestamps) - min(timestamps)
|
||||
print(f"Timestamp range: {min(timestamps):.1f} - {max(timestamps):.1f} ms (span: {time_span_ms:.1f} ms)")
|
||||
|
||||
# Create benchmark instance with context manager for cleanup
|
||||
with StorageBenchmark(
|
||||
storage_dir, bytes_per_token, max_blocks,
|
||||
block_size_tokens, fsync_mode, fsync_batch_size
|
||||
) as benchmark:
|
||||
|
||||
# Run benchmark
|
||||
start_time = time.perf_counter()
|
||||
total_io_time = 0.0 # Actual I/O time (excluding sleep)
|
||||
last_timestamp = None
|
||||
base_time = time.time() # Use wall time for replay synchronization
|
||||
|
||||
for i, req in enumerate(requests):
|
||||
# Replay by timestamps
|
||||
sleep_time = 0.0
|
||||
if replay_timestamps and last_timestamp is not None:
|
||||
# Calculate time interval from previous request
|
||||
delta_ms = req.timestamp - last_timestamp
|
||||
sleep_time = delta_ms / 1000.0 / time_scale # Apply time scaling
|
||||
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
# Process request (measure I/O time)
|
||||
req_start = time.perf_counter()
|
||||
benchmark.process_request(req)
|
||||
req_io_time = time.perf_counter() - req_start
|
||||
total_io_time += req_io_time
|
||||
|
||||
# Record current request timestamp
|
||||
last_timestamp = req.timestamp
|
||||
|
||||
# Progress output
|
||||
if (i + 1) % 100 == 0:
|
||||
if replay_timestamps:
|
||||
elapsed_wall_time = time.time() - base_time
|
||||
simulated_time = (req.timestamp - requests[0].timestamp) / 1000.0 / time_scale
|
||||
print(f" Processed {i + 1}/{len(requests)}... (wall: {elapsed_wall_time:.1f}s, simulated: {simulated_time:.1f}s, io: {total_io_time:.1f}s)")
|
||||
else:
|
||||
print(f" Processed {i + 1}/{len(requests)}...")
|
||||
|
||||
elapsed = time.perf_counter() - start_time
|
||||
|
||||
# Perform final sync to include it in stats
|
||||
benchmark.storage._finalize_sync()
|
||||
|
||||
# Get statistics (context manager will handle cleanup)
|
||||
stats = benchmark.get_stats()
|
||||
|
||||
# Calculate actual I/O time (excluding sleep)
|
||||
io_time = total_io_time if replay_timestamps else elapsed
|
||||
|
||||
return {
|
||||
'trace_file': Path(trace_path).name,
|
||||
'total_requests': len(requests),
|
||||
'simulation_time_s': elapsed,
|
||||
'io_time_s': io_time, # Actual I/O time
|
||||
'wall_time_s': elapsed, # Wall time (including sleep)
|
||||
'requests_per_second': len(requests) / io_time if io_time > 0 else 0, # Based on I/O time
|
||||
'timestamp_replay_enabled': replay_timestamps,
|
||||
'time_scale': time_scale,
|
||||
'bytes_per_token': bytes_per_token,
|
||||
'block_size_tokens': block_size_tokens,
|
||||
'fsync_mode': fsync_mode,
|
||||
**stats,
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Result Output
|
||||
# ============================================================================
|
||||
|
||||
def print_results(results: List[Dict]):
|
||||
"""Print benchmark results
|
||||
|
||||
Args:
|
||||
results: List of benchmark results
|
||||
"""
|
||||
for i, r in enumerate(results, 1):
|
||||
print(f"\n{'='*80}")
|
||||
print(f" [{i}/{len(results)}] {r['trace_file']}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
print(f"\n[Performance Overview]")
|
||||
print(f" Total Requests: {r['total_requests']:,}")
|
||||
print(f" Queries Per Second (QPS): {r['requests_per_second']:.2f}")
|
||||
print(f" Cache Hit Rate: {r['block_hit_rate']:.2%}")
|
||||
print(f" Write Ratio: {r['write_ratio']:.2%}")
|
||||
print(f" Total Blocks: {r['total_blocks']:,}")
|
||||
print(f" Read Blocks: {r['read_blocks']:,}")
|
||||
print(f" Write Blocks: {r['write_blocks']:,}")
|
||||
print(f" Prefix Hits: {r['prefix_hit_blocks']:,}")
|
||||
|
||||
print(f"\n[Latency Analysis]")
|
||||
req_lat = r['latency']
|
||||
print(f" Request Latency (End-to-End): Avg={req_lat['avg_ms']:.2f}ms, P50={req_lat['p50_ms']:.2f}ms, P95={req_lat['p95_ms']:.2f}ms, P99={req_lat['p99_ms']:.2f}ms")
|
||||
read_lat = r['storage']['read']
|
||||
write_lat = r['storage']['write']
|
||||
print(f" Single I/O Operation (Per Block):")
|
||||
print(f" Read: Avg={read_lat.get('avg_ms', 0):.3f}ms, P50={read_lat.get('p50_ms', 0):.3f}ms, P95={read_lat.get('p95_ms', 0):.3f}ms, P99={read_lat.get('p99_ms', 0):.3f}ms")
|
||||
print(f" Write: Avg={write_lat.get('avg_ms', 0):.3f}ms, P50={write_lat.get('p50_ms', 0):.3f}ms, P95={write_lat.get('p95_ms', 0):.3f}ms, P99={write_lat.get('p99_ms', 0):.3f}ms")
|
||||
|
||||
print(f"\n[I/O & Bandwidth]")
|
||||
print(f" Total Read I/O: {r['storage']['read']['mb']:>10.1f} MB ({r['storage']['read']['count']:,} ops)")
|
||||
print(f" Total Write I/O: {r['storage']['write']['mb']:>10.1f} MB ({r['storage']['write']['count']:,} ops)")
|
||||
io_time = r['io_time_s']
|
||||
bandwidth = (r['storage']['read']['mb'] + r['storage']['write']['mb']) / io_time
|
||||
print(f" Effective Bandwidth: {bandwidth:>10.1f} MB/s")
|
||||
|
||||
print(f"\n[Storage Details]")
|
||||
print(f" Blocks in Use: {r['storage']['total_blocks']:>10,}")
|
||||
print(f" Free Blocks: {r['storage']['free_blocks']:>10,}")
|
||||
print(f" Tokens per Block: {r['tokens_per_block']:>10,}")
|
||||
print(f" Block Size: {r['tokens_per_block'] * r.get('bytes_per_token', 2048) / 1024 / 1024:>10.2f} MB")
|
||||
if 'sync_count' in r['storage']:
|
||||
print(f" Fsync Operations: {r['storage']['sync_count']:>10,}")
|
||||
|
||||
print(f"\n[Execution Time]")
|
||||
if r.get('timestamp_replay_enabled'):
|
||||
print(f" Wall Time (Total): {r['wall_time_s']:>10.2f} s")
|
||||
print(f" I/O Time (Actual): {r['io_time_s']:>10.2f} s")
|
||||
print(f" Sleep Time (Replay): {r['wall_time_s'] - r['io_time_s']:>10.2f} s")
|
||||
else:
|
||||
print(f" Total Execution Time: {r['wall_time_s']:>10.2f} s")
|
||||
|
||||
print(f"\n{'='*80}\n")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Main Program
|
||||
# ============================================================================
|
||||
|
||||
def main():
|
||||
"""Main entry point"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Mooncake KVCache Storage Benchmark',
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Quick test (100 requests)
|
||||
python storage_benchmark.py --scenario=toolagent --max-requests=100
|
||||
|
||||
# Test with large model preset (Llama-3.1-405B)
|
||||
python storage_benchmark.py --scenario=toolagent --model=llama-3.1-405b --max-requests=100
|
||||
|
||||
# Test with Deepseek V3 (extra large model)
|
||||
python storage_benchmark.py --scenario=toolagent --model=deepseek-v3 --max-requests=100
|
||||
|
||||
# Realistic replay (with timestamps, 10x speed)
|
||||
python storage_benchmark.py --scenario=toolagent --max-requests=1000 \\
|
||||
--replay-timestamps --time-scale=0.1
|
||||
|
||||
# All scenarios with custom bytes_per_token
|
||||
python storage_benchmark.py --scenario=all --bytes-per-token=512
|
||||
|
||||
# Test with different block sizes and fsync modes
|
||||
python storage_benchmark.py --scenario=toolagent --block-size-tokens=256 --fsync-mode=always
|
||||
|
||||
# Test with custom fsync batch size
|
||||
python storage_benchmark.py --scenario=toolagent --fsync-mode=batch --fsync-batch-size=50
|
||||
|
||||
Performance Tuning:
|
||||
--fsync-mode=batch (default): Balance between performance and safety
|
||||
--fsync-mode=always: Safest but slowest, measures full persistence cost
|
||||
--fsync-mode=end: Fastest, only measures write I/O (not persistence)
|
||||
--fsync-mode=none: Testing only, no durability guarantees
|
||||
|
||||
Available model presets:
|
||||
llama-3.1-405b, qwen3-32b, deepseek-v3, glm-4.6, default
|
||||
|
||||
For more information: tools/STORAGE_BENCHMARK_README.md
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument('--trace-dir', type=str, default='../../FAST25-release/traces',
|
||||
help='Trace files directory')
|
||||
parser.add_argument('--scenario', type=str, choices=['conversation', 'synthetic', 'toolagent', 'all'],
|
||||
default='toolagent', help='Test scenario')
|
||||
parser.add_argument('--storage-dir', type=str, default='/tmp/mooncake_bench',
|
||||
help='Storage directory')
|
||||
parser.add_argument('--model', type=str, choices=list(MODEL_BYTES_PER_TOKEN.keys()),
|
||||
default='default',
|
||||
help=f'Model preset (overrides --bytes-per-token). Available: {", ".join(MODEL_BYTES_PER_TOKEN.keys())}')
|
||||
parser.add_argument('--bytes-per-token', type=int, default=DEFAULT_BYTES_PER_TOKEN,
|
||||
help='Bytes per token (default %d, overridden by --model if specified)' % DEFAULT_BYTES_PER_TOKEN)
|
||||
parser.add_argument('--max-requests', type=int, default=None,
|
||||
help='Maximum number of requests (default: unlimited)')
|
||||
parser.add_argument('--max-blocks', type=int, default=100000,
|
||||
help='Maximum number of blocks in storage file (determines file size)')
|
||||
parser.add_argument('--replay-timestamps', action='store_true',
|
||||
help='Enable timestamp replay (simulate realistic request timing)')
|
||||
parser.add_argument('--time-scale', type=float, default=1.0,
|
||||
help='Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed)')
|
||||
parser.add_argument('--block-size-tokens', type=int, default=512,
|
||||
help='Number of tokens per block (default: 512)')
|
||||
parser.add_argument('--fsync-mode', type=str, choices=['batch', 'always', 'end', 'none'],
|
||||
default='batch',
|
||||
help='When to fsync: batch=every N writes (default), always=after each write, end=only at close, none=never')
|
||||
parser.add_argument('--fsync-batch-size', type=int, default=100,
|
||||
help='Number of writes between fsync in batch mode (default: 100)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Print benchmark header
|
||||
print(f"\n{'='*80}")
|
||||
print(f"{'Mooncake KVCache Storage Benchmark':^80}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Determine bytes_per_token (model preset takes precedence)
|
||||
bytes_per_token = MODEL_BYTES_PER_TOKEN.get(args.model, args.bytes_per_token)
|
||||
if args.model != 'default':
|
||||
print(f"Using model preset: {args.model} ({bytes_per_token} bytes/token, ~{bytes_per_token/1024:.1f} KB/token)")
|
||||
else:
|
||||
print(f"Using custom bytes_per_token: {bytes_per_token}")
|
||||
|
||||
# Determine test scenarios
|
||||
scenarios = ['conversation', 'synthetic', 'toolagent'] if args.scenario == 'all' else [args.scenario]
|
||||
trace_files = {
|
||||
'conversation': 'conversation_trace.jsonl',
|
||||
'synthetic': 'synthetic_trace.jsonl',
|
||||
'toolagent': 'toolagent_trace.jsonl'
|
||||
}
|
||||
|
||||
# Run benchmarks
|
||||
results = []
|
||||
|
||||
for scenario in scenarios:
|
||||
trace_path = Path(args.trace_dir) / trace_files[scenario]
|
||||
if trace_path.exists():
|
||||
result = run_benchmark(
|
||||
str(trace_path),
|
||||
str(Path(args.storage_dir) / scenario),
|
||||
bytes_per_token,
|
||||
args.max_requests,
|
||||
args.max_blocks,
|
||||
args.replay_timestamps,
|
||||
args.time_scale,
|
||||
args.block_size_tokens,
|
||||
args.fsync_mode,
|
||||
args.fsync_batch_size
|
||||
)
|
||||
results.append(result)
|
||||
else:
|
||||
print(f"Warning: Trace file not found: {trace_path}")
|
||||
|
||||
# Print results
|
||||
if results:
|
||||
print_results(results)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
161
dependencies.sh
161
dependencies.sh
|
|
@ -23,7 +23,7 @@ NC="\033[0m" # No Color
|
|||
# Configuration
|
||||
REPO_ROOT=`pwd`
|
||||
GITHUB_PROXY=${GITHUB_PROXY:-"https://github.com"}
|
||||
GOVER=1.23.8
|
||||
GOVER=1.25.9
|
||||
|
||||
# Function to print section headers
|
||||
print_section() {
|
||||
|
|
@ -75,8 +75,7 @@ echo -e "${YELLOW}Mooncake Dependencies Installer${NC}"
|
|||
echo -e "This script will install all required dependencies for Mooncake."
|
||||
echo -e "The following components will be installed:"
|
||||
echo -e " - System packages (build tools, libraries)"
|
||||
echo -e " - yalantinglibs"
|
||||
echo -e " - Git submodules"
|
||||
echo -e " - Git submodules (including pybind11 and yalantinglibs)"
|
||||
echo -e " - Go $GOVER"
|
||||
echo
|
||||
|
||||
|
|
@ -102,8 +101,10 @@ echo -e "${YELLOW}This may take a few minutes...${NC}"
|
|||
|
||||
SYSTEM_PACKAGES="build-essential \
|
||||
cmake \
|
||||
ninja-build \
|
||||
git \
|
||||
wget \
|
||||
unzip \
|
||||
libibverbs-dev \
|
||||
libgoogle-glog-dev \
|
||||
libgtest-dev \
|
||||
|
|
@ -122,50 +123,47 @@ SYSTEM_PACKAGES="build-essential \
|
|||
libhiredis-dev \
|
||||
liburing-dev \
|
||||
libjemalloc-dev \
|
||||
libmsgpack-dev \
|
||||
libzstd-dev \
|
||||
libasio-dev \
|
||||
libxxhash-dev \
|
||||
pkg-config \
|
||||
patchelf"
|
||||
patchelf \
|
||||
libc6-dev \
|
||||
libc-bin"
|
||||
|
||||
apt-get install -y $SYSTEM_PACKAGES
|
||||
check_success "Failed to install system packages"
|
||||
print_success "System packages installed successfully"
|
||||
|
||||
# Install yalantinglibs
|
||||
# Initialize and update git submodules
|
||||
print_section "Initializing Git Submodules"
|
||||
|
||||
# Check if .gitmodules exists
|
||||
if [ -f "${REPO_ROOT}/.gitmodules" ]; then
|
||||
echo "Enter repository root: ${REPO_ROOT}"
|
||||
cd "${REPO_ROOT}"
|
||||
check_success "Failed to change to repository root directory"
|
||||
|
||||
echo "Initializing git submodules..."
|
||||
git submodule sync --recursive
|
||||
check_success "Failed to sync git submodules"
|
||||
git submodule update --init --recursive
|
||||
check_success "Failed to initialize git submodules"
|
||||
|
||||
print_success "Git submodules initialized and updated successfully"
|
||||
else
|
||||
echo -e "${YELLOW}No .gitmodules file found. Skipping...${NC}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Build and install yalantinglibs from submodule
|
||||
print_section "Installing yalantinglibs"
|
||||
|
||||
# Check if thirdparties directory exists
|
||||
if [ ! -d "${REPO_ROOT}/thirdparties" ]; then
|
||||
mkdir -p "${REPO_ROOT}/thirdparties"
|
||||
check_success "Failed to create thirdparties directory"
|
||||
fi
|
||||
|
||||
# Change to thirdparties directory
|
||||
cd "${REPO_ROOT}/thirdparties"
|
||||
check_success "Failed to change to thirdparties directory"
|
||||
|
||||
# Check if yalantinglibs is already installed
|
||||
if [ -d "yalantinglibs" ]; then
|
||||
echo -e "${YELLOW}yalantinglibs directory already exists. Removing for fresh install...${NC}"
|
||||
rm -rf yalantinglibs
|
||||
check_success "Failed to remove existing yalantinglibs directory"
|
||||
fi
|
||||
|
||||
# Clone yalantinglibs
|
||||
echo "Cloning yalantinglibs from ${GITHUB_PROXY}/alibaba/yalantinglibs.git"
|
||||
git clone ${GITHUB_PROXY}/alibaba/yalantinglibs.git
|
||||
check_success "Failed to clone yalantinglibs"
|
||||
|
||||
# Build and install yalantinglibs
|
||||
cd yalantinglibs
|
||||
check_success "Failed to change to yalantinglibs directory"
|
||||
|
||||
# Checkout version 0.5.6
|
||||
echo "Checking out yalantinglibs version 0.5.6..."
|
||||
git checkout 0.5.6
|
||||
check_success "Failed to checkout yalantinglibs version 0.5.6"
|
||||
cd "${REPO_ROOT}/extern/yalantinglibs"
|
||||
check_success "Failed to change to yalantinglibs submodule directory"
|
||||
|
||||
mkdir -p build
|
||||
check_success "Failed to create build directory"
|
||||
|
||||
cd build
|
||||
check_success "Failed to change to build directory"
|
||||
|
||||
|
|
@ -182,35 +180,25 @@ cmake --install .
|
|||
check_success "Failed to install yalantinglibs"
|
||||
|
||||
print_success "yalantinglibs installed successfully"
|
||||
cd "${REPO_ROOT}"
|
||||
|
||||
# Initialize and update git submodules
|
||||
print_section "Initializing Git Submodules"
|
||||
print_section "Verifying essential build tools"
|
||||
|
||||
# Check if .gitmodules exists
|
||||
if [ -f "${REPO_ROOT}/.gitmodules" ]; then
|
||||
# Check if submodules are already initialized by looking for the .git directory in the first submodule
|
||||
FIRST_SUBMODULE=$(grep "path" ${REPO_ROOT}/.gitmodules | head -1 | awk '{print $3}')
|
||||
|
||||
echo "Enter repository root: ${REPO_ROOT}"
|
||||
cd "${REPO_ROOT}"
|
||||
check_success "Failed to change to repository root directory"
|
||||
|
||||
if [ -d "${REPO_ROOT}/${FIRST_SUBMODULE}/.git" ] || [ -f "${REPO_ROOT}/${FIRST_SUBMODULE}/.git" ]; then
|
||||
echo -e "${YELLOW}Git submodules already initialized. Skipping...${NC}"
|
||||
else
|
||||
echo "Initializing git submodules..."
|
||||
git submodule update --init
|
||||
check_success "Failed to initialize git submodules"
|
||||
|
||||
print_success "Git submodules initialized and updated successfully"
|
||||
fi
|
||||
else
|
||||
echo -e "${YELLOW}No .gitmodules file found. Skipping...${NC}"
|
||||
exit 1
|
||||
# Verify getconf and ldd (required for glibc version detection in build_wheel.sh)
|
||||
# Both are provided by libc-bin, which is included in SYSTEM_PACKAGES
|
||||
if ! command -v getconf >/dev/null 2>&1; then
|
||||
print_error "getconf not found after installing system packages. This should not happen."
|
||||
fi
|
||||
if ! command -v ldd >/dev/null 2>&1; then
|
||||
print_error "ldd not found after installing system packages. This should not happen."
|
||||
fi
|
||||
print_success "getconf found: $(getconf --version 2>&1 | head -1)"
|
||||
print_success "ldd found: $(ldd --version 2>&1 | head -1)"
|
||||
|
||||
print_section "Installing Go $GOVER"
|
||||
|
||||
USED_CN_MIRROR=false
|
||||
|
||||
install_go() {
|
||||
ARCH=$(uname -m)
|
||||
if [ "$ARCH" = "aarch64" ]; then
|
||||
|
|
@ -221,18 +209,45 @@ install_go() {
|
|||
echo "Unsupported architecture: $ARCH"
|
||||
exit 1
|
||||
fi
|
||||
# Download Go
|
||||
echo "Downloading Go $GOVER..."
|
||||
wget -q --show-progress https://go.dev/dl/go$GOVER.linux-$ARCH.tar.gz
|
||||
check_success "Failed to download Go $GOVER"
|
||||
|
||||
GO_TARBALL="go$GOVER.linux-$ARCH.tar.gz"
|
||||
|
||||
# Try multiple download mirrors with fallback
|
||||
GO_DOWNLOAD_URLS=(
|
||||
"https://go.dev/dl/${GO_TARBALL}"
|
||||
"https://golang.google.cn/dl/${GO_TARBALL}"
|
||||
"https://mirrors.aliyun.com/golang/${GO_TARBALL}"
|
||||
)
|
||||
|
||||
DOWNLOAD_SUCCESS=false
|
||||
for url in "${GO_DOWNLOAD_URLS[@]}"; do
|
||||
echo "Downloading Go $GOVER from ${url}..."
|
||||
if wget -q --show-progress --timeout=30 --tries=2 -O "${GO_TARBALL}" "${url}"; then
|
||||
DOWNLOAD_SUCCESS=true
|
||||
# If the official source (go.dev) failed and we fell back to a CN mirror,
|
||||
# it likely means the network has restricted access to international sites.
|
||||
if [[ "$url" != "https://go.dev/dl/${GO_TARBALL}" ]]; then
|
||||
USED_CN_MIRROR=true
|
||||
fi
|
||||
print_success "Downloaded Go $GOVER from ${url}"
|
||||
break
|
||||
else
|
||||
echo -e "${YELLOW}Failed to download from ${url}, trying next mirror...${NC}"
|
||||
rm -f "${GO_TARBALL}"
|
||||
fi
|
||||
done
|
||||
|
||||
if [ "$DOWNLOAD_SUCCESS" = false ]; then
|
||||
print_error "Failed to download Go $GOVER from all mirrors"
|
||||
fi
|
||||
|
||||
# Install Go
|
||||
echo "Installing Go $GOVER..."
|
||||
tar -C /usr/local -xzf go$GOVER.linux-$ARCH.tar.gz
|
||||
tar -C /usr/local -xzf "${GO_TARBALL}"
|
||||
check_success "Failed to install Go $GOVER"
|
||||
|
||||
# Clean up downloaded file
|
||||
rm -f go$GOVER.linux-$ARCH.tar.gz
|
||||
rm -f "${GO_TARBALL}"
|
||||
check_success "Failed to clean up Go installation file"
|
||||
|
||||
print_success "Go $GOVER installed successfully"
|
||||
|
|
@ -258,6 +273,20 @@ if ! grep -q "export PATH=\$PATH:/usr/local/go/bin" ~/.bashrc; then
|
|||
echo -e "${YELLOW}Please run 'source ~/.bashrc' or start a new terminal to use Go${NC}"
|
||||
fi
|
||||
|
||||
# Set GOPROXY only if Go download fell back to a CN mirror, indicating restricted
|
||||
# network access to international sites. Skip if user already configured GOPROXY.
|
||||
if [ "$USED_CN_MIRROR" = true ] && [ -z "$GOPROXY" ]; then
|
||||
export GOPROXY=https://goproxy.cn,https://goproxy.io,direct
|
||||
echo -e "${YELLOW}Detected restricted network (Go was downloaded from a CN mirror).${NC}"
|
||||
echo -e "${YELLOW}GOPROXY set to: ${GOPROXY}${NC}"
|
||||
if ! grep -q "export GOPROXY=" ~/.bashrc; then
|
||||
echo 'export GOPROXY=https://goproxy.cn,https://goproxy.io,direct' >> ~/.bashrc
|
||||
echo -e "${YELLOW}GOPROXY added to ~/.bashrc for future sessions${NC}"
|
||||
fi
|
||||
elif [ -n "$GOPROXY" ]; then
|
||||
echo -e "${GREEN}GOPROXY already set to: ${GOPROXY}${NC}"
|
||||
fi
|
||||
|
||||
# Return to the repository root
|
||||
cd "${REPO_ROOT}"
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,121 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
###############################################################################
|
||||
# Stage 1: build Mooncake from source and produce a Python wheel
|
||||
###############################################################################
|
||||
ARG CUDA_VERSION=12.8.1
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
|
||||
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} AS builder
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1
|
||||
|
||||
ARG PYTHON_VERSION=3.10
|
||||
ARG PYPA_INDEX_URL=https://bootstrap.pypa.io
|
||||
ARG CMAKE_BUILD_TYPE=Release
|
||||
ARG EP_TORCH_VERSIONS="2.9.1"
|
||||
ARG TORCH_CUDA_ARCH_LIST="8.0;9.0"
|
||||
|
||||
ENV PYTHON_VERSION=${PYTHON_VERSION} \
|
||||
BUILD_WITH_EP=1 \
|
||||
EP_TORCH_VERSIONS=${EP_TORCH_VERSIONS} \
|
||||
TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
|
||||
PATH="/usr/local/go/bin:${PATH}"
|
||||
|
||||
# Install base build utilities and the requested Python version via deadsnakes PPA
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
curl \
|
||||
git \
|
||||
ninja-build \
|
||||
software-properties-common \
|
||||
pkg-config && \
|
||||
add-apt-repository -y ppa:deadsnakes/ppa && \
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
python${PYTHON_VERSION} \
|
||||
python${PYTHON_VERSION}-dev \
|
||||
python${PYTHON_VERSION}-venv && \
|
||||
curl -sS ${PYPA_INDEX_URL}/get-pip.py | python${PYTHON_VERSION} && \
|
||||
update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 && \
|
||||
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
|
||||
apt-get purge -y --auto-remove software-properties-common && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /workspace
|
||||
COPY . /workspace
|
||||
|
||||
# Install Mooncake dependencies (yalantinglibs, Go, etc.)
|
||||
RUN bash dependencies.sh -y
|
||||
|
||||
# Configure & build Mooncake
|
||||
RUN mkdir -p build && \
|
||||
cd build && \
|
||||
cmake -G Ninja .. \
|
||||
-DBUILD_UNIT_TESTS=OFF \
|
||||
-DUSE_HTTP=ON \
|
||||
-DUSE_ETCD=ON \
|
||||
-DUSE_CUDA=ON \
|
||||
-DWITH_EP=ON \
|
||||
-DSTORE_USE_ETCD=ON \
|
||||
-DPython3_EXECUTABLE=/usr/bin/python${PYTHON_VERSION} \
|
||||
-DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} && \
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH && \
|
||||
cmake --build .
|
||||
|
||||
# Build nvlink allocator to make wheel self-contained for CUDA paths
|
||||
RUN export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH && \
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH && \
|
||||
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH && \
|
||||
mkdir -p build/mooncake-transfer-engine/nvlink-allocator && \
|
||||
cd mooncake-transfer-engine/nvlink-allocator && \
|
||||
bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
|
||||
|
||||
# Build the Python wheel from local sources
|
||||
RUN OUTPUT_DIR=dist ./scripts/build_wheel.sh
|
||||
|
||||
###############################################################################
|
||||
# Stage 2: install the freshly built wheel into a runtime image
|
||||
###############################################################################
|
||||
FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu${UBUNTU_VERSION} AS runtime
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
# Inherit build-args so the runtime stage installs the matching interpreter
|
||||
ARG PYTHON_VERSION=3.10
|
||||
ARG PYPA_INDEX_URL=https://bootstrap.pypa.io
|
||||
ENV PYTHON_VERSION=${PYTHON_VERSION}
|
||||
|
||||
# Install runtime dependencies and the requested Python version
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
curl \
|
||||
software-properties-common \
|
||||
ibverbs-providers \
|
||||
rdma-core \
|
||||
libibverbs1 \
|
||||
librdmacm1 \
|
||||
libnuma1 \
|
||||
liburing2 \
|
||||
libyaml-0-2 \
|
||||
libcurl4 && \
|
||||
add-apt-repository -y ppa:deadsnakes/ppa && \
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
python${PYTHON_VERSION} && \
|
||||
curl -sS ${PYPA_INDEX_URL}/get-pip.py | python${PYTHON_VERSION} && \
|
||||
update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 && \
|
||||
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
|
||||
apt-get purge -y --auto-remove software-properties-common curl && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy wheels produced in builder stage and install them via pip
|
||||
COPY --from=builder /workspace/mooncake-wheel/dist /tmp/mooncake-wheel
|
||||
RUN python${PYTHON_VERSION} -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
|
@ -0,0 +1,87 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
###############################################################################
|
||||
# Stage 1: build Mooncake from source and produce a Python wheel
|
||||
###############################################################################
|
||||
ARG MUSA_VERSION=rc4.3.0
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
|
||||
FROM mthreads/musa:${MUSA_VERSION}-devel-ubuntu${UBUNTU_VERSION}-amd64 AS builder
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1
|
||||
|
||||
ARG PYTHON_VERSION=3.10
|
||||
ARG CMAKE_BUILD_TYPE=Release
|
||||
|
||||
ENV PYTHON_VERSION=${PYTHON_VERSION} \
|
||||
PATH="/usr/local/go/bin:${PATH}"
|
||||
|
||||
# Install base build utilities and python bindings
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
curl \
|
||||
git \
|
||||
python3 \
|
||||
python3-dev \
|
||||
python3-pip \
|
||||
python-is-python3 \
|
||||
pkg-config && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /workspace
|
||||
COPY . /workspace
|
||||
|
||||
# Install Mooncake dependencies (yalantinglibs, Go, etc.)
|
||||
RUN bash dependencies.sh -y
|
||||
|
||||
# Configure & build Mooncake
|
||||
RUN mkdir -p build && \
|
||||
cd build && \
|
||||
cmake .. \
|
||||
-DBUILD_UNIT_TESTS=OFF \
|
||||
-DUSE_HTTP=ON \
|
||||
-DUSE_ETCD=ON \
|
||||
-DUSE_MUSA=ON \
|
||||
-DSTORE_USE_ETCD=ON \
|
||||
-DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} && \
|
||||
cmake --build . -j"$(nproc)"
|
||||
|
||||
# Build nvlink allocator to make wheel self-contained for MUSA paths
|
||||
RUN mkdir -p build/mooncake-transfer-engine/nvlink-allocator && \
|
||||
cd mooncake-transfer-engine/nvlink-allocator && \
|
||||
bash build.sh --use-mcc ../../build/mooncake-transfer-engine/nvlink-allocator/
|
||||
|
||||
# Build the Python wheel from local sources
|
||||
RUN OUTPUT_DIR=dist ./scripts/build_wheel.sh
|
||||
|
||||
###############################################################################
|
||||
# Stage 2: install the freshly built wheel into a runtime image
|
||||
###############################################################################
|
||||
FROM mthreads/musa:${MUSA_VERSION}-devel-ubuntu${UBUNTU_VERSION}-amd64 AS runtime
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
# Install runtime dependencies required by Mooncake
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
python3 \
|
||||
python3-pip \
|
||||
ibverbs-providers \
|
||||
rdma-core \
|
||||
libibverbs1 \
|
||||
librdmacm1 \
|
||||
libnuma1 \
|
||||
liburing2 \
|
||||
libyaml-0-2 \
|
||||
libcurl4 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy wheels produced in builder stage and install them via pip
|
||||
COPY --from=builder /workspace/mooncake-wheel/dist /tmp/mooncake-wheel
|
||||
RUN python3 -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
# Governance
|
||||
|
||||
This document describes the governance model and code maintenance process for the Mooncake project.
|
||||
|
||||
Mooncake is an open-source project that follows a community-driven governance model. The project is maintained by a group of Collaborators who oversee the codebase and ensure code quality, while Contributors from the community help improve the project through code contributions, bug reports, documentation, and other valuable inputs.
|
||||
|
||||
## Collaborator
|
||||
|
||||
Collaborators are trusted members of the community who have been granted specific permissions to review, approve, and merge code changes. This model helps protect critical code paths while enabling efficient development workflows. Collaborators may hold one or more of the following roles:
|
||||
|
||||
- **Codeowner**: Codeowners are maintainers whose primary responsibility is to protect critical code. Each pull request needs at least one Codeowner approval if it modifies files protected by [CODEOWNERS](https://github.com/kvcache-ai/Mooncake/blob/main/.github/CODEOWNERS). When a pull request is submitted, Codeowners are responsible for reviewing the code in a timely manner, or delegating the review to appropriate reviewers when unavailable, and ultimately approving the pull request. This role is not just an honor but a significant responsibility, as pull requests cannot be merged without Codeowner approval. Current Codeowners are listed in the [CODEOWNERS](https://github.com/kvcache-ai/Mooncake/blob/main/.github/CODEOWNERS) file.
|
||||
|
||||
- **Write**: Members with Write permission are trusted contributors responsible for reviewing code in a timely manner and ensuring code quality. They actively participate in code reviews, providing feedback on code correctness, functionality, style adherence, test coverage, and documentation. After a Codeowner has approved a pull request and all required tests have passed, Write members have the permission to merge the pull request. Their role is crucial in maintaining both code quality and the project's development velocity.
|
||||
|
||||
## Development Process
|
||||
|
||||
### Technical Discussion
|
||||
|
||||
We encourage developers to initiate technical discussions in the community before starting formal development. This helps gather valuable feedback from other contributors and users, ensuring that proposed changes align with the project's goals and benefit the broader community.
|
||||
|
||||
- **Major Changes**: For significant architectural changes (typically >500 LOC excluding tests), an RFC (Request for Comments) must be submitted on GitHub to solicit comprehensive feedback from the community. This is a mandatory step to ensure all stakeholders have an opportunity to provide input and raise concerns before implementation begins.
|
||||
|
||||
- **Urgent Changes**: For urgent modifications or critical bug fixes that require immediate attention, developers can raise the issue on [Slack](https://join.slack.com/t/mooncake-project/shared_invite/zt-3ig4fjai8-KH1zIm3x8Vm8WqyH0i_JaA) and request relevant collaborators to participate quickly. This allows for rapid response while still maintaining communication with the team.
|
||||
|
||||
- **Breaking Changes**: Any changes that break backward compatibility require broader consensus and must be clearly documented with migration guides. These should always go through the RFC process to ensure all implications are thoroughly discussed.
|
||||
|
||||
### Pull Request Merge Process
|
||||
|
||||
The pull request merge process ensures code quality through collaborative review and automated checks:
|
||||
|
||||
1. **PR Submission**: Contributors submit pull requests with clear descriptions following the [PR template](https://github.com/kvcache-ai/Mooncake/blob/main/.github/pull_request_template.md). The description should explain the changes, their motivation, and any relevant context.
|
||||
|
||||
2. **Automatic Review Requests**: Once a PR is submitted, GitHub will automatically request reviews from Codeowners based on the CODEOWNERS file. Additionally, an AI assistant will automatically review the code to provide initial feedback and catch common issues.
|
||||
|
||||
3. **Code Review**: Codeowners conduct code reviews or invite appropriate reviewers with relevant expertise. Additionally, anyone in the community who is relevant or interested is welcome to help with reviews. Contributors should address all review comments and feedback.
|
||||
|
||||
4. **Approval and Merge**: Once approved by at least one Codeowner and all required tests have passed, a Write member can merge the pull request. The PR must meet all quality standards and have no blocking issues before merging.
|
||||
|
||||
If you encounter any issues during the merge process, you can discuss them on [Slack](https://join.slack.com/t/mooncake-project/shared_invite/zt-3ig4fjai8-KH1zIm3x8Vm8WqyH0i_JaA) to get help from Collaborators.
|
||||
|
||||
---
|
||||
|
||||
*This governance model is designed to be flexible and may evolve as the project grows. We welcome feedback and suggestions for improvement.*
|
||||
|
|
@ -8,6 +8,7 @@ This page summarizes useful flags, environment variables, and HTTP endpoints to
|
|||
- `--rpc_port` (int, default 50051): RPC listen port.
|
||||
- `--rpc_thread_num` (int, default min(4, CPU cores)): RPC worker threads. If not set, uses `--max_threads` (default 4) capped by CPU cores.
|
||||
- `--rpc_address` (str, default `0.0.0.0`): RPC bind address.
|
||||
- `--rpc_interface` (str, default empty): Network interface used to resolve the final RPC address. When set, Mooncake Master resolves the interface's current IPv4 address at startup and uses it as the final `rpc_address`. This overrides `--rpc_address`.
|
||||
- `--rpc_conn_timeout_seconds` (int, default `0`): RPC idle connection timeout; `0` disables.
|
||||
- `--rpc_enable_tcp_no_delay` (bool, default `true`): Enable TCP_NODELAY.
|
||||
|
||||
|
|
@ -20,9 +21,14 @@ This page summarizes useful flags, environment variables, and HTTP endpoints to
|
|||
- `--http_metadata_server_host` (str, default `0.0.0.0`): Metadata bind host.
|
||||
- `--http_metadata_server_port` (int, default `8080`): Metadata TCP port.
|
||||
|
||||
- Allocation Strategy
|
||||
- `--allocation_strategy` (str, default `random`): Memory allocation strategy for replica placement. Available options:
|
||||
- `random`: Pure random selection across segments (baseline, fastest).
|
||||
- `free_ratio_first`: Free-ratio-first strategy. Samples multiple candidates and selects those with highest free space ratio for better load balancing.
|
||||
|
||||
- Eviction and TTLs
|
||||
- `--default_kv_lease_ttl` (uint64, default `5000` ms): Default lease TTL for KV objects.
|
||||
- `--default_kv_soft_pin_ttl` (uint64, default `1800000` ms): Soft pin TTL (30 minutes).
|
||||
- `--default_kv_lease_ttl` (duration, default `5000` ms): Default lease TTL for KV objects. The default unit is milliseconds, so `5000` means `5000ms`. Duration strings such as `5000ms`, `5s`, `30m`, or `1h` are also supported.
|
||||
- `--default_kv_soft_pin_ttl` (duration, default `1800000` ms): Soft pin TTL (30 minutes). The default unit is milliseconds, so `1800000` means `1800000ms`. Duration strings such as `1800000ms`, `30m`, or `1h` are also supported.
|
||||
- `--allow_evict_soft_pinned_objects` (bool, default `true`): Allow evicting soft-pinned objects.
|
||||
- `--eviction_ratio` (double, default `0.05`): Fraction evicted when hitting high watermark.
|
||||
- `--eviction_high_watermark_ratio` (double, default `0.95`): Usage ratio to trigger eviction.
|
||||
|
|
@ -33,10 +39,30 @@ This page summarizes useful flags, environment variables, and HTTP endpoints to
|
|||
- `--client_ttl` (int64, default `10` s): Client alive TTL after last ping (HA mode).
|
||||
- `--cluster_id` (str, default `mooncake_cluster`): Cluster ID for persistence in HA mode.
|
||||
|
||||
- Task Manager (optional)
|
||||
- `--max_total_finished_tasks` (uint32, default `10000`): Maximum number of finished tasks to keep in memory. When this limit is reached, the oldest finished tasks will be pruned from memory.
|
||||
- `--max_total_pending_tasks` (uint32, default `10000`): Maximum number of pending tasks that can be queued in memory. When this limit is reached, new task submissions will fail with `TASK_PENDING_LIMIT_EXCEEDED` error.
|
||||
- `--max_total_processing_tasks` (uint32, default `10000`): Maximum number of tasks that can be processing simultaneously. When this limit is reached, no new tasks will be popped from the pending queue until some processing tasks complete.
|
||||
- `--max_retry_attempts` (uint32, default `10`): Maximum number of retry attempts for failed tasks. Tasks that fail with `NO_AVAILABLE_HANDLE` error will be retried up to this many times before being marked as failed.
|
||||
|
||||
- DFS Storage (optional)
|
||||
- `--root_fs_dir` (str, default empty): DFS mount directory for storage backend, used in Multi-layer Storage Support.
|
||||
- `--global_file_segment_size` (int64, default `int64_max`): Maximum available space for DFS segments.
|
||||
|
||||
- Snapshot / Restore (optional)
|
||||
- `--enable_snapshot` (bool, default `false`): Enable periodic snapshot of master metadata data (effective when using the `offset` memory allocator).
|
||||
- `--snapshot_interval_seconds` (uint64, default `600`): Interval in seconds between periodic snapshots of master data.
|
||||
- `--snapshot_child_timeout_seconds` (uint64, default `300`): Timeout in seconds for each snapshot child process.
|
||||
- `--snapshot_retention_count` (uint32, default `2`): Number of recent snapshots to keep. Older snapshots beyond this limit will be automatically deleted.
|
||||
- `--snapshot_backend_type` (str, required when snapshot enabled): Snapshot storage backend type: `local` for local filesystem, `s3` for S3 storage.
|
||||
- `--snapshot_backup_dir` (str, default empty): Optional local directory for snapshot backup. If empty (default), local backup is disabled. When set, it serves two purposes: (1) during snapshot persistence, data will be saved locally as a fallback if uploading to the backend fails; (2) during restore, downloaded metadata will also be saved to this directory as a local backup.
|
||||
- `--enable_snapshot_restore` (bool, default `false`): Enable restore from the latest snapshot at master startup.
|
||||
- **Environment variable** `MOONCAKE_SNAPSHOT_LOCAL_PATH` (**required** when `--snapshot_backend_type=local`): Persistent directory path for local snapshot storage. This variable **must** be set before starting the master; there is no default value. Example: `export MOONCAKE_SNAPSHOT_LOCAL_PATH=/data/mooncake_snapshots`.
|
||||
|
||||
> **Warning: Managed Directory**
|
||||
>
|
||||
> The snapshot storage path (`MOONCAKE_SNAPSHOT_LOCAL_PATH` for local backend, or S3 bucket for S3 backend) is a **managed directory** exclusively controlled by the Mooncake snapshot system. **DO NOT store other files or data in this directory.** Old snapshots exceeding `--snapshot_retention_count` will be automatically and permanently deleted during cleanup. Use a dedicated, isolated directory for snapshot storage to avoid accidental data loss.
|
||||
|
||||
Example (enable embedded HTTP metadata and metrics):
|
||||
|
||||
```bash
|
||||
|
|
@ -49,13 +75,41 @@ mooncake_master \
|
|||
--enable_metric_reporting=true
|
||||
```
|
||||
|
||||
Example (resolve the master RPC address from a stable interface name in a container):
|
||||
|
||||
```bash
|
||||
mooncake_master \
|
||||
--rpc_interface=eth0 \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_host=0.0.0.0 \
|
||||
--http_metadata_server_port=8080
|
||||
```
|
||||
|
||||
This resolves the current IPv4 address of `eth0` at startup and uses it as the final `rpc_address`.
|
||||
|
||||
Example (use free-ratio-first allocation strategy for better load balancing):
|
||||
|
||||
```bash
|
||||
mooncake_master \
|
||||
--allocation_strategy=free_ratio_first \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_port=8080
|
||||
```
|
||||
|
||||
**Tips:**
|
||||
|
||||
In addition to command-line flags, the Master also supports configuration via JSON and YAML files. For example:
|
||||
|
||||
```bash
|
||||
mooncake_master \
|
||||
--config_path=mooncake-store/conf/master.yaml
|
||||
--config_path=mooncake-store/conf/master.yaml
|
||||
```
|
||||
|
||||
For config files, the equivalent setting is:
|
||||
|
||||
```yaml
|
||||
rpc_interface: "eth0"
|
||||
rpc_port: 50051
|
||||
```
|
||||
|
||||
## Metrics Endpoints
|
||||
|
|
@ -80,7 +134,7 @@ curl -s http://<master_host>:9003/metrics/summary
|
|||
- If `MC_MS_AUTO_DISC=0`, pass `rdma_devices` (comma-separated) to the Python `setup(...)` call.
|
||||
|
||||
- Transfer Engine metrics (disabled by default)
|
||||
- `MC_TE_METRIC` (default `0`/unset): Set to `1` to enable periodic engine metrics logging.
|
||||
- `MC_TE_METRIC` (default `0`/unset): Set to `1` to enable periodic engine metrics logging. **Note:** Not supported when using Transfer Engine TENT.
|
||||
- `MC_TE_METRIC_INTERVAL_SECONDS` (default `5`): Positive integer seconds between reports (effective only if metrics enabled).
|
||||
|
||||
- Client metrics (enabled by default)
|
||||
|
|
@ -104,3 +158,13 @@ Available log levels: trace, debug, info, warn (or warning), error, and critical
|
|||
- Scale `--rpc_thread_num` with available CPU cores and workload.
|
||||
- Start with default eviction settings; adjust `--eviction_high_watermark_ratio` and `--eviction_ratio` based on memory pressure and object churn.
|
||||
- Use `/metrics/summary` during bring-up; integrate `/metrics` with Prometheus/Grafana for production.
|
||||
|
||||
|
||||
---
|
||||
|
||||
:::{toctree}
|
||||
:caption: Advanced Topics
|
||||
:maxdepth: 1
|
||||
|
||||
ssd-offload
|
||||
:::
|
||||
|
|
|
|||
|
|
@ -0,0 +1,280 @@
|
|||
# SSD Offload
|
||||
|
||||
## Overview
|
||||
|
||||
Mooncake Store supports offloading KV cache objects from distributed memory to local SSD. When memory pressure is high, the master instructs clients to persist selected objects to disk. On a cache miss, the client automatically falls back to reading from SSD.
|
||||
|
||||
SSD offload is currently **only available in Real Client mode**. The real client is a standalone process that communicates with the application (e.g., SGLang) via RPC. All SSD reads and writes happen within this process.
|
||||
|
||||
## Startup Steps
|
||||
|
||||
### Step 1: Create the SSD storage directory
|
||||
|
||||
```bash
|
||||
mkdir -p /nvme/mooncake_offload
|
||||
```
|
||||
|
||||
### Step 2: Start the master
|
||||
|
||||
```bash
|
||||
mooncake_master \
|
||||
--rpc_port=50051 \
|
||||
--enable-offload true
|
||||
```
|
||||
|
||||
### Step 3: Start the real client with SSD offload enabled
|
||||
|
||||
Use the `--enable_offload` flag to enable SSD offload, and set environment variables to specify the storage path and backend:
|
||||
|
||||
```bash
|
||||
export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
|
||||
export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
|
||||
|
||||
mooncake_client \
|
||||
--master_server_address=127.0.0.1:50051 \
|
||||
--host=<machine IP> \
|
||||
--protocol="rdma" \
|
||||
--device_names=<NIC name, e.g. eth0> \
|
||||
--port=50052 \
|
||||
--global_segment_size="4 GB" \
|
||||
--enable_offload=true \
|
||||
--metadata_server="P2PHANDSHAKE"
|
||||
```
|
||||
|
||||
> **Note:** On startup, the real client automatically scans existing SSD data and reports it to the master. No manual recovery is needed.
|
||||
|
||||
### Step 4: Connect the application to the real client
|
||||
|
||||
The application (e.g., SGLang) connects to the real client via the `MooncakeDistributedStore` Python SDK. SSD offload and fallback loading are handled transparently.
|
||||
|
||||
```python
|
||||
from mooncake.store import MooncakeDistributedStore
|
||||
|
||||
store = MooncakeDistributedStore()
|
||||
store.setup(
|
||||
local_hostname="<machine IP>",
|
||||
metadata_server="P2PHANDSHAKE",
|
||||
global_segment_size=4 * 1024 * 1024 * 1024, # 4 GB
|
||||
local_buffer_size=512 * 1024 * 1024, #512MB
|
||||
protocol="rdma",
|
||||
device_name="eth0",
|
||||
master_server_address="127.0.0.1:50051",
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Real Client Parameters
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--master_server_address` | `127.0.0.1:50051` | Master address |
|
||||
| `--host` | `0.0.0.0` | This machine's externally reachable IP |
|
||||
| `--port` | `50052` | Real client RPC listening port |
|
||||
| `--device_names` | ` ` | NIC name(s), e.g. `eth0` or `mlx5_0` |
|
||||
| `--protocol` | `tcp` | Transport protocol: `tcp` or `rdma` |
|
||||
| `--global_segment_size` | `4 GB` | Memory pool size allocated for this node |
|
||||
| `--enable_offload` | `false` | **Must be set to `true` to enable SSD offload** |
|
||||
| `--threads` | `1` | Number of RPC server threads |
|
||||
|
||||
---
|
||||
|
||||
## SSD Offload Configuration
|
||||
|
||||
### Core settings
|
||||
|
||||
| Environment Variable | Default | Description |
|
||||
|---|---|---|
|
||||
| `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | `/data/file_storage` | Absolute path to the SSD storage directory |
|
||||
| `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR` | `bucket_storage_backend` | Storage backend type (see below) |
|
||||
| `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` | `1342177280` (1.25 GB) | Client-side staging buffer size |
|
||||
| `MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT` | `20000` | Max keys processed per iteration when scanning existing SSD metadata on startup |
|
||||
| `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES` | `2199023255552` (2 TB) | Maximum disk usage |
|
||||
| `MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT` | `10000000` | Maximum number of objects on disk |
|
||||
| `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` | `10` | Interval for offload heartbeat to master (seconds) |
|
||||
| `MOONCAKE_OFFLOAD_USE_URING` | `false` | Enable io_uring for async file I/O |
|
||||
|
||||
### Bucket backend settings
|
||||
|
||||
Applies when `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend`.
|
||||
|
||||
| Environment Variable | Default | Description |
|
||||
|---|---|---|
|
||||
| `MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES` | `268435456` (256 MB) | Max size per bucket |
|
||||
| `MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT` | `500` | Max keys per bucket |
|
||||
| `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` | `0` | Eviction threshold in bytes. When set to `0`, the backend uses **90% of the physical disk capacity** as the quota — it does not mean unlimited. Set an explicit value to control disk usage precisely. |
|
||||
| `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY` | `none` | Eviction policy: `none` / `fifo` / `lru` |
|
||||
|
||||
---
|
||||
|
||||
## Storage Backends
|
||||
|
||||
### `bucket_storage_backend` (recommended)
|
||||
|
||||
Groups multiple objects into bucket files. Reduces filesystem overhead, supports efficient batch I/O, and supports FIFO and LRU eviction.
|
||||
|
||||
**File layout:**
|
||||
```
|
||||
/nvme/mooncake_offload/
|
||||
├── 1710000000000-0.bucket # data file (multiple KV pairs)
|
||||
├── 1710000000000-0.meta # metadata file
|
||||
├── 1710000000001-0.bucket
|
||||
└── ...
|
||||
```
|
||||
|
||||
Best for: general-purpose use, large-scale deployments.
|
||||
|
||||
### `file_per_key_storage_backend`
|
||||
|
||||
Stores each object in an individual file. Simple and easy to inspect, but generates many small files at scale.
|
||||
|
||||
| Environment Variable | Default | Description |
|
||||
|---|---|---|
|
||||
| `MOONCAKE_OFFLOAD_FSDIR` | `file_per_key_dir` | Subdirectory name under `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` where objects are stored |
|
||||
| `MOONCAKE_OFFLOAD_ENABLE_EVICTION` | `true` | Enable disk eviction when the total size exceeds the quota |
|
||||
|
||||
Best for: debugging or small-scale deployments.
|
||||
|
||||
### `offset_allocator_storage_backend`
|
||||
|
||||
Pre-allocates a single large file and manages offset-based allocation within it. Highest concurrency via 1024-shard metadata.
|
||||
|
||||
> **Warning:** This backend does **not** support metadata recovery on restart. On initialization, the data file is truncated and all in-memory metadata is cleared. Any previously offloaded objects become inaccessible after a process restart.
|
||||
|
||||
**Capacity:** `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES` is used directly as the pre-allocated file size (100%, no safety margin). Unlike `bucket_storage_backend`, there is no separate quota variable — this is the sole disk usage control. Set it below the physical disk capacity to avoid filling the disk; writes are rejected once usage reaches this limit.
|
||||
|
||||
Best for: high-concurrency scenarios with many small objects where restart durability is not required.
|
||||
|
||||
---
|
||||
|
||||
## Eviction (Bucket Backend Only)
|
||||
|
||||
When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend automatically evicts buckets before writing new ones if total disk usage would exceed the limit.
|
||||
|
||||
| Policy | Behavior |
|
||||
|--------|----------|
|
||||
| `none` | No eviction (default); writes fail when disk is full |
|
||||
| `fifo` | Evict the oldest bucket first |
|
||||
| `lru` | Evict the least recently read bucket first |
|
||||
|
||||
Eviction is two-phase: the bucket is removed from metadata and master is notified first, then in-flight reads are drained before files are deleted.
|
||||
|
||||
---
|
||||
|
||||
## Example
|
||||
|
||||
The following example starts a master and a real client on a single machine.
|
||||
|
||||
### Environment
|
||||
|
||||
- Machine IP: `192.168.1.10`
|
||||
- NIC: `eth0`
|
||||
- SSD mount point: `/nvme`
|
||||
- Memory pool size: 4 GB (smaller than the total data written, to trigger offload)
|
||||
|
||||
### Start the master
|
||||
|
||||
```bash
|
||||
mooncake_master \
|
||||
--rpc_port=50051
|
||||
```
|
||||
|
||||
### Start the real client (new terminal)
|
||||
|
||||
```bash
|
||||
export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
|
||||
export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
|
||||
export MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE=$((200 * 1024 * 1024 * 1024)) # 200 GB
|
||||
export MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY=lru
|
||||
|
||||
mooncake_client \
|
||||
--master_server_address="192.168.1.10:50051" \
|
||||
--host="192.168.1.10" \
|
||||
--device_names="eth0" \
|
||||
--port=50052 \
|
||||
--protocol="rdma" \
|
||||
--global_segment_size="4GB" \
|
||||
--enable_offload="true"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Notes
|
||||
|
||||
- `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` must be an absolute path to an existing, writable directory. Symbolic links and paths containing `..` are rejected.
|
||||
- On real client restart, the backend automatically scans existing SSD files and reports them to the master, so previously offloaded objects remain accessible.
|
||||
- Eviction only notifies the master and deletes local files; objects replicated on other nodes are unaffected.
|
||||
- Each machine requires its own real client process. In multi-node deployments, ensure `--host` and `--port` are correctly set so nodes can reach each other.
|
||||
|
||||
**2-node example:** suppose Node A (`192.168.1.10`) runs the master and Node B (`192.168.1.11`) is a second worker. Both real clients must point to the same master and advertise their own externally reachable IP:
|
||||
|
||||
```bash
|
||||
# Node A — runs the master and its own real client
|
||||
mooncake_master --rpc_port=50051 --enable-offload true &
|
||||
|
||||
export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
|
||||
mooncake_client \
|
||||
--master_server_address="192.168.1.10:50051" \
|
||||
--host="192.168.1.10" \ # externally reachable IP of Node A
|
||||
--device_names="eth0" \
|
||||
--protocol="rdma" \
|
||||
--metadata_server="P2PHANDSHAKE" \
|
||||
--port=50052 \
|
||||
--global_segment_size="4GB" \
|
||||
--enable_offload="true"
|
||||
```
|
||||
|
||||
```bash
|
||||
# Node B — real client only; points to the same master on Node A
|
||||
export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
|
||||
mooncake_client \
|
||||
--master_server_address="192.168.1.10:50051" \
|
||||
--host="192.168.1.11" \ # externally reachable IP of Node B, NOT 127.0.0.1
|
||||
--device_names="eth0" \
|
||||
--protocol="rdma" \
|
||||
--metadata_server="P2PHANDSHAKE" \
|
||||
--port=50052 \
|
||||
--global_segment_size="4GB" \
|
||||
--enable_offload="true"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### SSD offload is not triggering
|
||||
|
||||
- Confirm `--enable_offload=true` is passed to `mooncake_client` and `--enable-offload true` is passed to `mooncake_master`.
|
||||
- Check that `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` points to an existing, writable directory. The client will fail silently if the path is invalid.
|
||||
- Verify memory pressure is actually high enough for the master to trigger offload. If the memory pool (`--global_segment_size`) is large relative to the data written, offload may never activate.
|
||||
|
||||
### "Permission denied" or "No such file or directory" on the storage path
|
||||
|
||||
- Ensure the directory exists before starting the client: `mkdir -p <path>`.
|
||||
- Confirm the process user has read/write access to the directory.
|
||||
- Symbolic links and paths containing `..` are rejected — use an absolute, canonical path.
|
||||
|
||||
### "Failed to register buffer with UringFile" warning in logs
|
||||
|
||||
This warning appears when `MOONCAKE_OFFLOAD_USE_URING=true` and the io_uring fixed-buffer registration fails. The most common cause is that `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` exceeds the process's locked-memory limit (`RLIMIT_MEMLOCK`). io_uring requires the registered buffer to be pinned in physical memory, which counts against this limit.
|
||||
|
||||
Check the current limit:
|
||||
|
||||
```bash
|
||||
ulimit -l # in KB; "unlimited" means no cap
|
||||
```
|
||||
|
||||
To raise it for the current session:
|
||||
|
||||
```bash
|
||||
ulimit -l unlimited
|
||||
```
|
||||
|
||||
To raise it permanently, add the following to `/etc/security/limits.conf`:
|
||||
|
||||
```
|
||||
* soft memlock unlimited
|
||||
* hard memlock unlimited
|
||||
```
|
||||
|
||||
Alternatively, reduce `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` to a value within the existing limit. Note that the warning does not abort startup — the client falls back to non-fixed-buffer I/O — but performance may be lower than expected.
|
||||
|
|
@ -0,0 +1,224 @@
|
|||
# Mooncake Conductor Indexer
|
||||
|
||||
## Introduction
|
||||
The Mooncake Conductor Indexer is a specialized service designed to efficiently track and report token hit counts across various caching levels for different model instances. It provides a list of APIs that allow users to query token hit statistics based on token ID or chunked token hash, thereby facilitating optimized the performance of LLM inference.The figure below illustrates the architecture of Mooncake KVindexer: 
|
||||
|
||||
## tiered storage & pools
|
||||
We drew inspiration from the definition of [KVBM components](https://github.com/ai-dynamo/dynamo/blob/main/docs/kvbm/kvbm_components.md) and divided the KV cache into three levels: G1, G2, and G3. The detailed introduction is as follows:
|
||||
|
||||
- **Device Pool(G1)**: Device-resident KV block pool. Allocates mutable device blocks, registers completed blocks (immutable), serves lookups by sequence hash, and is the target for onboarding (Host→Device, Disk→Device).
|
||||
- **Host Pool(G2)**: Mooncake registered memory KV pool. Receives Device offloads (Device→Host), can onboard to Device (Host→Device), and offloads to Disk. For high-performance, zero-copy data transfers, it utilizes the Mooncake Transfer-Engine.
|
||||
- **Disk Pool(G3)**: SSD NVMe-backed KV pool. Receives Host offloads (Host→Disk), and provides large space for storing KV.
|
||||
|
||||
|
||||
## Indexer API
|
||||
|
||||
### `POST /query`
|
||||
query token hit count.
|
||||
- **Input**:
|
||||
- **Body** (JSON):
|
||||
```json
|
||||
{
|
||||
"model": "deepseek",
|
||||
"lora_name": "xx-adapter",
|
||||
"lora_id": 12, // defined for backward compatibility and should not be used together with `lora_name`
|
||||
"token_ids": [1, 15, 100],
|
||||
"tenant_id": None,
|
||||
"cache_salt": None,
|
||||
}
|
||||
```
|
||||
- **Parameter Description**:
|
||||
- `model`: (required, string) model name
|
||||
- `lora_name`: (optional, string) The name of the LoRA adapter, default is `None`(indicating no LoRA adapter is used)
|
||||
- `lora_id`: (optional, int) The ID of the LoRA adapter. This parameter is defined for backward compatibility and should not be used together with `lora_name`(Only one of them can be specified). Default is `-1`(indicating no LoRA adapter is used)
|
||||
- `token_ids`: (required, [int]) prompt token id list
|
||||
- `tenant_id`: (optional, int) In a multi-tenant architecture, tenant_id is the key identifier for distinguishing and isolating data from different tenants (such as different companies or user groups). All data operations are logically isolated based on this ID. If you provide it, the indexer will only return the token hit information for this tenant. Default is None(meaning there is only one tenant)
|
||||
- `cache_salt`: (optional, int) An optional salt value to ensure cached data blocks are kept separate for different customers. This prevents one customer's kv-index data from being served to another. Default is None, meaning no salt is used.
|
||||
- **example**:
|
||||
```json
|
||||
{
|
||||
"model": "deepseek-v3",
|
||||
"lora_name": "sql_adapter",
|
||||
"token_ids": [101, 15, 100, 55, 89],
|
||||
}
|
||||
```
|
||||
- **Output**:
|
||||
```json
|
||||
{
|
||||
"data": {
|
||||
"tenant_id": {
|
||||
"api_server_unique_name": {
|
||||
"longest_matched": 100, // the number of longest prefix matched token among multiple DPs(if there are)
|
||||
"GPU": 20,
|
||||
"DP": {
|
||||
0: 10,
|
||||
1: 20
|
||||
},
|
||||
"CPU": 60,
|
||||
"DISK": 10
|
||||
},
|
||||
... // other engine instance
|
||||
},
|
||||
... // other tenant
|
||||
}
|
||||
}
|
||||
```
|
||||
- **Parameter Description**
|
||||
- `tenant_id`: tenant id, only used in multi-tenant scenario.
|
||||
- `api_server_unique_name`: it is a unique name for a LLM API server endpoint in the engine side. For example, two service instances are currently started separately by running the `vllm server` command, and they are registered in the indexer with different names(such as vllm-1,vllm-2)
|
||||
- `longest_matched`: the number of longest prefix matched token among G1/G2/G3. Indexer will sequentially query the hit status of each token-block according to the prefix order. If it hits, count the situation of this token-block at each level; If it missed, terminate the query (ensuring prefix continuity).
|
||||
- `GPU`, `CPU`, `DISK`: token ids hit count for each tiered storage medium. The Indexer will track the storage status of KV-cache across various media. This requires different KV publishers to inform the Indexer of the actual storage medium type via kv-events. The following examples list several common names, such as using GPU or NPU to represent the Device Pool, using CPU to represent the Host Pool, and using DISK to represent the Disk Pool.
|
||||
- `DP`: token ids hit count for each DP rank.
|
||||
- **example**:
|
||||
|
||||
Assume the input token_ids are [101, 15, 100, 55, 89, 63], the block_size is 2, and the dp2 strategy is enabled. There are three block hashes to match [H1, H2, H3], where H1 hits in GPU (dp0, dp1), CPU, and DISK; H2 hits in GPU (dp0) and CPU; and H3 hits in DISK.
|
||||
```json
|
||||
{
|
||||
"vllm-1": {
|
||||
"longest_matched": 6,
|
||||
"GPU": 4,
|
||||
"DP": {
|
||||
0: 4,
|
||||
1: 2
|
||||
},
|
||||
"CPU": 4,
|
||||
"DISK": 4
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### `POST /query_by_hash`
|
||||
query token hit count by chunked_token hash key. Each model service uses its own independent page_size, `longest_matched = page_size * matched hash_key`
|
||||
- **Input**:
|
||||
- **Body** (JSON):
|
||||
```json
|
||||
{
|
||||
"model": "deepseek",
|
||||
"lora_name": "xx-adapter",
|
||||
"lora_id": 12, // defined for backward compatibility and should not be used together with `lora_name`
|
||||
"block_hash": ["hash_key_by_chunked_tokens"],
|
||||
"tenant_id": None,
|
||||
"cache_salt": None,
|
||||
}
|
||||
```
|
||||
- **Parameter Description**:
|
||||
- `model`: (required, string) model name
|
||||
- `lora_name`: (optional, string) The name of the LoRA adapter, default is `None`(indicating no LoRA adapter is used)
|
||||
- `lora_id`: (optional, int) The ID of the LoRA adapter. This parameter is defined for backward compatibility and should not be used together with `lora_name`(Only one of them can be specified). Default is `-1`(indicating no LoRA adapter is used)
|
||||
- `block_hash`: (required, [int]) chunk_token hash list
|
||||
- `tenant_id`: (optional, int) In a multi-tenant architecture, tenant_id is the key identifier for distinguishing and isolating data from different tenants (such as different companies or user groups). All data operations are logically isolated based on this ID. If you provide it, the indexer will only return the token hit information for this tenant. Default is None(meaning there is only one tenant)
|
||||
- `cache_salt`: (optional, int) An optional salt value to ensure cached data blocks are kept separate for different customers. This prevents one customer's kv-index data from being served to another. Default is None, meaning no salt is used.
|
||||
- **Output**:
|
||||
```json
|
||||
{
|
||||
"data": {
|
||||
"tenant_id": {
|
||||
"api_server_unique_name": {
|
||||
"longest_matched": 100, // the number of longest prefix matched token among multiple DPs(if there are)
|
||||
"GPU": 20,
|
||||
"DP": {
|
||||
0: 10,
|
||||
1: 20
|
||||
},
|
||||
"CPU": 60,
|
||||
"DISK": 10
|
||||
},
|
||||
... // other engine instance
|
||||
},
|
||||
... // other tenant
|
||||
}
|
||||
}
|
||||
```
|
||||
- **Parameter Description**: The output result is same as `/query` api.
|
||||
|
||||
|
||||
## Indexer KVEvents Structure
|
||||
Typically, the device pool is used for loading model weights, with the remaining space registered for KV blocks by the model inference service runtime, the host pool and disk pool are managed uniformly by the Mooncake Store. There is a difference in the management unit for KV data between the two: the device pool uses blocks as the smallest unit for KV data, while the host pool and disk pool use Mooncake Store Objects as the smallest unit for KV storage. In practice, users may split a complete KV block into multiple Mooncake Store Objects for maintenance according to parallel strategies such as tensor parallelism (tp) and context parallelism (cp).
|
||||
|
||||
### G1 KVEvents
|
||||
[vLLM](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/kv_events.py) and [SGLang](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/disaggregation/kv_events.py) publishes events using the `EventBatch` structure, with each batch containing three types of events:
|
||||
|
||||
- `BlockStored`:Adds a single KV block.
|
||||
- `BlockRemoved`:Removes a single KV block.
|
||||
- `AllBlocksCleared`:Clears all KV blocks.
|
||||
|
||||
```py
|
||||
EventBatch:
|
||||
{
|
||||
ts: float, # timestamp
|
||||
events:list[BlockStored | BlockRemoved | AllBlocksCleared],
|
||||
data_parallel_rank: int | None = None, # vLLM use this to indicate dp rank
|
||||
attn_dp_rank: int | None = None, # SGLang use this to indicate dp rank
|
||||
}
|
||||
|
||||
|
||||
BlockStored:
|
||||
{
|
||||
block_hashes: list[int]
|
||||
parent_block_hash: int | None
|
||||
token_ids: list[int]
|
||||
block_size: int
|
||||
|
||||
lora_id: int | None
|
||||
"""Deprecated: use `lora_name` for KV block key hash.
|
||||
Retained for backward compatibility.
|
||||
"""
|
||||
|
||||
medium: str | None
|
||||
"""KV cache is categorized by tier. Currently, the following types are supported:
|
||||
set "GPU", "NPU" for device pool(G1),
|
||||
set "CPU" for host pool(G2),
|
||||
set "DISK" for disk pool(G3).
|
||||
In the future, more medium types can be supported for each tier. For example, "TPU" and "AMD" could be added for device pool.
|
||||
"""
|
||||
lora_name: str | None
|
||||
}
|
||||
|
||||
BlockRemoved:
|
||||
{
|
||||
block_hashes: list[int]
|
||||
lora_name: str | None
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
### G2/G3 KVEvents
|
||||
Mooncake Store is a distributed key-value (KV) store. To ensure system consistency, a timestamp is assigned to each KVEvent for maintenance.
|
||||
|
||||
Mooncake publishes events using the `EventBatch` structure, with each batch containing three types of events:
|
||||
- `BlockStoreEvent`:Adds a single Mooncake Store Object.
|
||||
- `BlockUpdateEvent`:Updates a single Mooncake Store Object.
|
||||
- `RemoveAllEvent`:Removes all Mooncake Store Objects.
|
||||
|
||||
```cpp
|
||||
EventBatch {
|
||||
std::vector<std::variant<BlockStoreEvent, BlockUpdateEvent, RemoveAllEvent>> events,
|
||||
}
|
||||
|
||||
BlockStoreEvent {
|
||||
"BlockStoreEvent", // Event type identifier, string type
|
||||
float ts, // timestamp
|
||||
std::string mooncake_key,
|
||||
std::vector<std::string> addr_list, // Storage location of each replica
|
||||
|
||||
uint64_t block_hash,
|
||||
uint64_t parent_block_hash,
|
||||
std::vector<uint32_t> token_id,
|
||||
uint32_t block_size,
|
||||
|
||||
std::string model_name,
|
||||
std::string lora_name,
|
||||
uint32_t lora_id, // Retained for backward compatibility
|
||||
}
|
||||
|
||||
BlockUpdateEvent {
|
||||
"BlockUpdateEvent",
|
||||
float ts, // timestamp
|
||||
mooncake_key,
|
||||
std::vector<std::string> addr_list,
|
||||
}
|
||||
|
||||
RemoveAllEvent {
|
||||
"RemoveAllEvent"
|
||||
}
|
||||
```
|
||||
|
|
@ -36,13 +36,14 @@ It is possible to configure a `Client` instance to act in only one of its two ro
|
|||
* If `global_segment_size` is set to zero, the instance functions as a **pure client**, issuing requests but not contributing memory to the system.
|
||||
* If `local_buffer_size` is set to zero, it acts as a **pure server**, providing memory for storage. In this case, request operations such as `Get` or `Put` are not permitted from this instance.
|
||||
|
||||
The `Client` can be used in two modes:
|
||||
1. **Embedded mode**: Runs in the same process as the LLM inference program (e.g., a vLLM instance), by being imported as a shared library.
|
||||
2. **Standalone mode**: Runs as an independent process. In this mode, the `Client` is separated into two parts: a **dummy** `Client` and a **real** `Client`: The **real** `Client` is a full-featured implementation that runs as a standalone process and directly communicates with other Mooncake Store components. It handles all RPC communications, memory management, and data transfer operations. The **real** `Client` is typically deployed on nodes that contribute memory to the distributed cache pool; The **dummy** `Client` is a lightweight wrapper that forwards all operations to a local **real** `Client` via RPC calls, which is designed for scenarios where the client needs to be embedded in the same process as the application (such as vLLM), but the actual Mooncake Store operations should be handled by a standalone process. The **dummy** `Client` and the **real** `Client` communicate via RPC calls and shared memory to make sure that Zero-copy transfers are still possible.
|
||||
The `Client` can be used in three ways:
|
||||
1. **Embedded mode**: Runs in the same process as the LLM inference program (e.g., a vLLM instance), by being imported as a shared library. Embedded clients issue requests directly, and when configured with `global_segment_size > 0` they also contribute memory resources to the cluster.
|
||||
2. **Embedded mode with dummy-real clients**: Each LLM inference **rank** holds an embedded **dummy** client (which holds no resources). Each LLM inference **instance** has one resource-owning **real** client (for example, with TP=8 there can be 8 dummy clients and 1 real client). All dummy clients of the same inference instance forward requests to that one real client. The real client owns the global segment (optionally) and is responsible for RPC handling, memory management, and data transfer. Dummy and real clients communicate via RPC, and use shared memory/zero-copy mechanisms for data transfer, so that the data path remains efficient.
|
||||
3. **Standalone store service**: A standalone store service (e.g., `python -m mooncake.mooncake_store_service`) wraps a client and provides the global memory/SSD resource pool. With this service, embedded clients can be configured with `global_segment_size = 0` so they contribute network/NIC resources only, while the standalone store service owns memory and storage management. This service can be deployed on the same server as the inference engine or on separate servers.
|
||||
|
||||
Mooncake store supports two deployment methods to accommodate different availability requirements:
|
||||
1. **Default mode**: In this mode, the master service consists of a single master node, which simplifies deployment but introduces a single point of failure. If the master crashes or becomes unreachable, the system cannot continue to serve requests until it is restored.
|
||||
2. **High availability mode (unstable)**: This mode enhances fault tolerance by running the master service as a cluster of multiple master nodes coordinated through an etcd cluster. The master nodes use etcd to elect a leader, which is responsible for handling client requests.
|
||||
2. **High availability mode**: This mode enhances fault tolerance by running the master service as a cluster of multiple master nodes coordinated through an etcd cluster. The master nodes use etcd to elect a leader, which is responsible for handling client requests.
|
||||
If the current leader fails or becomes partitioned from the network, the remaining master nodes automatically perform a new leader election, ensuring continuous availability.
|
||||
|
||||
In both modes, the leader monitors the health of all client nodes through periodic heartbeats. If a client crashes or becomes unreachable, the leader quickly detects the failure and takes appropriate action. When a client node recovers or reconnects, it can automatically rejoin the cluster without manual intervention.
|
||||
|
|
@ -69,7 +70,7 @@ Initializes the Mooncake Store client. The parameters are as follows:
|
|||
### Get
|
||||
|
||||
```C++
|
||||
tl::expected<void, ErrorCode> Get(const std::string& object_key,
|
||||
tl::expected<void, ErrorCode> Get(const std::string& object_key,
|
||||
std::vector<Slice>& slices);
|
||||
```
|
||||
|
||||
|
|
@ -100,10 +101,30 @@ The data structure details of `ReplicateConfig` are as follows:
|
|||
struct ReplicateConfig {
|
||||
size_t replica_num{1}; // Total number of replicas for the object
|
||||
bool with_soft_pin{false}; // Whether to enable soft pin mechanism for this object
|
||||
bool with_hard_pin{false}; // Whether to enable hard pin (never evicted)
|
||||
std::string preferred_segment{}; // Preferred segment for allocation
|
||||
};
|
||||
```
|
||||
|
||||
### Upsert
|
||||
|
||||
```C++
|
||||
tl::expected<void, ErrorCode> Upsert(const ObjectKey& key,
|
||||
std::vector<Slice>& slices,
|
||||
const ReplicateConfig& config);
|
||||
|
||||
std::vector<tl::expected<void, ErrorCode>> BatchUpsert(
|
||||
const std::vector<ObjectKey>& keys,
|
||||
std::vector<std::vector<Slice>>& batched_slices,
|
||||
const ReplicateConfig& config);
|
||||
```
|
||||
|
||||
`Upsert` inserts `key` if it does not exist and updates the existing object if
|
||||
it does. It uses the same replication configuration model as `Put`, while
|
||||
allowing the store to reuse existing placement for in-place updates when the
|
||||
current layout permits it. `BatchUpsert` performs the same operation for
|
||||
multiple keys using a shared replication configuration.
|
||||
|
||||
### Remove
|
||||
|
||||
```C++
|
||||
|
|
@ -112,6 +133,69 @@ tl::expected<void, ErrorCode> Remove(const ObjectKey& key);
|
|||
|
||||
Used to delete the object corresponding to the specified key. This interface marks all data replicas associated with the key in the storage engine as deleted, without needing to communicate with the corresponding storage node (Client).
|
||||
|
||||
### CreateCopyTask
|
||||
|
||||
```C++
|
||||
tl::expected<UUID, ErrorCode> CreateCopyTask(
|
||||
const std::string& key,
|
||||
const std::vector<std::string>& targets);
|
||||
```
|
||||
|
||||

|
||||
|
||||
`CreateCopyTask` creates an asynchronous copy task that will be executed by the client's task execution system. This is useful when you want to submit multiple copy operations without waiting for each one to complete. The task is submitted to the master service, assigned a unique task ID, and executed asynchronously by an available client. The task status can be queried using `QueryTask`.
|
||||
|
||||
**Task Execution and Result Reporting:**
|
||||
1. **Task Assignment**: The master service assigns the task to an available client during the client's periodic ping operation
|
||||
2. **Task Execution**: The assigned client executes the copy operation asynchronously in a background thread pool
|
||||
3. **Result Reporting**: Upon completion (success or failure), the client automatically reports the result to the master service via `MarkTaskToComplete`:
|
||||
- On success: `status = SUCCESS`, `message = "Task completed successfully"`
|
||||
- On failure: `status = FAILED`, `message = <error description>`
|
||||
4. **Status Query**: You can query the task status at any time using `QueryTask` to monitor progress
|
||||
|
||||
### CreateMoveTask
|
||||
|
||||
```C++
|
||||
tl::expected<UUID, ErrorCode> CreateMoveTask(
|
||||
const std::string& key,
|
||||
const std::string& source,
|
||||
const std::string& target);
|
||||
```
|
||||
|
||||

|
||||
|
||||
`CreateMoveTask` creates an asynchronous move task that will be executed by the client's task execution system. This is useful when you want to submit multiple move operations without waiting for each one to complete. The task is submitted to the master service, assigned a unique task ID, and executed asynchronously by an available client. The task status can be queried using `QueryTask`.
|
||||
|
||||
**Task Execution and Result Reporting:**
|
||||
1. **Task Assignment**: The master service assigns the task to an available client during the client's periodic ping operation
|
||||
2. **Task Execution**: The assigned client executes the move operation asynchronously in a background thread pool
|
||||
3. **Result Reporting**: Upon completion (success or failure), the client automatically reports the result to the master service via `MarkTaskToComplete`:
|
||||
- On success: `status = SUCCESS`, `message = "Task completed successfully"`
|
||||
- On failure: `status = FAILED`, `message = <error description>`
|
||||
4. **Status Query**: You can query the task status at any time using `QueryTask` to monitor progress
|
||||
|
||||
### QueryTask
|
||||
|
||||
```C++
|
||||
tl::expected<QueryTaskResponse, ErrorCode> QueryTask(const UUID& task_id);
|
||||
```
|
||||
|
||||
`QueryTask` queries the status of an asynchronous task (copy or move). This allows you to monitor the progress of task-based operations. The response includes task status, type, creation time, last update time, assigned client, and status message.
|
||||
|
||||
The data structure details of `QueryTaskResponse` are as follows:
|
||||
|
||||
```C++
|
||||
struct QueryTaskResponse {
|
||||
UUID id; // Task UUID
|
||||
TaskType type; // Task type (REPLICA_COPY or REPLICA_MOVE)
|
||||
TaskStatus status; // Task status (PENDING, PROCESSING, SUCCESS, or FAILED)
|
||||
int64_t created_at_ms_epoch; // Task creation timestamp in milliseconds
|
||||
int64_t last_updated_at_ms_epoch; // Last update timestamp in milliseconds
|
||||
UUID assigned_client; // UUID of the client assigned to execute the task
|
||||
std::string message; // Status message or error description
|
||||
};
|
||||
```
|
||||
|
||||
### BatchQueryIp
|
||||
|
||||
```C++
|
||||
|
|
@ -151,10 +235,26 @@ Used to delete all objects from the store whose keys match the specified regular
|
|||
|
||||
### Master Service
|
||||
|
||||
The cluster's available resources are viewed as a large resource pool, managed centrally by a Master process for space allocation and guiding data replication
|
||||
The cluster's available resources are viewed as a large resource pool, managed centrally by a Master process for space allocation and guiding data replication
|
||||
|
||||
**Note: The Master Service does not take over any data flow, only providing corresponding metadata information.**
|
||||
|
||||
#### Snapshot & Restore
|
||||
|
||||
To reduce cache warm-up time after a master restart, the Master Service supports periodic snapshots of its in-memory metadata and recovery from these snapshots.
|
||||
|
||||
- Snapshot generation
|
||||
- A background snapshot thread periodically takes a consistent copy of the in-memory KV metadata, segment information, and allocator state using fork-based copy-on-write, without blocking normal RPC handling.
|
||||
- The child process serializes these structures into a compact binary format and writes them to the configured snapshot backend via the `SerializerBackend` abstraction.
|
||||
- Restore
|
||||
- On startup, when snapshot restore is enabled, the master reads the latest snapshot from the backend and reconstructs the Master Service's metadata state in memory.
|
||||
- Notes
|
||||
- Because snapshots are taken periodically rather than continuously, metadata changes after the last successful snapshot may be lost if the master fails before the next snapshot completes.
|
||||
|
||||
> **Warning: Managed Storage**
|
||||
>
|
||||
> The snapshot storage location is **exclusively managed** by the Mooncake snapshot system. Old snapshots are automatically deleted during cleanup. **DO NOT store other files in this location.** Use a dedicated, isolated storage for snapshots.
|
||||
|
||||
#### Master Service APIs
|
||||
|
||||
The protobuf definition between Master and Client is as follows:
|
||||
|
|
@ -225,7 +325,7 @@ service MasterService {
|
|||
|
||||
```protobuf
|
||||
message GetReplicaListRequest {
|
||||
required string key = 1;
|
||||
required string key = 1;
|
||||
};
|
||||
|
||||
message GetReplicaListResponse {
|
||||
|
|
@ -310,7 +410,7 @@ message PutStartRequest {
|
|||
};
|
||||
|
||||
message PutStartResponse {
|
||||
required int32 status_code = 1;
|
||||
required int32 status_code = 1;
|
||||
repeated ReplicaInfo replica_list = 2; // Replica information allocated by the Master Service
|
||||
};
|
||||
```
|
||||
|
|
@ -323,7 +423,7 @@ message PutStartResponse {
|
|||
|
||||
```protobuf
|
||||
message PutEndRequest {
|
||||
required string key = 1;
|
||||
required string key = 1;
|
||||
};
|
||||
|
||||
message PutEndResponse {
|
||||
|
|
@ -339,7 +439,7 @@ message PutEndResponse {
|
|||
|
||||
```protobuf
|
||||
message RemoveRequest {
|
||||
required string key = 1;
|
||||
required string key = 1;
|
||||
};
|
||||
|
||||
message RemoveResponse {
|
||||
|
|
@ -436,6 +536,40 @@ The Master Service handles object-related interfaces as follows:
|
|||
|
||||
Before writing an object, the Client calls PutStart to request storage space allocation from the Master Service. After completing data writing, the Client calls PutEnd to notify the Master Service to mark the object write as completed.
|
||||
|
||||
- Upsert
|
||||
|
||||
```C++
|
||||
tl::expected<std::vector<Replica::Descriptor>, ErrorCode> UpsertStart(
|
||||
const std::string& key,
|
||||
const std::vector<size_t>& slice_lengths,
|
||||
const ReplicateConfig& config);
|
||||
|
||||
std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
|
||||
BatchUpsertStart(const std::vector<std::string>& keys,
|
||||
const std::vector<std::vector<uint64_t>>& slice_lengths,
|
||||
const ReplicateConfig& config);
|
||||
|
||||
tl::expected<void, ErrorCode> UpsertEnd(
|
||||
const std::string& key, ReplicaType replica_type);
|
||||
|
||||
std::vector<tl::expected<void, ErrorCode>> BatchUpsertEnd(
|
||||
const std::vector<std::string>& keys);
|
||||
|
||||
tl::expected<void, ErrorCode> UpsertRevoke(
|
||||
const std::string& key, ReplicaType replica_type);
|
||||
|
||||
std::vector<tl::expected<void, ErrorCode>> BatchUpsertRevoke(
|
||||
const std::vector<std::string>& keys);
|
||||
```
|
||||
|
||||
`UpsertStart` / `UpsertEnd` / `UpsertRevoke` mirror the existing put lifecycle
|
||||
but operate on insert-or-update semantics. If the key does not exist, the flow
|
||||
behaves like `PutStart`. If the key already exists, the Master may reuse the
|
||||
current allocation for an in-place update or allocate new space when the object
|
||||
layout changes. The batch variants provide the same control flow for multiple
|
||||
keys and are the lower-level primitives used by the high-level `BatchUpsert`
|
||||
path.
|
||||
|
||||
- GetReplicaList
|
||||
|
||||
```C++
|
||||
|
|
@ -518,17 +652,68 @@ virtual tl::expected<std::vector<Replica>, ErrorCode> Allocate(
|
|||
- On success: vector of allocated replicas (may be fewer than requested due to resource constraints, but at least 1)
|
||||
- On failure: ErrorCode::NO_AVAILABLE_HANDLE if no replicas can be allocated, ErrorCode::INVALID_PARAMS for invalid configuration
|
||||
|
||||
#### Implementation Strategies
|
||||
#### Allocation Strategies
|
||||
|
||||
`RandomAllocationStrategy` is a subclass implementing `AllocationStrategy` that provides intelligent allocation with the following features:
|
||||
Mooncake Store provides multiple built-in allocation strategies to control how storage space is distributed across segments. Users can select a strategy via the `--allocation_strategy` flag when starting the master service:
|
||||
|
||||
1. **Preferred Segment Support**: If a preferred segment is specified in the `ReplicateConfig`, the strategy first attempts to allocate from that segment before falling back to random allocation.
|
||||
```bash
|
||||
./build/mooncake-store/src/mooncake_master --allocation_strategy=free_ratio_first
|
||||
```
|
||||
|
||||
2. **Random Allocation with Retry Logic**: When multiple allocators are available, it uses a randomized approach with up to 10 retry attempts to find a suitable allocator.
|
||||
Valid values are: `random` (default), `free_ratio_first`, `cxl` (case-sensitive).
|
||||
|
||||
3. **Deterministic Randomization**: Uses a Mersenne Twister random number generator with proper seeding for consistent behavior.
|
||||
##### How to Choose
|
||||
|
||||
The strategy automatically handles cases where the preferred segment is unavailable, full, or doesn't exist by gracefully falling back to random allocation among all available segments.
|
||||
| Strategy | Best For | Trade-off |
|
||||
|---|---|---|
|
||||
| `random` | Maximum throughput, stable clusters | Limited load balancing; slow convergence when new segments join |
|
||||
| `free_ratio_first` | Balanced utilization, dynamic scaling | Slightly lower throughput due to sampling and sorting overhead |
|
||||
| `cxl` | CXL memory hardware | CXL-specific; single-replica only |
|
||||
|
||||
**Use `random`** (default) when your cluster is relatively stable (segments rarely join or leave) and you want the highest possible allocation throughput.
|
||||
|
||||
**Use `free_ratio_first`** when you need better load balancing across segments, especially in scenarios where:
|
||||
- Segments have different capacities and you want even utilization ratios.
|
||||
- New segments are dynamically added at runtime and you need them to absorb load quickly. With `random`, convergence to a well-balanced state can be slow on large or dynamic clusters; `free_ratio_first` accelerates this by preferentially filling emptier segments, substantially increasing the likelihood that newly joined segments are selected for allocations (see details below).
|
||||
|
||||
**Use `cxl`** only when your hardware includes CXL (Compute Express Link) memory devices and you want to allocate data exclusively on CXL segments.
|
||||
|
||||
##### Strategy Details
|
||||
|
||||
**`random` — RandomAllocationStrategy**
|
||||
|
||||
Pure random allocation with preferred segment support. The allocation process for N replicas is:
|
||||
|
||||
1. **Preferred segment phase**: If preferred segments are specified in the `ReplicateConfig`, they are tried first in order. Each successful allocation consumes one replica slot. If all replicas are satisfied, the process finishes early.
|
||||
2. **Random phase**: For any remaining replicas, a random starting index is chosen among all available segments. The strategy then iterates consecutively from that index, attempting to allocate from each segment. Segments already used for a previous replica of the same slice, as well as explicitly excluded segments, are skipped to guarantee that each replica resides on a different segment.
|
||||
3. **Retry limit**: The iteration is capped at `min(100, total_segments)` to avoid excessive scanning when most segments are full.
|
||||
4. **Best-effort result**: If fewer than N replicas are allocated but at least one succeeds, the partial result is returned. The call fails only if zero replicas can be allocated.
|
||||
|
||||
All random state uses a thread-local Mersenne Twister (`std::mt19937`), so no locks or shared mutable data are involved.
|
||||
|
||||
**`free_ratio_first` — FreeRatioFirstAllocationStrategy (Best-of-N)**
|
||||
|
||||
An improved strategy built on top of `RandomAllocationStrategy`. Instead of picking segments purely at random, it samples a small pool of candidates and selects the ones with the most free space. The allocation process for N replicas is:
|
||||
|
||||
1. **Preferred segment phase**: Same as `random` — preferred segments are tried first in order.
|
||||
2. **Sampling phase**: Randomly picks a starting index and takes `min(6*remaining_replicas, total_segments)` consecutive segments as candidates. For each candidate, queries its free space ratio: `free_bytes / total_capacity`.
|
||||
3. **Sorting phase**: Sorts the candidates in descending order by free space ratio (most free first).
|
||||
4. **Allocation phase**: Iterates through the sorted candidates from top to bottom, attempting to allocate from each. Excluded and already-used segments are skipped.
|
||||
5. **Fallback phase**: If insufficient replicas are allocated from the sorted candidates, falls back to the base `RandomAllocationStrategy` random iteration logic for the remaining replicas.
|
||||
|
||||
The overhead is minimal: sampling is `O(K)` and sorting is `O(K log K)`, where K is the candidate count (at most `6*N`) — both small since `replica_num` is typically 1–3. The strategy is thread-safe, using `thread_local` random state with no shared mutable data.
|
||||
|
||||
The key insight behind Best-of-N is that if a new/empty segment is sampled, it will almost certainly be ranked first due to having the highest free ratio, which naturally accelerates convergence when new segments join the cluster.
|
||||
|
||||
**`cxl` — CxlAllocationStrategy**
|
||||
|
||||
Specialized for CXL (Compute Express Link) memory hardware. Unlike the other strategies, this one does not perform random or load-balanced selection — it always allocates from a specific CXL segment:
|
||||
|
||||
1. Requires `preferred_segments` to be non-empty; the first element is used as the target CXL segment name.
|
||||
2. Allocates a single replica from the specified CXL segment's allocator.
|
||||
3. Marks the allocated buffer as CXL type via `change_to_cxl()`, so downstream components can distinguish CXL-backed data from regular DRAM.
|
||||
|
||||
Limitations: This strategy only supports single-replica allocation (does not distribute across multiple segments) and does not support the `AllocateFrom()` interface.
|
||||
|
||||
### Eviction Policy
|
||||
|
||||
|
|
@ -558,6 +743,18 @@ There are two startup parameters in `master_service` related to the soft pin mec
|
|||
|
||||
Notably, soft pinned objects can still be removed using APIs such as `Remove` or `RemoveAll`.
|
||||
|
||||
### Hard Pin
|
||||
|
||||
For objects that must never be evicted under any circumstances (e.g., model weights, critical metadata), Mooncake Store provides a hard pin mechanism. Unlike soft pin, hard-pinned objects are permanently protected from eviction — they will never be selected as eviction candidates regardless of memory pressure.
|
||||
|
||||
Hard pin is set at object creation time through the `with_hard_pin` field in `ReplicateConfig` and cannot be changed afterward. Hard-pinned objects can only be removed explicitly via `Remove` (with force) or `RemoveAll`.
|
||||
|
||||
Key differences from soft pin:
|
||||
|
||||
- Hard pin never expires. Soft pin status is removed after a configurable TTL if the object is not accessed.
|
||||
- Hard-pinned objects are completely skipped during eviction. Soft-pinned objects may still be evicted when no other candidates are available.
|
||||
- Hard pin is immutable once set. Soft pin status is automatically refreshed on access.
|
||||
|
||||
### Zombie Object Cleanup
|
||||
|
||||
If a Client crashes or experiences a network failure after sending a `PutStart` request but before it can send the corresponding `PutEnd` or `PutRevoke` request to the Master, the object initiated by `PutStart` enters a "zombie" state—rendering it neither usable nor deletable. The existence of such "zombie objects" not only consumes storage space but also prevents subsequent `Put` operations on the same keys. To mitigate these issues, the Master records the start time of each `PutStart` request and employs two timeout thresholds—`put_start_discard_timeout` and `put_start_release_timeout`—to clean up zombie objects.
|
||||
|
|
@ -582,6 +779,7 @@ The preferred segment allocation feature is implemented through the `AllocationS
|
|||
struct ReplicateConfig {
|
||||
size_t replica_num{1}; // Total number of replicas for the object
|
||||
bool with_soft_pin{false}; // Whether to enable soft pin mechanism for this object
|
||||
bool with_hard_pin{false}; // Whether to enable hard pin (never evicted)
|
||||
std::string preferred_segment{}; // Preferred segment for allocation
|
||||
};
|
||||
```
|
||||
|
|
@ -608,7 +806,7 @@ When the user specifies `--root_fs_dir=/path/to/dir` when starting the master, a
|
|||
Note: When enabling this feature, the user must ensure that the DFS-mounted directory (`root_fs_dir=/path/to/dir`) is valid and consistent across all client hosts. If some clients have invalid or incorrect mount paths, it may cause abnormal behavior in Mooncake Store.
|
||||
|
||||
#### Persistent Storage Space Configuration
|
||||
Mooncake provides configurable DFS available space. Users can specify `--global_file_segment_size=1048576` when starting the master, indicating a maximum usable space of 1MB on DFS.
|
||||
Mooncake provides configurable DFS available space. Users can specify `--global_file_segment_size=1048576` when starting the master, indicating a maximum usable space of 1MB on DFS.
|
||||
The current default setting is the maximum value of int64 (as we generally do not restrict DFS storage usage), which is displayed as `infinite` in `mooncake_maseter`'s console logs.
|
||||
**Notice** The DFS cache space configuration must be used together with the `--root_fs_dir` parameter. Otherwise, you will observe that the `SSD Storage` usage consistently shows: `0 B / 0 B`
|
||||
**Notice** The capability for file eviction on DFS has not been provided yet
|
||||
|
|
@ -637,6 +835,8 @@ The HTTP metadata server can be configured using the following parameters:
|
|||
- MC_STORE_MEMCPY: Enables or disables local memcpy optimization, set to 1/true to enable, 0/false to disable.
|
||||
- MC_STORE_CLIENT_METRIC: Enables client metric reporting, enabled by default; set to 0/false to disable.
|
||||
- MC_STORE_CLIENT_METRIC_INTERVAL: Reporting interval in seconds, default 0 (collects but does not report).
|
||||
- MC_STORE_USE_HUGEPAGE: Enables huge page support, disabled by default.
|
||||
- MC_STORE_HUGEPAGE_SIZE: Specifies the page size of the huge page to use, default 2M.
|
||||
#### Usage Example
|
||||
To start the master service with the HTTP metadata server enabled:
|
||||
```bash
|
||||
|
|
@ -712,7 +912,7 @@ Mooncake Store provides various sample programs, including interface forms based
|
|||
`metadata_server`: the address of the Transfer Engine metadata service
|
||||
`master_server_address`: the address of the Master Service
|
||||
**Note**: The format of `master_server_address` depends on the deployment mode. In default mode, use the format `IP:Port`, specifying the address of a single master node. In HA mode, use the format `etcd://IP:Port;IP:Port;...;IP:Port`, specifying the addresses of the etcd cluster endpoints.
|
||||
For example:
|
||||
For example:
|
||||
```python
|
||||
import os
|
||||
import time
|
||||
|
|
@ -776,6 +976,8 @@ The **real** `Client` can be configured using the following parameters:
|
|||
|
||||
- **`host`**: (string, default: "0.0.0.0"): The hostname of the client.
|
||||
|
||||
- **`port`**: (int, default: 50052): The port number the client service listens on.
|
||||
|
||||
- **`global_segment_size`**: (string, default: "4GB"): The size of the global segment to be allocated by the client.
|
||||
|
||||
- **`master_server_address`**: (string, default: "localhost:50051"): The address of the Master Service.
|
||||
|
|
@ -841,3 +1043,13 @@ When to bump the version:
|
|||
* **Major version (X.0.0)**: For breaking API changes, major architectural changes, or significant new features that affect backward compatibility
|
||||
* **Minor version (0.X.0)**: For new features, API additions, or notable improvements that maintain backward compatibility
|
||||
* **Patch version (0.0.X)**: For bug fixes, performance optimizations, or minor improvements that don't affect the API
|
||||
|
||||
|
||||
---
|
||||
|
||||
:::{toctree}
|
||||
:caption: Related Design Docs
|
||||
:maxdepth: 1
|
||||
|
||||
ssd-offload
|
||||
:::
|
||||
|
|
|
|||
|
|
@ -0,0 +1,245 @@
|
|||
# SSD Offload Design
|
||||
|
||||
## Overview
|
||||
|
||||
Mooncake Store supports offloading KV cache objects from distributed memory to local SSD. This extends the effective cache capacity beyond DRAM limits at lower cost, while preserving the performance characteristics of the hot path through zero-copy RDMA-based memory transfers.
|
||||
|
||||
SSD offload is implemented as a background subsystem within the **real client** process. It is transparent to the application: a `Put` that would otherwise be evicted from memory is persisted to disk, and a `Get` that finds no memory replica automatically falls back to reading from SSD.
|
||||
|
||||
---
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────┐
|
||||
│ Application (vLLM, etc.) │
|
||||
└──────────────────────────┬──────────────────────────────┘
|
||||
│ MooncakeDistributedStore API
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────┐
|
||||
│ Real Client │
|
||||
│ │
|
||||
│ ┌──────────────────────────────────────────────────┐ │
|
||||
│ │ FileStorage │ │
|
||||
│ │ ┌────────────┐ ┌──────────────────────────┐ │ │
|
||||
│ │ │ Heartbeat │ │ ClientBuffer (staging) │ │ │
|
||||
│ │ │ Thread │ └──────────────────────────┘ │ │
|
||||
│ │ └─────┬──────┘ │ │
|
||||
│ │ │ offload / load │ │
|
||||
│ │ ▼ │ │
|
||||
│ │ ┌─────────────────────────────────────────┐ │ │
|
||||
│ │ │ StorageBackendInterface │ │ │
|
||||
│ │ │ ┌───────────┐ ┌──────────┐ ┌────────┐ │ │ │
|
||||
│ │ │ │ Bucket │ │FilePerKey│ │Offset │ │ │ │
|
||||
│ │ │ │ Backend │ │ Backend │ │Alloc. │ │ │ │
|
||||
│ │ │ └───────────┘ └──────────┘ └────────┘ │ │ │
|
||||
│ │ └─────────────────────────────────────────┘ │ │
|
||||
│ └──────────────────────────────────────────────────┘ │
|
||||
│ │
|
||||
│ ┌──────────────────────────────────────────────────┐ │
|
||||
│ │ In-memory distributed KV cache │ │
|
||||
│ │ (Transfer Engine / RDMA) │ │
|
||||
│ └──────────────────────────────────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────┘
|
||||
│
|
||||
Local SSD / NVMe
|
||||
```
|
||||
|
||||
The key components are:
|
||||
|
||||
- **FileStorage**: The top-level coordinator. It owns the storage backend, a staging buffer (`ClientBuffer`), and background threads for heartbeating and buffer garbage collection.
|
||||
- **StorageBackendInterface**: An abstract interface implemented by three backends (see below). Responsible for the actual on-disk layout and I/O.
|
||||
- **Heartbeat thread**: Periodically contacts the master. The master returns a list of objects to offload; the heartbeat thread writes them to SSD and notifies the master of completion.
|
||||
- **ClientBuffer**: A pre-registered, O_DIRECT-aligned staging area used for zero-copy reads from SSD back into application memory.
|
||||
|
||||
---
|
||||
|
||||
## Data Flow
|
||||
|
||||
### Offload (memory → SSD)
|
||||
|
||||
The offload path is driven entirely by the heartbeat thread inside `FileStorage`. No write path from the application is involved.
|
||||
|
||||
```
|
||||
Heartbeat Thread Master Local Memory Segment
|
||||
│ │ │
|
||||
│─OffloadObjectHB ───▶│ │
|
||||
│◀─ {key→size} map ───│ (objects to evict from │
|
||||
│ │ memory to SSD) │
|
||||
│ │ │
|
||||
│─ BatchQuery(keys) ───────────────────────────────▶│
|
||||
│◀─ {key→Slice} ────────────────────────────────────│
|
||||
│ │ │
|
||||
│ [PrepareEviction: remove old buckets, notify master via BatchEvictDiskReplica]
|
||||
│─ BatchEvictDiskReplica(evicted_keys) ────────────▶│ (master removes stale replicas)
|
||||
│ [FinalizeEviction: delete evicted files]
|
||||
│ │ │
|
||||
│ BatchOffload(slices) → StorageBackend → SSD │
|
||||
│ │ │
|
||||
│─ NotifyOffloadSuccess(keys, metadata) ───────────▶│
|
||||
│ │ (master adds LOCAL_DISK │
|
||||
│ │ replica to object entry) │
|
||||
```
|
||||
|
||||
Step by step:
|
||||
|
||||
1. **Heartbeat**: The heartbeat thread wakes up every `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` seconds and calls `client_->OffloadObjectHeartbeat(enable_offloading_, offloading_objects)`. The master replies with a map of `{key → size}` for objects it has selected to evict from memory.
|
||||
2. **Read slices from memory**: `FileStorage::OffloadObjects` groups the keys into buckets (for `BucketStorageBackend`) and calls `BatchQuerySegmentSlices` to obtain `{key → Slice}` from the local memory segment via `client_->BatchQuery`.
|
||||
3. **Eviction** (if capacity limit is set): Before writing, `PrepareEviction` removes old buckets from metadata under the exclusive lock and collects their keys. The `eviction_handler` callback calls `client_->BatchEvictDiskReplica` to notify the master in a single RPC. `FinalizeEviction` then deletes the corresponding files.
|
||||
4. **Write to SSD**: `StorageBackend::BatchOffload` serializes and writes the key-value data to disk.
|
||||
5. **Notify master**: On success, the `complete_handler` calls `client_->NotifyOffloadSuccess(keys, metadatas)`. The master adds a `LOCAL_DISK` replica entry (carrying the real client's RPC address as `transport_endpoint`) to the object's replica list.
|
||||
|
||||
### Load (SSD → memory)
|
||||
|
||||
The load path involves three parties: the **requesting client**, the **target client** that holds the SSD data, and the **Transfer Engine** for zero-copy data movement.
|
||||
|
||||
```
|
||||
Requesting Client Target Client Master
|
||||
│ │ │
|
||||
│─ BatchGet(keys) ──────────────────────────────────────────▶│
|
||||
│◀─ QueryResult {replicas: [LOCAL_DISK(rpc_addr)]} ──────────│
|
||||
│ │ │
|
||||
│ (no memory replica available) │ │
|
||||
│─ batch_get_offload_object(keys, sizes) ───────────────────▶│
|
||||
│ │ │
|
||||
│ FileStorage::BatchGet │
|
||||
│ → StorageBackend::BatchLoad │
|
||||
│ → read from SSD into ClientBuffer │
|
||||
│ │ │
|
||||
│◀─ BatchGetOffloadObjectResponse ───│ │
|
||||
│ {batch_id, pointers[], transfer_engine_addr, gc_ttl_ms} │
|
||||
│ │ │
|
||||
│─ Transfer Engine: BatchGetOffloadObject ──────────────────▶│
|
||||
│ (RDMA/TCP: pull data from ClientBuffer into app memory) │
|
||||
│◀─ done ────────────────────────────│ │
|
||||
│ │ │
|
||||
│─ release_offload_buffer(batch_id) ────────────────────────▶│
|
||||
│ │ (free ClientBuffer slot)│
|
||||
```
|
||||
|
||||
Step by step:
|
||||
|
||||
1. **Query master**: The requesting client calls `client_->BatchGet(keys, ...)` to query the master for replica locations. If the object has been offloaded, the master returns a `LOCAL_DISK` replica descriptor containing the target client's RPC address (`transport_endpoint`).
|
||||
2. **RPC to target client**: The requesting client calls `batch_get_offload_object(keys, sizes)` on the target client identified by `transport_endpoint`. The target client calls `FileStorage::BatchGet`, which allocates slots in `ClientBuffer` and reads the requested objects from SSD via `StorageBackend::BatchLoad`.
|
||||
3. **Response with buffer pointers**: The target client returns a `BatchGetOffloadObjectResponse` containing `batch_id`, a list of buffer `pointers` (addresses within `ClientBuffer`), the Transfer Engine address, and `gc_ttl_ms` (the buffer lease TTL).
|
||||
4. **Zero-copy transfer**: The requesting client invokes `client_->BatchGetOffloadObject(transfer_engine_addr, keys, pointers, slices)`, which uses the Transfer Engine (RDMA or TCP) to pull the data directly from the target client's `ClientBuffer` into the application's target memory (DRAM or VRAM). No intermediate copy is made on the requesting client side.
|
||||
5. **Release buffer**: After the transfer completes, the requesting client immediately calls `release_offload_buffer(batch_id)` on the target client to free the `ClientBuffer` slots. If the transfer takes longer than `gc_ttl_ms`, the buffer GC thread reclaims the slot automatically as a fallback.
|
||||
|
||||
---
|
||||
|
||||
## Storage Backends
|
||||
|
||||
### BucketStorageBackend (default)
|
||||
|
||||
Objects are grouped into **buckets** before being written to disk. Each bucket produces two files:
|
||||
|
||||
- **`.bucket`** — binary data file containing serialized key-value records
|
||||
- **`.meta`** — metadata file describing the keys and byte offsets within the data file
|
||||
|
||||
Bucket IDs are monotonically increasing timestamps with a sequence suffix, so `buckets_` (a `std::map<int64_t, BucketMetadata>`) is always ordered by creation time.
|
||||
|
||||
**Grouping strategy** (`GroupOffloadingKeysByBucket`): objects are accumulated into a bucket until either `bucket_size_limit` (default 256 MB) or `bucket_keys_limit` (default 500) is reached. Objects that do not fill a complete bucket are held in `ungrouped_offloading_objects_` and retried on the next heartbeat.
|
||||
|
||||
**In-flight read tracking**: A `BucketReadGuard` RAII object increments `BucketMetadata::inflight_reads_` on construction and decrements it on destruction. This allows safe deletion of bucket files even when concurrent reads are in progress.
|
||||
|
||||
### StorageBackendAdaptor (FilePerKey)
|
||||
|
||||
Each object is stored as an individual file. The file path is derived from the key via a two-level hash-sharded directory structure to avoid large flat directories. This backend is simple and easy to inspect but does not scale well to millions of objects.
|
||||
|
||||
### OffsetAllocatorStorageBackend
|
||||
|
||||
A single pre-allocated file (`kv_cache.data`) is shared by all objects. Space within the file is managed by an `OffsetAllocator`. Metadata is sharded across 1024 independent maps to reduce lock contention under high concurrency. Records follow the layout `[key_len: u32 | value_len: u32 | key | value]`.
|
||||
|
||||
---
|
||||
|
||||
## Eviction (BucketStorageBackend)
|
||||
|
||||
When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend evicts existing buckets to make room before writing a new one. Eviction is disabled by default (`BucketEvictionPolicy::NONE`).
|
||||
|
||||
### Policies
|
||||
|
||||
| Policy | Candidate selection |
|
||||
|--------|---------------------|
|
||||
| `FIFO` | `buckets_.begin()` — always the oldest bucket, since `buckets_` is ordered by bucket ID |
|
||||
| `LRU` | `std::min_element` over `BucketMetadata::last_access_ns_` — the bucket with the smallest last-read timestamp |
|
||||
|
||||
`last_access_ns_` is an atomic `int64_t` updated on every `BatchLoad` with relaxed ordering. Buckets that have never been read have `last_access_ns_ == 0` and are therefore always evicted first under LRU, giving FIFO-among-unread semantics.
|
||||
|
||||
### Two-phase eviction protocol
|
||||
|
||||
Eviction is split into two phases to ensure that the master is notified before files are deleted, and that no in-flight reads are interrupted.
|
||||
|
||||
**Phase 1 — `PrepareEviction(required_size)`** (called under exclusive lock):
|
||||
|
||||
1. Repeatedly call `SelectEvictionCandidate()` until `total_size_ + required_size <= max_total_size`.
|
||||
2. For each selected bucket: remove it from `buckets_` and `object_bucket_map_`, subtract its size from `total_size_`.
|
||||
3. Collect all evicted keys and bucket metadata into a `PendingEviction` struct and return it — no file I/O at this point.
|
||||
|
||||
**Between phases** — notify master:
|
||||
|
||||
The caller invokes the `eviction_handler` callback with the full list of evicted keys. The handler calls `MasterClient::BatchEvictDiskReplica`, which sends a single RPC to the master to remove the disk replicas for all evicted keys atomically.
|
||||
|
||||
**Phase 2 — `FinalizeEviction(pending)`** (called after master notification):
|
||||
|
||||
For each evicted bucket:
|
||||
1. Spin-wait (with a 10-second timeout) until `inflight_reads_ == 0`.
|
||||
2. Evict any stale file-handle cache entries.
|
||||
3. Delete the `.bucket` and `.meta` files.
|
||||
|
||||
This ordering guarantees:
|
||||
- The master never serves a stale disk-replica location for a file that has already been deleted.
|
||||
- Ongoing reads complete successfully before their files are removed.
|
||||
- Freed disk space is available for the incoming write before `WriteBucket` is called.
|
||||
|
||||
---
|
||||
|
||||
## io_uring File I/O
|
||||
|
||||
When `MOONCAKE_OFFLOAD_USE_URING=true`, the storage backends replace POSIX `pread`/`pwrite` calls with an io_uring-based implementation (`UringFile`). The design prioritizes eliminating inter-thread lock contention, which was the dominant latency source in the previous global-ring approach.
|
||||
|
||||
### Thread-local rings (`SharedUringRing`)
|
||||
|
||||
Each thread owns exactly one `io_uring` ring, stored in `thread_local` storage. This means:
|
||||
|
||||
- **No mutex between threads.** Each ring is accessed only by its owning thread, so concurrent I/O from multiple threads is fully parallel with zero synchronization overhead.
|
||||
- **Within-thread batching.** Multiple SQEs can be enqueued before calling `io_uring_submit_and_wait`, exposing NVMe queue depth > 1 within a single thread. `batch_read` exploits this to issue up to `QUEUE_DEPTH` (32) independent reads in one submission.
|
||||
- **File-descriptor registration is omitted.** The per-I/O `fdget()` overhead (~50 ns) is negligible compared to the lock contention (> 1 ms) the old global ring imposed, so `IOSQE_FIXED_FILE` is not used.
|
||||
|
||||
Rings are initialized lazily on first use and destroyed when the thread exits. If ring initialization fails (e.g., kernel too old), the backend falls back gracefully to POSIX I/O.
|
||||
|
||||
### Fixed-buffer registration
|
||||
|
||||
The `ClientBuffer` (the staging buffer used for SSD reads) is registered with io_uring as a **fixed buffer** via `io_uring_register_buffers`. When a read destination falls within the registered region, the backend uses `io_uring_prep_read_fixed` instead of `io_uring_prep_read`, which avoids a per-I/O `mmap`/`munmap` in the kernel and reduces system-call overhead.
|
||||
|
||||
Buffer registration is global but applied **lazily per thread**: `g_buf` stores the base address and length atomically; each thread-local ring calls `ensure_buf_registered()` on its first I/O and registers the buffer independently. This avoids a global barrier at startup.
|
||||
|
||||
To prevent `io_uring`'s `FOLL_LONGTERM` page pinning from failing on systems with Transparent Huge Pages (THP) enabled, `MADV_NOHUGEPAGE` is applied to the buffer region before registration. This forces the kernel to back the range with 4 KB pages, making long-term pinning reliable regardless of system THP policy.
|
||||
|
||||
### O_DIRECT and alignment
|
||||
|
||||
`UringFile` supports an optional `O_DIRECT` mode. When enabled:
|
||||
|
||||
- All file descriptors are opened with `O_DIRECT`.
|
||||
- Buffers, lengths, and offsets must be aligned to 4 KB (`ALIGNMENT_ = 4096`).
|
||||
- For unaligned writes (e.g., metadata serialized into a `std::string`), the backend allocates a temporary aligned bounce buffer via `posix_memalign`, copies the data, performs the aligned write, and frees the bounce buffer.
|
||||
- `read_aligned` and `write_aligned` are the primary I/O paths; they assert alignment constraints and delegate directly to the ring.
|
||||
|
||||
### I/O operations
|
||||
|
||||
| Method | Description |
|
||||
|--------|-------------|
|
||||
| `read` / `write` | Contiguous read or write, chunked into up to `QUEUE_DEPTH` SQEs per submission |
|
||||
| `read_aligned` / `write_aligned` | Same as above but with alignment preconditions for O_DIRECT |
|
||||
| `batch_read` | Submits multiple independent reads (different offsets) in batches of up to `QUEUE_DEPTH`, maximizing NVMe queue utilization |
|
||||
| `vector_read` / `vector_write` | Scatter/gather I/O: one SQE per `iovec`, submitted in batches |
|
||||
| `datasync` | Issues `IORING_FSYNC_DATASYNC` and waits for completion |
|
||||
|
||||
### Integration with storage backends
|
||||
|
||||
- **BucketStorageBackend**: uses `UringFile` for both bucket data files and metadata files when `use_uring_` is set. A file-handle cache (`file_cache_`) avoids repeated `open`/`close` for hot buckets. On eviction, the cache entry is explicitly removed before the file is deleted to prevent stale handles.
|
||||
- **OffsetAllocatorStorageBackend**: opens the single pre-allocated data file with `O_DIRECT` and `UringFile`, and uses `GetFileInstance()` to expose the file handle for external buffer registration.
|
||||
- **StorageBackendAdaptor** (FilePerKey): uses `UringFile` for reads when `use_uring_` is set; writes use POSIX paths.
|
||||
|
||||
## Metadata Recovery on Restart
|
||||
|
||||
On startup, `FileStorage::Init` calls `StorageBackend::ScanMeta`, which reads all on-disk metadata and invokes a callback for each discovered object. The callback calls `MasterClient::NotifyOffloadSuccess` to re-register the objects with the master. This restores the full disk-replica view without any application-level intervention.
|
||||
|
|
@ -0,0 +1,594 @@
|
|||
# TENT C++ API Reference
|
||||
|
||||
## Overview
|
||||
|
||||
This page summarizes the C++ APIs in `mooncake-transfer-engine/tent/include/tent/transfer_engine.h`.
|
||||
It follows the same structure as the Transfer Engine API documentation.
|
||||
|
||||
For conceptual background, see [TENT Overview](overview.md).
|
||||
|
||||
**Prerequisites and API modes**
|
||||
- **Build** with `-DUSE_TENT=ON` to enable TENT.
|
||||
- TENT provides two API surfaces:
|
||||
- **TENT-native API** (documented below).
|
||||
- **TE-compatible API** via the compatibility shim (set `MC_USE_TENT=1`). For TE-compatible API, see [TE C++ API Reference](../transfer-engine/cpp-api.md)
|
||||
|
||||
**Core APIs vs Advanced APIs**
|
||||
- Core APIs form the minimal path to move data: create the engine, register memory, open segments, submit transfers, and query status.
|
||||
- Advanced APIs are optional; they help with segment export/import, notifications, and engine introspection.
|
||||
|
||||
## Key Differences from Transfer Engine
|
||||
|
||||
TENT redesigns the API surface based on different design goals. The following table summarizes the key differences and the rationale behind them.
|
||||
|
||||
| Aspect | Transfer Engine (TE) | TENT | Rationale |
|
||||
|--------|---------------------|------|-----------|
|
||||
| **Initialization** | Two-step: constructor + `init(metadata_conn_string, server_name, ...)` | Single-step: constructor with `Config` object or config file path | Simplifies initialization; centralizes configuration in a single place |
|
||||
| **Transport Management** | Manual: `installTransport()`, `uninstallTransport()`, `getTransport()` | Removed from public API | TENT performs dynamic transport selection at runtime; applications should not manage transports directly |
|
||||
| **Memory Allocation** | Not provided; users allocate memory externally | `allocateLocalMemory()` / `freeLocalMemory()` | Provides integrated memory management with automatic registration |
|
||||
| **Memory Registration** | `registerLocalMemory(addr, len, location, remote_accessible, update_metadata)` | `registerLocalMemory(addr, size, Permission)` or `MemoryOptions` | Cleaner semantics; `Permission` replaces boolean flags; metadata updates are internal |
|
||||
| **Segment Discovery** | Via metadata service; manual cache sync with `syncSegmentCache()` | Internal control plane (metadata type = `p2p` or central); no manual sync | Simplifies metadata management; discovery handled inside runtime |
|
||||
| **Error Handling** | Mixed: some APIs return `int`, others return `Status` | Consistent: all APIs return `Status` | Uniform error handling across the API |
|
||||
| **Topology/Introspection** | Exposed: `getLocalTopology()`, `getMetadata()`, `checkOverlap()`, etc. | Internalized; not exposed | TENT handles topology and path selection internally; reduces API complexity |
|
||||
|
||||
### Design Philosophy
|
||||
|
||||
1. **Declarative over Imperative**: Applications describe *what* data to move, not *how* to move it. Transport selection, path optimization, and failure handling are delegated to the TENT runtime.
|
||||
|
||||
2. **Single Initialization**: Instead of a two-phase `constructor + init()` pattern, TENT uses configuration objects that fully describe the engine state at construction time.
|
||||
|
||||
3. **Internal Metadata Management**: TENT does not expose metadata internals (`TransferMetadata`, `Topology`). Segment discovery is handled by the control plane (P2P or central), and the cache is managed automatically.
|
||||
|
||||
4. **Consistent Error Model**: All TENT APIs return `Status` objects, eliminating the mixed `int` / `Status` return types in TE.
|
||||
|
||||
## API Mapping: TE to TENT
|
||||
|
||||
For users migrating from Transfer Engine, the following table shows how TE APIs map to TENT APIs. TENT provides a backward-compatible shim (`MC_USE_TENT=1`) that allows existing TE code to run on the TENT runtime.
|
||||
|
||||
| Transfer Engine API | TENT API | Notes |
|
||||
|---------------------|----------|-------|
|
||||
| **Initialization** |||
|
||||
| `TransferEngine(auto_discover)` | `TransferEngine()` | TENT ignores `auto_discover`; discovery is always automatic |
|
||||
| `init(conn_string, server_name, ip, port)` | Constructor with `Config` | Config sets `metadata_type`, `metadata_servers`, `local_segment_name` |
|
||||
| `freeEngine()` | Destructor | Resources released automatically on destruction |
|
||||
| **Transport** |||
|
||||
| `installTransport(proto, args)` | *Not available* | Transport selection is internal to TENT |
|
||||
| `uninstallTransport(proto)` | *Not available* | — |
|
||||
| `getTransport(proto)` | *Not available* | — |
|
||||
| **Memory** |||
|
||||
| `registerLocalMemory(addr, len, location, remote_accessible, update_metadata)` | `registerLocalMemory(addr, size, MemoryOptions)` | `MemoryOptions.location` replaces `location`; `MemoryOptions.perm` replaces `remote_accessible` |
|
||||
| `unregisterLocalMemory(addr, update_metadata)` | `unregisterLocalMemory(addr)` | Metadata update is internal |
|
||||
| `registerLocalMemoryBatch(buffer_list, location)` | `registerLocalMemory(addr_list, size_list, MemoryOptions)` | Batch API uses vectors instead of `BufferEntry` |
|
||||
| `unregisterLocalMemoryBatch(addr_list)` | `unregisterLocalMemory(addr_list)` | — |
|
||||
| *Not available* | `allocateLocalMemory(addr, size, location)` | TENT-only: allocates and registers in one call |
|
||||
| *Not available* | `freeLocalMemory(addr)` | TENT-only: frees allocated memory |
|
||||
| **Segment** |||
|
||||
| `openSegment(segment_name)` → returns handle | `openSegment(handle, segment_name)` → via output param | Return style differs; `Status` indicates success/failure |
|
||||
| `closeSegment(handle)` | `closeSegment(handle)` | Same semantics |
|
||||
| `removeLocalSegment(segment_name)` | *Not available* | Segment lifecycle managed internally |
|
||||
| `CheckSegmentStatus(sid)` | *Not available* | Status checking is internal |
|
||||
| `syncSegmentCache(segment_name)` | *Not available* | Cache sync is automatic |
|
||||
| *Not available* | `exportLocalSegment(shared_handle)` | Declared in TENT API but currently not implemented |
|
||||
| *Not available* | `importRemoteSegment(handle, shared_handle)` | Declared in TENT API but currently not implemented |
|
||||
| `getSegmentInfo(handle, info)` | `getSegmentInfo(handle, info)` | Same semantics |
|
||||
| **Batch & Transfer** |||
|
||||
| `allocateBatchID(batch_size)` | `allocateBatch(batch_size)` | Renamed |
|
||||
| `freeBatchID(batch_id)` | `freeBatch(batch_id)` | Renamed |
|
||||
| `submitTransfer(batch_id, entries)` | `submitTransfer(batch_id, request_list)` | `TransferRequest` → `Request` |
|
||||
| `submitTransferWithNotify(batch_id, entries, notify_msg)` | `submitTransfer(batch_id, request_list, notifi)` | Unified API with optional notification |
|
||||
| `getTransferStatus(batch_id, task_id, status)` | `getTransferStatus(batch_id, task_id, status)` | Same |
|
||||
| `getBatchTransferStatus(batch_id, status)` | `getTransferStatus(batch_id, status)` | Overloaded; single `TransferStatus` output = overall status |
|
||||
| *Not available* | `getTransferStatus(batch_id, status_list)` | TENT-only: get all task statuses at once |
|
||||
| **Notification** |||
|
||||
| `sendNotifyByID(target_id, notify_msg)` | `sendNotification(target_id, notifi)` | Renamed; uses `Notification` struct |
|
||||
| `sendNotifyByName(remote_agent, notify_msg)` | `openSegment()` + `sendNotification()` | TENT requires explicit segment handle |
|
||||
| `getNotifies(notifies)` | `receiveNotification(notifi_list)` | Renamed; uses `Notification` struct |
|
||||
| **Introspection** |||
|
||||
| `getLocalIpAndPort()` | `getRpcServerAddress()` + `getRpcServerPort()` | Split into two methods |
|
||||
| `getRpcPort()` | `getRpcServerPort()` | Same |
|
||||
| `getMetadata()` | *Not available* | Metadata internals not exposed |
|
||||
| `getLocalTopology()` | *Not available* | Topology internals not exposed |
|
||||
| `checkOverlap(addr, length)` | *Not available* | — |
|
||||
| `setAutoDiscover(auto_discover)` | *Not available* | Always enabled (ignored under `MC_USE_TENT`) |
|
||||
| `setWhitelistFilters(filters)` | *Not available* | Configure via `Config` |
|
||||
| `numContexts()` | *Not available* | — |
|
||||
| *Not available* | `available()` | TENT-only: check if engine initialized successfully |
|
||||
| *Not available* | `getSegmentName()` | TENT-only: get local segment name |
|
||||
|
||||
### Using the Backward-Compatible Shim
|
||||
|
||||
Existing Transfer Engine code can run on TENT by setting the environment variable:
|
||||
|
||||
```bash
|
||||
export MC_USE_TENT=1
|
||||
```
|
||||
|
||||
When this variable is set, the `mooncake::TransferEngine` class internally delegates to `mooncake::tent::TransferEngine`. Most TE APIs are translated automatically. APIs that have no TENT equivalent (e.g., `installTransport`, `getMetadata`) become no-ops or return placeholder values.
|
||||
|
||||
## Core APIs
|
||||
|
||||
### Core Usage Path (C++)
|
||||
|
||||
Most integrations follow a short, repeatable path: create the engine, register local memory, open a target segment, submit a batch transfer, poll status, and finally free resources. A minimal example is shown below.
|
||||
|
||||
```cpp
|
||||
#include "tent/transfer_engine.h"
|
||||
|
||||
using mooncake::tent::TransferEngine;
|
||||
using mooncake::tent::Request;
|
||||
using mooncake::tent::TransferStatus;
|
||||
using mooncake::tent::SegmentID;
|
||||
using mooncake::tent::BatchID;
|
||||
|
||||
// Create engine (loads config from default path or environment)
|
||||
TransferEngine engine;
|
||||
if (!engine.available()) {
|
||||
// handle initialization failure
|
||||
}
|
||||
|
||||
// Allocate and register local memory
|
||||
void* local_addr = nullptr;
|
||||
size_t length = 1024 * 1024; // 1MB
|
||||
engine.allocateLocalMemory(&local_addr, length, "cuda:0");
|
||||
|
||||
// Open remote segment
|
||||
SegmentID remote_segment;
|
||||
engine.openSegment(remote_segment, "remote_node");
|
||||
|
||||
// Prepare transfer request
|
||||
Request req{};
|
||||
req.opcode = Request::WRITE;
|
||||
req.source = local_addr;
|
||||
req.target_id = remote_segment;
|
||||
req.target_offset = 0;
|
||||
req.length = length;
|
||||
|
||||
// Allocate batch and submit transfer
|
||||
BatchID batch = engine.allocateBatch(1);
|
||||
engine.submitTransfer(batch, {req});
|
||||
|
||||
// Poll for completion
|
||||
TransferStatus status;
|
||||
do {
|
||||
engine.getTransferStatus(batch, status);
|
||||
} while (status.s == mooncake::tent::PENDING);
|
||||
|
||||
// Cleanup
|
||||
engine.freeBatch(batch);
|
||||
engine.closeSegment(remote_segment);
|
||||
engine.freeLocalMemory(local_addr);
|
||||
```
|
||||
|
||||
### Constructors
|
||||
|
||||
```cpp
|
||||
TransferEngine();
|
||||
TransferEngine(const std::string config_path);
|
||||
TransferEngine(std::shared_ptr<Config> config);
|
||||
```
|
||||
|
||||
Constructs a TENT Transfer Engine instance.
|
||||
|
||||
- Default constructor: Loads configuration from the default path or environment variables.
|
||||
- `config_path`: Path to a JSON configuration file.
|
||||
- `config`: A pre-constructed `Config` object for programmatic configuration.
|
||||
|
||||
The engine is ready to use after construction if `available()` returns `true`.
|
||||
|
||||
### Types
|
||||
|
||||
#### Request
|
||||
|
||||
The core API provided by TENT is submitting a group of asynchronous `Request` tasks through the `submitTransfer` interface, and querying their status through the `getTransferStatus` interface.
|
||||
|
||||
```cpp
|
||||
struct Request {
|
||||
enum OpCode { READ, WRITE };
|
||||
OpCode opcode;
|
||||
void* source;
|
||||
SegmentID target_id;
|
||||
uint64_t target_offset;
|
||||
size_t length;
|
||||
};
|
||||
```
|
||||
|
||||
- `opcode`: `READ` copies data from `<target_id, target_offset>` to `source`; `WRITE` copies data from `source` to `<target_id, target_offset>`.
|
||||
- `source`: Local buffer address, must be registered via `registerLocalMemory` or allocated via `allocateLocalMemory`.
|
||||
- `target_id`: Segment ID obtained from `openSegment`.
|
||||
- `target_offset`: Offset within the target segment.
|
||||
- `length`: Number of bytes to transfer.
|
||||
|
||||
#### TransferStatus
|
||||
|
||||
```cpp
|
||||
enum TransferStatusEnum {
|
||||
INITIAL, // Not yet started
|
||||
PENDING, // Transfer in progress
|
||||
INVALID, // Invalid parameters
|
||||
CANCELED, // Transfer canceled
|
||||
COMPLETED, // Transfer completed successfully
|
||||
TIMEOUT, // Transfer timed out
|
||||
FAILED // Transfer failed
|
||||
};
|
||||
|
||||
struct TransferStatus {
|
||||
TransferStatusEnum s;
|
||||
size_t transferred_bytes;
|
||||
};
|
||||
```
|
||||
|
||||
- `s`: Current status of the transfer.
|
||||
- `transferred_bytes`: Number of bytes successfully transferred (lower bound).
|
||||
|
||||
### Data Transfer
|
||||
|
||||
#### TransferEngine::allocateBatch
|
||||
|
||||
```cpp
|
||||
BatchID allocateBatch(size_t batch_size);
|
||||
```
|
||||
|
||||
Allocates a `BatchID` that can hold up to `batch_size` transfer requests.
|
||||
|
||||
- `batch_size`: Maximum number of requests that can be submitted under this batch.
|
||||
- Return value: A valid `BatchID` on success.
|
||||
|
||||
#### TransferEngine::submitTransfer
|
||||
|
||||
```cpp
|
||||
Status submitTransfer(BatchID batch_id,
|
||||
const std::vector<Request>& request_list);
|
||||
```
|
||||
|
||||
Submits transfer requests to the specified batch. Requests are executed asynchronously.
|
||||
|
||||
- `batch_id`: The batch to submit requests to.
|
||||
- `request_list`: Vector of `Request` objects.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
#### TransferEngine::getTransferStatus
|
||||
|
||||
```cpp
|
||||
// Get status of a single task
|
||||
Status getTransferStatus(BatchID batch_id, size_t task_id,
|
||||
TransferStatus& status);
|
||||
|
||||
// Get status of all tasks in a batch
|
||||
Status getTransferStatus(BatchID batch_id,
|
||||
std::vector<TransferStatus>& status_list);
|
||||
|
||||
// Get overall batch status
|
||||
Status getTransferStatus(BatchID batch_id, TransferStatus& overall_status);
|
||||
```
|
||||
|
||||
Queries the status of transfer requests.
|
||||
|
||||
- `batch_id`: The batch to query.
|
||||
- `task_id`: Index of the specific task (for single-task query).
|
||||
- `status` / `status_list` / `overall_status`: Output parameter(s) for status.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
#### TransferEngine::freeBatch
|
||||
|
||||
```cpp
|
||||
Status freeBatch(BatchID batch_id);
|
||||
```
|
||||
|
||||
Releases a batch. All transfers in the batch must be completed before calling this.
|
||||
|
||||
- `batch_id`: The batch to release.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
### Memory Management
|
||||
|
||||
#### TransferEngine::allocateLocalMemory
|
||||
|
||||
```cpp
|
||||
Status allocateLocalMemory(void** addr, size_t size,
|
||||
Location location = kWildcardLocation);
|
||||
|
||||
// Advanced version with MemoryOptions
|
||||
Status allocateLocalMemory(void** addr, size_t size,
|
||||
MemoryOptions& options);
|
||||
```
|
||||
|
||||
Allocates memory that is automatically registered for transfers.
|
||||
|
||||
- `addr`: Output pointer to the allocated memory.
|
||||
- `size`: Size in bytes to allocate.
|
||||
- `location`: Device location hint (e.g., `"cuda:0"`, `"cpu:0"`, or `"*"` for auto-detect).
|
||||
- `options`: Advanced options including location, permission, and transport type.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
#### TransferEngine::freeLocalMemory
|
||||
|
||||
```cpp
|
||||
Status freeLocalMemory(void* addr);
|
||||
```
|
||||
|
||||
Frees memory previously allocated with `allocateLocalMemory`.
|
||||
|
||||
- `addr`: Pointer to the memory to free.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
#### TransferEngine::registerLocalMemory
|
||||
|
||||
```cpp
|
||||
Status registerLocalMemory(void* addr, size_t size,
|
||||
Permission permission = kGlobalReadWrite);
|
||||
|
||||
// Batch registration
|
||||
Status registerLocalMemory(std::vector<void*> addr_list,
|
||||
std::vector<size_t> size_list,
|
||||
Permission permission = kGlobalReadWrite);
|
||||
|
||||
// Advanced version with MemoryOptions
|
||||
Status registerLocalMemory(void* addr, size_t size, MemoryOptions& options);
|
||||
```
|
||||
|
||||
Registers externally allocated memory for use in transfers.
|
||||
|
||||
- `addr`: Starting address of the memory region.
|
||||
- `size`: Size in bytes.
|
||||
- `permission`: Access permission (`kLocalReadWrite`, `kGlobalReadOnly`, `kGlobalReadWrite`).
|
||||
- `addr_list` / `size_list`: For batch registration of multiple buffers.
|
||||
- `options`: Advanced options for fine-grained control.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
#### TransferEngine::unregisterLocalMemory
|
||||
|
||||
```cpp
|
||||
Status unregisterLocalMemory(void* addr, size_t size = 0);
|
||||
|
||||
// Batch unregistration
|
||||
Status unregisterLocalMemory(std::vector<void*> addr_list,
|
||||
std::vector<size_t> size_list = {});
|
||||
```
|
||||
|
||||
Unregisters previously registered memory.
|
||||
|
||||
- `addr`: Starting address of the memory region.
|
||||
- `size`: Size in bytes (optional, can be 0 if the engine tracks it).
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
### Segment Management
|
||||
|
||||
#### TransferEngine::openSegment
|
||||
|
||||
```cpp
|
||||
Status openSegment(SegmentID& handle, const std::string& segment_name);
|
||||
```
|
||||
|
||||
Opens a segment by name and returns a handle for use in transfers.
|
||||
|
||||
- `handle`: Output parameter for the segment ID.
|
||||
- `segment_name`: Name of the segment to open (typically the remote node name).
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
#### TransferEngine::closeSegment
|
||||
|
||||
```cpp
|
||||
Status closeSegment(SegmentID handle);
|
||||
```
|
||||
|
||||
Closes a previously opened segment.
|
||||
|
||||
- `handle`: The segment ID to close.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
#### TransferEngine::getSegmentInfo
|
||||
|
||||
```cpp
|
||||
Status getSegmentInfo(SegmentID handle, SegmentInfo& info);
|
||||
```
|
||||
|
||||
Retrieves information about a segment.
|
||||
|
||||
- `handle`: The segment ID.
|
||||
- `info`: Output parameter containing segment details (type, buffers).
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
## Advanced APIs
|
||||
|
||||
These APIs are optional; use them for segment export/import, notifications, or engine introspection.
|
||||
|
||||
### Segment Export and Import
|
||||
|
||||
Reserved for future use. These APIs are declared but currently return
|
||||
`Status::NotImplemented`.
|
||||
|
||||
#### TransferEngine::exportLocalSegment
|
||||
|
||||
```cpp
|
||||
Status exportLocalSegment(std::string& shared_handle);
|
||||
```
|
||||
|
||||
Exports the local segment as a shareable handle string.
|
||||
|
||||
- `shared_handle`: Output string that can be passed to remote nodes.
|
||||
- Return value: Currently returns `Status::NotImplemented`.
|
||||
- Typical use: Share segment information through an external channel (e.g., gRPC, Redis).
|
||||
|
||||
#### TransferEngine::importRemoteSegment
|
||||
|
||||
```cpp
|
||||
Status importRemoteSegment(SegmentID& handle,
|
||||
const std::string& shared_handle);
|
||||
```
|
||||
|
||||
Imports a remote segment from a shared handle string.
|
||||
|
||||
- `handle`: Output segment ID.
|
||||
- `shared_handle`: The handle string obtained from `exportLocalSegment` on the remote side.
|
||||
- Return value: Currently returns `Status::NotImplemented`.
|
||||
|
||||
### Notifications
|
||||
|
||||
TENT supports lightweight notifications to coordinate data movement between nodes.
|
||||
|
||||
#### TransferEngine::submitTransfer (with notification)
|
||||
|
||||
```cpp
|
||||
Status submitTransfer(BatchID batch_id,
|
||||
const std::vector<Request>& request_list,
|
||||
const Notification& notifi);
|
||||
```
|
||||
|
||||
Submits transfers and sends a notification upon completion.
|
||||
|
||||
- `notifi`: A `{name, msg}` payload delivered to the receiver when the transfer completes.
|
||||
- Typical use: Signal that transferred data is ready for consumption.
|
||||
|
||||
#### TransferEngine::sendNotification
|
||||
|
||||
```cpp
|
||||
Status sendNotification(SegmentID target_id, const Notification& notifi);
|
||||
```
|
||||
|
||||
Sends a notification to a specific segment without data transfer.
|
||||
|
||||
- `target_id`: The segment to notify.
|
||||
- `notifi`: The notification payload.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
|
||||
#### TransferEngine::receiveNotification
|
||||
|
||||
```cpp
|
||||
Status receiveNotification(std::vector<Notification>& notifi_list);
|
||||
```
|
||||
|
||||
Receives pending notifications from peers.
|
||||
|
||||
- `notifi_list`: Output vector of received notifications.
|
||||
- Return value: `Status::OK()` on success; otherwise a non-OK status.
|
||||
- Typical use: Polling loop to trigger follow-up actions on received data.
|
||||
|
||||
### Engine Introspection
|
||||
|
||||
#### TransferEngine::available
|
||||
|
||||
```cpp
|
||||
bool available() const;
|
||||
```
|
||||
|
||||
Returns `true` if the engine was initialized successfully and is ready for use.
|
||||
|
||||
#### TransferEngine::getSegmentName
|
||||
|
||||
```cpp
|
||||
const std::string getSegmentName() const;
|
||||
```
|
||||
|
||||
Returns the local segment name (node identifier).
|
||||
|
||||
#### TransferEngine::getRpcServerAddress
|
||||
|
||||
```cpp
|
||||
const std::string getRpcServerAddress() const;
|
||||
```
|
||||
|
||||
Returns the RPC server address for this engine instance.
|
||||
|
||||
#### TransferEngine::getRpcServerPort
|
||||
|
||||
```cpp
|
||||
uint16_t getRpcServerPort() const;
|
||||
```
|
||||
|
||||
Returns the RPC server port for this engine instance.
|
||||
|
||||
## Type Reference
|
||||
|
||||
### Permission
|
||||
|
||||
```cpp
|
||||
enum Permission {
|
||||
kLocalReadWrite, // Only local access
|
||||
kGlobalReadOnly, // Remote read access
|
||||
kGlobalReadWrite, // Remote read/write access
|
||||
};
|
||||
```
|
||||
|
||||
### Location
|
||||
|
||||
```cpp
|
||||
using Location = std::string;
|
||||
const static std::string kWildcardLocation = "*";
|
||||
```
|
||||
|
||||
Location strings identify device affinity: `"cpu:0"`, `"cuda:0"`, `"cuda:1"`, etc. Use `"*"` for automatic detection.
|
||||
|
||||
### TransportType
|
||||
|
||||
```cpp
|
||||
enum TransportType {
|
||||
RDMA = 0,
|
||||
MNNVL,
|
||||
SHM,
|
||||
NVLINK,
|
||||
GDS,
|
||||
IOURING,
|
||||
TCP,
|
||||
AscendDirect,
|
||||
UNSPEC
|
||||
};
|
||||
```
|
||||
|
||||
Transport types used internally by TENT. Applications typically do not need to specify these directly.
|
||||
|
||||
### MemoryOptions
|
||||
|
||||
```cpp
|
||||
struct MemoryOptions {
|
||||
Location location = kWildcardLocation;
|
||||
Permission perm = kGlobalReadWrite;
|
||||
TransportType type = UNSPEC;
|
||||
std::string shm_path = "";
|
||||
size_t shm_offset = 0;
|
||||
bool internal = false;
|
||||
};
|
||||
```
|
||||
|
||||
Advanced options for memory allocation and registration.
|
||||
|
||||
### SegmentInfo
|
||||
|
||||
```cpp
|
||||
struct SegmentInfo {
|
||||
enum Type { Memory, File };
|
||||
struct Buffer {
|
||||
uint64_t base, length;
|
||||
Location location;
|
||||
};
|
||||
Type type;
|
||||
std::vector<Buffer> buffers;
|
||||
};
|
||||
```
|
||||
|
||||
Information about a segment, including its type and registered buffers.
|
||||
|
||||
### Notification
|
||||
|
||||
```cpp
|
||||
struct Notification {
|
||||
std::string name;
|
||||
std::string msg;
|
||||
};
|
||||
```
|
||||
|
||||
Lightweight notification payload for coordination between nodes.
|
||||
|
||||
### Status
|
||||
|
||||
The `Status` class is used for error handling throughout the API. Key methods:
|
||||
|
||||
```cpp
|
||||
bool ok() const; // Returns true if operation succeeded
|
||||
std::string ToString() const; // Human-readable error description
|
||||
|
||||
// Common status factory methods
|
||||
static Status OK();
|
||||
static Status InvalidArgument(std::string_view msg);
|
||||
static Status InternalError(std::string_view msg);
|
||||
// ... and more
|
||||
```
|
||||
|
|
@ -0,0 +1,388 @@
|
|||
# TENT Metrics System
|
||||
|
||||
TENT provides a built-in metrics system based on yalantinglibs, compatible with Prometheus for monitoring data transfer performance and system health.
|
||||
|
||||
## Overview
|
||||
|
||||
The metrics system supports two metric types:
|
||||
|
||||
- **Counter**: Monotonically increasing values (e.g., total bytes transferred, total requests)
|
||||
- **Histogram**: Distribution of values with configurable buckets (e.g., latency)
|
||||
|
||||
All metrics are thread-safe and designed for high-performance data paths.
|
||||
|
||||
## Performance Optimization
|
||||
|
||||
The metrics system provides two levels of control for performance optimization:
|
||||
|
||||
### Compile-time Disable (Zero Overhead)
|
||||
|
||||
By default, metrics are **disabled** at compile time for maximum performance. To enable metrics, build with:
|
||||
|
||||
```bash
|
||||
cmake -DTENT_METRICS_ENABLED=ON ..
|
||||
```
|
||||
|
||||
When disabled at compile time (`TENT_METRICS_ENABLED=OFF`, the default), all metrics macros expand to `((void)0)`, resulting in **zero runtime overhead**.
|
||||
|
||||
### Runtime Disable (Minimal Overhead)
|
||||
|
||||
When metrics are enabled at compile time, you can still disable them at runtime:
|
||||
|
||||
```cpp
|
||||
// Disable metrics collection at runtime
|
||||
TentMetrics::setEnabled(false);
|
||||
|
||||
// Re-enable metrics collection
|
||||
TentMetrics::setEnabled(true);
|
||||
|
||||
// Check current state
|
||||
bool enabled = TentMetrics::isEnabled();
|
||||
```
|
||||
|
||||
When disabled at runtime, record functions return immediately after a single atomic load (~1ns overhead).
|
||||
|
||||
## Configuration
|
||||
|
||||
### Configuration Sources (Priority Order)
|
||||
|
||||
1. **Config File** (highest priority)
|
||||
2. **Environment Variables** (medium priority)
|
||||
3. **Default Values** (lowest priority)
|
||||
|
||||
### Config File Format
|
||||
|
||||
TENT metrics configuration is integrated into the main `transfer-engine.json` configuration file:
|
||||
|
||||
```json
|
||||
{
|
||||
"local_segment_name": "",
|
||||
"metadata_type": "p2p",
|
||||
"metadata_servers": "127.0.0.1:2379",
|
||||
"log_level": "warning",
|
||||
"metrics": {
|
||||
"enabled": true,
|
||||
"http_port": 9100,
|
||||
"http_host": "0.0.0.0",
|
||||
"http_server_threads": 2,
|
||||
"report_interval_seconds": 30,
|
||||
"enable_prometheus": true,
|
||||
"enable_json": true,
|
||||
"latency_buckets": [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0],
|
||||
"size_buckets": [1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864, 268435456, 1073741824]
|
||||
},
|
||||
"transports": {
|
||||
// ... transport configuration
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Note**:
|
||||
- `report_interval_seconds`: Set to 0 to disable periodic logging
|
||||
- `latency_buckets`: Values are in **seconds** (e.g., 0.001 = 1ms). The system internally converts to microseconds for histogram storage.
|
||||
- `size_buckets`: Values are in bytes
|
||||
|
||||
### Environment Variables
|
||||
|
||||
```bash
|
||||
# Basic settings
|
||||
TENT_METRICS_ENABLED=true
|
||||
TENT_METRICS_HTTP_PORT=9100
|
||||
TENT_METRICS_HTTP_HOST=0.0.0.0
|
||||
TENT_METRICS_HTTP_SERVER_THREADS=2
|
||||
TENT_METRICS_REPORT_INTERVAL=30 # Set to 0 to disable periodic logging
|
||||
|
||||
# Output formats
|
||||
TENT_METRICS_ENABLE_PROMETHEUS=true
|
||||
TENT_METRICS_ENABLE_JSON=true
|
||||
|
||||
# Custom buckets (comma-separated, latency in seconds, size in bytes)
|
||||
TENT_METRICS_LATENCY_BUCKETS="0.0001,0.0005,0.001,0.005,0.01,0.05,0.1,0.5,1.0"
|
||||
TENT_METRICS_SIZE_BUCKETS="1024,4096,16384,65536,262144,1048576"
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Build with Metrics Enabled
|
||||
|
||||
```bash
|
||||
# Enable metrics at compile time (disabled by default)
|
||||
cmake -DTENT_METRICS_ENABLED=ON ..
|
||||
make
|
||||
```
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```cpp
|
||||
#include "tent/metrics/tent_metrics.h"
|
||||
#include "tent/metrics/config_loader.h"
|
||||
|
||||
// Load configuration from transfer-engine.json
|
||||
auto config = MetricsConfigLoader::loadWithDefaults();
|
||||
|
||||
// Initialize TENT metrics system
|
||||
auto& tent_metrics = TentMetrics::instance();
|
||||
tent_metrics.initialize(config);
|
||||
|
||||
// HTTP server starts automatically
|
||||
```
|
||||
|
||||
### Recording Transfer Metrics
|
||||
|
||||
```cpp
|
||||
// Using convenience macros (recommended)
|
||||
TENT_RECORD_READ_COMPLETED(1024*1024, 0.025); // 1MB read in 25ms
|
||||
TENT_RECORD_WRITE_COMPLETED(512*1024, 0.015); // 512KB write in 15ms
|
||||
TENT_RECORD_READ_FAILED(1024*1024); // 1MB read failed
|
||||
TENT_RECORD_WRITE_FAILED(512*1024); // 512KB write failed
|
||||
|
||||
// Direct API usage
|
||||
auto& tent_metrics = TentMetrics::instance();
|
||||
tent_metrics.recordReadCompleted(1024*1024, 0.025);
|
||||
tent_metrics.recordWriteCompleted(512*1024, 0.015);
|
||||
tent_metrics.recordReadFailed(1024*1024);
|
||||
tent_metrics.recordWriteFailed(512*1024);
|
||||
```
|
||||
|
||||
### RAII Latency Measurement
|
||||
|
||||
```cpp
|
||||
// Automatic latency measurement using RAII
|
||||
{
|
||||
TENT_SCOPED_READ_LATENCY(1024 * 1024); // e.g. 1MB
|
||||
// ... perform read operation ...
|
||||
} // latency automatically recorded when scope exits
|
||||
|
||||
{
|
||||
TENT_SCOPED_WRITE_LATENCY(512 * 1024); // e.g. 512KB
|
||||
// ... perform write operation ...
|
||||
}
|
||||
```
|
||||
|
||||
## HTTP Server Endpoints
|
||||
|
||||
The HTTP server provides multiple endpoints:
|
||||
|
||||
- **`/metrics`**: Prometheus format
|
||||
- **`/metrics/summary`**: Human-readable summary
|
||||
- **`/metrics/json`**: JSON format
|
||||
- **`/health`**: Health check endpoint
|
||||
|
||||
### Example Responses
|
||||
|
||||
**Prometheus Format (`/metrics`)**:
|
||||
```
|
||||
# HELP tent_read_bytes_total Total bytes read via TENT
|
||||
# TYPE tent_read_bytes_total counter
|
||||
tent_read_bytes_total 1048576
|
||||
|
||||
# HELP tent_write_bytes_total Total bytes written via TENT
|
||||
# TYPE tent_write_bytes_total counter
|
||||
tent_write_bytes_total 524288
|
||||
|
||||
# HELP tent_read_requests_total Total read requests via TENT
|
||||
# TYPE tent_read_requests_total counter
|
||||
tent_read_requests_total 100
|
||||
|
||||
# HELP tent_write_requests_total Total write requests via TENT
|
||||
# TYPE tent_write_requests_total counter
|
||||
tent_write_requests_total 50
|
||||
|
||||
# HELP tent_read_failures_total Total read failures via TENT
|
||||
# TYPE tent_read_failures_total counter
|
||||
tent_read_failures_total 2
|
||||
|
||||
# HELP tent_write_failures_total Total write failures via TENT
|
||||
# TYPE tent_write_failures_total counter
|
||||
tent_write_failures_total 1
|
||||
|
||||
# HELP tent_read_latency_us Read latency distribution in microseconds
|
||||
# TYPE tent_read_latency_us histogram
|
||||
tent_read_latency_us_bucket{le="100"} 10
|
||||
tent_read_latency_us_bucket{le="500"} 50
|
||||
...
|
||||
```
|
||||
|
||||
**JSON Format (`/metrics/json`)**:
|
||||
```json
|
||||
{
|
||||
"tent_read_bytes_total": 1048576,
|
||||
"tent_write_bytes_total": 524288,
|
||||
"tent_read_requests_total": 100,
|
||||
"tent_write_requests_total": 50,
|
||||
"tent_read_failures_total": 2,
|
||||
"tent_write_failures_total": 1
|
||||
}
|
||||
```
|
||||
|
||||
**Summary Format (`/metrics/summary`)**:
|
||||
```
|
||||
Read: 1.00 MB (100 reqs, 2 fails) | Write: 512.00 KB (50 reqs, 1 fails)
|
||||
```
|
||||
|
||||
## Available Metrics
|
||||
|
||||
| Metric Name | Type | Description |
|
||||
|-------------|------|-------------|
|
||||
| `tent_read_bytes_total` | Counter | Total bytes read via TENT |
|
||||
| `tent_write_bytes_total` | Counter | Total bytes written via TENT |
|
||||
| `tent_read_requests_total` | Counter | Total read requests via TENT |
|
||||
| `tent_write_requests_total` | Counter | Total write requests via TENT |
|
||||
| `tent_read_failures_total` | Counter | Total read failures via TENT |
|
||||
| `tent_write_failures_total` | Counter | Total write failures via TENT |
|
||||
| `tent_read_latency_us` | Histogram | Read latency distribution in microseconds |
|
||||
| `tent_write_latency_us` | Histogram | Write latency distribution in microseconds |
|
||||
| `tent_read_size_bytes` | Histogram | Read request size distribution in bytes |
|
||||
| `tent_write_size_bytes` | Histogram | Write request size distribution in bytes |
|
||||
|
||||
## Integration with TransferEngine
|
||||
|
||||
The metrics system is automatically integrated with TransferEngine. When TransferEngine starts, it initializes the metrics system:
|
||||
|
||||
```cpp
|
||||
#include "tent/metrics/tent_metrics.h"
|
||||
#include "tent/metrics/config_loader.h"
|
||||
|
||||
// Load configuration
|
||||
auto metrics_config = MetricsConfigLoader::loadWithDefaults();
|
||||
if (metrics_config.enabled) {
|
||||
TentMetrics::instance().initialize(metrics_config);
|
||||
}
|
||||
```
|
||||
|
||||
Metrics are automatically recorded at the TENT layer:
|
||||
|
||||
- **Latency tracking**: Start time is recorded when `submitTransfer` is called
|
||||
- **Metrics recording**: When `getTransferStatus` detects task completion, latency is calculated and metrics are recorded
|
||||
|
||||
This provides end-to-end latency measurement across all transport types (RDMA, TCP, NVLink, etc.).
|
||||
|
||||
**Note**: Remember to build with `-DTENT_METRICS_ENABLED=ON` to enable metrics collection.
|
||||
|
||||
## Adding New Metrics
|
||||
|
||||
To add new metrics to the TENT metrics system, follow these steps:
|
||||
|
||||
### Step 1: Declare the Metric
|
||||
|
||||
Add the metric member variable in `tent_metrics.h`:
|
||||
|
||||
```cpp
|
||||
// In TentMetrics class private section:
|
||||
|
||||
// For a new counter:
|
||||
ylt::metric::counter_t new_counter_{"tent_new_counter", "Description of the counter"};
|
||||
|
||||
// For a new histogram:
|
||||
ylt::metric::histogram_t new_histogram_{"tent_new_histogram", "Description",
|
||||
std::vector<double>{/* bucket boundaries */}};
|
||||
```
|
||||
|
||||
### Step 2: Register the Metric
|
||||
|
||||
Add the metric pointer to `registerMetrics()` in `tent_metrics.cpp`:
|
||||
|
||||
```cpp
|
||||
void TentMetrics::registerMetrics() {
|
||||
counters_ = {
|
||||
&read_bytes_total_,
|
||||
// ... existing counters ...
|
||||
&new_counter_, // Add new counter here
|
||||
};
|
||||
|
||||
histograms_ = {
|
||||
&read_latency_,
|
||||
// ... existing histograms ...
|
||||
&new_histogram_, // Add new histogram here
|
||||
};
|
||||
}
|
||||
```
|
||||
|
||||
### Step 3: Add Recording Methods (Optional)
|
||||
|
||||
If needed, add public methods to record the metric:
|
||||
|
||||
```cpp
|
||||
// In tent_metrics.h:
|
||||
void recordNewMetric(int64_t value);
|
||||
|
||||
// In tent_metrics.cpp:
|
||||
void TentMetrics::recordNewMetric(int64_t value) {
|
||||
if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) return;
|
||||
new_counter_.inc(value);
|
||||
// or for histogram:
|
||||
// new_histogram_.observe(value);
|
||||
}
|
||||
```
|
||||
|
||||
### Automatic Serialization
|
||||
|
||||
Once registered in `registerMetrics()`, the new metric will be **automatically included** in:
|
||||
- `/metrics` (Prometheus format)
|
||||
- `/metrics/json` (JSON format)
|
||||
|
||||
No changes to `getPrometheusMetrics()` or `getJsonMetrics()` are required.
|
||||
|
||||
## Advanced Configuration
|
||||
|
||||
### Custom Buckets
|
||||
|
||||
Define custom histogram buckets for specific use cases:
|
||||
|
||||
```cpp
|
||||
// Latency buckets (in seconds, converted to microseconds internally)
|
||||
std::vector<double> rdma_latency_buckets = {
|
||||
0.000001, 0.000005, 0.00001, 0.00005, 0.0001, // 1-100μs
|
||||
0.0005, 0.001, 0.005, 0.01, 0.05, 0.1 // 0.5-100ms
|
||||
};
|
||||
|
||||
// Size buckets for different data patterns (in bytes)
|
||||
std::vector<double> message_size_buckets = {
|
||||
64, 256, 1024, 4096, 16384, 65536, 262144 // 64B to 256KB
|
||||
};
|
||||
```
|
||||
|
||||
### Validation
|
||||
|
||||
```cpp
|
||||
MetricsConfig config = MetricsConfigLoader::loadWithDefaults();
|
||||
std::string error_msg;
|
||||
if (!MetricsConfigLoader::validateConfig(config, &error_msg)) {
|
||||
LOG(ERROR) << "Invalid metrics config: " << error_msg;
|
||||
return;
|
||||
}
|
||||
```
|
||||
|
||||
## Prometheus Integration
|
||||
|
||||
### Prometheus Configuration
|
||||
|
||||
```yaml
|
||||
# prometheus.yml
|
||||
scrape_configs:
|
||||
- job_name: 'tent-metrics'
|
||||
static_configs:
|
||||
- targets: ['localhost:9100']
|
||||
scrape_interval: 15s
|
||||
metrics_path: /metrics
|
||||
```
|
||||
|
||||
### Grafana Queries
|
||||
|
||||
```promql
|
||||
# Transfer throughput (MB/s)
|
||||
rate(tent_read_bytes_total[5m]) / 1024 / 1024
|
||||
rate(tent_write_bytes_total[5m]) / 1024 / 1024
|
||||
|
||||
# Request rate
|
||||
rate(tent_read_requests_total[5m])
|
||||
rate(tent_write_requests_total[5m])
|
||||
|
||||
# Failure rate
|
||||
rate(tent_read_failures_total[5m]) / rate(tent_read_requests_total[5m])
|
||||
|
||||
# P99 latency (note: latency is in microseconds, convert to seconds for display)
|
||||
histogram_quantile(0.99, rate(tent_read_latency_us_bucket[5m])) / 1000000
|
||||
histogram_quantile(0.99, rate(tent_write_latency_us_bucket[5m])) / 1000000
|
||||
```
|
||||
|
||||
|
|
@ -73,3 +73,18 @@ TENT extends the classic Mooncake Transfer Engine by moving transport selection,
|
|||
|
||||
The design favors predictable behavior and operational simplicity over manual tuning and static configuration.
|
||||
|
||||
## TENT C++ API Reference
|
||||
|
||||
:::{toctree}
|
||||
:maxdepth: 1
|
||||
|
||||
cpp-api
|
||||
:::
|
||||
|
||||
## TENT Metrics System
|
||||
|
||||
:::{toctree}
|
||||
:maxdepth: 1
|
||||
|
||||
metrics
|
||||
:::
|
||||
|
|
@ -125,7 +125,23 @@ Example:
|
|||
|
||||
> Consistency checking introduces CPU-side overhead and should be disabled for pure performance measurements.
|
||||
|
||||
### 5.3 Segment, Concurrency, and Memory Layout
|
||||
### 5.3 Notification Feature (`--notifi`)
|
||||
|
||||
When enabled, the benchmark sends a notification message along with each transfer batch:
|
||||
|
||||
* The notification contains the `target_addr` as the message payload
|
||||
* The peer can verify the notification was received correctly by checking the message
|
||||
* Useful for testing notification delivery and end-to-end communication
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
./tebench --target_seg_name=<SEG> --notifi=true
|
||||
```
|
||||
|
||||
> This feature is primarily for testing notification mechanisms. The notification message contains the target address for verification purposes.
|
||||
|
||||
### 5.4 Segment, Concurrency, and Memory Layout
|
||||
|
||||
**Segment and role**
|
||||
|
||||
|
|
@ -149,7 +165,7 @@ block_size × batch_size × num_threads > total_buffer_size
|
|||
|
||||
---
|
||||
|
||||
### 5.4 GPU Affinity
|
||||
### 5.5 GPU Affinity
|
||||
|
||||
* `--local_gpu_id` : initiator base GPU ID
|
||||
* `--target_gpu_id` : target base GPU ID
|
||||
|
|
@ -162,7 +178,7 @@ gpu_id + thread_id
|
|||
|
||||
---
|
||||
|
||||
### 5.5 Backend, Transport, and Metadata
|
||||
### 5.6 Backend, Transport, and Metadata
|
||||
|
||||
**Backend**
|
||||
|
||||
|
|
|
|||
|
|
@ -42,6 +42,33 @@ Complete command format is shown below:
|
|||
./transfer_engine_ascend_direct_perf --metadata_server=P2PHANDSHAKE --local_server_name=127.0.0.1:12346 --operation=write --device_logicid=1 --mode=initiator --block_size=16384 --batch_size=32 --block_iteration=10 --segment_id=127.0.0.1:real_port
|
||||
```
|
||||
|
||||
### Environment Variables Configuration
|
||||
|
||||
The following environment variables can be configured to control Ascend Direct Transport behavior:
|
||||
|
||||
| Variable | Description | Default Value | Example |
|
||||
|----------|-------------|---------------|---------|
|
||||
| `ASCEND_AUTO_CONNECT` | Enable automatic connection management | 0 (disabled) | `ASCEND_AUTO_CONNECT=1` |
|
||||
| `ASCEND_ENABLE_USE_FABRIC_MEM` | Enable fabric memory transfer mode in Mooncake Store (A3 only) | 0 (disabled) | `ASCEND_ENABLE_USE_FABRIC_MEM=1` |
|
||||
| `ASCEND_USE_ASYNC_TRANSFER` | Enable asynchronous transfer mode | 0 (disabled) | `ASCEND_USE_ASYNC_TRANSFER=1` |
|
||||
| `ASCEND_GLOBAL_RESOURCE_CONFIG` | Global resource configuration | - | `ASCEND_GLOBAL_RESOURCE_CONFIG="{\"fabric_memory.max_capacity\":32}"` |
|
||||
| `ASCEND_CONNECT_TIMEOUT` | Link establishment timeout in milliseconds | 3000 | `ASCEND_CONNECT_TIMEOUT=5000` |
|
||||
| `ASCEND_TRANSFER_TIMEOUT` | Data transfer timeout in milliseconds | 3000 | `ASCEND_TRANSFER_TIMEOUT=10000` |
|
||||
| `ASCEND_THREAD_POOL_SIZE` | Number of worker threads in the transfer thread pool | 8 (1 for buffer pool mode) | `ASCEND_THREAD_POOL_SIZE=16` |
|
||||
| `ASCEND_USE_SHORT_CONNECTION` | Enable short connection mode (disconnect after each transfer) | 0 (disabled) | `ASCEND_USE_SHORT_CONNECTION=1` |
|
||||
| `ASCEND_BUFFER_POOL` | Buffer pool configuration for intermediate transfer mode (BUFFER_NUM:BUFFER_SIZE_MB) | "0:0" (disabled) | `ASCEND_BUFFER_POOL=4:8` |
|
||||
| `ASCEND_BASE_PORT` | Base port for ADXL engine port allocation | 11000 | `ASCEND_BASE_PORT=20000` |
|
||||
| `HCCL_INTRA_ROCE_ENABLE` | Enable RDMA protocol for intra-node communication | 0 (disabled) | `HCCL_INTRA_ROCE_ENABLE=1` |
|
||||
| `HCCL_RDMA_TIMEOUT` | RDMA packet retransmission timeout coefficient | - | `HCCL_RDMA_TIMEOUT=14` |
|
||||
| `HCCL_RDMA_RETRY_CNT` | RDMA packet retransmission count | - | `HCCL_RDMA_RETRY_CNT=7` |
|
||||
|
||||
**Detailed Descriptions:**
|
||||
|
||||
- **ASCEND_AUTO_CONNECT**: Requires CANN 9.0 or later. Default is 0, recommended to enable on supported versions: link can be automatically disconnected when the remote end goes offline abnormally.
|
||||
- **ASCEND_ENABLE_USE_FABRIC_MEM**: Requires CANN 9.0+ and HDK 26.0+. Recommended when using Mooncake Store on supported A3 platforms: it can significantly improve transmission performance.
|
||||
- **ASCEND_USE_ASYNC_TRANSFER**: Requires CANN 8.5+. Enables HIXL asynchronous transfer mode, defaults to synchronous mode.
|
||||
- **ASCEND_GLOBAL_RESOURCE_CONFIG**: Configures HIXL global resources. Refer to HIXL documentation for `OPTION_GLOBAL_RESOURCE_CONFIG` settings.
|
||||
|
||||
### Important Notes
|
||||
|
||||
1. **Device Setup Required**: Before calling `TransferEngine.initialize()`, you must set the device (e.g., `torch.npu.set_device(0)`).
|
||||
|
|
@ -61,10 +88,16 @@ Complete command format is shown below:
|
|||
- Use `HCCL_RDMA_RETRY_CNT` to configure the RDMA NIC retransmission count
|
||||
- It is recommended to configure `ASCEND_TRANSFER_TIMEOUT` to be slightly larger than `retransmission_timeout * HCCL_RDMA_RETRY_CNT`
|
||||
|
||||
7. **Communication Protocol Selection**: Within A2 servers/A3 supernodes, the default communication protocol is HCCS. You can specify RDMA by setting `export HCCL_INTRA_ROCE_ENABLE=1`. For KV Cache transfer scenarios, RDMA transfer is recommended to avoid conflicts with model collective communication traffic that could impact inference performance.
|
||||
7. **Communication Protocol Selection**: Within A2 servers/A3 supernodes, the default communication protocol is HCCS. You can specify RDMA by setting `export HCCL_INTRA_ROCE_ENABLE=1`.
|
||||
|
||||
8. **RDMA Configuration**: When using the `RDMA` communication protocol, in scenarios where switch and NIC default configurations are inconsistent or traffic planning is required, you may need to modify the RDMA NIC's Traffic Class and Service Level configuration:
|
||||
- Set Traffic Class using the `ASCEND_RDMA_TC` environment variable
|
||||
- Set Service Level using the `ASCEND_RDMA_SL` environment variable
|
||||
|
||||
9. **Buffer Pool Configuration**: Under the default 4KB page table configuration, the registrable host memory is approximately 20GB using RDMA. Additionally, HCCS does not currently support host memory transmission. Under these two constrained scenarios, transmission can be achieved by using intermediate buffers. The specific way to enable this is by configuring the ASCEND_BUFFER_POOL environment variable in the format BUFFER_NUM:BUFFER_SIZE (in MB). The recommended size is 4:8, which can be adjusted to the most suitable configuration based on actual scenarios.
|
||||
9. **Buffer Pool Configuration**: Under constrained scenarios, transmission can be achieved by using intermediate buffers. The specific way to enable this is by configuring the ASCEND_BUFFER_POOL environment variable in the format BUFFER_NUM:BUFFER_SIZE (in MB). The recommended size is 4:8, which can be adjusted to the most suitable configuration based on actual scenarios.
|
||||
|
||||
10. **Async transfer**: The asynchronous transfer mode can be enabled by configuring the ASCEND_USE_ASYNC_TRANSFER environment variable.
|
||||
|
||||
12. **Fabric Memory mode**: On the A3, with the latest drivers and CANN installed, when using Mooncake store, the ASCEND_ENABLE_USE_FABRIC_MEM environment variable can be set to enable fabric memory transfer mode (which allows direct access remote HOST memory).
|
||||
|
||||
13. **Auto Connect**: The auto connect feature can be enabled by configuring the `ASCEND_AUTO_CONNECT` environment variable. The default value is 0 (disabled).
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
|
||||
The source code path for Ascend Transport is `Mooncake/mooncake-transfer-engine/src/transport/ascend_transport`, which also includes automated build scripts and the README file.
|
||||
|
||||
**ASCEND TRANSPORT is scheduled for deprecation, please use [ASCEND DIRECT TRANSPORT](./ascend_direct_transport.md) on ASCEND platform. **
|
||||
|
||||
## Overview
|
||||
|
||||
Ascend Transport is a high-performance zero-copy NPU data transfer library with one-sided semantics, directly compatible with Mooncake Transfer Engine. To compile and use the Ascend Transport library, please set the `USE_ASCEND` flag to `"ON"` in the `mooncake-common/common.cmake` file.
|
||||
|
|
@ -142,7 +144,7 @@ Therefore, in testing:
|
|||
Watch the log produced by `mooncake-transfer-engine/src/transfer_engine.cpp`; you should see a line similar to
|
||||
```
|
||||
Transfer Engine RPC using <protocol> listening on <IP>:<actual-port>
|
||||
```
|
||||
```
|
||||
Note the **actual port** the target node is listening on.
|
||||
|
||||
2. **Edit the initiator’s launch command**:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,556 @@
|
|||
# Transfer Engine C++ API Reference
|
||||
|
||||
## Overview
|
||||
This page summarizes the C++ APIs in `mooncake-transfer-engine/include/transfer_engine.h`.
|
||||
It follows the same narrative style as the design doc, while covering the full set of public C++ APIs.
|
||||
|
||||
For conceptual background, see [Transfer Engine](index.md).
|
||||
|
||||
**Core APIs vs Advanced APIs**
|
||||
- Core APIs form the minimal path to move data: initialize the engine, register memory, open segments, submit transfers, and query status.
|
||||
- Advanced APIs are optional; they help with transport control, notifications, metadata/cache maintenance, topology discovery, and debugging.
|
||||
|
||||
## Core APIs
|
||||
|
||||
### Core Usage Path (C++)
|
||||
Most integrations follow a short, repeatable path: initialize the engine, register local memory, open a target segment, submit a batch transfer, poll status, and finally free resources (batch, segment, memory). A minimal example is shown below.
|
||||
|
||||
```cpp
|
||||
#include "transfer_engine.h"
|
||||
|
||||
using mooncake::TransferEngine;
|
||||
using mooncake::TransferRequest;
|
||||
using mooncake::TransferStatus;
|
||||
|
||||
TransferEngine engine(true);
|
||||
engine.init("etcd://127.0.0.1:2379", "node0");
|
||||
|
||||
void *local_addr = /* local buffer address */;
|
||||
size_t length = /* bytes to transfer */;
|
||||
uint64_t remote_addr = /* remote address or file offset */;
|
||||
|
||||
engine.registerLocalMemory(local_addr, length, "cpu:0", true);
|
||||
|
||||
auto segment = engine.openSegment("node1");
|
||||
auto batch = engine.allocateBatchID(1);
|
||||
|
||||
TransferRequest req{};
|
||||
req.opcode = TransferRequest::WRITE;
|
||||
req.source = local_addr;
|
||||
req.target_id = segment;
|
||||
req.target_offset = remote_addr;
|
||||
req.length = length;
|
||||
|
||||
engine.submitTransfer(batch, {req});
|
||||
|
||||
TransferStatus status{};
|
||||
engine.getTransferStatus(batch, 0, status);
|
||||
|
||||
engine.freeBatchID(batch);
|
||||
engine.closeSegment(segment);
|
||||
engine.unregisterLocalMemory(local_addr);
|
||||
```
|
||||
|
||||
### Constructors
|
||||
|
||||
```cpp
|
||||
TransferEngine(bool auto_discover = false);
|
||||
TransferEngine(bool auto_discover, const std::vector<std::string>& filter);
|
||||
```
|
||||
|
||||
Constructs a Transfer Engine instance. `auto_discover` enables topology discovery; `filter` constrains discovery or device selection with a whitelist.
|
||||
|
||||
### Data Transfer
|
||||
|
||||
#### TransferEngine::TransferRequest
|
||||
|
||||
The core API provided by Mooncake Transfer Engine is submitting a group of asynchronous `TransferRequest` tasks through the `submitTransfer` interface, and querying their status through the `getTransferStatus` interface. Each `TransferRequest` specifies reading or writing a continuous data space of `length` starting from the local starting address `source`, to the position starting at `target_offset` in the segment corresponding to `target_id`.
|
||||
|
||||
The `TransferRequest` structure is defined as follows:
|
||||
|
||||
```cpp
|
||||
struct TransferRequest
|
||||
{
|
||||
enum OpCode { READ, WRITE };
|
||||
OpCode opcode;
|
||||
void *source;
|
||||
SegmentID target_id; // The ID of the target segment, which may correspond to local or remote DRAM/VRAM/NVMeof, with the specific routing logic hidden
|
||||
uint64_t target_offset;
|
||||
size_t length;
|
||||
int advise_retry_cnt = 0;
|
||||
};
|
||||
```
|
||||
|
||||
- `opcode` takes the values `READ` or `WRITE`. `READ` indicates that data is copied from the target address indicated by `<target_id, target_offset>` to the local starting address `source`; `WRITE` indicates that data is copied from `source` to the address indicated by `<target_id, target_offset>`.
|
||||
- `source` represents the DRAM/VRAM buffer managed by the current `TransferEngine`, which must have been registered in advance by the `registerLocalMemory` interface.
|
||||
- `target_id` represents the segment ID of the transfer target. The segment ID is obtained using the `openSegment` interface. Segments are divided into the following types:
|
||||
- RAM space type, covering DRAM/VRAM. As mentioned earlier, there is only one segment under the same process (or `TransferEngine` instance), which contains various types of Buffers (DRAM/VRAM). In this case, the segment name passed to the `openSegment` interface is equivalent to the server hostname. `target_offset` is the virtual address of the target server.
|
||||
- NVMeOF space type, where each file corresponds to a segment. In this case, the segment name passed to the `openSegment` interface is equivalent to the unique identifier of the file. `target_offset` is the offset of the target file.
|
||||
- `length` represents the amount of data transferred. TransferEngine may further split this into multiple read/write requests internally.
|
||||
|
||||
#### TransferEngine::allocateBatchID
|
||||
|
||||
```cpp
|
||||
BatchID allocateBatchID(size_t batch_size);
|
||||
```
|
||||
|
||||
Allocates a `BatchID`. A maximum of `batch_size` `TransferRequest`s can be submitted under the same `BatchID`.
|
||||
|
||||
- `batch_size`: The maximum number of `TransferRequest`s that can be submitted under the same `BatchID`;
|
||||
- Return value: If successful, returns a `BatchID`; on failure, returns an invalid handle (for example `INVALID_BATCH_ID`).
|
||||
|
||||
#### TransferEngine::submitTransfer
|
||||
|
||||
```cpp
|
||||
Status submitTransfer(BatchID batch_id,
|
||||
const std::vector<TransferRequest> &entries);
|
||||
```
|
||||
|
||||
Submits new `TransferRequest` tasks to `batch_id`. The task is asynchronously submitted to the background thread pool. The total number of `entries` accumulated under the same `batch_id` should not exceed the `batch_size` defined at creation.
|
||||
|
||||
- `batch_id`: The `BatchID` it belongs to;
|
||||
- `entries`: Array of `TransferRequest`;
|
||||
- Return value: If successful, returns an OK status; otherwise, returns a non-OK status.
|
||||
|
||||
#### TransferEngine::getTransferStatus
|
||||
|
||||
```cpp
|
||||
enum TransferStatusEnum
|
||||
{
|
||||
WAITING, // In the transfer phase
|
||||
PENDING, // Not supported
|
||||
INVALID, // Invalid parameters
|
||||
CANCELED, // Not supported
|
||||
COMPLETED, // Transfer completed
|
||||
TIMEOUT, // Not supported
|
||||
FAILED // Transfer failed even after retries
|
||||
};
|
||||
struct TransferStatus {
|
||||
TransferStatusEnum s;
|
||||
size_t transferred_bytes; // How much data has been successfully transferred (not necessarily an accurate value, but it is a lower bound)
|
||||
};
|
||||
Status getTransferStatus(BatchID batch_id, size_t task_id, TransferStatus &status)
|
||||
```
|
||||
|
||||
Obtains the running status of the `TransferRequest` with `task_id` in `batch_id`.
|
||||
|
||||
- `batch_id`: The `BatchID` it belongs to;
|
||||
- `task_id`: The sequence number of the `TransferRequest` to query;
|
||||
- `status`: Output Transfer status;
|
||||
- Return value: If successful, returns an OK status; otherwise, returns a non-OK status.
|
||||
|
||||
#### TransferEngine::getBatchTransferStatus
|
||||
|
||||
```cpp
|
||||
Status getBatchTransferStatus(BatchID batch_id, TransferStatus& status);
|
||||
```
|
||||
|
||||
Obtains the aggregated status of the batch and the total transferred bytes.
|
||||
|
||||
- `batch_id`: The `BatchID` it belongs to;
|
||||
- `status`: Output Transfer status;
|
||||
- Return value: If successful, returns an OK status; otherwise, returns a non-OK status.
|
||||
|
||||
#### TransferEngine::freeBatchID
|
||||
|
||||
```cpp
|
||||
Status freeBatchID(BatchID batch_id);
|
||||
```
|
||||
|
||||
Recycles `BatchID`, and subsequent operations on `submitTransfer` and `getTransferStatus` are undefined. If there are still `TransferRequest`s pending completion in the `BatchID`, the operation is refused.
|
||||
|
||||
- `batch_id`: The `BatchID` it belongs to;
|
||||
- Return value: If successful, returns an OK status; otherwise, returns a non-OK status.
|
||||
|
||||
### Space Registration
|
||||
|
||||
For the RDMA transfer process, the source pointer `TransferRequest::source` must be registered in advance as an RDMA readable/writable Memory Region space, that is, included as part of the RAM Segment of the current process. Therefore, the following functions are needed:
|
||||
|
||||
#### TransferEngine::registerLocalMemory
|
||||
|
||||
```cpp
|
||||
int registerLocalMemory(void *addr,
|
||||
size_t length,
|
||||
const std::string& location = kWildcardLocation,
|
||||
bool remote_accessible = true,
|
||||
bool update_metadata = true);
|
||||
```
|
||||
|
||||
Registers a space starting at address `addr` with a length of `length` on the local DRAM/VRAM.
|
||||
|
||||
- `addr`: The starting address of the registration space;
|
||||
- `length`: The length of the registration space;
|
||||
- `location`: The `device` corresponding to this memory segment, such as `cuda:0` indicating the GPU device, `cpu:0` indicating the CPU socket, by matching with the network card priority order table (see `installTransport`), the preferred network card is identified. You can also use `*`, Transfer Engine will try to automatically recognize the `device` corresponding to `addr`, if it fails to recognize the device, it will print a `WARNING` level log and use all network cards, no preferred network cards.
|
||||
- `remote_accessible`: Indicates whether this memory can be accessed by remote nodes.
|
||||
- `update_metadata`: Whether to publish the registration to the metadata service.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::unregisterLocalMemory
|
||||
|
||||
```cpp
|
||||
int unregisterLocalMemory(void *addr, bool update_metadata = true);
|
||||
```
|
||||
|
||||
Unregisters the region.
|
||||
|
||||
- `addr`: The starting address of the registration space;
|
||||
- `update_metadata`: Whether to publish the unregistration to the metadata service.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::registerLocalMemoryBatch
|
||||
|
||||
```cpp
|
||||
int registerLocalMemoryBatch(const std::vector<BufferEntry>& buffer_list,
|
||||
const std::string& location);
|
||||
```
|
||||
|
||||
Registers multiple buffers in one call to reduce registration overhead.
|
||||
|
||||
- `buffer_list`: A list of `{addr, length}` entries.
|
||||
- `location`: The `device` corresponding to these buffers.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::unregisterLocalMemoryBatch
|
||||
|
||||
```cpp
|
||||
int unregisterLocalMemoryBatch(const std::vector<void*>& addr_list);
|
||||
```
|
||||
|
||||
Unregisters multiple buffers in one call.
|
||||
|
||||
- `addr_list`: A list of buffer start addresses.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
### Segment Management and Metadata Format
|
||||
|
||||
TransferEngine provides the `openSegment` function, which obtains a `SegmentHandle` for subsequent `Transport` transfers.
|
||||
```cpp
|
||||
SegmentHandle openSegment(const std::string& segment_name);
|
||||
```
|
||||
|
||||
- `segment_name`: The unique identifier of the segment. For RAM Segment, this needs to be consistent with the `server_name` filled in by the peer process when initializing the TransferEngine object.
|
||||
- Return value: If successful, returns the corresponding `SegmentHandle`; otherwise, returns an invalid handle.
|
||||
|
||||
```cpp
|
||||
int closeSegment(SegmentHandle segment_id);
|
||||
```
|
||||
|
||||
- `segment_id`: The unique identifier of the segment.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::removeLocalSegment
|
||||
|
||||
```cpp
|
||||
int removeLocalSegment(const std::string& segment_name);
|
||||
```
|
||||
|
||||
Removes local segment metadata entries.
|
||||
|
||||
- `segment_name`: The local segment name to remove.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::CheckSegmentStatus
|
||||
|
||||
```cpp
|
||||
Status CheckSegmentStatus(SegmentID sid);
|
||||
```
|
||||
|
||||
Checks whether a segment is reachable and valid.
|
||||
|
||||
- `sid`: The segment identifier.
|
||||
- Return value: If successful, returns an OK status; otherwise, returns a non-OK status.
|
||||
|
||||
<details>
|
||||
<summary><strong>Metadata Format</strong></summary>
|
||||
|
||||
```
|
||||
// Used to find the communicable address and exposed rpc port based on server_name.
|
||||
// Created: when calling TransferEngine::init().
|
||||
// Deleted: when TransferEngine is destructed.
|
||||
Key = mooncake/rpc_meta/[server_name]
|
||||
Value = {
|
||||
'ip_or_host_name': 'node01'
|
||||
'rpc_port': 12345
|
||||
}
|
||||
|
||||
// For segments, the key naming method of mooncake/[proto]/[segment_name] is used, and the segment name can use the Server Name.
|
||||
// A segment corresponds to a machine, and a buffer corresponds to different segments of memory or different files or different disks on the machine. Different buffers of the same segment are in the same fault domain.
|
||||
|
||||
// RAM Segment, used by RDMA Transport to obtain transfer information.
|
||||
// Created: command line tool register.py, at this time buffers are empty, only fill in the information that can be known in advance.
|
||||
// Modified: TransferEngine at runtime through register / unregister to add or delete Buffer.
|
||||
Key = mooncake/ram/[segment_name]
|
||||
Value = {
|
||||
'server_name': server_name,
|
||||
'protocol': rdma,
|
||||
'devices': [
|
||||
{ 'name': 'mlx5_2', 'lid': 17, 'gid': 'fe:00:...' },
|
||||
{ 'name': 'mlx5_3', 'lid': 22, 'gid': 'fe:00:...' }
|
||||
],
|
||||
'priority_matrix': {
|
||||
"cpu:0": [["mlx5_2"], ["mlx5_3"]],
|
||||
"cpu:1": [["mlx5_3"], ["mlx5_2"]],
|
||||
"cuda:0": [["mlx5_2"], ["mlx5_3"]],
|
||||
},
|
||||
'buffers': [
|
||||
{
|
||||
'name': 'cpu:0',
|
||||
'addr': 0x7fa16bdf5000,
|
||||
'length': 1073741824,
|
||||
'rkey': [1fe000, 1fdf00, ...], // The length is the same as the number of elements in the 'devices' field
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
// Created: command line tool register.py, determine the file path that can be mounted.
|
||||
// Modified: command line tool mount.py, add a mapping of the machine mounting the file to the file path on the mounting machine to the buffers.local_path_map.
|
||||
Key = mooncake/nvmeof/[segment_name]
|
||||
Value = {
|
||||
'server_name': server_name,
|
||||
'protocol': nvmeof,
|
||||
'buffers':[{
|
||||
'length': 1073741824,
|
||||
'file_path': "/mnt/nvme0" // The file path on this machine
|
||||
'local_path_map': {
|
||||
"node01": "/mnt/transfer_engine/node01/nvme0", // The machine mounting the file -> The file path on the mounting machine
|
||||
...
|
||||
},
|
||||
},
|
||||
{
|
||||
'length': 1073741824,
|
||||
'file_path': "/mnt/nvme1",
|
||||
'local_path_map': {
|
||||
"node02": "/mnt/transfer_engine/node02/nvme1",
|
||||
...
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
</details>
|
||||
|
||||
### HTTP Metadata Server
|
||||
|
||||
The HTTP server should implement three following RESTful APIs, while the metadata server configured to `http://host:port/metadata` as an example:
|
||||
|
||||
1. `GET /metadata?key=$KEY`: Get the metadata corresponding to `$KEY`.
|
||||
2. `PUT /metadata?key=$KEY`: Update the metadata corresponding to `$KEY` to the value of the request body.
|
||||
3. `DELETE /metadata?key=$KEY`: Delete the metadata corresponding to `$KEY`.
|
||||
|
||||
For specific implementation, refer to the demo service implemented in Golang at [mooncake-transfer-engine/example/http-metadata-server](../../../mooncake-transfer-engine/example/http-metadata-server).
|
||||
|
||||
### Initialization
|
||||
|
||||
TransferEngine needs to be initialized by calling the `init` method before further actions:
|
||||
```cpp
|
||||
TransferEngine();
|
||||
|
||||
int init(const std::string &metadata_conn_string,
|
||||
const std::string &local_server_name);
|
||||
```
|
||||
There is also an extended form to override RPC binding:
|
||||
```cpp
|
||||
int init(const std::string &metadata_conn_string,
|
||||
const std::string &local_server_name,
|
||||
const std::string &ip_or_host_name,
|
||||
uint64_t rpc_port);
|
||||
```
|
||||
- `metadata_conn_string`: Connecting string of metadata storage servers, i.e., the IP address/hostname of `etcd`/`redis` or the URI of the http service.
|
||||
The general form is `[proto]://[hostname:port]`. For example, the following metadata server addresses are legal:
|
||||
- Using `etcd` as a metadata storage service: `"10.0.0.1:2379"` or `"etcd://10.0.0.1:2379"`.
|
||||
- Using `redis` as a metadata storage service: `"redis://10.0.0.1:6379"`
|
||||
- Using `http` as a metadata storage service: `"http://10.0.0.1:8080/metadata"`
|
||||
|
||||
- `local_server_name`: The local server name, ensuring uniqueness within the cluster. It also serves as the name of the RAM Segment that other nodes refer to the current instance (i.e., Segment Name).
|
||||
- `ip_or_host_name`: Optional explicit bind address or hostname for the RPC service.
|
||||
- `rpc_port`: Optional explicit RPC port.
|
||||
|
||||
```cpp
|
||||
~TransferEngine();
|
||||
```
|
||||
|
||||
Reclaims all allocated resources and also deletes the global metadata server information.
|
||||
|
||||
#### TransferEngine::freeEngine
|
||||
|
||||
```cpp
|
||||
int freeEngine();
|
||||
```
|
||||
|
||||
Releases resources early without waiting for destruction.
|
||||
|
||||
## Advanced APIs
|
||||
These APIs are optional; use them when you need manual transport control, coordination notifications, metadata/cache refresh, topology discovery, or debugging. They are not required for the basic data path.
|
||||
|
||||
### Multi-Transport Management
|
||||
|
||||
The `TransferEngine` class internally manages multiple backend `Transport` classes.
|
||||
And it will discover the topology between CPU/CUDA and RDMA devices automatically
|
||||
(more device types are working in progress, feedbacks are welcome when the automatic discovery mechanism is not accurate),
|
||||
and it will install `Transport` automatically based on the topology.
|
||||
|
||||
#### TransferEngine::installTransport
|
||||
|
||||
```cpp
|
||||
Transport* installTransport(const std::string& proto, void** args);
|
||||
```
|
||||
|
||||
Installs a transport backend explicitly.
|
||||
|
||||
- `proto`: Transport protocol name, such as `rdma`, `tcp`, or `nvmeof`.
|
||||
- `args`: Transport-specific arguments.
|
||||
> Note: In TENT, `installTransport` is not exposed (removed from the public API, including compatibility surfaces). Transport selection is internal to TENT.
|
||||
|
||||
#### TransferEngine::uninstallTransport
|
||||
|
||||
```cpp
|
||||
int uninstallTransport(const std::string& proto);
|
||||
```
|
||||
|
||||
Uninstalls a transport backend.
|
||||
|
||||
- `proto`: Transport protocol name.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::getTransport
|
||||
|
||||
```cpp
|
||||
Transport* getTransport(const std::string& proto);
|
||||
```
|
||||
|
||||
Returns a transport instance by protocol name, mainly for advanced inspection or debugging.
|
||||
|
||||
### Notifications
|
||||
|
||||
TransferEngine can send and receive lightweight notifications across segments to coordinate data movement.
|
||||
|
||||
#### TransferEngine::submitTransferWithNotify
|
||||
|
||||
```cpp
|
||||
Status submitTransferWithNotify(BatchID batch_id,
|
||||
const std::vector<TransferRequest>& entries,
|
||||
TransferMetadata::NotifyDesc notify_msg);
|
||||
```
|
||||
|
||||
Submits a batch transfer and delivers a notification upon completion.
|
||||
|
||||
- `notify_msg`: A `{name, msg}` payload delivered to the receiver.
|
||||
- Typical use: signal that a transferred buffer or KV-cache slice is ready to consume on the receiver side.
|
||||
|
||||
#### TransferEngine::getNotifies
|
||||
|
||||
```cpp
|
||||
int getNotifies(std::vector<TransferMetadata::NotifyDesc>& notifies);
|
||||
```
|
||||
|
||||
Gets pending notifications from peers.
|
||||
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
- Typical use: receiver-side polling loop to trigger follow-up actions (e.g., attaching the transferred buffer to a scheduler or cache).
|
||||
|
||||
#### TransferEngine::sendNotifyByID
|
||||
|
||||
```cpp
|
||||
int sendNotifyByID(SegmentID target_id, TransferMetadata::NotifyDesc notify_msg);
|
||||
```
|
||||
|
||||
Sends a notification to a specific segment by ID.
|
||||
- Typical use: control-plane message such as "ready", "invalidate", or "retry" tied to a known segment handle.
|
||||
|
||||
#### TransferEngine::sendNotifyByName
|
||||
|
||||
```cpp
|
||||
int sendNotifyByName(std::string remote_agent, TransferMetadata::NotifyDesc notify_msg);
|
||||
```
|
||||
|
||||
Sends a notification to a remote agent by name.
|
||||
- Typical use: best-effort coordination when you only know the peer's name (segment name), not the handle.
|
||||
|
||||
### Segment Cache and Metadata Access
|
||||
|
||||
#### TransferEngine::syncSegmentCache
|
||||
|
||||
```cpp
|
||||
int syncSegmentCache(const std::string& segment_name = "");
|
||||
```
|
||||
|
||||
Synchronizes local segment cache from the metadata service.
|
||||
|
||||
- `segment_name`: If empty, refreshes all segments; otherwise, refreshes the specified segment.
|
||||
- Typical use: when peers dynamically register/unregister buffers and you need to refresh before opening or submitting transfers.
|
||||
|
||||
#### TransferEngine::getMetadata
|
||||
|
||||
```cpp
|
||||
std::shared_ptr<TransferMetadata> getMetadata();
|
||||
```
|
||||
|
||||
Returns the metadata subsystem instance for advanced workflows.
|
||||
- Typical use: advanced integration or debugging of metadata content beyond the standard APIs.
|
||||
|
||||
### Topology and Discovery
|
||||
|
||||
#### TransferEngine::setAutoDiscover
|
||||
|
||||
```cpp
|
||||
void setAutoDiscover(bool auto_discover);
|
||||
```
|
||||
|
||||
Enables or disables topology discovery.
|
||||
- Typical use: controlled environments or debugging where auto discovery is not desired.
|
||||
|
||||
#### TransferEngine::setWhitelistFilters
|
||||
|
||||
```cpp
|
||||
void setWhitelistFilters(std::vector<std::string>&& filters);
|
||||
```
|
||||
|
||||
Sets the device whitelist used during topology discovery.
|
||||
- Typical use: restrict transfers to a subset of NICs or GPUs for performance isolation or testing.
|
||||
|
||||
#### TransferEngine::numContexts
|
||||
|
||||
```cpp
|
||||
int numContexts() const;
|
||||
```
|
||||
|
||||
Returns the number of active transport contexts.
|
||||
- Typical use: introspection and resource monitoring.
|
||||
|
||||
#### TransferEngine::getLocalTopology
|
||||
|
||||
```cpp
|
||||
std::shared_ptr<Topology> getLocalTopology();
|
||||
```
|
||||
|
||||
Returns the local topology model used for path selection.
|
||||
- Typical use: diagnose path selection or export topology to a scheduler.
|
||||
|
||||
### Introspection and Safety
|
||||
|
||||
#### TransferEngine::getLocalIpAndPort
|
||||
|
||||
```cpp
|
||||
std::string getLocalIpAndPort();
|
||||
```
|
||||
|
||||
Returns the resolved local RPC address.
|
||||
- Typical use: publish the resolved address to external components or logs.
|
||||
|
||||
#### TransferEngine::getRpcPort
|
||||
|
||||
```cpp
|
||||
int getRpcPort();
|
||||
```
|
||||
|
||||
Returns the active RPC port.
|
||||
- Typical use: verify port binding when dynamic ports are used.
|
||||
|
||||
#### TransferEngine::checkOverlap
|
||||
|
||||
```cpp
|
||||
bool checkOverlap(void* addr, uint64_t length);
|
||||
```
|
||||
|
||||
Checks whether a given address range overlaps with existing registered buffers.
|
||||
- Typical use: validate registration plans before calling `registerLocalMemory` in complex buffer managers.
|
||||
|
|
@ -0,0 +1,545 @@
|
|||
# AWS EFA Transport for Mooncake
|
||||
|
||||
This document describes how to build and use Mooncake with AWS Elastic Fabric Adapter (EFA) support using libfabric.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### 1. AWS EFA Driver and libfabric
|
||||
|
||||
EFA driver and libfabric should be pre-installed on AWS instances with EFA support (e.g., p6-b300.48xlarge, p6-b200.48xlarge, p5en.48xlarge, p5e.48xlarge, p5.48xlarge).
|
||||
|
||||
Verify installation:
|
||||
```bash
|
||||
# Check EFA devices
|
||||
fi_info -p efa
|
||||
|
||||
# Verify libfabric location
|
||||
ls /opt/amazon/efa/lib/libfabric.so
|
||||
ls /opt/amazon/efa/include/rdma/fabric.h
|
||||
```
|
||||
|
||||
If not installed, follow [AWS EFA documentation](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start.html).
|
||||
|
||||
### 2. Build Dependencies
|
||||
|
||||
Clone the repository and install all dependencies:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/kvcache-ai/Mooncake.git
|
||||
cd Mooncake
|
||||
sudo ./dependencies.sh -y
|
||||
```
|
||||
|
||||
This installs all system packages, git submodules (including pybind11 and yalantinglibs), and Go.
|
||||
|
||||
**Additional EFA-specific dependencies** (not covered by `dependencies.sh`):
|
||||
|
||||
```bash
|
||||
# gflags is needed by transfer_engine_bench and EFA unit tests
|
||||
sudo apt-get install -y libgflags-dev
|
||||
```
|
||||
|
||||
> **Note:** The EFA driver and libfabric are **not** installed by `dependencies.sh`. They must be pre-installed on the instance (see section 1 above).
|
||||
|
||||
## Building Mooncake with EFA Support
|
||||
|
||||
### 1. Build with EFA Enabled
|
||||
|
||||
**GPU memory transfers (e.g., KV cache in vLLM):**
|
||||
|
||||
```bash
|
||||
mkdir build && cd build
|
||||
|
||||
cmake .. \
|
||||
-DUSE_EFA=ON \
|
||||
-DUSE_CUDA=ON \
|
||||
-DCMAKE_BUILD_TYPE=RelWithDebInfo
|
||||
|
||||
make -j$(nproc)
|
||||
```
|
||||
|
||||
> **Note:** `-DUSE_CUDA=ON` is required when transferring GPU memory (e.g., KV cache in vLLM). Without it, the TCP transport (used as fallback when `mooncake_protocol` is set to `"tcp"`) cannot detect GPU memory and will fail with "Bad address" (EFAULT) errors.
|
||||
|
||||
**CPU memory transfers only (no GPU dependency):**
|
||||
|
||||
```bash
|
||||
mkdir build && cd build
|
||||
|
||||
cmake .. \
|
||||
-DUSE_EFA=ON \
|
||||
-DUSE_CUDA=OFF \
|
||||
-DCMAKE_BUILD_TYPE=RelWithDebInfo
|
||||
|
||||
make -j$(nproc)
|
||||
```
|
||||
|
||||
> **Note:** With `-DUSE_CUDA=OFF`, the benchmark tool uses DRAM buffers allocated via `numa_alloc_onnode`. This is useful for measuring EFA transport throughput independently of GPU hardware.
|
||||
|
||||
### 2. Install Python Package
|
||||
|
||||
```bash
|
||||
# Copy built modules to wheel directory
|
||||
cp mooncake-integration/engine.cpython-*.so ../mooncake-wheel/mooncake/
|
||||
cp mooncake-integration/store.cpython-*.so ../mooncake-wheel/mooncake/
|
||||
cp mooncake-common/libasio.so ../mooncake-wheel/mooncake/
|
||||
|
||||
# Install with pip
|
||||
pip install -e ../mooncake-wheel --no-build-isolation
|
||||
```
|
||||
|
||||
## Verification
|
||||
|
||||
Test EFA transport initialization:
|
||||
|
||||
```python
|
||||
from mooncake.engine import TransferEngine
|
||||
|
||||
te = TransferEngine()
|
||||
result = te.initialize('127.0.0.1', 'P2PHANDSHAKE', 'efa', '')
|
||||
print(f'Initialize result: {result}') # Should be 0
|
||||
|
||||
# You should see logs like:
|
||||
# EFA device (libfabric): rdmap79s0, domain: rdmap79s0-rdm, provider: efa
|
||||
```
|
||||
|
||||
## Unit Tests
|
||||
|
||||
Run the EFA transport unit tests (requires EFA hardware):
|
||||
|
||||
```bash
|
||||
./build/mooncake-transfer-engine/tests/efa_transport_test
|
||||
```
|
||||
|
||||
The test suite includes:
|
||||
|
||||
| Test | Description |
|
||||
|------|-------------|
|
||||
| `InstallTransport` | Verify EFA transport installation |
|
||||
| `LoopbackWrite` | Loopback write operation |
|
||||
| `WriteAndRead` | Write then read with data integrity check |
|
||||
| `MultiWrite` | Batch write (16 requests) |
|
||||
| `StressMultipleBatches` | Stress test (20 batches x 8 requests) |
|
||||
|
||||
You can also run all unit tests via CTest:
|
||||
|
||||
```bash
|
||||
cd build && ctest --output-on-failure
|
||||
```
|
||||
|
||||
Environment variables for test configuration:
|
||||
|
||||
```bash
|
||||
export MC_METADATA_SERVER=P2PHANDSHAKE # default
|
||||
export MC_LOCAL_SERVER_NAME=127.0.0.1:12345 # default
|
||||
```
|
||||
|
||||
## Performance Benchmark
|
||||
|
||||
Use `transfer_engine_bench` to measure EFA transport throughput between two nodes.
|
||||
|
||||
### Target Node (receiver)
|
||||
|
||||
```bash
|
||||
./build/mooncake-transfer-engine/example/transfer_engine_bench \
|
||||
--mode=target \
|
||||
--protocol=efa \
|
||||
--metadata_server=P2PHANDSHAKE
|
||||
```
|
||||
|
||||
### Initiator Node (sender)
|
||||
|
||||
```bash
|
||||
./build/mooncake-transfer-engine/example/transfer_engine_bench \
|
||||
--mode=initiator \
|
||||
--protocol=efa \
|
||||
--metadata_server=P2PHANDSHAKE \
|
||||
--segment_id=<target_hostname>:<target_port> \
|
||||
--operation=write \
|
||||
--duration=10 \
|
||||
--threads=8 \
|
||||
--block_size=65536 \
|
||||
--batch_size=128 \
|
||||
--buffer_size=1073741824 \
|
||||
--report_unit=GB
|
||||
```
|
||||
|
||||
> **Tip:** For CPU-to-CPU benchmarks, prepend `CUDA_VISIBLE_DEVICES=""` to prevent the CUDA runtime from being initialized. Without it, `nvidia-smi` may show GPU memory usage (due to CUDA context initialization) even though the benchmark only uses DRAM.
|
||||
|
||||
Replace `<target_hostname>:<target_port>` with the target node's address shown in the target's startup log (e.g., `ip-172-31-29-226:12345`).
|
||||
|
||||
### Key Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--block_size` | 65536 | Bytes per transfer request |
|
||||
| `--batch_size` | 128 | Requests per batch |
|
||||
| `--threads` | 12 | Concurrent submission threads |
|
||||
| `--buffer_size` | 1 GB | Total buffer size (per GPU when `--gpu_id=-1`) |
|
||||
| `--duration` | 10 | Test duration in seconds |
|
||||
| `--operation` | write | `read` or `write` |
|
||||
| `--report_unit` | GB | `GB\|GiB\|Gb\|MB\|MiB\|Mb` |
|
||||
| `--gpu_id` | 0 | GPU device ID; `-1` to use all GPUs (requires `-DUSE_CUDA=ON`) |
|
||||
|
||||
| Environment Variable | Default | Description |
|
||||
|---------------------|---------|-------------|
|
||||
| `MC_SLICE_SIZE` | 65536 | Slice size for RDMA transport. **Not used by EFA transport** (see note below). |
|
||||
| `MC_EFA_STRIPING_THRESHOLD` | 2097152 | Transfers larger than this (bytes) are striped across all NICs |
|
||||
|
||||
> **Note on EFA slicing:** Unlike RDMA transport which splits every transfer into fixed `MC_SLICE_SIZE` chunks, EFA transport uses a different strategy: transfers ≤ `MC_EFA_STRIPING_THRESHOLD` (default 2MB) are sent as a **single `fi_write`/`fi_read`** whose size equals `block_size`; transfers larger than the threshold are striped across all NICs (one chunk per NIC). This means **`block_size` directly determines per-operation size** and is the key tuning parameter for EFA, while `MC_SLICE_SIZE` has no effect.
|
||||
|
||||
> **Note:** `buffer_size` must be >= `block_size * batch_size * threads`. The benchmark auto-adjusts if too small.
|
||||
|
||||
### Benchmark Results
|
||||
|
||||
#### p6-b200.48xlarge (B200, 8 EFA × 400 Gbps)
|
||||
|
||||
Tested on two p6-b200.48xlarge instances in the same AWS placement group.
|
||||
|
||||
**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs):
|
||||
|
||||
| Configuration | Write | Read |
|
||||
|---------------|-------|------|
|
||||
| block=1MB, threads=32, batch=64, buf=2GB/GPU | 285-296 GB/s | 312 GB/s |
|
||||
| **block=1MB, threads=16, batch=128, buf=2GB/GPU** | **302 GB/s** | **313 GB/s** |
|
||||
|
||||
**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
|
||||
|
||||
| Configuration | Write | Read |
|
||||
|---------------|-------|------|
|
||||
| block=1MB, threads=32, batch=128, buf=4GB | **222 GB/s** (stable over 6 runs) | **226 GB/s** |
|
||||
|
||||
<details>
|
||||
<summary>CPU Parameter Tuning History (p6-b200)</summary>
|
||||
|
||||
Earlier CPU-to-CPU tuning results (before EFA striping optimization, when `MC_SLICE_SIZE` was still used by EFA):
|
||||
|
||||
| block_size | threads | batch_size | MC_SLICE_SIZE | Throughput |
|
||||
|-----------|---------|------------|---------------|-----------|
|
||||
| 64KB | 8 | 128 | default (64KB) | 69.47 GB/s |
|
||||
| 128KB | 32 | 128 | default | 92.33 GB/s |
|
||||
| 128KB | 32 | 128 | 256KB | 156.18 GB/s |
|
||||
| 128KB | 48 | 128 | 256KB | 160.34 GB/s |
|
||||
|
||||
> **Note:** These results predate the EFA striping optimization. With the current code, `MC_SLICE_SIZE` no longer affects EFA performance. Use `--block_size=1048576` (1MB) instead, which achieves 222 GB/s.
|
||||
|
||||
</details>
|
||||
|
||||
#### p6-b300.48xlarge (B300, 16 EFA × 400 Gbps)
|
||||
|
||||
Tested on two p6-b300.48xlarge instances (Intel Xeon Platinum 8559C, 8× B300, 16 EFA devices) in the same AWS placement group.
|
||||
|
||||
**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs, `--buffer_size=2147483648`):
|
||||
|
||||
| Configuration | Write | Read |
|
||||
|---------------|-------|------|
|
||||
| block=1MB, threads=16, batch=128 | 701 GB/s | **697 GB/s** |
|
||||
| **block=1MB, threads=32, batch=64** | **752 GB/s** | 713 GB/s |
|
||||
| block=1MB, threads=32, batch=32 | 751 GB/s | - |
|
||||
| block=1MB, threads=64, batch=32 | 728 GB/s | - |
|
||||
|
||||
> **Peak: 752 GB/s write**, reaching ~94% of the 800 GB/s theoretical line rate (16×400 Gbps). GPUDirect RDMA bypasses DRAM entirely (HBM3e → PCIe switch → NIC), so performance is not bottlenecked by CPU memory bandwidth.
|
||||
|
||||
**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
|
||||
|
||||
| Configuration | Write | Read |
|
||||
|---------------|-------|------|
|
||||
| **block=1MB, threads=32, batch=128, buf=4GB** | **230 GB/s** | 180 GB/s |
|
||||
| block=16MB, threads=32, batch=8, buf=8GB (striping off) | 233 GB/s | - |
|
||||
|
||||
> CPU-to-CPU is bounded by DRAM bandwidth (~250 GB/s/socket on Xeon 8559C). Per-NIC sampling shows NUMA-0 NICs at 90 Gbps and NUMA-1 NICs at 53 Gbps, confirming DRAM controller saturation rather than NIC limit.
|
||||
|
||||
#### p5en.48xlarge (H200, 16 EFA × 200 Gbps)
|
||||
|
||||
Tested on two p5en.48xlarge instances (Intel Xeon 8488C, 8× H200 141GB, 16 EFA devices) in the same AWS placement group.
|
||||
|
||||
**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs):
|
||||
|
||||
| Configuration | Write | Read |
|
||||
|---------------|-------|------|
|
||||
| block=1MB, threads=8, batch=128, buf=1GB/GPU | 236 GB/s | 271 GB/s |
|
||||
| block=1MB, threads=16, batch=128, buf=2GB/GPU | 271 GB/s | **297-308 GB/s** |
|
||||
| **block=1MB, threads=32, batch=64, buf=2GB/GPU** | **337-347 GB/s** | 274 GB/s |
|
||||
|
||||
> GPU HBM bandwidth (>3 TB/s) eliminates the memory bottleneck, allowing full EFA utilization. Write and read have different optimal thread counts: write peaks at 32 threads, read peaks at 16 threads.
|
||||
|
||||
> **Note:** EFA memory region registration (fi_mr_reg) for GPU memory segfaults at 4GB+ per GPU. Use `--buffer_size=2147483648` (2GB) as the maximum per-GPU buffer.
|
||||
|
||||
**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
|
||||
|
||||
| Configuration | Write | Read |
|
||||
|---------------|-------|------|
|
||||
| Single instance (block=1MB, threads=32, batch=128, buf=4GB) | 179 GB/s | 185 GB/s |
|
||||
| NUMA-split (block=1MB, 2 instances, 8 NICs each, threads=16, buf=2GB) | **192 GB/s** | **182 GB/s** |
|
||||
|
||||
> CPU-to-CPU throughput is bottlenecked by DRAM bandwidth (~155 GB/s per NUMA node, measured with STREAM Copy).
|
||||
|
||||
#### Cross-Transport Comparison
|
||||
|
||||
| Transport | Throughput | Notes |
|
||||
|-----------|-----------|-------|
|
||||
| **EFA GPU-to-GPU (B300)** | **752 GB/s** | p6-b300.48xlarge, 16×400G, block=1MB, ~94% line rate |
|
||||
| **EFA GPU-to-GPU (H200)** | **347 GB/s** | p5en.48xlarge, 16×200G, block=1MB |
|
||||
| **EFA GPU-to-GPU (B200)** | **313 GB/s** | p6-b200.48xlarge, 8×400G, block=1MB |
|
||||
| **EFA CPU-to-CPU (B300)** | **230 GB/s** | p6-b300.48xlarge, 16×400G, block=1MB, DRAM-limited |
|
||||
| **EFA CPU-to-CPU (B200)** | **222 GB/s** | p6-b200.48xlarge, 8×400G, block=1MB, DRAM-limited |
|
||||
| **EFA CPU-to-CPU (H200)** | **192 GB/s** | p5en.48xlarge, block=1MB, NUMA-split, DRAM-limited |
|
||||
| EFA (default params) | 69.47 GB/s | Default block=64KB |
|
||||
| TCP (iperf3 baseline) | 9.5 GB/s | Kernel TCP stack, 8 parallel streams |
|
||||
|
||||
**EFA vs RoCE RDMA**: On comparable 8×400 Gbps RoCE networks, Mooncake's RDMA transport achieves ~190 GB/s. Tuned EFA **exceeds** RoCE performance with GPU memory (313-347 GB/s) and on CPU-to-CPU (222 GB/s).
|
||||
|
||||
### Tuning Tips
|
||||
|
||||
- **Use `--block_size=1048576` (1MB)** — this is the most important tuning parameter for EFA. Each `block_size`-sized transfer becomes a single `fi_write`/`fi_read` call, so larger blocks amortize per-operation overhead. 1MB gives ~2× throughput over the 64KB default.
|
||||
- `MC_SLICE_SIZE` has **no effect** on EFA transport (it only applies to RDMA transport). Use `block_size` instead.
|
||||
- Increase `--threads` to 32-48 to saturate multiple EFA devices (2-4 threads per device is a good starting point)
|
||||
- For **CPU-to-CPU**: use `--block_size=1048576` (1MB) with NUMA-split (separate instances per NUMA node) for best results
|
||||
- For **GPU-to-GPU**: use `--block_size=1048576` (1MB), `--gpu_id=-1` (all GPUs), and `--buffer_size=2147483648` (2GB max per GPU). Write peaks at threads=32, read at threads=16
|
||||
- Keep `--batch_size` such that `block_size * batch_size * threads <= buffer_size`
|
||||
- Allocate buffers on both NUMA nodes for balanced NIC utilization (the bench tool does this by default for CPU mode)
|
||||
- On 16-NIC instances (p5en), writes are NUMA-sensitive: 8 local-NUMA NICs reach 90 Gbps each, while 8 cross-NUMA NICs only reach ~20 Gbps without NUMA-split
|
||||
|
||||
### Eager endpoint warmup (first-request latency)
|
||||
|
||||
libfabric `FI_EP_RDM` endpoints resolve peer addresses lazily: `fi_av_insert()` and the metadata handshake fire on the first send to each `(local_ctx, peer_nic)` pair. On 16-NIC instances that gives `16 × N_peer_NICs` serial handshakes inside the first `submitTransfer`, which shows up as a single-digit-second first-batch stall (measured ~4 s on p6-B300 for a 100 × 0.5 MB batch; the first batch runs at <0.1 GB/s while the CQ drains, steady-state afterwards is unaffected).
|
||||
|
||||
Mooncake exposes an explicit eager-warmup API to eliminate the stall:
|
||||
|
||||
- C++: `EfaTransport::warmupSegment(const std::string& segment_name)`
|
||||
- C: `int warmupEfaSegment(transfer_engine_t engine, const char *segment_name)`
|
||||
- Rust: `TransferEngine::warmup_efa_segment(name: &str)`
|
||||
|
||||
Call it once per peer segment, right after `openSegment` (or after any metadata change that adds a new peer). Every `(local_ctx, peer_nic)` endpoint is connected concurrently via `std::async`; the critical path becomes `max(handshake RTT)` instead of `sum(handshake RTT)`. The call is idempotent — safe to re-run.
|
||||
|
||||
Measured on p6-B300 (16 local NICs × 16 peer NICs, dual-NUMA initiator, 100 × 0.5 MB batch):
|
||||
|
||||
| | first-batch latency | steady-state |
|
||||
|---|---:|---:|
|
||||
| No warmup | 4,043 ms | 141 GB/s |
|
||||
| `warmup_efa_segment` (256 endpoints connected in 4.1 s) | **13.5 ms** (~300×) | 230 GB/s |
|
||||
|
||||
The warmup call itself takes roughly the same wall time as the stall it replaces — the win is that it's a one-time setup cost decoupled from the critical path of the first real transfer, not paid inside your latency budget.
|
||||
|
||||
## Usage with vLLM
|
||||
|
||||
### Prefill Instance
|
||||
|
||||
```bash
|
||||
VLLM_MOONCAKE_BOOTSTRAP_PORT=8998 \
|
||||
vllm serve <model_path> -tp 8 \
|
||||
--port 8010 \
|
||||
--trust-remote-code \
|
||||
--kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}'
|
||||
```
|
||||
|
||||
### Decode Instance
|
||||
|
||||
```bash
|
||||
vllm serve <model_path> -tp 8 \
|
||||
--port 8020 \
|
||||
--trust-remote-code \
|
||||
--kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}'
|
||||
```
|
||||
|
||||
## Usage with SGLang
|
||||
|
||||
SGLang's Mooncake integration currently hardcodes the `"rdma"` protocol. To use EFA transport, apply the provided patch and set environment variables.
|
||||
|
||||
### 1. Apply EFA Patch
|
||||
|
||||
SGLang's transfer engine initialization needs to be patched to read the protocol from an environment variable instead of using hardcoded `"rdma"`. Use the [patch script](https://github.com/whn09/kimi-k2-sglang):
|
||||
|
||||
```bash
|
||||
bash patch_sglang_efa.sh
|
||||
```
|
||||
|
||||
This is idempotent and safe to rerun.
|
||||
|
||||
### 2. Environment Variables
|
||||
|
||||
```bash
|
||||
export MOONCAKE_PROTOCOL=efa
|
||||
export FI_PROVIDER=efa
|
||||
export FI_EFA_USE_DEVICE_RDMA=1
|
||||
export GLOO_SOCKET_IFNAME=enp71s0 # adjust to your instance's primary interface
|
||||
```
|
||||
|
||||
For multi-node expert parallelism (EP) deployments, also set:
|
||||
|
||||
```bash
|
||||
export NVSHMEM_REMOTE_TRANSPORT=libfabric
|
||||
export NVSHMEM_LIBFABRIC_PROVIDER=efa
|
||||
```
|
||||
|
||||
> **Warning:** Do **not** set NVSHMEM variables on single-node deployments — doing so causes segmentation faults.
|
||||
|
||||
### 3. Docker Launch Example
|
||||
|
||||
```bash
|
||||
docker run -d --name sglang \
|
||||
--runtime=nvidia --gpus all --network host \
|
||||
--privileged --shm-size=600g \
|
||||
--device=/dev/infiniband \
|
||||
-e MOONCAKE_PROTOCOL=efa \
|
||||
-e FI_PROVIDER=efa \
|
||||
-e FI_EFA_USE_DEVICE_RDMA=1 \
|
||||
<image> bash start.sh
|
||||
```
|
||||
|
||||
> **Note:** Ensure the Docker image's libfabric version matches the host's EFA driver. If not, mount the host's EFA libraries into the container (see [Troubleshooting](#libfabric-version-mismatch-in-docker)).
|
||||
|
||||
## Technical Details
|
||||
|
||||
### Why libfabric instead of ibverbs?
|
||||
|
||||
AWS EFA exposes RDMA-like devices through the ibverbs interface, but does not support the full ibverbs API. Specifically:
|
||||
- Queue Pair (QP) creation fails with "Operation not supported" (error 95)
|
||||
- EFA requires using libfabric's `FI_EP_RDM` (Reliable Datagram Message) endpoint type
|
||||
|
||||
### EFA Transport Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────┐
|
||||
│ EfaTransport │
|
||||
├─────────────────────────────────────────────────────┤
|
||||
│ EfaContext (per device) │
|
||||
│ ├── fi_info (fabric info) │
|
||||
│ ├── fid_fabric (fabric handle) │
|
||||
│ ├── fid_domain (protection domain) │
|
||||
│ ├── fid_av (address vector for peer lookup) │
|
||||
│ ├── fid_cq (completion queues) │
|
||||
│ └── fid_mr (memory regions) │
|
||||
├─────────────────────────────────────────────────────┤
|
||||
│ EfaEndpoint (per connection) │
|
||||
│ ├── fid_ep (RDM endpoint) │
|
||||
│ ├── fi_addr_t (peer address) │
|
||||
│ └── local_addr (local endpoint address) │
|
||||
└─────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Thread Safety
|
||||
|
||||
The EFA transport requests `FI_THREAD_SAFE` from the libfabric provider and adds per-endpoint spinlocks to serialize `fi_write`/`fi_read` calls. This is necessary because:
|
||||
|
||||
- Multiple submission threads may route slices to the same endpoint concurrently
|
||||
- libfabric RDM endpoints default to `FI_THREAD_UNSPEC` (no thread safety guarantees)
|
||||
- Concurrent `fi_write`/`fi_read` without serialization corrupts provider internals, causing completions to silently vanish
|
||||
|
||||
CQ completion queues are polled by dedicated worker threads (one per EFA device) that run independently of submission threads.
|
||||
|
||||
### EFA vs RoCE RDMA
|
||||
|
||||
| Feature | EFA (libfabric SRD) | RoCE (ibverbs) |
|
||||
|---------|--------------------|--------------------|
|
||||
| Protocol | Scalable Reliable Datagram | RDMA over Converged Ethernet |
|
||||
| Endpoint type | `FI_EP_RDM` (message-based) | Queue Pairs (true RDMA) |
|
||||
| Write operation | Software-emulated via messages + ACKs | Hardware-offloaded one-sided RDMA |
|
||||
| CPU overhead | Moderate (provider processes ACKs) | Minimal (NIC handles everything) |
|
||||
| Throughput CPU-to-CPU (8×400G) | 222 GB/s (tuned) | ~190 GB/s |
|
||||
| Throughput GPU-to-GPU (16×200G) | 347 GB/s (tuned) | N/A |
|
||||
| Throughput GPU-to-GPU (8×400G) | 313 GB/s (tuned) | N/A |
|
||||
| AWS availability | All EFA-enabled instances | Not available on AWS |
|
||||
|
||||
### Supported AWS Instance Types
|
||||
|
||||
- p6-b300.48xlarge (16 EFA devices × 400 Gbps = 6,400 Gbps, `rdmap*` naming)
|
||||
- p6-b200.48xlarge (8 EFA devices × 400 Gbps = 3,200 Gbps, `rdmap*` naming)
|
||||
- p5en.48xlarge (16 EFA devices × 200 Gbps = 3,200 Gbps, `rdmap*` naming)
|
||||
- p5e.48xlarge (32 EFA devices × 100 Gbps = 3,200 Gbps, `rdmap*` naming)
|
||||
- p5.48xlarge (32 EFA devices × 100 Gbps = 3,200 Gbps, `rdmap*` naming)
|
||||
- Other EFA-enabled instances
|
||||
|
||||
Use `fi_info -p efa` to list available EFA devices on your instance.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### No EFA devices found
|
||||
|
||||
```
|
||||
EfaTransport: No EFA devices found
|
||||
```
|
||||
|
||||
Solution: Verify EFA is available with `fi_info -p efa`
|
||||
|
||||
### Permission denied
|
||||
|
||||
```
|
||||
fi_fabric failed: Permission denied
|
||||
```
|
||||
|
||||
Solution: Ensure proper permissions or run with sudo for testing
|
||||
|
||||
### libfabric not found
|
||||
|
||||
```
|
||||
cannot find -lfabric
|
||||
```
|
||||
|
||||
Solution: Verify `/opt/amazon/efa/lib` is in the library path:
|
||||
```bash
|
||||
export LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH
|
||||
```
|
||||
|
||||
### Workers hang under high concurrency
|
||||
|
||||
If `transfer_engine_bench` hangs with some workers never completing:
|
||||
|
||||
1. **Ensure both nodes are running the same build** — the CQ backpressure and thread-safety fixes must be present on both sides
|
||||
2. **Reduce concurrency** to verify basic connectivity: `--threads=1 --batch_size=16`
|
||||
3. **Check CQ poller threads**: logs should show "Started N CQ polling worker threads" where N matches the number of EFA devices
|
||||
|
||||
### Building on AWS Deep Learning AMI
|
||||
|
||||
On AWS Deep Learning AMI (e.g., Ubuntu 24.04), the system Python and CUDA toolkit are bundled inside the `/opt/pytorch` virtual environment. You must activate it and set CUDA paths before building:
|
||||
|
||||
```bash
|
||||
# Activate the PyTorch environment (provides Python 3.13 + CUDA toolkit)
|
||||
source /opt/pytorch/bin/activate
|
||||
|
||||
# Set CUDA paths (nvcc, headers and libs are inside the pip-installed nvidia packages)
|
||||
export CUDA_HOME=/opt/pytorch/lib/python3.13/site-packages/nvidia/cu13
|
||||
export PATH=$CUDA_HOME/bin:$PATH
|
||||
export CPLUS_INCLUDE_PATH=$CUDA_HOME/include:$CPLUS_INCLUDE_PATH
|
||||
export LD_LIBRARY_PATH=$CUDA_HOME/lib:$LD_LIBRARY_PATH
|
||||
export LIBRARY_PATH=$CUDA_HOME/lib:$LIBRARY_PATH
|
||||
|
||||
# Build with CUDA support
|
||||
cd ~/Mooncake
|
||||
mkdir -p build && cd build
|
||||
cmake .. -DUSE_EFA=ON -DUSE_CUDA=ON -DCMAKE_BUILD_TYPE=RelWithDebInfo
|
||||
make -j$(nproc)
|
||||
```
|
||||
|
||||
Without activating the environment, you may encounter:
|
||||
- `Could not find nvcc, please set CUDAToolkit_ROOT` — nvcc is not in PATH
|
||||
- `fatal error: cuda.h: No such file or directory` — CUDA headers not in include path, set `CPLUS_INCLUDE_PATH`
|
||||
- `cannot find -lcudart: No such file or directory` — CUDA libs not in library path, set `LIBRARY_PATH` and `LD_LIBRARY_PATH`
|
||||
- `ModuleNotFoundError: No module named 'mooncake.engine'` — `.so` built against wrong Python version (e.g., 3.12 vs 3.13)
|
||||
|
||||
### libfabric version mismatch in Docker
|
||||
|
||||
```
|
||||
fi_ep_bind (av) failed: Function not implemented
|
||||
```
|
||||
|
||||
or:
|
||||
|
||||
```
|
||||
undefined reference to `efadv_query_qp_wqs@EFA_1.4'
|
||||
```
|
||||
|
||||
This happens when the Docker container's libfabric version is older than the host's EFA driver. Check with `fi_info --version` on both host and container.
|
||||
|
||||
Solution: Mount the host's EFA libraries into the container:
|
||||
|
||||
```bash
|
||||
docker run --gpus all --device=/dev/infiniband --net=host --privileged \
|
||||
-v /opt/amazon/efa:/opt/amazon/efa \
|
||||
-v /lib/x86_64-linux-gnu/libefa.so.1:/lib/x86_64-linux-gnu/libefa.so.1 \
|
||||
-v /lib/x86_64-linux-gnu/libefa.so:/lib/x86_64-linux-gnu/libefa.so \
|
||||
-v /lib/x86_64-linux-gnu/libibverbs.so.1:/lib/x86_64-linux-gnu/libibverbs.so.1 \
|
||||
-e LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH \
|
||||
-it <image>
|
||||
```
|
||||
|
||||
Then rebuild Mooncake inside the container to link against the host's libfabric.
|
||||
|
|
@ -1,17 +1,17 @@
|
|||
# Transfer Engine
|
||||
|
||||
## Overview
|
||||
## Overview
|
||||
Mooncake Transfer Engine is a high-performance, zero-copy data transfer library designed around two core abstractions: Segment and BatchTransfer.
|
||||
|
||||
- [**Segment**](#segment) represents a contiguous address space that can be remotely read and written, which can be either non-persistent storage provided by DRAM or VRAM, known as **RAM Segment**, or persistent storage provided by NVMeof, known as **NVMeof Segment**.
|
||||
|
||||
- [**BatchTransfer**](#batchtransfer) encapsulates operation requests, specifically responsible for synchronizing data between a set of non-contiguous data spaces in one Segment and the corresponding spaces in another set of Segments, supporting Read/Write in both directions, thus acting like an asynchronous and more flexible AllScatter/AllGather.
|
||||
|
||||

|
||||

|
||||
|
||||
As shown in the diagram, each specific client corresponds to a `TransferEngine`, which not only includes a RAM Segment but also integrates management for high-speed transfers across multiple threads and network cards. The RAM Segment, in principle, corresponds to the entire virtual address space of this `TransferEngine`, but in reality, only parts of it (known as a `Buffer`) are registered for (GPUDirect) RDMA Read/Write. Each Buffer can have separate permissions (corresponding to RDMA `rkey`, etc.) and network card affinity (e.g., preferred NICs for different types of memory).
|
||||
|
||||
Mooncake Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `NVMeoFTransport`, `NvlinkTransport`, and `HipTransport`.
|
||||
Mooncake Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `EfaTransport`, `NVMeoFTransport`, `NvlinkTransport`, `IntraNodeNvlinkTransport`, and `HipTransport`.
|
||||
|
||||
### Segment
|
||||
Segment represents a collection of source address ranges and target address ranges available during the data transfer process in Transfer Engine. That is, all local and remote addresses involved in `BatchTransfer` requests must be within the valid segment range. Transfer Engine supports the following two types of Segments.
|
||||
|
|
@ -49,17 +49,17 @@ The BatchTransfer API uses an array of requests, which specify the operation typ
|
|||
### Topology Aware Path Selection
|
||||
Modern inference servers often consist of multiple CPU sockets, DRAM, GPUs, and RDMA NIC devices. Although it's technically possible to transfer data from local DRAM or VRAM to a remote location using any RDMA NIC, these transfers can be limited by the bandwidth constraints of the Ultra Path Interconnect (UPI) or PCIe Switch. To overcome these limitations, Transfer Engine implements a topology-aware path selection algorithm.
|
||||
|
||||
Before processing requests, each server generates a topology matrix and broadcasts it across the cluster.
|
||||
This matrix categorizes network interface cards (NICs) into preferred and secondary lists for various types of memory, which types are specified during memory registration.
|
||||
Under normal conditions, a NIC from the preferred list is selected for transfers, facilitating RDMA operations within the local NUMA or GPU Direct RDMA through the local PCIe switch only.
|
||||
Before processing requests, each server generates a topology matrix and broadcasts it across the cluster.
|
||||
This matrix categorizes network interface cards (NICs) into preferred and secondary lists for various types of memory, which types are specified during memory registration.
|
||||
Under normal conditions, a NIC from the preferred list is selected for transfers, facilitating RDMA operations within the local NUMA or GPU Direct RDMA through the local PCIe switch only.
|
||||
In case of failures, NICs from both lists may be utilized.
|
||||
The process involves identifying the appropriate local and target NICs based on the memory addresses, establishing a connection, and executing the data transfer.
|
||||
|
||||

|
||||

|
||||
|
||||
For instance, as illustrated in figure above, to transfer data from buffer 0 (assigned to cpu:0) in the local node to buffer 1 (assigned to cpu:1) in the target node, the engine first identifies the preferred NICs for cpu:0 using the local server's topology matrix and selects one, such as mlx5_1, as the local NIC. Similarly, the target NIC, such as mlx5_3, is selected based on the target memory address. This setup enables establishing an RDMA connection from mlx5_1@local to mlx5_3@target to carry out RDMA read and write operations.
|
||||
|
||||
To further maximize bandwidth utilization, if a single request's transfer is internally divided into multiple slices if its length exceeds 64KB.
|
||||
To further maximize bandwidth utilization, if a single request's transfer is internally divided into multiple slices if its length exceeds 64KB.
|
||||
Each slice might use a different path, enabling collaborative work among all RDMA NICs.
|
||||
|
||||
### Endpoint Management
|
||||
|
|
@ -140,7 +140,7 @@ After successfully compiling Transfer Engine, the test program `transfer_engine_
|
|||
```
|
||||
The meanings of the various parameters are as follows (the rest are the same as before):
|
||||
- `--segment_id` is the segment name of target node. It needs to be consistent with the value passed to `--local_server_name` when starting the target node (if any).
|
||||
|
||||
|
||||
Under normal circumstances, the initiator node will start the transfer operation, wait for 10 seconds, and then display the "Test completed" message, indicating that the test is complete.
|
||||
|
||||
The initiator node can also configure the following test parameters: `--operation` (can be `"read"` or `"write"`), `batch_size`, `block_size`, `duration`, `threads`, etc.
|
||||
|
|
@ -152,99 +152,15 @@ After successfully compiling Transfer Engine, the test program `transfer_engine_
|
|||
|
||||
The following video shows a normal run as described above, with the Target on the right and the Initiator on the left, at the end of the test the Initiator reports the test duration (10 seconds), IOPS (379008 requests/s), and throughput (19.87 GiB/s). The throughput here exceeds the maximum throughput supported by a single card on the host computer used.
|
||||
|
||||

|
||||

|
||||
|
||||
## Transfer Engine C/C++ API
|
||||
Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `NVMeoFTransport`, `NvlinkTransport` (for NVIDIA GPUs), and `HipTransport` (for AMD GPUs).
|
||||
Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `EfaTransport` (for AWS EFA), `NVMeoFTransport`, `NvlinkTransport` (for NVIDIA GPUs), `IntraNodeNvlinkTransport` (for NVIDIA GPUs), and `HipTransport` (for AMD GPUs).
|
||||
|
||||
For a complete C++ API reference, see [Transfer Engine C++ API Reference](cpp-api.md).
|
||||
|
||||
### Data Transfer
|
||||
|
||||
#### TransferEngine::TransferRequest
|
||||
|
||||
The core API provided by Mooncake Transfer Engine is submitting a group of asynchronous `TransferRequest` tasks through the `submitTransfer` interface, and querying their status through the `getTransferStatus` interface. Each `TransferRequest` specifies reading or writing a continuous data space of `length` starting from the local starting address `source`, to the position starting at `target_offset` in the segment corresponding to `target_id`.
|
||||
|
||||
The `TransferRequest` structure is defined as follows:
|
||||
|
||||
```cpp
|
||||
using SegmentID = int32_t;
|
||||
struct TransferRequest
|
||||
{
|
||||
enum OpCode { READ, WRITE };
|
||||
OpCode opcode;
|
||||
void *source;
|
||||
SegmentID target_id; // The ID of the target segment, which may correspond to local or remote DRAM/VRAM/NVMeof, with the specific routing logic hidden
|
||||
size_t target_offset;
|
||||
size_t length;
|
||||
};
|
||||
```
|
||||
|
||||
- `opcode` takes the values `READ` or `WRITE`. `READ` indicates that data is copied from the target address indicated by `<target_id, target_offset>` to the local starting address `source`; `WRITE` indicates that data is copied from `source` to the address indicated by `<target_id, target_offset>`.
|
||||
- `source` represents the DRAM/VRAM buffer managed by the current `TransferEngine`, which must have been registered in advance by the `registerLocalMemory` interface.
|
||||
- `target_id` represents the segment ID of the transfer target. The segment ID is obtained using the `openSegment` interface. Segments are divided into the following types:
|
||||
- RAM space type, covering DRAM/VRAM. As mentioned earlier, there is only one segment under the same process (or `TransferEngine` instance), which contains various types of Buffers (DRAM/VRAM). In this case, the segment name passed to the `openSegment` interface is equivalent to the server hostname. `target_offset` is the virtual address of the target server.
|
||||
- NVMeOF space type, where each file corresponds to a segment. In this case, the segment name passed to the `openSegment` interface is equivalent to the unique identifier of the file. `target_offset` is the offset of the target file.
|
||||
- `length` represents the amount of data transferred. TransferEngine may further split this into multiple read/write requests internally.
|
||||
|
||||
#### TransferEngine::allocateBatchID
|
||||
|
||||
```cpp
|
||||
BatchID allocateBatchID(size_t batch_size);
|
||||
```
|
||||
|
||||
Allocates a `BatchID`. A maximum of `batch_size` `TransferRequest`s can be submitted under the same `BatchID`.
|
||||
|
||||
- `batch_size`: The maximum number of `TransferRequest`s that can be submitted under the same `BatchID`;
|
||||
- Return value: If successful, returns `BatchID` (non-negative); otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::submitTransfer
|
||||
|
||||
```cpp
|
||||
int submitTransfer(BatchID batch_id, const std::vector<TransferRequest> &entries);
|
||||
```
|
||||
|
||||
Submits new `TransferRequest` tasks to `batch_id`. The task is asynchronously submitted to the background thread pool. The total number of `entries` accumulated under the same `batch_id` should not exceed the `batch_size` defined at creation.
|
||||
|
||||
- `batch_id`: The `BatchID` it belongs to;
|
||||
- `entries`: Array of `TransferRequest`;
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::getTransferStatus
|
||||
|
||||
```cpp
|
||||
enum TaskStatus
|
||||
{
|
||||
WAITING, // In the transfer phase
|
||||
PENDING, // Not supported
|
||||
INVALID, // Ilvalid parameters
|
||||
CANCELED, // Not supported
|
||||
COMPLETED, // Transfer completed
|
||||
TIMEOUT, // Not supported
|
||||
FAILED // Transfer failed even after retries
|
||||
};
|
||||
struct TransferStatus {
|
||||
TaskStatus s;
|
||||
size_t transferred; // How much data has been successfully transferred (not necessarily an accurate value, but it is a lower bound)
|
||||
};
|
||||
int getTransferStatus(BatchID batch_id, size_t task_id, TransferStatus &status)
|
||||
```
|
||||
|
||||
Obtains the running status of the `TransferRequest` with `task_id` in `batch_id`.
|
||||
|
||||
- `batch_id`: The `BatchID` it belongs to;
|
||||
- `task_id`: The sequence number of the `TransferRequest` to query;
|
||||
- `status`: Output Transfer status;
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::freeBatchID
|
||||
|
||||
```cpp
|
||||
int freeBatchID(BatchID batch_id);
|
||||
```
|
||||
|
||||
Recycles `BatchID`, and subsequent operations on `submitTransfer` and `getTransferStatus` are undefined. If there are still `TransferRequest`s pending completion in the `BatchID`, the operation is refused.
|
||||
|
||||
- `batch_id`: The `BatchID` it belongs to;
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
Transfer Engine provides batch-based read/write transfers between segments (DRAM/VRAM/NVMeof). A typical flow is: register local memory, open a target segment, submit a batch, and poll status. Detailed function signatures and usage are documented in the C++ API reference.
|
||||
|
||||
### Multi-Transport Management
|
||||
|
||||
|
|
@ -253,51 +169,8 @@ And it will discover the toplogy between CPU/CUDA and RDMA devices automatically
|
|||
(more device types are working in progress, feedbacks are welcome when the automatic discovery mechanism is not accurate),
|
||||
and it will install `Transport` automatically based on the topology.
|
||||
|
||||
### Space Registration
|
||||
|
||||
For the RDMA transfer process, the source pointer `TransferRequest::source` must be registered in advance as an RDMA readable/writable Memory Region space, that is, included as part of the RAM Segment of the current process. Therefore, the following functions are needed:
|
||||
|
||||
#### TransferEngine::registerLocalMemory
|
||||
|
||||
```cpp
|
||||
int registerLocalMemory(void *addr, size_t size, string location, bool remote_accessible);
|
||||
```
|
||||
|
||||
Registers a space starting at address `addr` with a length of `size` on the local DRAM/VRAM.
|
||||
|
||||
- `addr`: The starting address of the registration space;
|
||||
- `size`: The length of the registration space;
|
||||
- `location`: The `device` corresponding to this memory segment, such as `cuda:0` indicating the GPU device, `cpu:0` indicating the CPU socket, by matching with the network card priority order table (see `installTransport`), the preferred network card is identified. You can also use `*`, Transfer Engine will try to automatically recognize the `device` corresponding to `addr`, if it fails to recognize the device, it will print a `WARNING` level log and use all network cards, no preferred network cards.
|
||||
- `remote_accessible`: Indicates whether this memory can be accessed by remote nodes.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
#### TransferEngine::unregisterLocalMemory
|
||||
|
||||
```cpp
|
||||
int unregisterLocalMemory(void *addr);
|
||||
```
|
||||
|
||||
Unregisters the region.
|
||||
|
||||
- `addr`: The starting address of the registration space;
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
|
||||
### Segment Management and Metadata Format
|
||||
|
||||
TransferEngine provides the `openSegment` function, which obtains a `SegmentHandle` for subsequent `Transport` transfers.
|
||||
```cpp
|
||||
SegmentHandle openSegment(const std::string& segment_name);
|
||||
```
|
||||
|
||||
- `segment_name`: The unique identifier of the segment. For RAM Segment, this needs to be consistent with the `server_name` filled in by the peer process when initializing the TransferEngine object.
|
||||
- Return value: If successful, returns the corresponding `SegmentHandle`; otherwise, returns a negative value.
|
||||
|
||||
```cpp
|
||||
int closeSegment(SegmentHandle segment_id);
|
||||
```
|
||||
|
||||
- `segment_id`: The unique identifier of the segment.
|
||||
- Return value: If successful, returns 0; otherwise, returns a negative value.
|
||||
Segment metadata is stored in the metadata service. The following format is provided for reference.
|
||||
|
||||
<details>
|
||||
<summary><strong>Metadata Format</strong></summary>
|
||||
|
|
@ -347,7 +220,7 @@ Key = mooncake/nvmeof/[segment_name]
|
|||
Value = {
|
||||
'server_name': server_name,
|
||||
'protocol': nvmeof,
|
||||
'buffers':[
|
||||
'buffers':[
|
||||
{
|
||||
'length': 1073741824,
|
||||
'file_path': "/mnt/nvme0" // The file path on this machine
|
||||
|
|
@ -358,7 +231,7 @@ Value = {
|
|||
},
|
||||
{
|
||||
'length': 1073741824,
|
||||
'file_path': "/mnt/nvme1",
|
||||
'file_path': "/mnt/nvme1",
|
||||
'local_path_map': {
|
||||
"node02": "/mnt/transfer_engine/node02/nvme1",
|
||||
.....
|
||||
|
|
@ -379,29 +252,6 @@ The HTTP server should implement three following RESTful APIs, while the metadat
|
|||
|
||||
For specific implementation, refer to the demo service implemented in Golang at [mooncake-transfer-engine/example/http-metadata-server](../../../mooncake-transfer-engine/example/http-metadata-server).
|
||||
|
||||
### Initialization
|
||||
|
||||
TransferEngine needs to initializing by calling the `init` method before further actions:
|
||||
```cpp
|
||||
TransferEngine();
|
||||
|
||||
int init(const std::string &metadata_conn_string,
|
||||
const std::string &local_server_name);
|
||||
```
|
||||
- `metadata_conn_string`: Connecting string of metadata storage servers, i.e., the IP address/hostname of `etcd`/`redis` or the URI of the http service.
|
||||
The general form is `[proto]://[hostname:port]`. For example, the following metadata server addresses are legal:
|
||||
- Using `etcd` as a metadata storage service: `“10.0.0.1:2379”` or `“etcd://10.0.0.1:2379”`.
|
||||
- Using `redis` as a metadata storage service: `“redis://10.0.0.1:6379”`
|
||||
- Using `http` as a metadata storage service: `“http://10.0.0.1:8080/metadata”`
|
||||
|
||||
- `local_server_name`: The local server name, ensuring uniqueness within the cluster. It also serves as the name of the RAM Segment that other nodes refer to the current instance (i.e., Segment Name).
|
||||
|
||||
```cpp
|
||||
~TransferEngine();
|
||||
```
|
||||
|
||||
Reclaims all allocated resources and also deletes the global meta data server information.
|
||||
|
||||
## Using Transfer Engine to Your Projects
|
||||
|
||||
### Using C/C++ Interface
|
||||
|
|
@ -437,6 +287,7 @@ For advanced users, TransferEngine provides the following advanced runtime optio
|
|||
- `MC_LOG_LEVEL` This option can be set as `TRACE`/`INFO`/`WARNING`/`ERROR` (see [glog doc](https://github.com/google/glog/blob/master/docs/logging.md)), and more detailed logs will be output during runtime
|
||||
- `MC_DISABLE_METACACHE` Disable local meta cache to prevent transfer failure due to dynamic memory registrations, which may downgrades the performance
|
||||
- `MC_HANDSHAKE_LISTEN_BACKLOG` The backlog size of socket listening for handshaking, default value is 128
|
||||
- `MC_HANDSHAKE_MAX_LENGTH` The maximum handshake message length in bytes for P2P mode. Valid range: 1MB to 128MB. Default value is 1MB (1048576 bytes). Increase this value when using a single RDMA instance with many registered memory buffers (>10,000) to avoid handshake failures. Example: set to 10485760 for 10MB
|
||||
- `MC_LOG_DIR` Specify the directory path for log redirection files. If invalid, log to stderr instead.
|
||||
- `MC_REDIS_PASSWORD` The password for Redis storage plugin, only takes effect when Redis is specified as the metadata server. If not set, no authentication will be attempted to log in to the Redis.
|
||||
- `MC_REDIS_DB_INDEX` The database index for Redis storage plugin, must be an integer between 0 and 255. Only takes effect when Redis is specified as the metadata server. If not set or invalid, the default value is 0.
|
||||
|
|
@ -445,11 +296,29 @@ For advanced users, TransferEngine provides the following advanced runtime optio
|
|||
- `MC_ENABLE_PARALLEL_REG_MR` Control parallel memory region registration across multiple RDMA NICs. Valid values: -1 (auto, default), 0 (disabled), 1 (enabled). When set to -1, parallel registration is automatically enabled when multiple RNICs exist and memory has been pre-touched. Note: If memory hasn't been touched before registration, parallel registration can be slower than sequential registration
|
||||
- `MC_FORCE_HCA` Force to use RDMA as the active transport, return error if no HCA has been found.
|
||||
- `MC_FORCE_MNNVL` Force to use Multi-Node NVLink as the active transport regardless whether RDMA devices are installed.
|
||||
- `MC_INTRA_NVLINK` Enable intra-node NVLINK transport, and cannot be used together with MC_FORCE_MNNVL.
|
||||
- `MC_FORCE_TCP` Force to use TCP as the active transport regardless whether RDMA devices are installed.
|
||||
- `MC_MIN_PRC_PORT` Specifies the minimum port number for RPC service. The default value is 15000.
|
||||
- `MC_MAX_PRC_PORT` Specifies the maximum port number for RPC service. The default value is 17000.
|
||||
- `MC_PATH_ROUNDROBIN` Use round-robin mode in the RDMA path selection. This may be beneficial for transferring large bulks.
|
||||
- `MC_ENDPOINT_STORE_TYPE` Choose FIFO Endpoint Store (`FIFO`) or Sieve Endpoint Store (`SIEVE`), default is `SIEVE`.
|
||||
- `MC_TCP_ENABLE_CONNECTION_POOL` Enable TCP Connection Pool to avoid excessive sockets.
|
||||
|
||||
## C++ API Reference
|
||||
|
||||
::::{toctree}
|
||||
:maxdepth: 1
|
||||
|
||||
cpp-api
|
||||
::::
|
||||
|
||||
## EFA Transport (AWS)
|
||||
|
||||
:::{toctree}
|
||||
:maxdepth: 1
|
||||
|
||||
efa_transport
|
||||
:::
|
||||
|
||||
## Ascend Transport Component
|
||||
|
||||
|
|
@ -467,4 +336,4 @@ heterogeneous_ascend
|
|||
:maxdepth: 1
|
||||
|
||||
transfer-engine-bench-tuning
|
||||
:::
|
||||
:::
|
||||
|
|
|
|||
|
|
@ -0,0 +1,293 @@
|
|||
# Kunpeng UB Transport for Mooncake
|
||||
|
||||
This document describes how to build and use Mooncake with Kunpeng UB (Unified Bus) transport support using URMA (Unified Remote Memory Access).
|
||||
|
||||
## Overview
|
||||
|
||||
UB (Unified Bus) is a transport protocol at the same abstraction layer as RDMA, CXL, NVLink, and TCP, providing a flexible transport solution that can be selected at the application layer. Currently, UB protocol has two open-source implementations:
|
||||
|
||||
- **URMA (Unified Remote Memory Access)**: Provides a unified programming abstraction and core semantic layer for upper-layer applications. It offers unified APIs and semantic interfaces for remote shared memory access and operations, leveraging the low-latency, high-bandwidth characteristics of the UB protocol.
|
||||
- URMA open-source repository: https://atomgit.com/openeuler/umdk
|
||||
|
||||
- **OBMM (Ownership Based Memory Management)**: A kernel memory management system for supernode environments, supporting cross-node physical memory sharing. It provides efficient remote memory access capabilities through a kernel module (obmm.ko) and a user-space library (libobmm.so).
|
||||
- OBMM open-source repository: https://atomgit.com/openeuler/obmm
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### 1. Hardware and Operating System
|
||||
|
||||
- **Hardware Platform**: Kunpeng 950 CPU with native UB interconnect architecture
|
||||
- **OS Version**: openEuler 24.03 (LTS-SP3) [Download link](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
|
||||
|
||||
### 2. URMA Dependencies
|
||||
|
||||
Install UMDK (URMA development package):
|
||||
|
||||
```bash
|
||||
# Install via yum
|
||||
yum install umdk-urma-devel
|
||||
|
||||
# Or build from source
|
||||
git clone https://atomgit.com/openeuler/umdk.git
|
||||
cd umdk
|
||||
mkdir build && cd build
|
||||
cmake ..
|
||||
make -j$(nproc)
|
||||
sudo make install
|
||||
```
|
||||
|
||||
### 3. Build Dependencies
|
||||
|
||||
```bash
|
||||
# Ubuntu/Debian
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y \
|
||||
build-essential \
|
||||
cmake \
|
||||
git \
|
||||
libgflags-dev \
|
||||
libgoogle-glog-dev \
|
||||
libjsoncpp-dev \
|
||||
libnuma-dev \
|
||||
libibverbs-dev \
|
||||
libboost-all-dev \
|
||||
libcurl4-openssl-dev \
|
||||
libgtest-dev \
|
||||
libmsgpack-dev \
|
||||
libxxhash-dev \
|
||||
libyaml-cpp-dev \
|
||||
pybind11-dev \
|
||||
python3-dev
|
||||
|
||||
# Install yalantinglibs (required)
|
||||
cd /tmp
|
||||
git clone https://github.com/alibaba/yalantinglibs.git
|
||||
cd yalantinglibs
|
||||
mkdir build && cd build
|
||||
cmake .. -DCMAKE_INSTALL_PREFIX=/usr/local
|
||||
make -j$(nproc)
|
||||
sudo make install
|
||||
```
|
||||
|
||||
## Building Mooncake with UB Support
|
||||
|
||||
### 1. Clone the Repository
|
||||
|
||||
```bash
|
||||
git clone https://github.com/kvcache-ai/Mooncake.git
|
||||
cd Mooncake
|
||||
git submodule update --init --recursive
|
||||
```
|
||||
|
||||
### 2. Build with UB Enabled
|
||||
|
||||
```bash
|
||||
mkdir build && cd build
|
||||
|
||||
cmake .. \
|
||||
-DUSE_UB=ON \
|
||||
-DURMA_INCLUDE_DIR=/usr/include \
|
||||
-DURMA_LIBRARY=/usr/lib64/liburma.so \
|
||||
-DCMAKE_BUILD_TYPE=RelWithDebInfo
|
||||
|
||||
make -j$(nproc)
|
||||
```
|
||||
|
||||
### 3. Install Python Package
|
||||
|
||||
```bash
|
||||
# Copy built modules to wheel directory
|
||||
cp mooncake-integration/engine.cpython-*.so ../mooncake-wheel/mooncake/
|
||||
cp mooncake-integration/store.cpython-*.so ../mooncake-wheel/mooncake/
|
||||
cp mooncake-common/libasio.so ../mooncake-wheel/mooncake/
|
||||
|
||||
# Install with pip
|
||||
pip install -e ../mooncake-wheel --no-build-isolation
|
||||
```
|
||||
|
||||
## Verification
|
||||
|
||||
### Check UB Transport Registration
|
||||
|
||||
```bash
|
||||
# Check if UB transport is registered
|
||||
./mooncake_server --list-transports
|
||||
# Expected output: rdma, tcp, nvlink, ub
|
||||
```
|
||||
|
||||
### Test UB Transport Initialization
|
||||
|
||||
```python
|
||||
from mooncake.engine import TransferEngine
|
||||
|
||||
te = TransferEngine()
|
||||
result = te.initialize('127.0.0.1', 'P2PHANDSHAKE', 'ub', '')
|
||||
print(f'Initialize result: {result}') # Should be 0
|
||||
|
||||
# You should see logs like:
|
||||
# URMA module init success
|
||||
# found 1 devices.
|
||||
# device_name : urma0 EID : 01:02:03:04:05:06:07:08:09:0a:0b:0c:0d:0e:0f:10
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Single Node Benchmark Test
|
||||
|
||||
```bash
|
||||
# Terminal 1: Target (receiver)
|
||||
./transfer_engine_bench \
|
||||
--mode=target \
|
||||
--protocol=ub \
|
||||
--device_name=urma0 \
|
||||
--local_server_name=127.0.0.1 \
|
||||
--metadata_server=P2PHANDSHAKE
|
||||
|
||||
# Terminal 2: Initiator (sender)
|
||||
./transfer_engine_bench \
|
||||
--mode=initiator \
|
||||
--protocol=ub \
|
||||
--device_name=urma0 \
|
||||
--metadata_server=P2PHANDSHAKE \
|
||||
--segment_size=8388608 \
|
||||
--batch_size=1 \
|
||||
--segment_id=127.0.0.1:$PORT
|
||||
```
|
||||
|
||||
### Multi-device Benchmark Test
|
||||
|
||||
```bash
|
||||
# Auto-discovery of multiple URMA devices
|
||||
./transfer_engine_bench \
|
||||
--protocol=ub \
|
||||
--device_name=urma0,urma1,urma2,urma3
|
||||
```
|
||||
|
||||
## Unit Tests
|
||||
|
||||
Run the UB transport unit tests:
|
||||
|
||||
```bash
|
||||
./build/mooncake-transfer-engine/tests/ub_transport_test
|
||||
```
|
||||
|
||||
The test suite includes:
|
||||
|
||||
| Test | Description |
|
||||
|------|-------------|
|
||||
| `MultiWrite` | Multiple write operations |
|
||||
| `MultipleRead` | Multiple read operations with data integrity check |
|
||||
|
||||
You can also run all unit tests via CTest:
|
||||
|
||||
```bash
|
||||
cd build && ctest --output-on-failure
|
||||
```
|
||||
|
||||
Environment variables for test configuration:
|
||||
|
||||
```bash
|
||||
export MC_METADATA_SERVER=P2PHANDSHAKE # default
|
||||
export MC_LOCAL_SERVER_NAME=127.0.0.1:12345 # default
|
||||
```
|
||||
|
||||
## Technical Details
|
||||
|
||||
### UB Transport Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────┐
|
||||
│ UbTransport │
|
||||
├─────────────────────────────────────────────────────┤
|
||||
│ UrmaContext (per device) │
|
||||
│ ├── urma_device (URMA device handle) │
|
||||
│ ├── urma_context (URMA context) │
|
||||
│ ├── urma_jfce (URMA jetty factory create) │
|
||||
│ ├── urma_jfc (URMA jetty factory send) │
|
||||
│ └── urma_jfr (URMA jetty factory receive) │
|
||||
├─────────────────────────────────────────────────────┤
|
||||
│ UrmaEndpoint (per connection) │
|
||||
│ ├── urma_jetty (URMA jetty for communication) │
|
||||
│ ├── local_jetty (local jetty ID) │
|
||||
│ └── remote_jetty (remote jetty ID) │
|
||||
└─────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Key Components
|
||||
|
||||
1. **UbTransport**: The main transport class that manages URMA resources and endpoints
|
||||
2. **UrmaContext**: Represents a URMA device context, handling device initialization and resource management
|
||||
3. **UrmaEndpoint**: Represents a connection to a remote peer, handling data transfer operations
|
||||
4. **mock_urma_api.cpp**: Mock implementation of URMA API for testing without real URMA hardware
|
||||
|
||||
### Protocol Advantages
|
||||
|
||||
- **Optimized for Kunpeng**: URMA is specifically optimized for Kunpeng chip on-chip interconnect
|
||||
- **RDMA-like Semantics**: Provides similar memory semantics to RDMA
|
||||
- **High Performance**: Leverages UB's low-latency, high-bandwidth characteristics
|
||||
- **Unified Abstraction**: Offers a unified programming model for remote memory access
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### No URMA devices found
|
||||
|
||||
```
|
||||
UbTransport: No URMA devices found
|
||||
```
|
||||
|
||||
Solution: Verify URMA is properly installed and devices are available:
|
||||
```bash
|
||||
# Check URMA installation
|
||||
ls /usr/lib64/liburma.so
|
||||
ls /usr/include/ub/umdk/urma/urma_api.h
|
||||
|
||||
# Check for URMA devices
|
||||
urma_admin -l
|
||||
```
|
||||
|
||||
### URMA initialization failed
|
||||
|
||||
```
|
||||
URMA module init failed
|
||||
```
|
||||
|
||||
Solution: Ensure the URMA kernel module is loaded and the device is properly configured:
|
||||
```bash
|
||||
# Load URMA module
|
||||
sudo modprobe urma
|
||||
|
||||
# Check module status
|
||||
sudo lsmod | grep urma
|
||||
|
||||
# Check device status
|
||||
urma_admin -l
|
||||
```
|
||||
|
||||
### Device port inactive
|
||||
|
||||
```
|
||||
Device urma0 port not active
|
||||
```
|
||||
|
||||
Solution: Ensure the UB port is properly configured and active:
|
||||
```bash
|
||||
# Check port status
|
||||
urma_admin -p urma0
|
||||
```
|
||||
|
||||
### Missing liburma.so
|
||||
|
||||
```
|
||||
cannot find -lurma
|
||||
```
|
||||
|
||||
Solution: Verify URMA library is installed and in the library path:
|
||||
```bash
|
||||
export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH
|
||||
```
|
||||
|
||||
## Conclusion
|
||||
|
||||
Kunpeng UB Transport provides a high-performance, optimized transport solution for Mooncake on Kunpeng 950 CPU platforms. By leveraging the UB protocol's low-latency and high-bandwidth characteristics, it offers comparable performance to RDMA while being specifically tailored for Kunpeng chip architectures.
|
||||
|
||||
With proper configuration and tuning, UB Transport can significantly improve the performance of distributed AI workloads, particularly for scenarios involving large-scale parameter transfers and distributed training.
|
||||
|
|
@ -18,6 +18,7 @@ pip install mooncake-transfer-engine-non-cuda
|
|||
📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine-non-cuda/](https://pypi.org/project/mooncake-transfer-engine-non-cuda/)
|
||||
|
||||
> **Note**: The CUDA version includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+. The non-CUDA version is for environments without CUDA dependencies.
|
||||
> **Note**: MLU support is currently source-build only. If you need Cambricon MLU memory support, install Neuware and build with `-DUSE_MLU=ON`.
|
||||
|
||||
## Automatic
|
||||
|
||||
|
|
@ -112,8 +113,43 @@ pip install mooncake-transfer-engine-non-cuda
|
|||
```bash
|
||||
export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/musa/lib
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/musa/lib
|
||||
```
|
||||
|
||||
4. Install yalantinglibs
|
||||
4. If you want to compile Cambricon MLU support, first install the Cambricon Neuware SDK. After that:
|
||||
1) Export `NEUWARE_HOME` or pass `-DNEUWARE_ROOT=/path/to/neuware` to CMake
|
||||
2) Configure `LIBRARY_PATH` and `LD_LIBRARY_PATH` to ensure linking of `cnrt`, `cndrv`, and other Neuware libraries during compilation:
|
||||
```bash
|
||||
export NEUWARE_HOME=/usr/local/neuware
|
||||
export LIBRARY_PATH=$LIBRARY_PATH:${NEUWARE_HOME}/lib64
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${NEUWARE_HOME}/lib64
|
||||
```
|
||||
|
||||
If your Neuware installation lives outside the default include/library layout, you can also pass:
|
||||
```bash
|
||||
cmake .. -DUSE_MLU=ON \
|
||||
-DMLU_INCLUDE_DIR=/path/to/neuware/include \
|
||||
-DMLU_LIB_DIR=/path/to/neuware/lib64
|
||||
```
|
||||
|
||||
For Cambricon MLU builds, enable the MLU backend explicitly:
|
||||
```bash
|
||||
cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=${NEUWARE_HOME:-/usr/local/neuware}
|
||||
make -j
|
||||
```
|
||||
|
||||
5. If you want to compile MetaX (Muxi) MACA support (e.g. C500), install the MACA SDK so headers and libraries are available under `MACA_ROOT` (defaults to `MACA_HOME` env var if set, otherwise `/opt/maca`). SDK layouts vary; include both `lib` and `lib64` in runtime paths when needed:
|
||||
```bash
|
||||
export MACA_HOME=/opt/maca
|
||||
export LIBRARY_PATH=$LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
|
||||
```
|
||||
Build with `-DUSE_MACA=ON`. Optional overrides:
|
||||
- `-DMACA_ROOT=/path/to/maca`
|
||||
- `-DMACA_INCLUDE_DIR=/path/to/maca/include`
|
||||
- `-DMACA_LIB_DIR=/path/to/maca/lib64`
|
||||
- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"` (semicolon-separated CMake list)
|
||||
|
||||
6. Install yalantinglibs
|
||||
```bash
|
||||
git clone https://github.com/alibaba/yalantinglibs.git
|
||||
cd yalantinglibs
|
||||
|
|
@ -123,7 +159,7 @@ pip install mooncake-transfer-engine-non-cuda
|
|||
make install
|
||||
```
|
||||
|
||||
5. In the root directory of this project, run the following commands:
|
||||
7. In the root directory of this project, run the following commands:
|
||||
```bash
|
||||
mkdir build
|
||||
cd build
|
||||
|
|
@ -131,7 +167,7 @@ pip install mooncake-transfer-engine-non-cuda
|
|||
make -j
|
||||
```
|
||||
|
||||
6. Install Mooncake python package and mooncake_master executable
|
||||
8. Install Mooncake python package and mooncake_master executable
|
||||
```bash
|
||||
make install
|
||||
```
|
||||
|
|
@ -150,13 +186,26 @@ cd /Mooncake-main/build/mooncake-transfer-engine/example
|
|||
|
||||
## Advanced Compile Options
|
||||
The following options can be used during `cmake ..` to specify whether to compile certain components of Mooncake.
|
||||
- `-DUSE_CUDA=[ON|OFF]`: Enable GPU Direct RDMA and NVMe-of support
|
||||
- `-DUSE_CUDA=[ON|OFF]`: Enable GPU memory support (GPUDirect RDMA, NVMe-oF, and GPU-aware TCP transport). **Default: OFF.** Required when transferring GPU memory (e.g., KV cache in vLLM disaggregated serving), even when using TCP protocol.
|
||||
- `-DUSE_MNNVL=[ON|OFF]`: Enable Multi-Node NVLink transport support, default is OFF. **Note:** `-DUSE_CUDA` is required when `-DUSE_MNNVL` is on (not used when building with `-DUSE_MUSA=ON`, `-DUSE_HIP=ON`, or `-DUSE_MACA=ON`).
|
||||
- `-DUSE_MUSA=[ON|OFF]`: Enable Moore Threads GPU support via MUSA
|
||||
- `-DUSE_MACA=[ON|OFF]`: Enable MetaX (Muxi) GPU support via MACA.
|
||||
- `-DMACA_ROOT=/path/to/maca`: Override the MACA SDK root (`MACA_HOME` env var is also honored; default `/opt/maca`).
|
||||
- `-DMACA_INCLUDE_DIR=/path/to/include`: Override MACA include directory when `-DUSE_MACA=ON`.
|
||||
- `-DMACA_LIB_DIR=/path/to/lib64`: Override MACA library directory when `-DUSE_MACA=ON`.
|
||||
- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`: Override MACA runtime libraries linked by `transfer_engine`.
|
||||
- `-DUSE_HIP=[ON|OFF]`: Enable AMD GPU support via HIP/ROCm
|
||||
- `-DUSE_MLU=[ON|OFF]`: Enable Cambricon MLU memory support via Neuware. **Default: OFF.** Supports MLU memory detection, topology discovery, and RDMA registration for Transfer Engine.
|
||||
- `-DNEUWARE_ROOT=/path/to/neuware`: Override the default Neuware SDK root used when `-DUSE_MLU=ON`. If unset, Mooncake uses `NEUWARE_HOME` or `/usr/local/neuware`.
|
||||
- `-DMLU_INCLUDE_DIR=/path/to/include`: Override the Neuware include directory when `-DUSE_MLU=ON`.
|
||||
- `-DMLU_LIB_DIR=/path/to/lib64`: Override the Neuware library directory when `-DUSE_MLU=ON`.
|
||||
- `-DUSE_EFA=[ON|OFF]`: Enable AWS Elastic Fabric Adapter transport via libfabric. **Default: OFF.** See [EFA Transport](../design/transfer-engine/efa_transport.md) for details.
|
||||
- `-DUSE_INTRA_NVLINK=[ON|OFF]`: Enable intranode nvlink transport
|
||||
- `-DUSE_CXL=[ON|OFF]`: Enable CXL support
|
||||
- `-DWITH_STORE=[ON|OFF]`: Build Mooncake Store component
|
||||
- `-DWITH_P2P_STORE=[ON|OFF]`: Enable Golang support and build P2P Store component, require go 1.23+
|
||||
- `-DWITH_WITH_RUST_EXAMPLE=[ON|OFF]`: Enable Rust support
|
||||
- `-DWITH_EP=[ON|OFF]`: Build the EP (Expert Parallelism) and PG Python extensions for CUDA. Requires CUDA toolkit and PyTorch. Use `-DEP_TORCH_VERSIONS="2.9.1"` (semicolon-separated) to build for specific PyTorch versions, or leave empty to use the currently-installed torch. The CUDA version is detected automatically. **Default: OFF.**
|
||||
- `-DUSE_REDIS=[ON|OFF]`: Enable Redis-based metadata service
|
||||
- `-DUSE_HTTP=[ON|OFF]`: Enable Http-based metadata service
|
||||
- `-DUSE_ETCD=[ON|OFF]`: Enable etcd-based metadata service, require go 1.23+
|
||||
|
|
|
|||
|
|
@ -144,3 +144,32 @@ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disag
|
|||
```
|
||||
- Set `--elastic-ep-backend` and `--moe-a2a-backend` to "mooncake" to enable Mooncake EP Backend.
|
||||
- The value of `--mooncake-ib-device` should be the same as `--disaggregation-ib-device`.
|
||||
|
||||
|
||||
### To enable Mooncake EPD Backend
|
||||
|
||||
Encoder:
|
||||
```bash
|
||||
python -m sglang.launch_server \
|
||||
--model-path $MODEL \
|
||||
--encoder-only \
|
||||
--encoder-transfer-backend mooncake \
|
||||
--port $PORT
|
||||
```
|
||||
|
||||
Prefill:
|
||||
```bash
|
||||
python -m sglang.launch_server \
|
||||
--model-path $MODEL \
|
||||
--disaggregation-mode prefill \
|
||||
--disaggregation-transfer-backend mooncake \
|
||||
--encoder-transfer-backend mooncake \
|
||||
--tp $TP \
|
||||
--mem-fraction-static $MEM_FRACTION \
|
||||
--chunked-prefill-size $CHUNK_SIZE \
|
||||
--language-only \
|
||||
--encoder-urls http://127.0.0.1:30002 http://127.0.0.1:30003 http://127.0.0.1:30004 http://127.0.0.1:30005 http://127.0.0.1:30006 http://127.0.0.1:30007 \
|
||||
--port $PORT
|
||||
```
|
||||
|
||||
- Set `--encoder-transfer-backend` to "mooncake" to enable Mooncake Backend.
|
||||
|
|
|
|||
|
|
@ -280,6 +280,89 @@ Distributed deployment of Mooncake is straightforward. Similar to the single-nod
|
|||
|
||||
Mooncake also supports high availability mode. This mode enhances fault tolerance by running the `master service` as a cluster of multiple master nodes coordinated through an `etcd` cluster. The master nodes use `etcd` to elect a leader, which is responsible for handling client requests. For more details about how to deploy in this mode, please refer to our [documents](https://kvcache-ai.github.io/Mooncake/).
|
||||
|
||||
### Deployment with Dummy Client (Experimental)
|
||||
|
||||
In addition to the standard deployment where SGLang acts as a full Mooncake node, you can use the **Dummy Client** mode. In this mode, SGLang connects to a local **Mooncake Store Service** (Real Client) via RPC/IPC. This decouples the SGLang process from the heavy RDMA and memory management, potentially improving stability and allowing the cache to persist even if the SGLang process restarts.
|
||||
|
||||
**Architecture:**
|
||||
* **Mooncake Master**: Manages the cluster topology (same as standard).
|
||||
* **Mooncake Store Service (Real Client)**: Manages the actual memory pool and RDMA connections. Must be running locally.
|
||||
* **SGLang Server (Dummy Client)**: Connects to the local Store Service to access the cache.
|
||||
|
||||
#### 1. Launch Services (Master & Store)
|
||||
|
||||
First, start the `master service` and the `store service`. The `store service` acts as the Real Client.
|
||||
|
||||
**Start Master:**
|
||||
```bash
|
||||
mooncake_master --eviction_high_watermark_ratio=0.95
|
||||
```
|
||||
|
||||
**Start Store Service (Real Client):** Crucially, the default port (50052) is used for internal RPC, which the Dummy Client will connect to.
|
||||
```bash
|
||||
mooncake_client --global_segment_size=4GB
|
||||
```
|
||||
|
||||
**Parameter Explanation:**
|
||||
|
||||
- **`host`**: (string, default: "0.0.0.0"): The hostname of the client.
|
||||
|
||||
- **`port`**: (int, default: 50052): The port number the client service listens on.
|
||||
|
||||
- **`global_segment_size`**: (string, default: "4GB"): The size of the global segment to be allocated by the client.
|
||||
|
||||
- **`master_server_address`**: (string, default: "localhost:50051"): The address of the Master Service.
|
||||
|
||||
- **`metadata_server`**: (string, default: "http://localhost:8080/metadata"): The address of the metadata service.
|
||||
|
||||
- **`protocol`**: (string, default: "tcp"): The protocol used by the Transfer Engine.
|
||||
|
||||
- **`device_name`**: (string, default: ""): The device name used by the Transfer Engine.
|
||||
|
||||
- **`threads`**: (int, default: 1): The number of threads used by the client.
|
||||
|
||||
#### 2. Launch SGLang (Dummy Client)
|
||||
Configure SGLang to connect to the Real Client using the client_server_address parameter.
|
||||
|
||||
**Using extra-config of sglang arguments to configure Mooncake**
|
||||
|
||||
```bash
|
||||
python -m sglang.launch_server \
|
||||
--enable-hierarchical-cache \
|
||||
--hicache-storage-backend mooncake \
|
||||
--model-path [model_path] \
|
||||
--hicache-storage-backend-extra-config '{"standalone_storage": true, "client_server_address": "127.0.0.1:50052"}'
|
||||
```
|
||||
|
||||
**Using JSON file to configure Mooncake**
|
||||
|
||||
SGLang server can load Mooncake config from `SGLANG_HICACHE_MOONCAKE_CONFIG_PATH`.
|
||||
|
||||
```bash
|
||||
export SGLANG_HICACHE_MOONCAKE_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hicache/mooncake_config.json
|
||||
|
||||
echo '{
|
||||
"standalone_storage": true,
|
||||
"client_server_address": "127.0.0.1:50052"
|
||||
}' > ${SGLANG_HICACHE_MOONCAKE_CONFIG_PATH}
|
||||
|
||||
python -m sglang.launch_server \
|
||||
--enable-hierarchical-cache \
|
||||
--hicache-storage-backend mooncake \
|
||||
--model-path [model_path]
|
||||
```
|
||||
|
||||
**Using env variables to configure Mooncake**
|
||||
|
||||
```bash
|
||||
MOONCAKE_STANDALONE_STORAGE=1
|
||||
MOONCAKE_CLIENT="127.0.0.1:50052"
|
||||
python -m sglang.launch_server \
|
||||
--enable-hierarchical-cache \
|
||||
--hicache-storage-backend mooncake \
|
||||
--model-path [model_path]
|
||||
```
|
||||
|
||||
### Prefill/Decode Disaggregation
|
||||
|
||||
In **PD disaggregation**, the configurations for the `metadata service`, `mooncake master`, and the optional `store service` remain the same as described above. The difference is that SGLang introduces three distinct roles: `prefill worker`, `decode worker`, and `router`.
|
||||
|
|
|
|||
|
|
@ -5,7 +5,8 @@
|
|||
|
||||
vllmv1-lmcache-integration
|
||||
vllm-integration-v0.2
|
||||
vllm-integration-v1
|
||||
vllm-integration-v0.3
|
||||
vllm-integration-v1.0
|
||||
::::
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -22,9 +22,12 @@ pip install mooncake-transfer-engine-non-cuda
|
|||
|
||||
## Transfer Engine Quick Start
|
||||
|
||||
> **Note**: When using RDMA protocol, you may need to run with `sudo` for proper permissions.
|
||||
|
||||
### Start Transfer Engine Receiver (Server)
|
||||
|
||||
```python
|
||||
|
||||
import numpy as np
|
||||
import zmq
|
||||
from mooncake.engine import TransferEngine
|
||||
|
|
@ -37,7 +40,7 @@ def main():
|
|||
|
||||
HOSTNAME = "localhost" # localhost for simple demo
|
||||
METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
|
||||
PROTOCOL = "tcp" # [rdma, tcp, ...]
|
||||
PROTOCOL = "rdma" # [rdma, tcp, ...]
|
||||
DEVICE_NAME = "" # auto discovery if empty
|
||||
|
||||
# Initialize server engine
|
||||
|
|
@ -93,12 +96,14 @@ def main():
|
|||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
```
|
||||
|
||||
### Start Transfer Engine Sender (Client)
|
||||
|
||||
```python
|
||||
|
||||
|
||||
import numpy as np
|
||||
import zmq
|
||||
from mooncake.engine import TransferEngine
|
||||
|
|
@ -121,7 +126,7 @@ def main():
|
|||
# Initialize client engine
|
||||
HOSTNAME = "localhost" # localhost for simple demo
|
||||
METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
|
||||
PROTOCOL = "tcp" # [rdma, tcp, ...]
|
||||
PROTOCOL = "rdma" # [rdma, tcp, ...]
|
||||
DEVICE_NAME = "" # auto discovery if empty
|
||||
|
||||
client_engine = TransferEngine()
|
||||
|
|
@ -139,7 +144,8 @@ def main():
|
|||
client_len = client_buffer.nbytes
|
||||
|
||||
# Register memory with Mooncake
|
||||
ret_value = client_engine.register_memory(client_ptr, client_len)
|
||||
if PROTOCOL == "rdma":
|
||||
ret_value = client_engine.register_memory(client_ptr, client_len)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory registration failed.")
|
||||
raise RuntimeError("Mooncake memory registration failed.")
|
||||
|
|
@ -162,16 +168,18 @@ def main():
|
|||
print("Transfer failed!")
|
||||
|
||||
# Cleanup
|
||||
ret_value = client_engine.unregister_memory(client_ptr)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory deregistration failed.")
|
||||
raise RuntimeError("Mooncake memory deregistration failed.")
|
||||
if PROTOCOL == "rdma":
|
||||
ret_value = client_engine.unregister_memory(client_ptr)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory deregistration failed.")
|
||||
raise RuntimeError("Mooncake memory deregistration failed.")
|
||||
|
||||
socket.close()
|
||||
context.term()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
```
|
||||
|
||||
### More Examples and Documentation
|
||||
|
|
@ -192,6 +200,19 @@ mooncake_master \
|
|||
```
|
||||
This exposes the metadata endpoint at `http://<host>:<port>/metadata`.
|
||||
|
||||
If the master runs in a container and its IP is dynamic, set `--rpc_interface=<ifname>` such as `--rpc_interface=eth0`. Mooncake Master will resolve the current IPv4 address from that interface at startup instead of relying on a fixed `--rpc_address`.
|
||||
|
||||
Optional: Use the free-ratio-first allocation strategy for better load balancing across segments with different sizes or utilization:
|
||||
|
||||
```bash
|
||||
mooncake_master \
|
||||
--allocation_strategy=free_ratio_first \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_port=8080
|
||||
```
|
||||
|
||||
The free-ratio-first strategy balances memory utilization ratio across segments by sampling multiple candidates and preferentially allocating to those with higher free space ratios, leading to more even utilization.
|
||||
|
||||
### Hello World Example
|
||||
|
||||
```python
|
||||
|
|
@ -224,4 +245,4 @@ store.close()
|
|||
|
||||
### More Examples and Documentation
|
||||
|
||||
Please refer to the [Mooncake Store Python API](../python-api-reference/mooncake-store.md), [Mooncake Store](../design/mooncake-store.md) and [Mooncake Store Deployment & Operations Guide](../deployment/mooncake-store-deployment-guide.md) for more examples and documentation.
|
||||
Please refer to the [Mooncake Store Python API](../python-api-reference/mooncake-store.md), [Mooncake Store](../design/mooncake-store.md) and [Mooncake Store Deployment & Operations Guide](../deployment/mooncake-store-deployment-guide.md) for more examples and documentation.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,370 @@
|
|||
# Supported Communication Protocols
|
||||
|
||||
Mooncake Transfer Engine supports multiple communication protocols for data transfer between nodes in a cluster. The protocol selection depends on your hardware capabilities and performance requirements.
|
||||
|
||||
## Quick Reference
|
||||
|
||||
| Protocol | Hardware Required | Use Case | Python API Support |
|
||||
|----------|-------------------|----------|-------------------|
|
||||
| **tcp** | Standard network | General purpose, works everywhere | ✅ Primary |
|
||||
| **rdma** | RDMA-capable NIC | High-performance, low-latency | ✅ Primary |
|
||||
| **efa** | AWS EFA-capable instance | High-performance on AWS (libfabric SRD) | ✅ Primary |
|
||||
| **nvmeof** | NVMe-oF capable storage | Direct NVMe storage access | ⚠️ Advanced |
|
||||
| **nvlink** | NVIDIA MNNVL | Inter-node GPU communication | ⚠️ Advanced |
|
||||
| **nvlink_intra** | NVIDIA NVLink | Intra-node GPU communication | ⚠️ Advanced |
|
||||
| **hip** | AMD ROCm/HIP | AMD GPU communication | ⚠️ Advanced |
|
||||
| **barex** | RDMA-capable NIC | Bare-metal RDMA extension | ⚠️ Advanced |
|
||||
| **cxl** | CXL-capable hardware | Memory pooling and sharing | ⚠️ Advanced |
|
||||
| **ascend** | Huawei Ascend NPU | Ascend NPU communication | ⚠️ Advanced |
|
||||
|
||||
## Commonly Used Protocols (Python API)
|
||||
|
||||
### TCP (Default)
|
||||
|
||||
**Description:** Standard TCP/IP network protocol.
|
||||
|
||||
**Use When:**
|
||||
- No special hardware is available
|
||||
- Testing or development environments
|
||||
- Compatibility is more important than performance
|
||||
|
||||
**Configuration:**
|
||||
```python
|
||||
# Python API
|
||||
engine.initialize(
|
||||
hostname="localhost",
|
||||
metadata_server="P2PHANDSHAKE",
|
||||
protocol="tcp", # No device_name needed
|
||||
device_name=""
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
# Environment variables
|
||||
export MOONCAKE_PROTOCOL="tcp"
|
||||
```
|
||||
|
||||
**Advantages:**
|
||||
- Works in all environments
|
||||
- No special hardware required
|
||||
- Simple setup
|
||||
|
||||
**Limitations:**
|
||||
- Lower throughput compared to RDMA
|
||||
- Higher CPU overhead
|
||||
- Higher latency
|
||||
|
||||
### RDMA (Recommended for Production)
|
||||
|
||||
**Description:** Remote Direct Memory Access protocol providing high-performance, low-latency data transfer with minimal CPU overhead. Supports accelerator-aware memory registration, including NVIDIA GPUDirect RDMA for CUDA buffers and Cambricon MLU buffers when built with Neuware.
|
||||
|
||||
**Hardware Support:**
|
||||
- InfiniBand
|
||||
- RoCE (RDMA over Converged Ethernet)
|
||||
- eRDMA (Elastic RDMA)
|
||||
- NVIDIA GPUDirect RDMA
|
||||
- Non-NVIDAI GPUDirect RDMA (e.g., Intel E810 RDMA NIC)
|
||||
- Cambricon MLU memory via Neuware (`-DUSE_MLU=ON`)
|
||||
|
||||
**Use When:**
|
||||
- High-performance networking is required
|
||||
- RDMA-capable NICs are available
|
||||
- Low latency is critical (e.g., distributed inference, KV cache transfer)
|
||||
|
||||
**Note:** If no RDMA HCA (Host Channel Adapter) is detected on the system, the Transfer Engine will automatically fall back to TCP protocol for compatibility.
|
||||
|
||||
**MLU Note:** Cambricon MLU support uses the standard `rdma` data path. There is no separate `mlu` protocol string. To enable MLU memory detection, topology discovery, and DMA-BUF based registration, build Transfer Engine with `-DUSE_MLU=ON` and make Neuware available through `NEUWARE_HOME` or `NEUWARE_ROOT`.
|
||||
|
||||
**Configuration:**
|
||||
```python
|
||||
# Python API - With specific device
|
||||
engine.initialize(
|
||||
hostname="node1",
|
||||
metadata_server="etcd://10.0.0.1:2379",
|
||||
protocol="rdma",
|
||||
device_name="mlx5_0" # Specify your RDMA device
|
||||
)
|
||||
|
||||
# Python API - With auto-discovery
|
||||
engine.initialize(
|
||||
hostname="node1",
|
||||
metadata_server="P2PHANDSHAKE",
|
||||
protocol="rdma",
|
||||
device_name="auto-discovery" # Automatically detect optimal device
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
# Environment variables
|
||||
export MOONCAKE_PROTOCOL="rdma"
|
||||
export MOONCAKE_DEVICE="mlx5_0" # or "auto-discovery"
|
||||
```
|
||||
|
||||
**Device Discovery:**
|
||||
To find available RDMA devices on your system:
|
||||
```bash
|
||||
ibv_devices # List InfiniBand/RDMA devices
|
||||
# Example output: mlx5_0, mlx5_1, erdma_0, etc.
|
||||
```
|
||||
|
||||
**Advantages:**
|
||||
- Very high throughput (up to 200 Gbps per NIC)
|
||||
- Ultra-low latency (sub-microsecond)
|
||||
- Minimal CPU overhead
|
||||
- Supports GPUDirect RDMA for zero-copy GPU transfers
|
||||
- Multi-NIC bandwidth aggregation
|
||||
- Topology-aware path selection
|
||||
|
||||
**Limitations:**
|
||||
- Requires RDMA-capable hardware
|
||||
- May require elevated permissions (sudo)
|
||||
- More complex network configuration
|
||||
|
||||
**Performance Tips:**
|
||||
- Use multiple RDMA NICs for bandwidth aggregation
|
||||
- Enable GPUDirect RDMA for GPU memory transfers
|
||||
- Configure proper NUMA affinity for optimal performance
|
||||
- See [Transfer Engine Benchmark Tuning](../design/transfer-engine/transfer-engine-bench-tuning.md) for detailed optimization
|
||||
|
||||
### EFA (AWS Elastic Fabric Adapter)
|
||||
|
||||
**Description:** AWS EFA transport using libfabric's Scalable Reliable Datagram (SRD) protocol, providing high-bandwidth RDMA-like performance on AWS instances without traditional RDMA support.
|
||||
|
||||
**Use When:**
|
||||
- Running on AWS EFA-enabled instances (e.g., p5e.48xlarge, p6-b200.48xlarge, p4d.24xlarge)
|
||||
- High-performance networking is required on AWS
|
||||
- Traditional RDMA (ibverbs QP) is not supported by the hardware
|
||||
|
||||
**Configuration:**
|
||||
```python
|
||||
# Python API
|
||||
engine.initialize(
|
||||
hostname="localhost",
|
||||
metadata_server="P2PHANDSHAKE",
|
||||
protocol="efa",
|
||||
device_name=""
|
||||
)
|
||||
```
|
||||
|
||||
**Build Requirements:**
|
||||
```bash
|
||||
cmake .. -DUSE_EFA=ON -DUSE_CUDA=ON
|
||||
```
|
||||
|
||||
> **Note:** `-DUSE_CUDA=ON` is required when transferring GPU memory. Without it, fallback to TCP protocol will fail with "Bad address" errors on GPU buffers.
|
||||
|
||||
**Advantages:**
|
||||
- High throughput (~170 GB/s with 8 EFA devices, tuned)
|
||||
- Bypasses kernel network stack
|
||||
- Available on all AWS EFA-enabled instances
|
||||
|
||||
**Limitations:**
|
||||
- AWS-only
|
||||
- Software-emulated RDMA writes (higher CPU overhead than true RDMA)
|
||||
- ~88% of RoCE RDMA throughput
|
||||
|
||||
**Documentation:** See [EFA Transport](../design/transfer-engine/efa_transport.md) for build instructions, benchmarks, and tuning.
|
||||
|
||||
## Advanced Protocols (C++ Transfer Engine)
|
||||
|
||||
The following protocols are available at the C++ Transfer Engine level for specialized use cases. They are not commonly used through the Python API.
|
||||
|
||||
### NVMe over Fabric (nvmeof)
|
||||
|
||||
**Description:** Direct data transfer between NVMe storage and DRAM/VRAM using GPUDirect Storage, bypassing the CPU for zero-copy operations.
|
||||
|
||||
**Use When:**
|
||||
- Direct NVMe storage access is needed
|
||||
- Implementing multi-tier storage (DRAM/VRAM/NVMe)
|
||||
- Working with large datasets that don't fit in memory
|
||||
|
||||
**Requirements:**
|
||||
- NVMe-oF capable storage
|
||||
- Properly mounted remote storage nodes
|
||||
|
||||
### NVLink (nvlink)
|
||||
|
||||
**Description:** NVIDIA MNNVL (Multi-Node NVLink) protocol for high-bandwidth, low-latency GPU-to-GPU communication across nodes.
|
||||
|
||||
**Use When:**
|
||||
- Inter-node GPU communication is required
|
||||
- Using NVIDIA MNNVL (Multi-Node NVLink)
|
||||
- Maximum GPU bandwidth is needed
|
||||
|
||||
**Requirements:**
|
||||
- NVIDIA MNNVL hardware
|
||||
- Compiled with `USE_MNNVL=ON`
|
||||
|
||||
**Configuration:**
|
||||
```bash
|
||||
# Set MC_FORCE_MNNVL=true to use MNNVL even when RDMA NICs are present
|
||||
export MC_FORCE_MNNVL=true
|
||||
```
|
||||
|
||||
**Note:** When `protocol="rdma"` is set and RDMA NICs exist, you must explicitly set `MC_FORCE_MNNVL=true` to use MNNVL instead of RDMA. If no RDMA HCA is detected, MNNVL will be used automatically.
|
||||
|
||||
### Intra-Node NVLink (nvlink_intra)
|
||||
|
||||
**Description:** NVIDIA NVLink for GPU-to-GPU communication within a single node.
|
||||
|
||||
**Use When:**
|
||||
- Local GPU-to-GPU transfers are needed
|
||||
- Maximizing intra-node GPU bandwidth
|
||||
|
||||
**Requirements:**
|
||||
- NVIDIA NVLink hardware
|
||||
- Compiled with `USE_INTRA_NVLINK=ON`
|
||||
|
||||
### HIP Transport (hip)
|
||||
|
||||
**Description:** AMD ROCm/HIP transport for GPU communication using IPC handles or Shareable handles.
|
||||
|
||||
**Use When:**
|
||||
- Working with AMD GPUs
|
||||
- Need intra-node GPU communication on AMD hardware
|
||||
|
||||
**Requirements:**
|
||||
- AMD ROCm/HIP runtime
|
||||
- AMD GPUs
|
||||
|
||||
### Barex Transport (barex)
|
||||
|
||||
**Description:** Bare-metal RDMA extension protocol for specialized RDMA configurations.
|
||||
|
||||
**Use When:**
|
||||
- Advanced RDMA features are required
|
||||
- Custom RDMA configurations
|
||||
|
||||
**Requirements:**
|
||||
- RDMA-capable hardware
|
||||
- Specialized configuration
|
||||
|
||||
### CXL Transport (cxl)
|
||||
|
||||
**Description:** Compute Express Link for memory pooling and sharing across devices.
|
||||
|
||||
**Use When:**
|
||||
- CXL memory pooling is available
|
||||
- Memory disaggregation is needed
|
||||
|
||||
**Requirements:**
|
||||
- CXL-capable hardware
|
||||
|
||||
### Ascend Transport (ascend)
|
||||
|
||||
**Description:** Huawei Ascend NPU communication using HCCL (Huawei Collective Communication Library) or direct transport.
|
||||
|
||||
**Use When:**
|
||||
- Working with Huawei Ascend NPUs
|
||||
- Distributed inference on Ascend hardware
|
||||
|
||||
**Requirements:**
|
||||
- Huawei Ascend NPU hardware
|
||||
- HCCL runtime
|
||||
|
||||
**Documentation:**
|
||||
- [Heterogeneous Ascend](../design/transfer-engine/heterogeneous_ascend.md)
|
||||
- [Ascend Transport](../design/transfer-engine/ascend_transport.md)
|
||||
|
||||
## Configuration Examples
|
||||
|
||||
### Configuration File (JSON)
|
||||
|
||||
**TCP Configuration:**
|
||||
```json
|
||||
{
|
||||
"local_hostname": "localhost",
|
||||
"metadata_server": "localhost:8080",
|
||||
"protocol": "tcp",
|
||||
"device_name": "",
|
||||
"master_server_address": "localhost:8081"
|
||||
}
|
||||
```
|
||||
|
||||
**RDMA Configuration:**
|
||||
```json
|
||||
{
|
||||
"local_hostname": "node1",
|
||||
"metadata_server": "etcd://10.0.0.1:2379",
|
||||
"global_segment_size": "3GB",
|
||||
"local_buffer_size": "1GB",
|
||||
"protocol": "rdma",
|
||||
"device_name": "mlx5_0",
|
||||
"master_server_address": "10.0.0.1:8081"
|
||||
}
|
||||
```
|
||||
|
||||
### Environment Variables
|
||||
|
||||
```bash
|
||||
# TCP (Default)
|
||||
export MOONCAKE_PROTOCOL="tcp"
|
||||
|
||||
# RDMA with specific device
|
||||
export MOONCAKE_PROTOCOL="rdma"
|
||||
export MOONCAKE_DEVICE="mlx5_0"
|
||||
|
||||
# RDMA with auto-discovery
|
||||
export MOONCAKE_PROTOCOL="rdma"
|
||||
export MOONCAKE_DEVICE="auto-discovery"
|
||||
|
||||
# Other configuration
|
||||
export MOONCAKE_MASTER="10.0.0.1:50051"
|
||||
export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE"
|
||||
export MOONCAKE_LOCAL_HOSTNAME="node1"
|
||||
```
|
||||
|
||||
## Choosing the Right Protocol
|
||||
|
||||
| Scenario | Recommended Protocol | Notes |
|
||||
|----------|---------------------|-------|
|
||||
| Development/Testing | tcp | Simple setup, no special hardware |
|
||||
| Production Inference | rdma | Best performance and latency |
|
||||
| AWS Cloud (EFA instances) | efa | High performance on p5e, p6-b200, p4d, etc. |
|
||||
| Cloud Environments | tcp or rdma (if available) | Check cloud provider support |
|
||||
| Multi-tier Storage | rdma + nvmeof | Combine protocols for different layers |
|
||||
| AMD GPU Clusters | rdma + hip | Use HIP for local GPU communication |
|
||||
| Cambricon MLU Clusters | rdma | Build with `-DUSE_MLU=ON`; MLU uses the normal RDMA protocol |
|
||||
| Ascend NPU Clusters | rdma + ascend | Use Ascend for NPU-specific operations |
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### RDMA Connection Issues
|
||||
|
||||
1. **Check RDMA devices:**
|
||||
```bash
|
||||
ibv_devices
|
||||
ibv_devinfo
|
||||
```
|
||||
|
||||
2. **Verify network connectivity:**
|
||||
```bash
|
||||
# Test RDMA connectivity (requires rdma-core tools)
|
||||
rping -s # On server
|
||||
rping -c -a <server_ip> -v # On client
|
||||
```
|
||||
|
||||
3. **Check permissions:**
|
||||
- RDMA may require elevated permissions
|
||||
- Run with `sudo` if necessary
|
||||
- Configure proper udev rules for non-root access
|
||||
|
||||
4. **Firewall configuration:**
|
||||
- Ensure RDMA ports are not blocked
|
||||
- Check InfiniBand subnet manager is running
|
||||
|
||||
### Protocol Selection
|
||||
|
||||
If a protocol fails to initialize:
|
||||
1. Verify hardware support
|
||||
2. Check that required drivers are installed
|
||||
3. Ensure compile-time flags are set correctly (for C++ protocols)
|
||||
4. Fall back to TCP for basic functionality
|
||||
|
||||
## See Also
|
||||
|
||||
- [Quick Start Guide](quick-start.md) - Getting started with Mooncake
|
||||
- [Transfer Engine Design](../design/transfer-engine/index.md) - Detailed architecture
|
||||
- [Transfer Engine Benchmark](../design/transfer-engine/transfer-engine-bench-tuning.md) - Performance tuning
|
||||
- [Python API Reference](../python-api-reference/transfer-engine.md) - API documentation
|
||||
- [Deployment Guide](../deployment/mooncake-store-deployment-guide.md) - Production deployment
|
||||
|
|
@ -61,6 +61,41 @@ curl "http://localhost:8080/query_key?key=my_object"
|
|||
}
|
||||
```
|
||||
|
||||
#### `/batch_query_keys`
|
||||
Retrieve replica information for multiple keys in a single request, including memory locations and transport endpoints for each key.
|
||||
|
||||
**Method**: `GET`
|
||||
**Parameters**: `keys` (query parameter) - Comma-separated list of object keys to query (format: key1,key2,key3)
|
||||
**Content-Type**: `application/json; charset=utf-8`
|
||||
**Response**: JSON-formatted mapping of keys to their respective replica descriptors
|
||||
|
||||
**Example**:
|
||||
```bash
|
||||
curl "http://localhost:8080/batch_query_keys?keys=key1,key2,key3"
|
||||
```
|
||||
|
||||
**Response Format**:
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"data": {
|
||||
"key1": {
|
||||
"ok": true,
|
||||
"values": [
|
||||
{
|
||||
"transport_endpoint_": "hostname:port",
|
||||
"buffer_descriptor": {...}
|
||||
}
|
||||
]
|
||||
},
|
||||
"key2": {
|
||||
"ok": false,
|
||||
"error": "error message"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### `/get_all_keys`
|
||||
List all keys currently stored in the distributed system.
|
||||
|
||||
|
|
|
|||
Binary file not shown.
|
After Width: | Height: | Size: 42 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 118 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 139 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 73 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 140 KiB |
|
|
@ -21,31 +21,38 @@
|
|||
</p>
|
||||
:::
|
||||
|
||||
Mooncake is the serving platform for <a href="https://kimi.ai/">Kimi</a>, a leading LLM service provided by <a href="https://www.moonshot.cn/">Moonshot AI</a>.
|
||||
Mooncake is the serving platform for <a href="https://kimi.ai/">Kimi</a>, a leading LLM service provided by <a href="https://www.moonshot.cn/">Moonshot AI</a>.
|
||||
Now both the Transfer Engine and Mooncake Store are open-sourced!
|
||||
This repository also hosts its technical report and the open sourced traces.
|
||||
This repository also hosts its technical report and the open-sourced traces.
|
||||
|
||||
<h2 id="updates">🔄 Updates</h2>
|
||||
|
||||
- **Dec 18, 2025**: Mooncake has now implemented a vLLM connector, enabling direct support for the Prefill-Decode (PD) separation architecture in vLLM v1.
|
||||
- **Mar 19, 2026**: [TorchSpec: Speculative Decoding Training at Scale](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale) is [open sourced](https://github.com/torchspec-project/TorchSpec), using Mooncake to decouple inference and training via efficient hidden states management.
|
||||
- **Feb 12, 2026**: [Mooncake Joins PyTorch Ecosystem](https://pytorch.org/blog/mooncake-joins-pytorch-ecosystem/) We are thrilled to announce that Mooncake has officially joined the PyTorch Ecosystem!
|
||||
- **Jan 28, 2026**: [FlexKV](https://github.com/taco-project/FlexKV), a distributed KV store and cache system from Tencent and NVIDIA in collaboration with the community, now supports [distributed KVCache reuse](https://github.com/taco-project/FlexKV/blob/main/docs/dist_reuse/README_en.md) with the Mooncake Transfer Engine.
|
||||
- **Dec 23, 2025**: SGLang introduces [Encode-Prefill-Decode (EPD) Disaggregation](https://lmsys.org/blog/2026-01-12-epd/) with Mooncake as a transfer backend. This integration allows decoupling compute-intensive multimodal encoders (e.g., Vision Transformers) from language model nodes, utilizing Mooncake's RDMA engine for zero-copy transfer of large multimodal embeddings.
|
||||
- **Dec 19, 2025**: Mooncake Transfer Engine has been [integrated into TensorRT LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils) for KVCache transfer in PD-disaggregated inference.
|
||||
- **Dec 19, 2025**: Mooncake Transfer Engine has been directly integrated into vLLM v1 as a [KV Connector](https://docs.vllm.ai/en/latest/features/mooncake_connector_usage/) in PD-disaggregated setups.
|
||||
- **Nov 07, 2025**: [RBG + SGLang HiCache + Mooncake](https://github.com/sgl-project/rbg/blob/main/keps/74-mooncake-integration/README.md), a role-based out-of-the-box solution for cloud native deployment, which is elastic, scalable, and high-performance.
|
||||
- **Sept 18, 2025**: Mooncake Store empowers vLLM Ascend by serving as [the distributed KV cache pool backend](https://docs.vllm.ai/projects/ascend/zh-cn/main/user_guide/feature_guide/kv_pool.html).
|
||||
- **Sept 10, 2025**: SGLang officially supports Mooncake Store as a [hierarchical KV caching storage backend](https://lmsys.org/blog/2025-09-10-sglang-hicache/). The integration extends RadixAttention with multi-tier KV cache storage across device, host, and remote storage layers.
|
||||
- **Sept 10, 2025**: The official & high-performance version of Mooncake P2P Store is open-sourced as [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine/). It has been successfully applied in K1.5 and K2 production training, updating Kimi-K2 model (1T parameters) across thousands of GPUs in ~20s.
|
||||
- **Aug 23, 2025**: [xLLM](https://github.com/jd-opensource/xllm) high-performance inference engine builds hybrid KV cache management based on Mooncake, supporting global KV cache management with intelligent offloading and prefetching.
|
||||
- **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://github.com/vllm-project/vllm-ascend/blob/main/examples/disaggregated_prefill_v1/mooncake_connector_deployment_guide.md) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs.
|
||||
- **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/feature_guide/disaggregated_prefill.html) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs.
|
||||
- **Jul 20, 2025**: Mooncake powers [the deployment of Kimi K2](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/) on 128 H200 GPUs with PD disaggregation and large-scale expert parallelism, achieving 224k tokens/sec prefill throughput and 288k tokens/sec decode throughput.
|
||||
- **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/lmdeploy-integration-v0.9.md) for LMDeploy.
|
||||
- **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](https://kvcache-ai.github.io/Mooncake/getting_started/examples/lmdeploy-integration-v0.9.html) for LMDeploy.
|
||||
- **May 9, 2025**: NIXL officially supports Mooncake Transfer Engine as [a backend plugin](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/mooncake/README.md).
|
||||
- **May 8, 2025**: Mooncake x LMCache <a href="https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/lmcache-integration.md" target="_blank">unite</a> to pioneer KVCache-centric LLM serving system.
|
||||
- **May 5, 2025**: Supported by Mooncake Team, SGLang release <a href="https://lmsys.org/blog/2025-05-05-large-scale-ep/" target="_blank">guidance</a> to deploy DeepSeek with PD Disaggregation on 96 H100 GPUs.
|
||||
- **Apr 22, 2025**: LMCache officially supports Mooncake Store as a <a href="https://blog.lmcache.ai/2025-04-22-tencent/" target="_blank">remote connector</a>.
|
||||
- **Apr 10, 2025**: SGLang officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer.
|
||||
- **Mar 7, 2025**: We open sourced the Mooncake Store, a distributed KVCache based on Transfer Engine. vLLM's xPyD disaggregated prefilling & decoding based on Mooncake Store will be released soon.
|
||||
- **Mar 7, 2025**: We open-sourced the Mooncake Store, a distributed KVCache based on Transfer Engine. vLLM's xPyD disaggregated prefilling & decoding based on Mooncake Store will be released soon.
|
||||
- **Feb 25, 2025**: Mooncake receives the **Best Paper Award** at **FAST 2025**!
|
||||
- **Feb 21, 2025**: The updated <a href="https://github.com/kvcache-ai/Mooncake/tree/main/FAST25-release/traces" target="_blank">traces</a> used in our FAST'25 paper have been released.
|
||||
- **Dec 16, 2024**: vLLM officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer.
|
||||
- **Nov 28, 2024**: We open sourced the Transfer Engine, the central component of Mooncake. We also provide two demonstrations of Transfer Engine: a P2P Store and vLLM integration.
|
||||
- **July 9, 2024**: We open sourced the trace as a <a href="https://github.com/kvcache-ai/Mooncake/blob/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl" target="_blank">jsonl file</a>.
|
||||
- **June 27, 2024**: We present a series of Chinese blogs with more discussions on <a href="https://zhuanlan.zhihu.com/p/705754254">zhihu 1</a>, <a href="https://zhuanlan.zhihu.com/p/705910725">2</a>, <a href="https://zhuanlan.zhihu.com/p/706204757">3</a>, <a href="https://zhuanlan.zhihu.com/p/707997501">4</a>.
|
||||
- **Nov 28, 2024**: We open-sourced the Transfer Engine, the central component of Mooncake. We also provide two demonstrations of Transfer Engine: a P2P Store and vLLM integration.
|
||||
- **July 9, 2024**: We open-sourced the trace as a <a href="https://github.com/kvcache-ai/Mooncake/blob/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl" target="_blank">JSONL file</a>.
|
||||
- **June 27, 2024**: We present a series of Chinese blogs with more discussions on <a href="https://zhuanlan.zhihu.com/p/705754254">zhihu 1</a>, <a href="https://zhuanlan.zhihu.com/p/705910725">2</a>, <a href="https://zhuanlan.zhihu.com/p/706204757">3</a>, <a href="https://zhuanlan.zhihu.com/p/707997501">4</a>, <a href="https://zhuanlan.zhihu.com/p/9461861451">5</a>, <a href="https://zhuanlan.zhihu.com/p/1939988652114580803">6</a>, <a href="https://zhuanlan.zhihu.com/p/1959366095443064318">7</a>.
|
||||
- **June 26, 2024**: Initial technical report release.
|
||||
|
||||
## Documentation
|
||||
|
|
@ -58,11 +65,11 @@ This repository also hosts its technical report and the open sourced traces.
|
|||
|
||||
getting_started/build
|
||||
getting_started/quick-start
|
||||
getting_started/supported-protocols
|
||||
getting_started/plugin-usage/3FS-USRBIO-Plugin
|
||||
getting_started/examples/lmcache-integration
|
||||
getting_started/examples/lmdeploy-integration-v0.9
|
||||
getting_started/examples/sglang-integration-v1
|
||||
getting_started/examples/vllm-integration-v1
|
||||
getting_started/examples/sglang-integration/index
|
||||
getting_started/examples/vllm-integration/index
|
||||
:::
|
||||
|
|
@ -78,7 +85,8 @@ performance/vllm-benchmark-results-v0.2
|
|||
performance/vllm-benchmark-results-v1
|
||||
performance/sglang-hicache-benchmark-results-v1
|
||||
performance/vllm-v1-support-benchmark
|
||||
performance/allocator-benchmark-result.md
|
||||
performance/allocator-benchmark-result
|
||||
performance/ssd-offload-benchmark-results
|
||||
:::
|
||||
|
||||
% API Documentation
|
||||
|
|
@ -103,9 +111,9 @@ design/architecture
|
|||
design/mooncake-store
|
||||
design/p2p-store
|
||||
design/transfer-engine/index
|
||||
design/hicache-design
|
||||
design/tent/overview
|
||||
design/tent/tebench
|
||||
design/hicache-design
|
||||
:::
|
||||
|
||||
% Q&A for Mooncake
|
||||
|
|
@ -122,7 +130,16 @@ troubleshooting/troubleshooting
|
|||
|
||||
:::{toctree}
|
||||
:caption: Deployment
|
||||
:maxdepth: 1
|
||||
:maxdepth: 2
|
||||
|
||||
deployment/mooncake-store-deployment-guide
|
||||
:::
|
||||
|
||||
% Community
|
||||
|
||||
:::{toctree}
|
||||
:caption: Community
|
||||
:maxdepth: 1
|
||||
|
||||
community/governance
|
||||
:::
|
||||
|
|
|
|||
|
|
@ -0,0 +1,157 @@
|
|||
# Mooncake SSD Offload Benchmark
|
||||
|
||||
This benchmark measures the performance benefit of Mooncake's SSD offload feature in multi-turn conversation scenarios. In the test, multiple clients send requests concurrently, each simulating a multi-round dialogue where every new round appends the previous context.
|
||||
|
||||
We compare four storage configurations for the KV cache:
|
||||
|
||||
* **GPU only**: KV cache resides entirely in GPU memory.
|
||||
* **(HiCache L1) + L2**: KV cache spans GPU and host memory via HiCache's two-level hierarchy.
|
||||
* **(HiCache L1 + L2) + Mooncake**: KV cache is further extended into an 80GB Mooncake distributed memory pool.
|
||||
* **(HiCache L1 + L2) + Mooncake + SSD**: On top of the above, SSD offload is enabled so that evicted cache entries are written to local NVMe storage rather than discarded.
|
||||
|
||||
The benchmark targets the prefill stage and reports two primary metrics: Time-To-First-Token (TTFT) and input token throughput.
|
||||
|
||||
## Benchmark Result
|
||||
|
||||

|
||||
|
||||
The figure above summarizes the end-to-end results on a single DGX node (8 × A100-SXM4-40GB, dual RDMA NICs). Enabling SSD offload cuts average TTFT by **57%** relative to GPU only and by **34%** relative to Mooncake without SSD, while delivering a **2.4×** improvement in input token throughput.
|
||||
|
||||

|
||||
|
||||
To better understand where the gains come from, we break down TTFT and cache hit rate by conversation round. The output length is fixed to 1 token so that decode overhead does not obscure prefill differences.
|
||||
|
||||
During the first six rounds the 80GB memory pool is large enough, so `+ Mooncake` and `+ Mooncake + SSD` behave identically — both sustain hit rates above 80%.
|
||||
|
||||
The divergence appears in round 7. Once the accumulated KV cache exceeds memory capacity, `+ Mooncake` must evict entries and its hit rate plunges from 83% to 36%, pushing TTFT from 6s to 16s. With SSD offload, those evicted entries survive on disk and remain retrievable; the hit rate stays above 84% through round 8, and TTFT remains at 9.4s — roughly half the latency of Mooncake without SSD.
|
||||
|
||||
Note that a slight increase in TTFT is visible in round 8 with SSD offload (9.4s vs 7.4s in round 7), reflecting the additional latency of reading evicted entries from NVMe storage rather than RDMA memory. This overhead is modest compared to the alternative of re-computing evicted KV cache from scratch.
|
||||
|
||||
This demonstrates that SSD offload turns local NVMe drives into a cost-effective extension of the cache hierarchy. In production, where long conversations and high concurrency are common, this prevents the sharp performance cliff that occurs when DRAM-based caching alone is exhausted.
|
||||
|
||||
## Benchmark Setup
|
||||
|
||||
### DGX Server
|
||||
|
||||
**Experimental Environment**
|
||||
|
||||
- GPU: 8 × NVIDIA A100-SXM4-40GB
|
||||
- Network: Dual RDMA NICs (ibp12s0, ibp75s0), InfiniBand 4X HDR 200 Gb/s each
|
||||
- Storage: 5 × Samsung NVMe SSDs in RAID0 — 3 × PM1733 3.84TB (PCIe Gen4, 7,000 MB/s seq read each) + 2 × PM983 1.92TB (PCIe Gen3, 3,000 MB/s seq read each). Aggregate theoretical sequential read bandwidth: ~27 GB/s. Mounted at /mnt/data (~14TB usable), used as the SSD offload target.
|
||||
- Model: Qwen3-8B
|
||||
|
||||
**Benchmark Script:**
|
||||
|
||||
We used SGLang's [multiturn benchmark](https://github.com/sgl-project/sglang/blob/main/benchmark/hicache/bench_multiturn.py) for the evaluation.
|
||||
|
||||
```bash
|
||||
python3 benchmark/hicache/bench_multiturn.py \
|
||||
--model-path $MODEL_PATH \
|
||||
--host 127.0.0.1 \
|
||||
--port 8189 \
|
||||
--disable-random-sample \
|
||||
--output-length 1 \
|
||||
--request-length 4096 \
|
||||
--num-clients 20 \
|
||||
--num-rounds 10 \
|
||||
--max-parallel 4 \
|
||||
--request-rate 16 \
|
||||
--ready-queue-policy random \
|
||||
--disable-auto-run \
|
||||
--enable-round-barrier
|
||||
```
|
||||
|
||||
**GPU Only:**
|
||||
|
||||
```bash
|
||||
python3 -m sglang.launch_server \
|
||||
--model-path $MODEL_PATH \
|
||||
--tp 1 \
|
||||
--page-size 64 \
|
||||
--attention-backend triton
|
||||
```
|
||||
|
||||
**HiCache L1 + L2:**
|
||||
|
||||
```bash
|
||||
python3 -m sglang.launch_server \
|
||||
--model-path $MODEL_PATH \
|
||||
--tp 1 \
|
||||
--page-size 64 \
|
||||
--attention-backend triton \
|
||||
--enable-hierarchical-cache \
|
||||
--hicache-ratio 2
|
||||
```
|
||||
|
||||
**L1 + L2 + Mooncake:**
|
||||
|
||||
Mooncake master and client must be started before launching the SGLang server.
|
||||
|
||||
```bash
|
||||
# Start Mooncake master
|
||||
mooncake_master \
|
||||
-http_metadata_server_port=8081 \
|
||||
-metrics_port=9004 \
|
||||
-logtostderr
|
||||
|
||||
# Start Mooncake client (requires root)
|
||||
# Total Distributed Memory Pool: 80GB
|
||||
mooncake_client \
|
||||
--host=127.0.0.1 \
|
||||
--global_segment_size=80GB \
|
||||
--master_server_address=localhost:50051 \
|
||||
--metadata_server=P2PHANDSHAKE \
|
||||
--protocol=rdma \
|
||||
--device_names=ibp12s0,ibp75s0 \
|
||||
--port=50052 \
|
||||
--logtostderr
|
||||
```
|
||||
|
||||
```bash
|
||||
MOONCAKE_MASTER="127.0.0.1:50051" \
|
||||
MOONCAKE_GLOBAL_SEGMENT_SIZE=0 \
|
||||
MOONCAKE_PROTOCOL="rdma" \
|
||||
MOONCAKE_DEVICE="ibp12s0,ibp75s0" \
|
||||
python3 -m sglang.launch_server \
|
||||
--model-path $MODEL_PATH \
|
||||
--tp 1 \
|
||||
--page-size 64 \
|
||||
--attention-backend triton \
|
||||
--enable-hierarchical-cache \
|
||||
--hicache-ratio 2 \
|
||||
--hicache-storage-prefetch-policy wait_complete \
|
||||
--hicache-mem-layout page_first_direct \
|
||||
--hicache-storage-backend mooncake
|
||||
```
|
||||
|
||||
**L1 + L2 + Mooncake + SSD:**
|
||||
|
||||
Compared to the previous configuration, the only change is enabling SSD offload on both master and client. A 20GB local buffer absorbs write bursts before flushing to SSD.
|
||||
|
||||
```bash
|
||||
# Start Mooncake master with offload enabled
|
||||
mooncake_master \
|
||||
-enable_offload=true \
|
||||
-http_metadata_server_port=8081 \
|
||||
-metrics_port=9004 \
|
||||
-logtostderr
|
||||
|
||||
# Start Mooncake client with offload enabled (requires root)
|
||||
# Total Distributed Memory Pool: 80GB
|
||||
# SSD Offload Buffer: 20GB
|
||||
MOONCAKE_OFFLOAD_FILE_STORAGE_PATH="/mnt/data/file_storage" \
|
||||
MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES=21474836480 \
|
||||
MOONCAKE_OFFLOAD_USE_URING=1 \
|
||||
mooncake_client \
|
||||
--host=127.0.0.1 \
|
||||
--global_segment_size=80GB \
|
||||
--master_server_address=localhost:50051 \
|
||||
--metadata_server=P2PHANDSHAKE \
|
||||
--protocol=rdma \
|
||||
--device_names=ibp12s0,ibp75s0 \
|
||||
--enable_offload=true \
|
||||
--port=50052 \
|
||||
--logtostderr
|
||||
```
|
||||
|
||||
The SGLang server launch command is identical to `L1 + L2 + Mooncake`.
|
||||
|
|
@ -0,0 +1,163 @@
|
|||
# Mooncake KVCache Storage Benchmark
|
||||
|
||||
High-performance KVCache storage benchmark tool based on Mooncake Store architecture.
|
||||
|
||||
## Overview
|
||||
|
||||
Evaluates I/O performance of KVCache storage systems using:
|
||||
- Single large file (100GB) with offset-based block management
|
||||
- Prefix caching simulation with hash-based block lookup
|
||||
- Timestamp-based request replay for realistic testing
|
||||
- Comprehensive metrics: latency, bandwidth, hit rates
|
||||
|
||||
## Test Flow
|
||||
|
||||
1. **Load Traces**: Read request sequences from JSONL files (`FAST25-release/traces`)
|
||||
2. **Process Requests**: For each request, check hash_id prefix cache hits/misses
|
||||
3. **Perform I/O**: Read cached blocks from disk, write new blocks to storage
|
||||
4. **Collect Metrics**: Track latency, bandwidth, and cache hit rates
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Quick test (100 requests, no timestamp replay)
|
||||
python storage_benchmark.py --scenario=toolagent --max-requests=100
|
||||
|
||||
# Test with large model preset (Llama-3.1-405B)
|
||||
python storage_benchmark.py --scenario=toolagent --model=llama-3.1-405b --max-requests=100
|
||||
|
||||
# Test with Deepseek V3 (extra large model)
|
||||
python storage_benchmark.py --scenario=toolagent --model=deepseek-v3 --max-requests=100
|
||||
|
||||
# Realistic replay (with timestamps, 10x speed)
|
||||
python storage_benchmark.py --scenario=toolagent --max-requests=1000 \
|
||||
--replay-timestamps --time-scale=0.1
|
||||
|
||||
# Test all scenarios with replay
|
||||
python storage_benchmark.py --scenario=all --time-scale=1.0
|
||||
```
|
||||
|
||||
## Command-Line Options
|
||||
|
||||
| Option | Description | Default |
|
||||
|--------|-------------|---------|
|
||||
| `--trace-dir` | Trace files directory | `../FAST25-release/traces` |
|
||||
| `--scenario` | Test scenario: `conversation`, `synthetic`, `toolagent`, `all` | `toolagent` |
|
||||
| `--storage-dir` | Storage directory | `/tmp/mooncake_bench` |
|
||||
| `--model` | Model preset (overrides `--bytes-per-token`) | `default` |
|
||||
| `--bytes-per-token` | Bytes per token (2048 for 7B FP16) | `2048` |
|
||||
| `--max-requests` | Maximum requests per scenario (unlimited if not specified) | `None` |
|
||||
| `--max-blocks` | Maximum number of blocks | `100000` |
|
||||
| `--replay-timestamps` | Enable timestamp replay | `False` |
|
||||
| `--time-scale` | Time scaling factor (1.0 = real-time, 0.1 = 10x faster) | `1.0` |
|
||||
|
||||
## Model Presets
|
||||
|
||||
The tool includes presets for popular LLM models with accurate KVCache sizes based on the [LMCache KVCache Calculator](https://lmcache.ai/kv_cache_calculator.html).
|
||||
|
||||
| Model | Bytes/Token | Size | Notes |
|
||||
|-------|-------------|------|-------|
|
||||
| **Small Models (7B-13B)** |
|
||||
| `llama-3-8b` | 128 | 128 B/token | GQA optimized |
|
||||
| `mistral-7b` | 128 | 128 B/token | GQA optimized |
|
||||
| `qwen-14b` | 40 | 40 B/token | GQA optimized |
|
||||
| `gemma-7b` | 224 | 224 B/token | |
|
||||
| `llama-2-7b` | 512 | 512 B/token | |
|
||||
| `llama-2-13b` | 800 | 800 B/token | |
|
||||
| **Large Models (70B-405B)** |
|
||||
| `llama-2-70b` | 320 | 320 B/token | GQA optimized |
|
||||
| `llama-3-70b` | 320 | 320 B/token | GQA optimized |
|
||||
| `mixtral-8x7b` | 128 | 128 B/token | GQA optimized |
|
||||
| `mixtral-8x22b` | 224 | 224 B/token | GQA optimized |
|
||||
| `qwen-72b` | 320 | 320 B/token | GQA optimized |
|
||||
| `qwen-110b` | 320 | 320 B/token | GQA optimized |
|
||||
| `llama-3.1-405b` | 516018 | ~504 KB/token | Very large KVCache |
|
||||
| **Extra Large Models** |
|
||||
| `glm-4.6` | 156991 | ~153 KB/token | |
|
||||
| `deepseek-v3` | 1749384 | ~1.67 MB/token | Largest KVCache |
|
||||
| **Default** |
|
||||
| `default` | 2048 | 2 KB/token | Legacy 7B FP16 |
|
||||
|
||||
**Usage**: `--model=llama-3.1-405b` (overrides `--bytes-per-token`)
|
||||
|
||||
## Test Scenarios
|
||||
|
||||
- **`conversation`**: Write-intensive workload (dialogue patterns)
|
||||
- **`synthetic`**: Read-intensive workload (cached patterns)
|
||||
- **`toolagent`**: Balanced read/write mix (tool use patterns)
|
||||
|
||||
## Output Example
|
||||
|
||||
```
|
||||
================================================================================
|
||||
Mooncake KVCache Storage Benchmark
|
||||
================================================================================
|
||||
Using model preset: llama-3.1-405b (516018 bytes/token, ~504.0 KB/token)
|
||||
|
||||
[1/1] toolagent_trace.jsonl
|
||||
================================================================================
|
||||
|
||||
[Performance Overview]
|
||||
Total Requests: 100
|
||||
Queries Per Second (QPS): 14.45
|
||||
Cache Hit Rate: 24.27%
|
||||
Write Ratio: 75.73%
|
||||
Total Blocks: 1,949
|
||||
Read Blocks: 473
|
||||
Write Blocks: 1,476
|
||||
Prefix Hits: 376
|
||||
|
||||
[Latency Analysis]
|
||||
Request Latency (End-to-End): Avg=69.18ms, P50=15.49ms, P95=239.99ms, P99=310.58ms
|
||||
Single I/O Operation (Per Block):
|
||||
Read: Avg=14.572ms, P50=0.280ms, P95=0.280ms, P99=0.280ms
|
||||
Write: Avg=5.120ms, P50=5.120ms, P95=5.120ms, P99=5.120ms
|
||||
|
||||
[I/O & Bandwidth]
|
||||
Total Read I/O: 473.0 MB (473 ops)
|
||||
Total Write I/O: 1476.0 MB (1,476 ops)
|
||||
Effective Bandwidth: 280.8 MB/s
|
||||
|
||||
[Storage Details]
|
||||
Blocks in Use: 1,476
|
||||
Free Blocks: 0
|
||||
Tokens per Block: 512
|
||||
Block Size: 1.00 MB
|
||||
|
||||
[Execution Time]
|
||||
Total Execution Time: 8.42 s
|
||||
|
||||
================================================================================
|
||||
```
|
||||
|
||||
## Metrics
|
||||
|
||||
| Metric | Description |
|
||||
|--------|-------------|
|
||||
| **QPS** | Queries per second (based on I/O time, excluding sleep) |
|
||||
| **Request Latency** | End-to-end latency for entire request (all I/O operations) |
|
||||
| **Single I/O Latency** | Latency for individual block read/write operations (512 tokens) |
|
||||
| **P50/P95/P99** | Latency percentiles (milliseconds) using linear interpolation |
|
||||
| **Hit Rate** | Cache hit ratio for blocks |
|
||||
| **Write Ratio** | Percentage of blocks that needed to be written |
|
||||
| **Bandwidth** | Effective throughput based on I/O time only |
|
||||
| **Prefix Hits** | Number of blocks served from prefix cache |
|
||||
|
||||
**Note**: Request Latency measures the total time to process all blocks in a request, while Single I/O Latency measures the time for one block operation (512 tokens).
|
||||
|
||||
## Trace Data Format
|
||||
|
||||
```json
|
||||
{
|
||||
"timestamp": 1234.567,
|
||||
"hash_ids": [1, 2, 4, 7],
|
||||
"input_length": 2048,
|
||||
"output_length": 512
|
||||
}
|
||||
```
|
||||
|
||||
Each `hash_id` corresponds to a 512-token block. The tool simulates prefix caching by checking if blocks are already in storage before writing.
|
||||
|
||||
## Requirements
|
||||
|
||||
- Python 3.10+
|
||||
|
|
@ -52,14 +52,14 @@ Basic usage:
|
|||
```python
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from mooncake import ep
|
||||
from mooncake import pg
|
||||
|
||||
active_ranks = torch.ones((world_size,), dtype=torch.int32, device="cuda")
|
||||
dist.init_process_group(
|
||||
backend="mooncake",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
pg_options=ep.MooncakeBackendOptions(active_ranks),
|
||||
pg_options=pg.MooncakeBackendOptions(active_ranks),
|
||||
)
|
||||
|
||||
dist.all_gather(...) # Standard API usage
|
||||
|
|
@ -76,16 +76,16 @@ Recover usage (e.g., wants to recover rank #2):
|
|||
# For the healthy processes, execute:
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from mooncake import ep
|
||||
from mooncake import pg
|
||||
|
||||
...
|
||||
|
||||
broken_rank = 2
|
||||
backend = dist.group.WORLD._get_backend(torch.device("cpu"))
|
||||
while True:
|
||||
(peer_state,) = ep.get_peer_state(backend, [broken_rank])
|
||||
(peer_state,) = pg.get_peer_state(backend, [broken_rank])
|
||||
if peer_state:
|
||||
ep.recover_ranks(backend, [broken_rank])
|
||||
pg.recover_ranks(backend, [broken_rank])
|
||||
break
|
||||
else:
|
||||
# Handle ongoing logic, like inference
|
||||
|
|
@ -96,7 +96,7 @@ dist.init_process_group(
|
|||
backend="mooncake-cpu",
|
||||
rank=broken_rank,
|
||||
world_size=num_processes,
|
||||
pg_options=ep.MooncakeBackendOptions(
|
||||
pg_options=pg.MooncakeBackendOptions(
|
||||
torch.ones((num_processes,), dtype=torch.int32),
|
||||
is_extension=True, # Must set this option to True
|
||||
),
|
||||
|
|
|
|||
|
|
@ -264,6 +264,151 @@ def get_into(self, key: str, buffer_ptr: int, size: int) -> int
|
|||
|
||||
**Returns:** Number of bytes read, or negative on error
|
||||
|
||||
#### get_into_ranges()
|
||||
Retrieve multiple byte ranges from multiple objects into registered buffers (zero-copy).
|
||||
|
||||
```python
|
||||
def get_into_ranges(self, buffer_ptrs: List[int], all_keys: List[List[str]], all_dst_offsets: List[List[List[int]]], all_src_offsets: List[List[List[int]]], all_sizes: List[List[List[int]]]) -> List[List[List[int]]]
|
||||
```
|
||||
|
||||
This API is **buffer-major** and supports **multiple fragments per key**.
|
||||
|
||||
Think of the input shape as:
|
||||
- `buffer_ptrs[i]`: the `i`-th destination buffer
|
||||
- `all_keys[i][j]`: the `j`-th key that writes into buffer `i`
|
||||
- `all_dst_offsets[i][j][k]`: destination offset of fragment `k` for key `j` in buffer `i`
|
||||
- `all_src_offsets[i][j][k]`: source offset of fragment `k` inside key `j` for buffer `i`
|
||||
- `all_sizes[i][j][k]`: byte size of fragment `k`
|
||||
|
||||
For each triple `(i, j, k)`, Mooncake reads the source range
|
||||
`[all_src_offsets[i][j][k], all_src_offsets[i][j][k] + all_sizes[i][j][k])`
|
||||
from object `all_keys[i][j]`, then writes it into destination buffer
|
||||
`buffer_ptrs[i]` at offset `all_dst_offsets[i][j][k]`.
|
||||
|
||||
This lets one buffer gather interleaved fragments from multiple keys, and lets one key contribute multiple disjoint fragments to the same buffer in a single call.
|
||||
|
||||
**Parameters:**
|
||||
- `buffer_ptrs`: Memory addresses of pre-allocated destination buffers. Every buffer must be registered with `register_buffer()` before calling this API.
|
||||
- `all_keys`: For each buffer, the ordered list of source object keys to read from.
|
||||
- `all_dst_offsets`: For each buffer and key, the destination offsets of that key's fragments.
|
||||
- `all_src_offsets`: For each buffer and key, the source offsets of that key's fragments inside the object.
|
||||
- `all_sizes`: For each buffer and key, the byte lengths of that key's fragments.
|
||||
|
||||
**Shape rules:**
|
||||
- `len(buffer_ptrs) == len(all_keys) == len(all_dst_offsets) == len(all_src_offsets) == len(all_sizes)`
|
||||
- For each buffer `i`, `len(all_keys[i]) == len(all_dst_offsets[i]) == len(all_src_offsets[i]) == len(all_sizes[i])`
|
||||
- For each `(buffer i, key j)`, `len(all_dst_offsets[i][j]) == len(all_src_offsets[i][j]) == len(all_sizes[i][j])`
|
||||
|
||||
If a top-level shape or per-key fragment shape does not match, the corresponding result entries are negative error codes.
|
||||
|
||||
**Returns:** A nested list of per-buffer, per-key, per-fragment results. `results[i][j][k]` is the number of bytes read for fragment `k`, or a negative value on error.
|
||||
|
||||
A successful call can still contain per-fragment failures. For example, if one key is missing but another key in the same buffer is valid, the missing key's fragment result will be negative while the valid fragment can still succeed.
|
||||
|
||||
**Typical scenarios:**
|
||||
- **Partial read from one object:** You only need a slice of a large value, such as a header, metadata block, or a small subrange of a tensor shard. In this case, use one buffer, one key, and one or more fragments under that key.
|
||||
- **Stitch multiple fragments from one object into one buffer:** You need several non-contiguous ranges from the same object and want to pack them into one destination buffer. In this case, keep a single key entry and place multiple fragments under that key.
|
||||
- **Stitch data from multiple objects into one buffer:** You want to assemble one logical payload from several keys. In this case, use one destination buffer and list multiple keys under that buffer, with each key contributing one or more fragments.
|
||||
- **Fill multiple output buffers in one call:** You have several destination buffers, each with its own read plan. In this case, each top-level entry in `buffer_ptrs` and the parallel nested arrays describes one independent destination buffer.
|
||||
|
||||
**How to use it for partial reads:**
|
||||
If you only want part of an object, do not call `get_into()` with the full object buffer size. Instead:
|
||||
1. Allocate and register a destination buffer sized for the bytes you actually want to materialize.
|
||||
2. Put that buffer pointer into `buffer_ptrs`.
|
||||
3. Put the source key into `all_keys`.
|
||||
4. Set `all_src_offsets` to the start offsets of the object ranges you want.
|
||||
5. Set `all_sizes` to the lengths of those ranges.
|
||||
6. Set `all_dst_offsets` to where those ranges should land in your destination buffer.
|
||||
|
||||
A useful way to think about the arguments is:
|
||||
- `buffer_ptrs` answers **where does the data land**
|
||||
- `all_keys` answers **which object does it come from**
|
||||
- `all_src_offsets` and `all_sizes` answer **which bytes should be read**
|
||||
- `all_dst_offsets` answers **where each fragment should be placed in the destination buffer**
|
||||
|
||||
If you are extracting a single contiguous slice from one object, the minimal shape is:
|
||||
|
||||
```python
|
||||
results = store.get_into_ranges(
|
||||
[buffer_ptr],
|
||||
[["my_key"]],
|
||||
[[[0]]],
|
||||
[[[src_offset]]],
|
||||
[[[size]]],
|
||||
)
|
||||
```
|
||||
|
||||
This means:
|
||||
- one destination buffer
|
||||
- one source key for that buffer
|
||||
- one fragment for that key
|
||||
- read `size` bytes from `my_key[src_offset:src_offset + size]`
|
||||
- write them into `buffer_ptr[0:size]`
|
||||
|
||||
If you want to read several disjoint ranges from the same object and pack them together, keep the same key and add more fragments under it. For example:
|
||||
|
||||
```python
|
||||
results = store.get_into_ranges(
|
||||
[buffer_ptr],
|
||||
[["my_key"]],
|
||||
[[[0, 16, 40]]],
|
||||
[[[128, 4096, 8192]]],
|
||||
[[[8, 12, 4]]],
|
||||
)
|
||||
```
|
||||
|
||||
This reads three fragments from `my_key` and places them into the same destination buffer at offsets `0`, `16`, and `40`. This pattern is useful when you want to assemble only the needed pieces of a large object without reading the whole value.
|
||||
|
||||
If you want to assemble one output buffer from multiple objects, keep one top-level buffer entry and add multiple keys under it. Each key can still contribute one or more fragments. For example, you might put a header from `meta_key` at the front of the buffer, then place a payload slice from `data_key` after it.
|
||||
|
||||
**Usage example:**
|
||||
|
||||
```python
|
||||
import ctypes
|
||||
|
||||
buffer_size = 32
|
||||
buffer0 = (ctypes.c_ubyte * buffer_size)()
|
||||
buffer1 = (ctypes.c_ubyte * buffer_size)()
|
||||
buffer_ptr0 = ctypes.addressof(buffer0)
|
||||
buffer_ptr1 = ctypes.addressof(buffer1)
|
||||
|
||||
store.register_buffer(buffer_ptr0, buffer_size)
|
||||
store.register_buffer(buffer_ptr1, buffer_size)
|
||||
|
||||
# Buffer 0 reads:
|
||||
# - from key1: two fragments -> src[1:5] -> dst[0:4], src[30:33] -> dst[20:23]
|
||||
# - from key2: one fragment -> src[2:7] -> dst[8:13]
|
||||
# Buffer 1 reads:
|
||||
# - from key2: one fragment -> src[0:6] -> dst[4:10]
|
||||
# - from key1: one fragment -> src[10:14] -> dst[16:20]
|
||||
results = store.get_into_ranges(
|
||||
[buffer_ptr0, buffer_ptr1],
|
||||
[["key1", "key2"], ["key2", "key1"]],
|
||||
[[[0, 20], [8]], [[4], [16]]],
|
||||
[[[1, 30], [2]], [[0], [10]]],
|
||||
[[[4, 3], [5]], [[6], [4]]],
|
||||
)
|
||||
|
||||
# results == [
|
||||
# [[4, 3], [5]],
|
||||
# [[6], [4]],
|
||||
# ]
|
||||
```
|
||||
|
||||
In the example above:
|
||||
- `results[0][0][0] == 4`: buffer 0, key 0 (`"key1"`), fragment 0 succeeded with 4 bytes
|
||||
- `results[0][0][1] == 3`: buffer 0, key 0 (`"key1"`), fragment 1 succeeded with 3 bytes
|
||||
- `results[0][1][0] == 5`: buffer 0, key 1 (`"key2"`), fragment 0 succeeded with 5 bytes
|
||||
|
||||
**Common pitfalls:**
|
||||
- Do not flatten all fragments for a buffer into one list. Fragments must be grouped under their corresponding key.
|
||||
- `all_dst_offsets`, `all_src_offsets`, and `all_sizes` are 3D, but `all_keys` is 2D.
|
||||
- Buffer overflow is checked against the registered destination buffer size.
|
||||
- Source overflow is checked against the source object's size.
|
||||
- Full-object `get_into()` and ranged `get_into_ranges()` are different APIs; use `get_into()` when you want the whole object into one buffer.
|
||||
|
||||
**Current limitation:** true ranged items currently require the selected source replica to be memory-backed. Whole-object reads still follow the normal full-read path, but partial reads through `get_into_ranges()` do not support non-memory replicas.
|
||||
|
||||
---
|
||||
|
||||
## ReplicateConfig Configuration
|
||||
|
|
@ -301,6 +446,16 @@ config = ReplicateConfig()
|
|||
config.with_soft_pin = True # Keep this object in memory longer
|
||||
```
|
||||
|
||||
#### with_hard_pin
|
||||
**Type:** `bool`
|
||||
**Default:** `False`
|
||||
**Description:** Enables hard pinning for the stored object. Hard pinned objects will not be evicted. This grants user to manually control the life time of stored objects.
|
||||
|
||||
```python
|
||||
config = ReplicateConfig()
|
||||
config.with_hard_pin = True # Keep this object in memory that will not be evicted
|
||||
```
|
||||
|
||||
#### preferred_segment
|
||||
**Type:** `str`
|
||||
**Default:** `""` (empty string)
|
||||
|
|
@ -629,6 +784,120 @@ result = store.put_batch(keys, values)
|
|||
|
||||
---
|
||||
|
||||
#### upsert()
|
||||
|
||||
Insert a new object if the key does not exist, or update the existing object in place when possible. They use the same replication configuration model as `put()`.
|
||||
|
||||
Upsert binary data in the distributed storage.
|
||||
|
||||
```python
|
||||
def upsert(self, key: str, value: bytes, config: ReplicateConfig = None) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `key` (str): Unique object identifier
|
||||
- `value` (bytes): Binary data to insert or update
|
||||
- `config` (ReplicateConfig, optional): Replication configuration
|
||||
|
||||
**Returns:**
|
||||
- `int`: Status code (0 = success, non-zero = error code)
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = 2
|
||||
|
||||
rc = store.upsert("weights", b"new-bytes", config)
|
||||
if rc == 0:
|
||||
print("Upsert succeeded")
|
||||
```
|
||||
|
||||
#### upsert_from()
|
||||
|
||||
Upsert object data directly from a pre-allocated buffer (zero-copy).
|
||||
|
||||
```python
|
||||
def upsert_from(self, key: str, buffer_ptr: int, size: int, config: ReplicateConfig = None) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `key` (str): Object identifier
|
||||
- `buffer_ptr` (int): Memory address of the source buffer
|
||||
- `size` (int): Number of bytes to insert or update
|
||||
- `config` (ReplicateConfig, optional): Replication configuration
|
||||
|
||||
**Returns:**
|
||||
- `int`: Status code (0 = success, non-zero = error code)
|
||||
|
||||
**Note:** This is the zero-copy counterpart of `upsert()`. As with
|
||||
`put_from()`, register the buffer before issuing the request.
|
||||
|
||||
#### batch_upsert_from()
|
||||
|
||||
Upsert multiple objects directly from pre-allocated buffers.
|
||||
|
||||
```python
|
||||
def batch_upsert_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int],
|
||||
config: ReplicateConfig = None) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `keys` (List[str]): List of object identifiers
|
||||
- `buffer_ptrs` (List[int]): List of source buffer addresses
|
||||
- `sizes` (List[int]): List of byte lengths for each buffer
|
||||
- `config` (ReplicateConfig, optional): Replication configuration shared by all objects
|
||||
|
||||
**Returns:**
|
||||
- `List[int]`: List of status codes for each upsert
|
||||
|
||||
#### upsert_parts()
|
||||
|
||||
Upsert data from multiple buffer parts as a single object (insert or update).
|
||||
|
||||
```python
|
||||
def upsert_parts(self, key: str, *parts, config: ReplicateConfig = None) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `key` (str): Object identifier
|
||||
- `*parts`: Variable number of bytes-like objects to concatenate
|
||||
- `config` (ReplicateConfig, optional): Replication configuration
|
||||
|
||||
**Returns:**
|
||||
- `int`: Status code (0 = success, non-zero = error code)
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
part1 = b"Hello, "
|
||||
part2 = b"World!"
|
||||
result = store.upsert_parts("greeting", part1, part2)
|
||||
```
|
||||
|
||||
#### upsert_batch()
|
||||
|
||||
Upsert multiple objects in a single batch operation.
|
||||
|
||||
```python
|
||||
def upsert_batch(self, keys: List[str], values: List[bytes], config: ReplicateConfig = None) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `keys` (List[str]): List of object identifiers
|
||||
- `values` (List[bytes]): List of binary data to insert or update
|
||||
- `config` (ReplicateConfig, optional): Replication configuration for all objects
|
||||
|
||||
**Returns:**
|
||||
- `int`: Status code (0 = success, non-zero = error code)
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
keys = ["key1", "key2", "key3"]
|
||||
values = [b"value1", b"value2", b"value3"]
|
||||
result = store.upsert_batch(keys, values)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
#### get_batch()
|
||||
Retrieve multiple objects in a single batch operation.
|
||||
|
||||
|
|
@ -721,6 +990,39 @@ print(f"Removed {count} objects")
|
|||
|
||||
---
|
||||
|
||||
#### batch_remove()
|
||||
Remove multiple objects by their keys in a single batch operation.
|
||||
|
||||
```python
|
||||
def batch_remove(self, keys: List[str], force: bool = False) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `keys` (List[str]): List of object identifiers to remove
|
||||
- `force` (bool): If True, skip lease and replication task checks (default: False)
|
||||
|
||||
**Returns:**
|
||||
- `List[int]`: List of status codes for each key (0 = success, negative = error code)
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
# Remove multiple keys in one batch
|
||||
keys = ["key1", "key2", "key3", "key4", "key5"]
|
||||
results = store.batch_remove(keys)
|
||||
|
||||
# Check results
|
||||
for key, result in zip(keys, results):
|
||||
if result == 0:
|
||||
print(f"✓ {key} removed successfully")
|
||||
else:
|
||||
print(f"✗ {key} failed with error code: {result}")
|
||||
|
||||
# Force remove (bypass lease checks)
|
||||
results = store.batch_remove(keys, force=True)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
#### is_exist()
|
||||
Check if an object exists in the storage system.
|
||||
|
||||
|
|
@ -977,6 +1279,132 @@ for key, desc_list in descriptors_map.items():
|
|||
|
||||
---
|
||||
|
||||
#### create_copy_task()
|
||||
|
||||
Creates an asynchronous copy task to replicate an object to target segments.
|
||||
|
||||
```python
|
||||
def create_copy_task(self, key: str, targets: List[str]) -> Tuple[UUID, int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `key` (str): Object key to copy
|
||||
- `targets` (List[str]): List of target segment names where replicas should be created
|
||||
|
||||
**Returns:**
|
||||
- `Tuple[UUID, int]`: (task UUID, error code)
|
||||
- If successful: (task UUID, 0)
|
||||
- If failed: (UUID{0, 0}, error code)
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
# Create an asynchronous copy task
|
||||
task_id, error_code = store.create_copy_task("my_key", ["segment1", "segment2"])
|
||||
if error_code == 0:
|
||||
print(f"Copy task created with ID: {task_id}")
|
||||
# Query task status later
|
||||
response, status = store.query_task(task_id)
|
||||
if status == 0:
|
||||
print(f"Task status: {response.status}")
|
||||
else:
|
||||
print(f"Failed to create copy task: {error_code}")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
#### create_move_task()
|
||||
|
||||
Creates an asynchronous move task to move an object from source segment to target segment.
|
||||
|
||||
```python
|
||||
def create_move_task(self, key: str, source: str, target: str) -> Tuple[UUID, int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `key` (str): Object key to move
|
||||
- `source` (str): Source segment name where the replica currently exists
|
||||
- `target` (str): Target segment name where the replica should be moved to
|
||||
|
||||
**Returns:**
|
||||
- `Tuple[UUID, int]`: (task UUID, error code)
|
||||
- If successful: (task UUID, 0)
|
||||
- If failed: (UUID{0, 0}, error code)
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
# Create an asynchronous move task
|
||||
task_id, error_code = store.create_move_task("my_key", "old_segment", "new_segment")
|
||||
if error_code == 0:
|
||||
print(f"Move task created with ID: {task_id}")
|
||||
# Query task status later
|
||||
response, status = store.query_task(task_id)
|
||||
if status == 0:
|
||||
print(f"Task status: {response.status}")
|
||||
else:
|
||||
print(f"Failed to create move task: {error_code}")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
#### query_task()
|
||||
|
||||
Queries the status of an asynchronous task (copy or move).
|
||||
|
||||
```python
|
||||
def query_task(self, task_id: UUID) -> Tuple[QueryTaskResponse | None, int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `task_id` (UUID): UUID of the task to query
|
||||
|
||||
**Returns:**
|
||||
- `Tuple[QueryTaskResponse | None, int]`: (QueryTaskResponse if success, error code)
|
||||
- If successful: (QueryTaskResponse, 0)
|
||||
- If failed: (None, error code)
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
from mooncake.store import MooncakeDistributedStore, TaskStatus
|
||||
import time
|
||||
|
||||
# Initialize store
|
||||
store = MooncakeDistributedStore()
|
||||
store.setup("localhost", "http://localhost:8080/metadata",
|
||||
512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051")
|
||||
|
||||
# Submit multiple copy tasks
|
||||
tasks = []
|
||||
for key in ["key1", "key2", "key3"]:
|
||||
task_id, error = store.create_copy_task(key, ["segment1", "segment2"])
|
||||
if error == 0:
|
||||
tasks.append(task_id)
|
||||
print(f"Created copy task {task_id} for {key}")
|
||||
|
||||
# Monitor task progress
|
||||
while tasks:
|
||||
completed = []
|
||||
for task_id in tasks:
|
||||
response, status = store.query_task(task_id)
|
||||
if status == 0 and response:
|
||||
if response.status == TaskStatus.SUCCESS:
|
||||
print(f"Task {task_id} succeeded")
|
||||
completed.append(task_id)
|
||||
elif response.status == TaskStatus.FAILED:
|
||||
print(f"Task {task_id} failed: {response.message}")
|
||||
completed.append(task_id)
|
||||
|
||||
# Remove completed tasks
|
||||
tasks = [t for t in tasks if t not in completed]
|
||||
|
||||
if tasks:
|
||||
time.sleep(1) # Wait before next check
|
||||
|
||||
print("All tasks completed")
|
||||
store.close()
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
#### close()
|
||||
Clean up all resources and terminate connections.
|
||||
|
||||
|
|
@ -1104,6 +1532,28 @@ def put_tensor_with_tp(self, key: str, tensor: torch.Tensor, tp_rank: int = 0, t
|
|||
|
||||
- `int`: Status code (0 = success, non-zero = error code).
|
||||
|
||||
#### pub_tensor_with_tp()
|
||||
|
||||
Publish a PyTorch tensor into the store with configurable replication settings, optionally splitting it into shards for tensor parallelism.
|
||||
The tensor is chunked immediately and stored as separate keys (e.g., `key_tp_0`, `key_tp_1`...).
|
||||
|
||||
```python
|
||||
def pub_tensor_with_tp(self, key: str, tensor: torch.Tensor, config: ReplicateConfig, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `key` (str): Base identifier for the tensor.
|
||||
- `tensor` (torch.Tensor): The PyTorch tensor to store.
|
||||
- `config` (ReplicateConfig): Optional replication configuration.
|
||||
- `tp_rank` (int): Current tensor parallel rank (default: 0). *Note: The method splits and stores all chunks for all ranks regardless of this value.*
|
||||
- `tp_size` (int): Total tensor parallel size (default: 1). If \> 1, the tensor is split into `tp_size` chunks.
|
||||
- `split_dim` (int): The dimension to split the tensor along (default: 0).
|
||||
|
||||
**Returns:**
|
||||
|
||||
- `int`: Status code (0 = success, non-zero = error code).
|
||||
|
||||
#### get_tensor_with_tp()
|
||||
|
||||
Get a PyTorch tensor from the store, specifically retrieving the shard corresponding to the given Tensor Parallel rank.
|
||||
|
|
@ -1143,6 +1593,27 @@ def batch_put_tensor_with_tp(self, base_keys: List[str], tensors_list: List[torc
|
|||
|
||||
- `List[int]`: List of status codes for each tensor operation.
|
||||
|
||||
#### batch_pub_tensor_with_tp()
|
||||
|
||||
Publish a batch of PyTorch tensors into the store with configurable replication settings, splitting each into shards for tensor parallelism.
|
||||
|
||||
```python
|
||||
def batch_pub_tensor_with_tp(self, base_keys: List[str], tensors_list: List[torch.Tensor], config: ReplicateConfig, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `base_keys` (List[str]): List of base identifiers.
|
||||
- `tensors_list` (List[torch.Tensor]): List of tensors to store.
|
||||
- `config` (ReplicateConfig): Optional replication configuration.
|
||||
- `tp_rank` (int): Current rank (default: 0).
|
||||
- `tp_size` (int): Total tp size (default: 1).
|
||||
- `split_dim` (int): Split dimension (default: 0).
|
||||
|
||||
**Returns:**
|
||||
|
||||
- `List[int]`: List of status codes for each tensor operation.
|
||||
|
||||
#### batch_get_tensor_with_tp()
|
||||
|
||||
Get a batch of PyTorch tensor shards from the store for a given Tensor Parallel rank.
|
||||
|
|
@ -1311,6 +1782,151 @@ for i, result in enumerate(results):
|
|||
print(f"Tensor {i} failed to store with code: {result}")
|
||||
```
|
||||
|
||||
#### batch_pub_tensor()
|
||||
|
||||
Pub a batch of PyTorch tensors into the store with configurable replication settings.
|
||||
|
||||
```python
|
||||
def batch_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tensor], config: ReplicateConfig) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `keys` (List[str]): List of object identifiers
|
||||
- `tensors_list` (List[torch.Tensor]): List of tensors to store
|
||||
- `config` (ReplicateConfig): Optional replication configuration.
|
||||
|
||||
**Returns:**
|
||||
- `List[int]`: List of status codes for each tensor operation.
|
||||
|
||||
**Note:** This function requires `torch` to be installed and available in the environment.
|
||||
|
||||
---
|
||||
|
||||
#### upsert_tensor()
|
||||
|
||||
Insert a tensor if its key is missing, or update the existing tensor if the key already exists. The current tensor upsert helpers use the default `ReplicateConfig` and therefore do not take a `config` parameter.
|
||||
|
||||
Upsert a PyTorch tensor into the store.
|
||||
|
||||
```python
|
||||
def upsert_tensor(self, key: str, tensor: torch.Tensor) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `key` (str): Object identifier
|
||||
- `tensor` (torch.Tensor): The PyTorch tensor to insert or update
|
||||
|
||||
**Returns:**
|
||||
- `int`: Status code (0 = success, non-zero = error code)
|
||||
|
||||
**Note:** This function requires `torch` to be installed and available in the environment.
|
||||
|
||||
#### upsert_tensor_from()
|
||||
|
||||
Upsert a tensor directly from a pre-allocated buffer. The buffer layout must be
|
||||
`[TensorMetadata][tensor data]`, matching the layout used by
|
||||
`get_tensor_into()`.
|
||||
|
||||
```python
|
||||
def upsert_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `key` (str): Object identifier
|
||||
- `buffer_ptr` (int): Buffer pointer containing serialized tensor metadata and payload
|
||||
- `size` (int): Actual serialized byte length of the tensor buffer
|
||||
|
||||
**Returns:**
|
||||
- `int`: Status code (0 = success, non-zero = error code)
|
||||
|
||||
**Note:** This function is not supported for dummy client.
|
||||
|
||||
#### batch_upsert_tensor_from()
|
||||
|
||||
Upsert multiple tensors directly from pre-allocated buffers. Each buffer must
|
||||
use layout `[TensorMetadata][tensor data]`.
|
||||
|
||||
```python
|
||||
def batch_upsert_tensor_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int]) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `keys` (List[str]): List of object identifiers
|
||||
- `buffer_ptrs` (List[int]): List of serialized tensor buffer pointers
|
||||
- `sizes` (List[int]): List of actual serialized byte lengths
|
||||
|
||||
**Returns:**
|
||||
- `List[int]`: List of status codes for each tensor upsert
|
||||
|
||||
#### batch_upsert_tensor()
|
||||
|
||||
Upsert a batch of PyTorch tensors into the store (insert or update).
|
||||
|
||||
```python
|
||||
def batch_upsert_tensor(self, keys: List[str], tensors_list: List[torch.Tensor]) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `keys` (List[str]): List of object identifiers
|
||||
- `tensors_list` (List[torch.Tensor]): List of tensors to insert or update
|
||||
|
||||
**Returns:**
|
||||
- `List[int]`: List of status codes for each tensor operation.
|
||||
|
||||
**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
|
||||
|
||||
#### upsert_pub_tensor()
|
||||
|
||||
Upsert a PyTorch tensor with configurable replication settings (insert or update).
|
||||
|
||||
```python
|
||||
def upsert_pub_tensor(self, key: str, tensor: torch.Tensor, config: ReplicateConfig = None) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `key` (str): Unique object identifier
|
||||
- `tensor` (torch.Tensor): PyTorch tensor to insert or update
|
||||
- `config` (ReplicateConfig, optional): Replication configuration
|
||||
|
||||
**Returns:**
|
||||
- `int`: Status code (0 = success, non-zero = error code)
|
||||
|
||||
**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
import torch
|
||||
from mooncake.store import ReplicateConfig
|
||||
|
||||
tensor = torch.randn(100, 100)
|
||||
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = 2
|
||||
config.with_soft_pin = True
|
||||
|
||||
result = store.upsert_pub_tensor("my_tensor", tensor, config)
|
||||
if result == 0:
|
||||
print("Tensor upserted successfully")
|
||||
```
|
||||
|
||||
#### batch_upsert_pub_tensor()
|
||||
|
||||
Batch upsert PyTorch tensors with configurable replication settings (insert or update).
|
||||
|
||||
```python
|
||||
def batch_upsert_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tensor], config: ReplicateConfig = None) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `keys` (List[str]): List of object identifiers
|
||||
- `tensors_list` (List[torch.Tensor]): List of tensors to insert or update
|
||||
- `config` (ReplicateConfig, optional): Replication configuration
|
||||
|
||||
**Returns:**
|
||||
- `List[int]`: List of status codes for each tensor operation.
|
||||
|
||||
**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
|
||||
|
||||
---
|
||||
|
||||
### PyTorch Tensor Operations (Zero Copy)
|
||||
|
|
@ -1348,8 +1964,8 @@ def batch_get_tensor_into(self, base_keys: List[str], buffer_ptrs: List[int], si
|
|||
**Parameters:**
|
||||
|
||||
- `base_keys` (List[str]): List of base identifiers.
|
||||
- `buffer_ptrs` (List[int]): List of the buffers pointer pre-allocated for tensor, and the buffers should be registered.
|
||||
- `sizes` (List[int]): List of the size of buffers.
|
||||
- `buffer_ptrs` (List[int]): List of buffer pointers pre-allocated for tensor; buffers should be registered.
|
||||
- `sizes` (List[int]): List of buffer sizes.
|
||||
|
||||
**Returns:**
|
||||
|
||||
|
|
@ -1387,8 +2003,8 @@ def batch_get_tensor_with_tp_into(self, base_keys: List[str], buffer_ptrs: List[
|
|||
**Parameters:**
|
||||
|
||||
- `base_keys` (List[str]): List of base identifiers.
|
||||
- `buffer_ptrs` (List[int]): List of the buffers pointer pre-allocated for tensor, and the buffers should be registered.
|
||||
- `sizes` (List[int]): List of the size of buffers.
|
||||
- `buffer_ptrs` (List[int]): List of buffer pointers pre-allocated for tensor; buffers should be registered.
|
||||
- `sizes` (List[int]): List of buffer sizes.
|
||||
- `tp_rank` (int): The tensor parallel rank to retrieve (default: 0).
|
||||
- `tp_size` (int): Total tensor parallel size (default: 1).
|
||||
|
||||
|
|
@ -1396,6 +2012,84 @@ def batch_get_tensor_with_tp_into(self, base_keys: List[str], buffer_ptrs: List[
|
|||
|
||||
- `List[torch.Tensor]`: List of retrieved tensors (or shards). Contains `None` for missing keys.
|
||||
|
||||
#### put_tensor_from()
|
||||
|
||||
Put a PyTorch tensor into the store directly from a pre-allocated buffer (zero-copy). The buffer must contain data in the same layout as produced by `get_tensor_into`: **\[TensorMetadata\]\[tensor data\]**. The buffer is only read during this call; no Python object references it.
|
||||
|
||||
```python
|
||||
def put_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `key` (str): Object identifier for the tensor.
|
||||
- `buffer_ptr` (int): The buffer pointer; the buffer should be registered. Layout must be \[TensorMetadata\]\[tensor data\].
|
||||
- `size` (int): **Actual serialized byte length** of the data in the buffer (metadata + tensor bytes), not the buffer capacity.
|
||||
|
||||
**Returns:**
|
||||
|
||||
- `int`: Status code (0 = success, non-zero = error code).
|
||||
|
||||
#### batch_put_tensor_from()
|
||||
|
||||
Put a batch of PyTorch tensors into the store directly from pre-allocated buffers (zero-copy). Each buffer must contain data in the layout **\[TensorMetadata\]\[tensor data\]**, same as `get_tensor_into`.
|
||||
|
||||
```python
|
||||
def batch_put_tensor_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int]) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `keys` (List[str]): List of object identifiers.
|
||||
- `buffer_ptrs` (List[int]): List of buffer pointers; buffers should be registered.
|
||||
- `sizes` (List[int]): List of **actual serialized byte lengths** for each buffer (metadata + tensor bytes), not buffer capacities.
|
||||
|
||||
**Returns:**
|
||||
|
||||
- `List[int]`: List of status codes for each tensor operation (0 = success, non-zero = error code).
|
||||
|
||||
#### put_tensor_with_tp_from()
|
||||
|
||||
Put a **full tensor** into the store directly from a pre-allocated buffer (zero-copy), for use with Tensor Parallelism. This is the zero-copy counterpart of `put_tensor_with_tp()`: the buffer must contain the complete tensor in layout **\[TensorMetadata\]\[tensor data\]**, and Mooncake will split it internally and store all shards under `key_tp_<rank>`.
|
||||
|
||||
```python
|
||||
def put_tensor_with_tp_from(self, key: str, buffer_ptr: int, size: int, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> int
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `key` (str): Base identifier for the tensor.
|
||||
- `buffer_ptr` (int): The buffer pointer; the buffer should be registered.
|
||||
- `size` (int): **Actual serialized byte length** of the full tensor in the buffer.
|
||||
- `tp_rank` (int): Kept for signature compatibility with `put_tensor_with_tp()` (default: 0). It does **not** mean "only write one shard".
|
||||
- `tp_size` (int): Total tensor parallel size (default: 1). If 1, equivalent to `put_tensor_from(key, buffer_ptr, size)`.
|
||||
- `split_dim` (int): Dimension along which the full tensor is split before storing shards.
|
||||
|
||||
**Returns:**
|
||||
|
||||
- `int`: Status code (0 = success, non-zero = error code).
|
||||
|
||||
#### batch_put_tensor_with_tp_from()
|
||||
|
||||
Put a batch of **full tensors** into the store directly from pre-allocated buffers (zero-copy). This is the zero-copy counterpart of `batch_put_tensor_with_tp()`: each buffer contains one full tensor in layout **\[TensorMetadata\]\[tensor data\]**, and Mooncake splits each tensor internally and stores all TP shards.
|
||||
|
||||
```python
|
||||
def batch_put_tensor_with_tp_from(self, base_keys: List[str], buffer_ptrs: List[int], sizes: List[int], tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> List[int]
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `base_keys` (List[str]): List of base identifiers.
|
||||
- `buffer_ptrs` (List[int]): List of buffer pointers; buffers should be registered.
|
||||
- `sizes` (List[int]): List of **actual serialized byte lengths** for each full-tensor buffer.
|
||||
- `tp_rank` (int): Kept for signature compatibility with `batch_put_tensor_with_tp()` (default: 0). It does **not** select a single shard to write.
|
||||
- `tp_size` (int): Total tensor parallel size (default: 1). If 1, equivalent to `batch_put_tensor_from(base_keys, buffer_ptrs, sizes)`.
|
||||
- `split_dim` (int): Dimension along which each full tensor is split before storing shards.
|
||||
|
||||
**Returns:**
|
||||
|
||||
- `List[int]`: List of status codes for each tensor operation (0 = success, non-zero = error code).
|
||||
|
||||
---
|
||||
|
||||
### Batch Zero-Copy Operations
|
||||
|
|
|
|||
|
|
@ -18,155 +18,7 @@ pip install mooncake-transfer-engine
|
|||
|
||||
## Quick Start
|
||||
|
||||
### Start Transfer Engine Receiver (Server)
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
import zmq
|
||||
from mooncake.engine import TransferEngine
|
||||
|
||||
def main():
|
||||
# Initialize ZMQ context and socket
|
||||
context = zmq.Context()
|
||||
socket = context.socket(zmq.PUSH)
|
||||
socket.bind("tcp://*:5555") # Bind to port 5555 for buffer info
|
||||
|
||||
HOSTNAME = "localhost" # localhost for simple demo
|
||||
METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
|
||||
PROTOCOL = "tcp" # [rdma, tcp, ...]
|
||||
DEVICE_NAME = "" # auto discovery if empty
|
||||
|
||||
# Initialize server engine
|
||||
server_engine = TransferEngine()
|
||||
server_engine.initialize(
|
||||
HOSTNAME,
|
||||
METADATA_SERVER,
|
||||
PROTOCOL,
|
||||
DEVICE_NAME
|
||||
)
|
||||
session_id = f"{HOSTNAME}:{server_engine.get_rpc_port()}"
|
||||
|
||||
# Allocate memory on server side (1MB buffer)
|
||||
server_buffer = np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
server_ptr = server_buffer.ctypes.data
|
||||
server_len = server_buffer.nbytes
|
||||
|
||||
# Register memory with Mooncake
|
||||
ret_value = server_engine.register_memory(server_ptr, server_len)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory registration failed.")
|
||||
raise RuntimeError("Mooncake memory registration failed.")
|
||||
|
||||
print(f"Server initialized with session ID: {session_id}")
|
||||
print(f"Server buffer address: {server_ptr}, length: {server_len}")
|
||||
|
||||
# Send buffer info to client
|
||||
buffer_info = {
|
||||
"session_id": session_id,
|
||||
"ptr": server_ptr,
|
||||
"len": server_len
|
||||
}
|
||||
socket.send_json(buffer_info)
|
||||
print("Buffer information sent to client")
|
||||
|
||||
# Keep server running
|
||||
try:
|
||||
while True:
|
||||
input("Press Ctrl+C to exit...")
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down server...")
|
||||
finally:
|
||||
# Cleanup
|
||||
ret_value = server_engine.unregister_memory(server_ptr)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory deregistration failed.")
|
||||
raise RuntimeError("Mooncake memory deregistration failed.")
|
||||
|
||||
socket.close()
|
||||
context.term()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
### Start Transfer Engine Sender (Client)
|
||||
|
||||
```python
|
||||
|
||||
import numpy as np
|
||||
import zmq
|
||||
from mooncake.engine import TransferEngine
|
||||
|
||||
def main():
|
||||
# Initialize ZMQ context and socket
|
||||
context = zmq.Context()
|
||||
socket = context.socket(zmq.PULL)
|
||||
socket.connect(f"tcp://localhost:5555")
|
||||
|
||||
# Wait for buffer info from server
|
||||
print("Waiting for server buffer information...")
|
||||
buffer_info = socket.recv_json()
|
||||
server_session_id = buffer_info["session_id"]
|
||||
server_ptr = buffer_info["ptr"]
|
||||
server_len = buffer_info["len"]
|
||||
print(f"Received server info - Session ID: {server_session_id}")
|
||||
print(f"Server buffer address: {server_ptr}, length: {server_len}")
|
||||
|
||||
# Initialize client engine
|
||||
HOSTNAME = "localhost" # localhost for simple demo
|
||||
METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
|
||||
PROTOCOL = "tcp" # [rdma, tcp, ...]
|
||||
DEVICE_NAME = "" # auto discovery if empty
|
||||
|
||||
client_engine = TransferEngine()
|
||||
client_engine.initialize(
|
||||
HOSTNAME,
|
||||
METADATA_SERVER,
|
||||
PROTOCOL,
|
||||
DEVICE_NAME
|
||||
)
|
||||
session_id = f"{HOSTNAME}:{client_engine.get_rpc_port()}"
|
||||
|
||||
# Allocate and initialize client buffer (1MB)
|
||||
client_buffer = np.ones(1024 * 1024, dtype=np.uint8) # Fill with ones
|
||||
client_ptr = client_buffer.ctypes.data
|
||||
client_len = client_buffer.nbytes
|
||||
|
||||
# Register memory with Mooncake
|
||||
ret_value = client_engine.register_memory(client_ptr, client_len)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory registration failed.")
|
||||
raise RuntimeError("Mooncake memory registration failed.")
|
||||
|
||||
print(f"Client initialized with session ID: {session_id}")
|
||||
|
||||
# Transfer data from client to server
|
||||
print("Transferring data to server...")
|
||||
for _ in range(10):
|
||||
ret = client_engine.transfer_sync_write(
|
||||
server_session_id,
|
||||
client_ptr,
|
||||
server_ptr,
|
||||
min(client_len, server_len) # Transfer minimum of both lengths
|
||||
)
|
||||
|
||||
if ret >= 0:
|
||||
print("Transfer successful!")
|
||||
else:
|
||||
print("Transfer failed!")
|
||||
|
||||
# Cleanup
|
||||
ret_value = client_engine.unregister_memory(client_ptr)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory deregistration failed.")
|
||||
raise RuntimeError("Mooncake memory deregistration failed.")
|
||||
|
||||
socket.close()
|
||||
context.term()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
See the [Transfer Engine Quick Start](../getting_started/quick-start.md#transfer-engine-quick-start) guide for a complete example of setting up and using the Transfer Engine.
|
||||
|
||||
## API Reference
|
||||
|
||||
|
|
@ -174,7 +26,7 @@ if __name__ == "__main__":
|
|||
|
||||
The main class that provides all transfer engine functionality.
|
||||
|
||||
### Constructor
|
||||
#### Constructor
|
||||
|
||||
```python
|
||||
TransferEngine()
|
||||
|
|
@ -182,6 +34,28 @@ TransferEngine()
|
|||
|
||||
Creates a new TransferEngine instance with default settings.
|
||||
|
||||
### Class: TransferNotify
|
||||
|
||||
A class representing a transfer notification message.
|
||||
|
||||
#### Constructor
|
||||
|
||||
```python
|
||||
TransferNotify()
|
||||
TransferNotify(name, msg)
|
||||
```
|
||||
|
||||
**Constructor Parameters:**
|
||||
- `name` (str): The notification name/identifier
|
||||
- `msg` (str): The notification message content
|
||||
|
||||
### Enums: TransferOpcode
|
||||
|
||||
```python
|
||||
TransferOpcode.READ # Read operation
|
||||
TransferOpcode.WRITE # Write operation
|
||||
```
|
||||
|
||||
### Initialization Methods
|
||||
|
||||
#### initialize()
|
||||
|
|
@ -232,8 +106,6 @@ Gets the inner transfer engine instance, which can be reused for mooncake store.
|
|||
**Returns:**
|
||||
- `InnerTransferEngine`: The inner transfer engine
|
||||
|
||||
### Network Information
|
||||
|
||||
#### get_rpc_port()
|
||||
|
||||
```python
|
||||
|
|
@ -288,7 +160,7 @@ Gets the address of the first buffer in a specified segment.
|
|||
- `segment_name` (str): The name of the segment
|
||||
|
||||
**Returns:**
|
||||
- `int`: The memory address of the first buffer in the segment
|
||||
- `int`: The memory address of the first buffer in the segment, or 0 if the segment is not found or has no registered buffers
|
||||
|
||||
### Data Transfer Operations
|
||||
|
||||
|
|
@ -329,10 +201,10 @@ Performs a synchronous read operation to transfer data from remote buffer to loc
|
|||
#### transfer_sync()
|
||||
|
||||
```python
|
||||
transfer_sync(target_hostname, buffer, peer_buffer_address, length, opcode)
|
||||
transfer_sync(target_hostname, buffer, peer_buffer_address, length, opcode, notify=None)
|
||||
```
|
||||
|
||||
Performs a synchronous transfer operation with specified opcode.
|
||||
Performs a synchronous transfer operation with specified opcode and optional notification.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
|
|
@ -340,6 +212,7 @@ Performs a synchronous transfer operation with specified opcode.
|
|||
- `peer_buffer_address` (int): The remote buffer address
|
||||
- `length` (int): The number of bytes to transfer
|
||||
- `opcode` (TransferOpcode): The transfer operation type (READ or WRITE)
|
||||
- `notify` (TransferNotify, optional): Notification object to send after transfer completion
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
|
@ -373,12 +246,228 @@ Checks the status of an asynchronous transfer operation.
|
|||
- `batch_id` (int): The batch ID returned from transfer_submit_write()
|
||||
|
||||
**Returns:**
|
||||
- `int`:
|
||||
- `int`:
|
||||
- 1: Transfer completed successfully
|
||||
- 0: Transfer still in progress
|
||||
- -1: Transfer failed
|
||||
- -2: Transfer timed out
|
||||
|
||||
#### transfer_write_on_cuda()
|
||||
|
||||
```python
|
||||
transfer_write_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr)
|
||||
```
|
||||
|
||||
Performs a write operation to transfer data from local buffer to remote buffer on a given cuda stream.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffer` (int): The local buffer address
|
||||
- `peer_buffer_address` (int): The remote buffer address
|
||||
- `length` (int): The number of bytes to transfer
|
||||
- `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`.
|
||||
|
||||
**Returns:**
|
||||
- `None`: The function returns immediately after successfully scheduling the transfer callback.
|
||||
|
||||
**Raises:**
|
||||
- `RuntimeError`: If the segment cannot be opened or if the `cudaLaunchHostFunc` call fails.
|
||||
|
||||
**Warning:**
|
||||
- `Unrecoverable Error`: If an error occurs during the asynchronous execution inside the CUDA callback, the process will terminate immediately via _exit(1).
|
||||
|
||||
#### transfer_read_on_cuda()
|
||||
|
||||
```python
|
||||
transfer_read_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr)
|
||||
```
|
||||
|
||||
Performs a read operation to transfer data from remote buffer to local buffer on a given cuda stream.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffer` (int): The local buffer address
|
||||
- `peer_buffer_address` (int): The remote buffer address
|
||||
- `length` (int): The number of bytes to transfer
|
||||
- `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`.
|
||||
|
||||
**Returns:**
|
||||
- `None`: The function returns immediately after successfully scheduling the transfer callback.
|
||||
|
||||
**Raises:**
|
||||
- `RuntimeError`: If the segment cannot be opened or if the `cudaLaunchHostFunc` call fails.
|
||||
|
||||
**Warning:**
|
||||
- `Unrecoverable Error`: If an error occurs during the asynchronous execution inside the CUDA callback, the process will terminate immediately via _exit(1).
|
||||
|
||||
### Batch Data Transfer Operations
|
||||
|
||||
**Note:** In a few inference engines and benchmarks, accuracy may be affected when using batch transfer APIs. This issue has been found only in multi-node NVLink transfers.
|
||||
|
||||
#### batch_transfer_sync_write()
|
||||
|
||||
```python
|
||||
batch_transfer_sync_write(target_hostname, buffers, peer_buffer_addresses, lengths)
|
||||
```
|
||||
|
||||
Performs a batch synchronous write operation to transfer multiple data chunks from local buffers to remote buffers.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffers` (List[int]): List of local buffer addresses
|
||||
- `peer_buffer_addresses` (List[int]): List of remote buffer addresses
|
||||
- `lengths` (List[int]): List of byte lengths for each transfer
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### batch_transfer_sync_read()
|
||||
|
||||
```python
|
||||
batch_transfer_sync_read(target_hostname, buffers, peer_buffer_addresses, lengths)
|
||||
```
|
||||
|
||||
Performs a batch synchronous read operation to transfer multiple data chunks from remote buffers to local buffers.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffers` (List[int]): List of local buffer addresses
|
||||
- `peer_buffer_addresses` (List[int]): List of remote buffer addresses
|
||||
- `lengths` (List[int]): List of byte lengths for each transfer
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### batch_transfer_sync()
|
||||
|
||||
```python
|
||||
batch_transfer_sync(target_hostname, buffers, peer_buffer_addresses, lengths, opcode, notify=None)
|
||||
```
|
||||
|
||||
Performs a batch synchronous transfer operation with specified opcode and optional notification.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffers` (List[int]): List of local buffer addresses
|
||||
- `peer_buffer_addresses` (List[int]): List of remote buffer addresses
|
||||
- `lengths` (List[int]): List of byte lengths for each transfer
|
||||
- `opcode` (TransferOpcode): The transfer operation type (READ or WRITE)
|
||||
- `notify` (TransferNotify, optional): Notification object to send after transfer completion
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### batch_transfer_async_write()
|
||||
|
||||
```python
|
||||
batch_transfer_async_write(target_hostname, buffers, peer_buffer_addresses, lengths)
|
||||
```
|
||||
|
||||
Submits a batch asynchronous write operation and returns immediately.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffers` (List[int]): List of local buffer addresses
|
||||
- `peer_buffer_addresses` (List[int]): List of remote buffer addresses
|
||||
- `lengths` (List[int]): List of byte lengths for each transfer
|
||||
|
||||
**Returns:**
|
||||
- `int`: Batch ID for tracking the operation, or 0 on failure
|
||||
|
||||
#### batch_transfer_async_read()
|
||||
|
||||
```python
|
||||
batch_transfer_async_read(target_hostname, buffers, peer_buffer_addresses, lengths)
|
||||
```
|
||||
|
||||
Submits a batch asynchronous read operation and returns immediately.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffers` (List[int]): List of local buffer addresses
|
||||
- `peer_buffer_addresses` (List[int]): List of remote buffer addresses
|
||||
- `lengths` (List[int]): List of byte lengths for each transfer
|
||||
|
||||
**Returns:**
|
||||
- `int`: Batch ID for tracking the operation, or 0 on failure
|
||||
|
||||
#### batch_transfer_async()
|
||||
|
||||
```python
|
||||
batch_transfer_async(target_hostname, buffers, peer_buffer_addresses, lengths, opcode)
|
||||
```
|
||||
|
||||
Submits a batch asynchronous transfer operation with specified opcode and returns immediately.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffers` (List[int]): List of local buffer addresses
|
||||
- `peer_buffer_addresses` (List[int]): List of remote buffer addresses
|
||||
- `lengths` (List[int]): List of byte lengths for each transfer
|
||||
- `opcode` (TransferOpcode): The transfer operation type (READ or WRITE)
|
||||
|
||||
**Returns:**
|
||||
- `int`: Batch ID for tracking the operation, or 0 on failure
|
||||
|
||||
#### get_batch_transfer_status()
|
||||
|
||||
```python
|
||||
get_batch_transfer_status(batch_ids)
|
||||
```
|
||||
|
||||
Waits for multiple batch asynchronous transfer operations to complete.
|
||||
|
||||
**Parameters:**
|
||||
- `batch_ids` (List[int]): List of batch IDs returned from batch async transfer operations
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 if all transfers completed successfully, -1 if any transfer failed or timed out
|
||||
|
||||
#### batch_transfer_write_on_cuda()
|
||||
|
||||
```python
|
||||
batch_transfer_write_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr)
|
||||
```
|
||||
|
||||
Performs a batch write operation to transfer multiple data chunks from local buffers to remote buffers on a given cuda stream.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffers` (List[int]): List of local buffer addresses
|
||||
- `peer_buffer_addresses` (List[int]): List of remote buffer addresses
|
||||
- `lengths` (List[int]): List of byte lengths for each transfer
|
||||
- `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`.
|
||||
|
||||
**Returns:**
|
||||
- `None`: The function returns immediately after successfully scheduling the transfer callback.
|
||||
|
||||
**Raises:**
|
||||
- `RuntimeError`: If the segment cannot be opened or if the `cudaLaunchHostFunc` call fails.
|
||||
|
||||
**Warning:**
|
||||
- `Unrecoverable Error`: If an error occurs during the asynchronous execution inside the CUDA callback, the process will terminate immediately via _exit(1).
|
||||
|
||||
#### batch_transfer_read_on_cuda()
|
||||
|
||||
```python
|
||||
batch_transfer_read_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr)
|
||||
```
|
||||
|
||||
Performs a batch read operation to transfer multiple data chunks from remote buffers to local buffers on a given cuda stream.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffers` (List[int]): List of local buffer addresses
|
||||
- `peer_buffer_addresses` (List[int]): List of remote buffer addresses
|
||||
- `lengths` (List[int]): List of byte lengths for each transfer
|
||||
- `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`.
|
||||
|
||||
**Returns:**
|
||||
- `None`: The function returns immediately after successfully scheduling the transfer callback.
|
||||
|
||||
**Raises:**
|
||||
- `RuntimeError`: If the segment cannot be opened or if the `cudaLaunchHostFunc` call fails.
|
||||
|
||||
### Buffer I/O Operations
|
||||
|
||||
#### write_bytes_to_buffer()
|
||||
|
|
@ -412,7 +501,7 @@ Reads bytes from a buffer at the specified address and returns them as a Python
|
|||
**Returns:**
|
||||
- `bytes`: The bytes read from the buffer
|
||||
|
||||
### Memory Registration (Experimental)
|
||||
### Memory Registration
|
||||
|
||||
#### register_memory()
|
||||
|
||||
|
|
@ -443,15 +532,62 @@ Unregisters a previously registered memory region.
|
|||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
### Enums
|
||||
|
||||
#### TransferOpcode
|
||||
#### batch_register_memory()
|
||||
|
||||
```python
|
||||
TransferOpcode.READ # Read operation
|
||||
TransferOpcode.WRITE # Write operation
|
||||
batch_register_memory(buffer_addresses, capacities)
|
||||
```
|
||||
|
||||
Registers multiple memory regions for RDMA access in a single batch operation.
|
||||
|
||||
**Parameters:**
|
||||
- `buffer_addresses` (List[int]): List of memory addresses to register
|
||||
- `capacities` (List[int]): List of sizes in bytes for each memory region
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### batch_unregister_memory()
|
||||
|
||||
```python
|
||||
batch_unregister_memory(buffer_addresses)
|
||||
```
|
||||
|
||||
Unregisters multiple previously registered memory regions in a single batch operation.
|
||||
|
||||
**Parameters:**
|
||||
- `buffer_addresses` (List[int]): List of memory addresses to unregister
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
### Topology and Notification
|
||||
|
||||
#### get_local_topology()
|
||||
|
||||
```python
|
||||
get_local_topology(device_name=None)
|
||||
```
|
||||
|
||||
Gets the local network topology information as a JSON string.
|
||||
|
||||
**Parameters:**
|
||||
- `device_name` (str, optional): Comma-separated list of device names to filter, or None for all devices
|
||||
|
||||
**Returns:**
|
||||
- `str`: JSON string representing the local network topology
|
||||
|
||||
#### get_notifies()
|
||||
|
||||
```python
|
||||
get_notifies()
|
||||
```
|
||||
|
||||
Gets the list of pending transfer notifications received from other nodes.
|
||||
|
||||
**Returns:**
|
||||
- `List[TransferNotify]`: List of notification objects containing name and message
|
||||
|
||||
## Environment Variables
|
||||
|
||||
The Transfer Engine respects the following environment variables:
|
||||
|
|
@ -461,7 +597,7 @@ The Transfer Engine respects the following environment variables:
|
|||
- `MC_LEGACY_RPC_PORT_BINDING`: Enables legacy RPC port binding behavior
|
||||
- `MC_TCP_BIND_ADDRESS`: Specifies the TCP bind address
|
||||
- `MC_CUSTOM_TOPO_JSON`: Path to custom topology JSON file
|
||||
- `MC_TE_METRIC`: Enables metrics reporting (set to "1", "true", "yes", or "on")
|
||||
- `MC_TE_METRIC`: Enables metrics reporting (set to "1", "true", "yes", or "on"). **Note:** Not supported when using Transfer Engine TENT.
|
||||
- `MC_TE_METRIC_INTERVAL_SECONDS`: Sets metrics reporting interval in seconds
|
||||
|
||||
## Usage Examples
|
||||
|
|
@ -557,15 +693,74 @@ else:
|
|||
# Use the buffer
|
||||
test_data = b"Test data for managed buffer"
|
||||
engine.write_bytes_to_buffer(buffer_addr, test_data, len(test_data))
|
||||
|
||||
|
||||
# Read back
|
||||
read_data = engine.read_bytes_from_buffer(buffer_addr, len(test_data))
|
||||
print(f"Read data: {read_data}")
|
||||
|
||||
|
||||
# Free the buffer when done
|
||||
engine.free_managed_buffer(buffer_addr, buffer_size)
|
||||
```
|
||||
|
||||
### Batch Transfer Operations
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
from mooncake.engine import TransferEngine, TransferOpcode
|
||||
|
||||
# Prepare multiple buffers
|
||||
num_chunks = 4
|
||||
chunk_size = 256 * 1024 # 256KB each
|
||||
|
||||
# Create local buffers
|
||||
local_buffers = [np.ones(chunk_size, dtype=np.uint8) for _ in range(num_chunks)]
|
||||
local_addrs = [buf.ctypes.data for buf in local_buffers]
|
||||
lengths = [chunk_size] * num_chunks
|
||||
|
||||
# Register all buffers in batch
|
||||
engine.batch_register_memory(local_addrs, lengths)
|
||||
|
||||
# Assume remote_addrs are obtained from the remote node
|
||||
remote_addrs = [...] # List of remote buffer addresses
|
||||
|
||||
# Synchronous batch write
|
||||
ret = engine.batch_transfer_sync_write(
|
||||
"target_host:port",
|
||||
local_addrs,
|
||||
remote_addrs,
|
||||
lengths
|
||||
)
|
||||
if ret == 0:
|
||||
print("Batch transfer completed successfully")
|
||||
|
||||
# Cleanup
|
||||
engine.batch_unregister_memory(local_addrs)
|
||||
```
|
||||
|
||||
### Transfer with Notification
|
||||
|
||||
```python
|
||||
from mooncake.engine import TransferEngine, TransferOpcode, TransferNotify
|
||||
|
||||
# Create a notification
|
||||
notify = TransferNotify("transfer_complete", "chunk_1_done")
|
||||
|
||||
# Transfer with notification - the receiver will get this notification
|
||||
ret = engine.transfer_sync(
|
||||
"target_host:port",
|
||||
local_addr,
|
||||
remote_addr,
|
||||
length,
|
||||
TransferOpcode.WRITE,
|
||||
notify
|
||||
)
|
||||
|
||||
# On the receiving side, get notifications
|
||||
notifications = engine.get_notifies()
|
||||
for n in notifications:
|
||||
print(f"Received notification: name={n.name}, msg={n.msg}")
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
|
||||
All methods return integer status codes:
|
||||
|
|
@ -582,9 +777,11 @@ Common error scenarios:
|
|||
## Performance Considerations
|
||||
|
||||
1. **Buffer Reuse**: Reuse allocated buffers when possible to avoid frequent allocation/deallocation overhead
|
||||
2. **Batch Operations**: Use `transfer_submit_write()` and `transfer_check_status()` for better throughput when multiple transfers are needed
|
||||
3. **Memory Alignment**: Ensure buffers are properly aligned for optimal RDMA performance
|
||||
4. **Timeout Configuration**: Adjust `MC_TRANSFER_TIMEOUT` based on your network characteristics and data sizes
|
||||
2. **Batch Operations**: Use batch transfer APIs (`batch_transfer_sync_write()`, `batch_transfer_async_write()`, etc.) for better throughput when transferring multiple chunks to the same target
|
||||
3. **Batch Memory Registration**: Use `batch_register_memory()` and `batch_unregister_memory()` when working with multiple buffers to reduce overhead
|
||||
4. **Asynchronous Transfers**: Use asynchronous APIs (`transfer_submit_write()`, `batch_transfer_async_*()`) with `transfer_check_status()` or `get_batch_transfer_status()` to overlap computation with data transfer
|
||||
5. **Memory Alignment**: Ensure buffers are properly aligned for optimal RDMA performance
|
||||
6. **Timeout Configuration**: Adjust `MC_TRANSFER_TIMEOUT` based on your network characteristics and data sizes
|
||||
|
||||
## Thread Safety
|
||||
|
||||
|
|
@ -598,4 +795,4 @@ The Transfer Engine Python API is thread-safe for most operations. However, it's
|
|||
1. **Initialization Failures**: Check metadata server connectivity and network configuration
|
||||
2. **Transfer Failures**: Verify target hostname is correct and network connectivity is established
|
||||
3. **Memory Issues**: Ensure sufficient system memory and proper buffer alignment
|
||||
4. **Performance Issues**: Check RDMA device configuration and network topology
|
||||
4. **Performance Issues**: Check RDMA device configuration and network topology
|
||||
|
|
|
|||
|
|
@ -63,7 +63,54 @@ Errors in this part usually indicate that the error occurred within the `mooncak
|
|||
**Solution:**
|
||||
Ensure that the total memory registration does not exceed the device's upper limit. You may need to reduce the amount of memory being registered or split large memory regions into smaller chunks that fit within the device's `max_mr_size` limit.
|
||||
|
||||
5. If the error `Failed to create QP: Cannot allocate memory` is displayed, it typically caused by too many QP have been created, reaching the driver limit. You can use `rdma resource` to trace how many QP is created. One possible way to resolve this issue:
|
||||
5. If you encounter `Failed to register memory 0x...: Resource temporarily unavailable [11]` and kernel logs show `CREATE_MKEY failed, status no resources(0xf)`, this indicates that the RDMA NIC has exhausted its internal Memory Key (MKEY) resources, even though `ulimit -l` and `vm.max_map_count` may appear sufficient.
|
||||
|
||||
This typically happens when:
|
||||
- Applications that use RDMA (e.g., SGLang with HiCache + Mooncake) have crashed or been killed multiple times without cleanly releasing RDMA resources.
|
||||
- The leaked MKEY entries accumulate in the NIC firmware and are not reclaimed by the kernel, eventually hitting the hardware limit.
|
||||
- Large memory regions (e.g., NSA indexer buffers at ~4.68 GB each across multiple TP ranks) amplify the problem since each registration consumes more internal NIC resources.
|
||||
|
||||
**Diagnostic Commands:**
|
||||
```bash
|
||||
# Check current RDMA resource usage per device
|
||||
rdma resource show
|
||||
|
||||
# Check kernel logs for CREATE_MKEY failures
|
||||
dmesg | grep -i "CREATE_MKEY\|no resources\|mlx5_cmd_out_err"
|
||||
# Example output:
|
||||
# mlx5_core 0000:65:01.0: mlx5_cmd_out_err:829:(pid 3958462): CREATE_MKEY(0x200) op_mod(0x0) failed, status no resources(0xf), syndrome (0x2aac7c), err(-11)
|
||||
|
||||
# Ensure vm.max_map_count is large enough (default 65530 may be too small)
|
||||
sysctl vm.max_map_count
|
||||
```
|
||||
|
||||
**Solutions:**
|
||||
- **Reboot the node** to fully reset NIC firmware state and reclaim all leaked MKEY resources. This is the most reliable fix.
|
||||
- Increase `vm.max_map_count` if it is at the default value: `sysctl -w vm.max_map_count=16777216`
|
||||
- Ensure applications shut down cleanly (avoid `kill -9` when possible) so RDMA resources are properly deregistered.
|
||||
- If rebooting is not feasible, try unloading and reloading the mlx5 kernel modules (may disrupt other services):
|
||||
```bash
|
||||
modprobe -r mlx5_ib mlx5_core && modprobe mlx5_core mlx5_ib
|
||||
```
|
||||
|
||||
6. If you encounter errors indicating inability to allocate memory space when requesting large memory regions, this may be due to ulimit restrictions. When the total memory requirement (number of registered RDMA devices × requested space) exceeds the ulimit, the system will display errors about failing to allocate space.
|
||||
|
||||
**Diagnostic Commands:**
|
||||
- Use `ulimit -a` to check current limits, particularly the `max locked memory` value
|
||||
- Calculate total memory requirement: number of RDMA devices × requested space per device
|
||||
- Verify if the total requirement exceeds the ulimit
|
||||
|
||||
**Solutions:**
|
||||
- Switch to a higher privilege level (root) to bypass ulimit restrictions
|
||||
- Modify ulimit settings: use `ulimit -l unlimited` to remove locked memory limits (may require root privileges)
|
||||
- Start multiple store instances with smaller memory allocations that stay within ulimit constraints
|
||||
- Add permanent ulimit configuration in `/etc/security/limits.conf`:
|
||||
```
|
||||
* soft memlock unlimited
|
||||
* hard memlock unlimited
|
||||
```
|
||||
|
||||
7. If the error `Failed to create QP: Cannot allocate memory` is displayed, it is typically caused by too many QP have been created, reaching the driver limit. You can use `rdma resource` to trace how many QP is created. One possible way to resolve this issue:
|
||||
- Update Mooncake to version v0.3.5 or later
|
||||
- Set the environment variable `MC_ENABLE_DEST_DEVICE_AFFINITY=1` before starting the application
|
||||
|
||||
|
|
|
|||
|
|
@ -30,6 +30,33 @@ Transfer Engine RPC using <协议> listening on <IP>:<实际端口>,记录目
|
|||
./transfer_engine_ascend_direct_perf --metadata_server=P2PHANDSHAKE --local_server_name=127.0.0.1:12346 --operation=write --device_logicid=1 --mode=initiator --block_size=16384 --batch_size=32 --block_iteration=10 --segment_id=127.0.0.1:real_port
|
||||
```
|
||||
|
||||
### 环境变量配置
|
||||
|
||||
以下环境变量可用于控制Ascend Direct Transport的行为:
|
||||
|
||||
| 变量名 | 描述 | 默认值 | 示例 |
|
||||
|--------|--------------------------------------|--------|-----------------------------------------------------------------------|
|
||||
| `ASCEND_AUTO_CONNECT` | 启用自动连接管理 | 0(禁用) | `ASCEND_AUTO_CONNECT=1` |
|
||||
| `ASCEND_ENABLE_USE_FABRIC_MEM` | 在Mooncake Store中启用fabric内存传输模式(仅A3) | 0(禁用) | `ASCEND_ENABLE_USE_FABRIC_MEM=1` |
|
||||
| `ASCEND_USE_ASYNC_TRANSFER` | 启用异步传输模式 | 0(禁用) | `ASCEND_USE_ASYNC_TRANSFER=1` |
|
||||
| `ASCEND_GLOBAL_RESOURCE_CONFIG` | 全局资源配置 | - | `ASCEND_GLOBAL_RESOURCE_CONFIG="{\"fabric_memory.max_capacity\":32}"` |
|
||||
| `ASCEND_CONNECT_TIMEOUT` | 链路建链超时时间(毫秒) | 3000 | `ASCEND_CONNECT_TIMEOUT=5000` |
|
||||
| `ASCEND_TRANSFER_TIMEOUT` | 数据传输超时时间(毫秒) | 3000 | `ASCEND_TRANSFER_TIMEOUT=10000` |
|
||||
| `ASCEND_THREAD_POOL_SIZE` | 传输线程池的工作线程数 | 8(缓冲池模式下为1) | `ASCEND_THREAD_POOL_SIZE=16` |
|
||||
| `ASCEND_USE_SHORT_CONNECTION` | 启用短连接模式(每次传输后断开) | 0(禁用) | `ASCEND_USE_SHORT_CONNECTION=1` |
|
||||
| `ASCEND_BUFFER_POOL` | 中转模式缓冲池配置(BUFFER_NUM:BUFFER_SIZE_MB) | "0:0"(禁用) | `ASCEND_BUFFER_POOL=4:8` |
|
||||
| `ASCEND_BASE_PORT` | ADXL引擎端口分配的基础端口 | 11000 | `ASCEND_BASE_PORT=20000` |
|
||||
| `HCCL_INTRA_ROCE_ENABLE` | 启用节点内RDMA通信协议 | 0(禁用) | `HCCL_INTRA_ROCE_ENABLE=1` |
|
||||
| `HCCL_RDMA_TIMEOUT` | RDMA数据包重传超时时间系数 | - | `HCCL_RDMA_TIMEOUT=14` |
|
||||
| `HCCL_RDMA_RETRY_CNT` | RDMA数据包重传次数 | - | `HCCL_RDMA_RETRY_CNT=7` |
|
||||
|
||||
详细说明:
|
||||
ASCEND_AUTO_CONNECT: 需要CANN升级到9.0之后的版本,所以默认值为0,在支持该功能的版本推荐启用:当对端异常下线后可以自动断链。
|
||||
ASCEND_ENABLE_USE_FABRIC_MEM:需要CANN升级到9.0之后的版本,HDK升级到26.0之后的版本,在支持该功能的版本使用Mooncake Store时推荐启用:可显著提升传输性能。
|
||||
ASCEND_USE_ASYNC_TRANSFER: 需要CANN升级到8.5之后的版本,用于开启Hixl异步传输模式,默认为同步模式。
|
||||
ASCEND_GLOBAL_RESOURCE_CONFIG:配置Hixl的全局资源,具体查看hixl的文档关于OPTION_GLOBAL_RESOURCE_CONFIG的配置。
|
||||
|
||||
|
||||
### 注意事项(必看)
|
||||
1. 调用TransferEngine initialize前需要set device, 比如`torch.npu.set_device(0)`。
|
||||
|
||||
|
|
@ -43,9 +70,13 @@ Transfer Engine RPC using <协议> listening on <IP>:<实际端口>,记录目
|
|||
|
||||
6. 通过`HCCL_RDMA_TIMEOUT` 用于配置RDMA网卡数据包重传超时时间系数,真实的数据包重传超时时间为`4.096us * 2 ^ $HCCL_RDMA_TIMEOUT`,通过`HCCL_RDMA_RETRY_CNT`来配置RDMA网卡的重传次数,建议配置`ASCEND_TRANSFER_TIMEOUT`略大于`重传时间 * HCCL_RDMA_RETRY_CNT`。
|
||||
|
||||
7. A2 server内/A3超节点内默认通信协议为`HCCS`,可以通过设置`export HCCL_INTRA_ROCE_ENABLE=1`来指定走`RDMA`。通常在KV Cache传输场景,为避免与模型集合通信流量冲突,影响推理性能,建议走RDMA传输。
|
||||
7. A2 server内/A3超节点内默认通信协议为`HCCS`,可以通过设置`export HCCL_INTRA_ROCE_ENABLE=1`来指定走`RDMA`。
|
||||
|
||||
8. 当使用`RDMA`通信协议时,在交换机和网卡默认配置不一致场景/需要流量规划场景下,可能需要修改RDMA网卡的Traffic Class和Service Level配置,通过`ASCEND_RDMA_TC`环境变量来设置Traffic Class, 通过`ASCEND_RDMA_SL`环境变量来设置Service Level。
|
||||
|
||||
9. 使用`RDMA`注册Host内存会按页表大小消耗device系统内存,在默认4KB页表情况下,可注册的Host的内存大约20GB, 另外`HCCS`也暂不支持Host内存传输,在这两种受约束的场景下,可通过中传Buffer的方式进行传输,具体开启方式是配置`ASCEND_BUFFER_POOL`环境变量,格式为`BUFFER_NUM:BUFFER_SIZE(单位MB)`, 推荐大小为`4:8`, 可根据实际场景调试出最合适的配置。
|
||||
9. 在向Host内存直接传输不通的场景下,可通过中传Buffer的方式进行传输,具体开启方式是配置`ASCEND_BUFFER_POOL`环境变量,格式为`BUFFER_NUM:BUFFER_SIZE(单位MB)`, 推荐大小为`4:8`, 可根据实际场景调试出最合适的配置。
|
||||
|
||||
10. 可以通过配置`ASCEND_USE_ASYNC_TRANSFER`环境变量来开启异步传输。
|
||||
|
||||
11. 在A3上,在获取最新驱动和CANN的前提下,在使用Mooncake store时,可以设置`ASCEND_ENABLE_USE_FABRIC_MEM`环境变量来开启fabric mem传输模式(能直接访问远端的HOST内存)。
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,8 @@
|
|||
# Ascend Transport
|
||||
Ascend Transport源代码路径为Mooncake/mooncake-transfer-engine/src/transport/ascend_transport,该路径下还包含自动化编译脚本、README文件。
|
||||
|
||||
**Ascend Transport 已不再维护,昇腾平台推荐使用 [Ascend Direct Transport](./ascend_direct_transport.md). **
|
||||
|
||||
## 概述
|
||||
Ascend Transport是一个单边语义的高性能零拷贝NPU数据传输库,直接兼容Mooncake Transfer Engine。要编译使用Ascend Transport库,请在mooncake-common\common.cmake文件中将USE_ASCEND开关置于"ON"。
|
||||
|
||||
|
|
|
|||
|
|
@ -110,7 +110,41 @@
|
|||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/musa/lib
|
||||
```
|
||||
|
||||
4. 安装 yalantinglibs
|
||||
4. 若需编译寒武纪 MLU 支持,请先安装寒武纪 Neuware SDK。之后:
|
||||
1) 导出 `NEUWARE_HOME`,或在 CMake 中传入 `-DNEUWARE_ROOT=/path/to/neuware`
|
||||
2) 配置 `LIBRARY_PATH` 与 `LD_LIBRARY_PATH`,确保编译时能链接 `cnrt`、`cndrv` 等 Neuware 库:
|
||||
```bash
|
||||
export NEUWARE_HOME=/usr/local/neuware
|
||||
export LIBRARY_PATH=$LIBRARY_PATH:${NEUWARE_HOME}/lib64
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${NEUWARE_HOME}/lib64
|
||||
```
|
||||
|
||||
若 Neuware 安装路径与默认头文件/库布局不一致,还可显式指定:
|
||||
```bash
|
||||
cmake .. -DUSE_MLU=ON \
|
||||
-DMLU_INCLUDE_DIR=/path/to/neuware/include \
|
||||
-DMLU_LIB_DIR=/path/to/neuware/lib64
|
||||
```
|
||||
|
||||
启用 MLU 后端示例:
|
||||
```bash
|
||||
cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=${NEUWARE_HOME:-/usr/local/neuware}
|
||||
make -j
|
||||
```
|
||||
|
||||
5. 若需编译沐曦 MetaX MACA 支持(如 C500),请安装 MACA SDK,使头文件与库位于 `MACA_ROOT`(优先取 `MACA_HOME` 环境变量,未设置时默认 `/opt/maca`)。不同安装包可能把库放在 `lib` 或 `lib64`,建议在环境变量中同时加入两者,避免链接或运行时找不到共享库:
|
||||
```bash
|
||||
export MACA_HOME=/opt/maca
|
||||
export LIBRARY_PATH=$LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
|
||||
```
|
||||
使用 `-DUSE_MACA=ON` 配置构建。可选覆盖项:
|
||||
- `-DMACA_ROOT=/path/to/maca`
|
||||
- `-DMACA_INCLUDE_DIR=/path/to/maca/include`
|
||||
- `-DMACA_LIB_DIR=/path/to/maca/lib64`
|
||||
- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`(分号分隔的 CMake 列表)
|
||||
|
||||
6. 安装 yalantinglibs
|
||||
```bash
|
||||
git clone https://github.com/alibaba/yalantinglibs.git
|
||||
cd yalantinglibs
|
||||
|
|
@ -120,7 +154,7 @@
|
|||
make install
|
||||
```
|
||||
|
||||
5. 进入项目根目录,运行下列命令进行编译
|
||||
7. 进入项目根目录,运行下列命令进行编译
|
||||
```bash
|
||||
mkdir build
|
||||
cd build
|
||||
|
|
@ -128,7 +162,7 @@
|
|||
make -j
|
||||
```
|
||||
|
||||
6. 安装 Mooncake python 包和 mooncake_master 可执行文件
|
||||
8. 安装 Mooncake python 包和 mooncake_master 可执行文件
|
||||
```bash
|
||||
make install
|
||||
```
|
||||
|
|
@ -137,6 +171,14 @@
|
|||
在执行 `cmake ..` 期间可以使用下列选项指定是否编译 Mooncake 的某些组件。
|
||||
- `-DUSE_CUDA=[ON|OFF]`: 启用 GPU Direct RDMA 及 NVMe-of 支持
|
||||
- `-DUSE_MUSA=[ON|OFF]`: 通过 MUSA 启用对摩尔线程 GPU 的支持
|
||||
- `-DUSE_MACA=[ON|OFF]`: 通过 MACA 启用对沐曦 MetaX GPU 的支持。
|
||||
- `-DMACA_ROOT=/path/to/maca`: 覆盖 MACA SDK 根路径(也支持 `MACA_HOME` 环境变量,默认 `/opt/maca`)。
|
||||
- `-DMACA_INCLUDE_DIR=/path/to/include`: 在 `-DUSE_MACA=ON` 时覆盖 MACA 头文件目录。
|
||||
- `-DMACA_LIB_DIR=/path/to/lib64`: 在 `-DUSE_MACA=ON` 时覆盖 MACA 库目录。
|
||||
- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`: 覆盖 `transfer_engine` 链接的 MACA 运行时库列表。
|
||||
- `-DUSE_MLU=[ON|OFF]`: 通过 Neuware 启用寒武纪 MLU 显存支持。默认 OFF;支持 MLU 显存探测、拓扑发现及 Transfer Engine 的 RDMA 注册。
|
||||
- `-DNEUWARE_ROOT=/path/to/neuware`: 在 `-DUSE_MLU=ON` 时覆盖默认 Neuware SDK 根路径;未设置时使用 `NEUWARE_HOME` 或 `/usr/local/neuware`。
|
||||
- `-DMLU_INCLUDE_DIR=/path/to/include` / `-DMLU_LIB_DIR=/path/to/lib64`: 在 `-DUSE_MLU=ON` 时覆盖 Neuware 头文件与库目录。
|
||||
- `-DUSE_HIP=[ON|OFF]`: 通过 HIP/ROCm 启用对 AMD GPU 的支持
|
||||
- `-DUSE_CXL=[ON|OFF]`: 启用 CXL 支持
|
||||
- `-DWITH_STORE=[ON|OFF]`: 编译 Mooncake Store 组件
|
||||
|
|
@ -150,3 +192,44 @@
|
|||
- `-DBUILD_UNIT_TESTS=[ON|OFF]`: 编译单元测试,默认为 ON
|
||||
- `-DBUILD_EXAMPLES=[ON|OFF]`: 编译示例程序,默认为 ON
|
||||
- `-DUSE_ASCEND_DIRECT=[ON|OFF]`: 启用 Ascend Direct RDMA 及 HCCS 支持
|
||||
- `-DUSE_MUSA=[ON|OFF]`: 启用Moore Threads GPUDirect RDMA
|
||||
|
||||
## 在 Docker 容器中使用 Mooncake
|
||||
|
||||
Mooncake 支持基于 Docker 的部署。您可以通过以下命令获取 Docker 镜像:
|
||||
|
||||
```bash
|
||||
docker pull alogfans/mooncake
|
||||
```
|
||||
|
||||
为了让容器能够使用主机的网络资源(特别是 InfiniBand RDMA),您需要在启动容器时添加 --device 选项。以下是使用示例:
|
||||
|
||||
## 在宿主机中运行容器
|
||||
```bash
|
||||
sudo docker run --net=host \
|
||||
--device=/dev/infiniband/uverbs0 \
|
||||
--device=/dev/infiniband/rdma_cm \
|
||||
--ulimit memlock=-1 \
|
||||
-t -i mooncake:v0.9.0 /bin/bash
|
||||
```
|
||||
|
||||
|
||||
## 进入容器后,运行 transfer engine 示例
|
||||
|
||||
```bash
|
||||
cd /app/build/mooncake-transfer-engine/example
|
||||
./transfer_engine_bench --device_name=ibp6s0 \
|
||||
--metadata_server=10.1.101.3:2379 \
|
||||
--mode=target \
|
||||
--local_server_name=10.1.100.3
|
||||
|
||||
```
|
||||
|
||||
注意事项:
|
||||
|
||||
--device 参数将宿主机的 RDMA 设备映射到容器内
|
||||
|
||||
--ulimit memlock=-1 解除内存锁定限制,RDMA 操作需要
|
||||
|
||||
--net=host 让容器使用宿主机的网络命名空间
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,90 @@
|
|||
# Kunpeng UB Transport
|
||||
Kunpeng UbTransport源代码路径为Mooncake/mooncake-transfer-engine/src/transport/kunpneg_transport,该路径下有UB协议的Transport对接代码和实现逻辑。
|
||||
|
||||
## 概述
|
||||
UB(Unified Bus,统一总线) 是与RDMA、CXL、NVLink 和TCP处于同一抽象层的传输协议,属于可在应用层灵活选择的传输方案。目前 UB 协议有两个开源实现:URMA(远程内存访问语义)和 OBMM(Load/Store 语义)。
|
||||
|
||||
URMA(Unified Remote Memory Access,统一远程内存访问)是UB协议为上层应用提供的统一编程抽象与核心语义层。它基于 UB 协议低延迟、高带宽的底层特性,为远程共享内存的访问与操作提供统一的 API 和语义接口。
|
||||
|
||||
URMA 开源代码仓库:https://atomgit.com/openeuler/umdk
|
||||
|
||||
OBMM (Ownership Based Memory Management) 是面向超节点环境的内核内存管理系统,支持跨节点的物理内存共享。该系统通过内核模块 (obmm.ko) 和用户态库 (libobmm.so) 提供高效的远程内存访问能力。
|
||||
|
||||
OBMM 开源代码仓库:https://atomgit.com/openeuler/obmm
|
||||
|
||||
## 新增依赖
|
||||
Kunpeng UbTransport在Mooncake本身依赖的基础上,新增了一部分URMA和OBMM的依赖:
|
||||
|
||||
- **硬件平台**: 支持原生UB互联架构的鲲鹏950 CPU
|
||||
- **OS版本**: openEuler 24.03 (LTS-SP3) [下载链接](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
|
||||
- **URMA依赖**: UMDK: `yum install umdk-urma-devel` 或从[源码](https://atomgit.com/openeuler/umdk)构建。
|
||||
- **协议优势**: URMA 提供类似 RDMA 的内存语义,针对鲲鹏芯片片上互联进行了优化
|
||||
|
||||
---
|
||||
|
||||
## 构建与编译
|
||||
|
||||
**前置条件**
|
||||
|
||||
- openEuler 24.03 (LTS-SP3) [下载链接](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
|
||||
- 已安装 UMDK: `yum install umdk-urma-devel` 或从[源码](https://atomgit.com/openeuler/umdk)构建
|
||||
|
||||
**CMake 配置**
|
||||
|
||||
```bash
|
||||
# 克隆 Mooncake 仓库
|
||||
git clone https://github.com/kvcache-ai/Mooncake.git
|
||||
cd Mooncake
|
||||
|
||||
# 启用 UB 传输层进行配置
|
||||
mkdir build && cd build
|
||||
cmake .. -DUSE_UB=ON \
|
||||
-DURMA_INCLUDE_DIR=/usr/include \
|
||||
-DURMA_LIBRARY=/usr/lib64/liburma.so
|
||||
|
||||
# 编译
|
||||
make -j$(nproc)
|
||||
```
|
||||
|
||||
**验证**
|
||||
|
||||
```bash
|
||||
# 检查 UB 传输层是否已注册
|
||||
./mooncake_server --list-transports
|
||||
# 预期输出: rdma, tcp, nvlink, ub
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 运行与测试
|
||||
|
||||
**单节点基准测试**
|
||||
|
||||
```bash
|
||||
# 终端 1: 目标端(Target)
|
||||
./transfer_engine_bench \
|
||||
--mode=target \
|
||||
--protocol=ub \
|
||||
--device_name=urma0 \
|
||||
--local_server_name=127.0.0.1 \
|
||||
--metadata_server=P2PHANDSHAKE
|
||||
|
||||
# 终端 2: 发起端(Initiator)
|
||||
./transfer_engine_bench \
|
||||
--mode=initiator \
|
||||
--protocol=ub \
|
||||
--device_name=urma0 \
|
||||
--metadata_server=P2PHANDSHAKE \
|
||||
--segment_size=8388608 \
|
||||
--batch_size=1\
|
||||
--segment_id=127.0.0.1:$PORT
|
||||
```
|
||||
|
||||
**多设备基准测试**
|
||||
|
||||
```bash
|
||||
# 自动发现多个 URMA 设备
|
||||
./transfer_engine_bench \
|
||||
--protocol=ub \
|
||||
--device_name=urma0,urma1,urma2,urma3
|
||||
```
|
||||
|
|
@ -73,7 +73,7 @@ ErrorCode Init(const std::string& local_hostname,
|
|||
### Get 接口
|
||||
|
||||
```C++
|
||||
tl::expected<void, ErrorCode> Get(const std::string& object_key,
|
||||
tl::expected<void, ErrorCode> Get(const std::string& object_key,
|
||||
std::vector<Slice>& slices);
|
||||
```
|
||||
|
||||
|
|
@ -116,6 +116,65 @@ tl::expected<void, ErrorCode> Remove(const ObjectKey& key);
|
|||
|
||||
用于删除指定 key 对应的对象。该接口标记存储引擎中与 key 关联的所有数据副本已被删除,不需要与对应存储节点(Client)通信。
|
||||
|
||||
### CreateCopyTask 接口
|
||||
|
||||
```C++
|
||||
tl::expected<UUID, ErrorCode> CreateCopyTask(
|
||||
const std::string& key,
|
||||
const std::vector<std::string>& targets);
|
||||
```
|
||||
|
||||
`CreateCopyTask` 创建一个异步复制任务,将由客户端的任务执行系统执行。当您需要提交多个复制操作而不等待每个操作完成时,这很有用。任务被提交到 master 服务,分配唯一的任务 ID,并由可用的客户端异步执行。可以使用 `QueryTask` 查询任务状态。
|
||||
|
||||
**任务执行和结果反馈:**
|
||||
1. **任务分配**:master 服务在客户端定期通讯中将任务分配给可用的客户端
|
||||
2. **任务执行**:分配的客户端在后台线程池中异步执行复制操作
|
||||
3. **结果反馈**:执行完成后(成功或失败),客户端通过 `MarkTaskToComplete` 自动向 master 服务报告结果:
|
||||
- 成功时:`status = SUCCESS`,`message = "Task completed successfully"`
|
||||
- 失败时:`status = FAILED`,`message = <错误描述>`
|
||||
4. **状态查询**:您可以随时使用 `QueryTask` 查询任务状态以监控进度
|
||||
|
||||
### CreateMoveTask 接口
|
||||
|
||||
```C++
|
||||
tl::expected<UUID, ErrorCode> CreateMoveTask(
|
||||
const std::string& key,
|
||||
const std::string& source,
|
||||
const std::string& target);
|
||||
```
|
||||
|
||||
`CreateMoveTask` 创建一个异步移动任务,将由客户端的任务执行系统执行。当您需要提交多个移动操作而不等待每个操作完成时,这很有用。任务被提交到 master 服务,分配唯一的任务 ID,并由可用的客户端异步执行。可以使用 `QueryTask` 查询任务状态。
|
||||
|
||||
**任务执行和结果反馈:**
|
||||
1. **任务分配**:master 服务在客户端定期通讯中将任务分配给可用的客户端
|
||||
2. **任务执行**:分配的客户端在后台线程池中异步执行移动操作
|
||||
3. **结果反馈**:执行完成后(成功或失败),客户端通过 `MarkTaskToComplete` 自动向 master 服务报告结果:
|
||||
- 成功时:`status = SUCCESS`,`message = "Task completed successfully"`
|
||||
- 失败时:`status = FAILED`,`message = <错误描述>`
|
||||
4. **状态查询**:您可以随时使用 `QueryTask` 查询任务状态以监控进度
|
||||
|
||||
### QueryTask 接口
|
||||
|
||||
```C++
|
||||
tl::expected<QueryTaskResponse, ErrorCode> QueryTask(const UUID& task_id);
|
||||
```
|
||||
|
||||
`QueryTask` 查询异步任务(复制或移动)的状态。这允许您监控基于任务的操作进度。响应包括任务状态、类型、创建时间、最后更新时间、分配的客户端和状态消息。
|
||||
|
||||
其中`QueryTaskResponse` 的数据结构细节如下:
|
||||
|
||||
```C++
|
||||
struct QueryTaskResponse {
|
||||
UUID id; // 任务 UUID
|
||||
TaskType type; // 任务类型 (REPLICA_COPY 或 REPLICA_MOVE)
|
||||
TaskStatus status; // 任务状态 (PENDING, PROCESSING, SUCCESS, 或 FAILED)
|
||||
int64_t created_at_ms_epoch; // 任务创建时间戳(毫秒)
|
||||
int64_t last_updated_at_ms_epoch; // 最后更新时间戳(毫秒)
|
||||
UUID assigned_client; // 分配给执行任务的客户端 UUID
|
||||
std::string message; // 状态消息或错误描述
|
||||
};
|
||||
```
|
||||
|
||||
### RemoveByRegex
|
||||
|
||||
```C++
|
||||
|
|
@ -156,6 +215,22 @@ QueryByRegex(const std::string& str);
|
|||
|
||||
将集群中所有可用的资源看做一个巨大的资源池,由一个中心化的 Master 进程进行空间分配,并指导实现数据复制(**注意 Master Service 不接管任何的数据流,只是提供对应的元数据信息**)。
|
||||
|
||||
#### Snapshot 与 Restore
|
||||
|
||||
为减少 master 重启后的缓存预热时间,Master Service 支持对自身元数据进行周期性快照(snapshot),并在启动时从快照中恢复(restore)。
|
||||
|
||||
- Snapshot 生成
|
||||
- 后台快照线程会定期在不阻塞正常 RPC 请求的情况下,基于 fork 的写时复制(copy-on-write)机制,获取一份一致性的内存快照,其中包含 KV 元数据、segment 信息以及分配器状态。
|
||||
- 子进程将这些结构序列化为紧凑的二进制格式,并通过 `SerializerBackend` 抽象写入配置的 snapshot 后端。
|
||||
- Restore
|
||||
- 在启动阶段,当启用 snapshot restore 时,master 会从后端读取最新的快照,在内存中重建 Master Service 的元数据状态。
|
||||
- 提示
|
||||
- 由于快照是周期性生成而不是实时更新,如果 master 在两次快照之间发生故障,自上一次成功快照以来的部分元数据变更可能无法恢复。
|
||||
|
||||
> **警告:受管存储**
|
||||
>
|
||||
> 快照存储位置由 Mooncake 快照系统**完全管理**,过期快照会被自动删除。**请勿在此位置存放其他文件**,请使用独立、隔离的存储用于快照。
|
||||
|
||||
#### Master Service 接口
|
||||
|
||||
Master与Client的通信协议如下:
|
||||
|
|
@ -226,7 +301,7 @@ service MasterService {
|
|||
|
||||
```protobuf
|
||||
message GetReplicaListRequest {
|
||||
required string key = 1;
|
||||
required string key = 1;
|
||||
};
|
||||
|
||||
message GetReplicaListResponse {
|
||||
|
|
@ -315,7 +390,7 @@ message PutStartRequest {
|
|||
};
|
||||
|
||||
message PutStartResponse {
|
||||
required int32 status_code = 1;
|
||||
required int32 status_code = 1;
|
||||
repeated ReplicaInfo replica_list = 2; // Master Service 分配好的副本信息
|
||||
};
|
||||
```
|
||||
|
|
@ -329,7 +404,7 @@ message PutStartResponse {
|
|||
|
||||
```protobuf
|
||||
message PutEndRequest {
|
||||
required string key = 1;
|
||||
required string key = 1;
|
||||
};
|
||||
|
||||
message PutEndResponse {
|
||||
|
|
@ -346,7 +421,7 @@ Client 完成数据写入后,调用 PutEnd 通知 `Master Service`。`Master S
|
|||
|
||||
```protobuf
|
||||
message RemoveRequest {
|
||||
required string key = 1;
|
||||
required string key = 1;
|
||||
};
|
||||
|
||||
message RemoveResponse {
|
||||
|
|
@ -355,7 +430,7 @@ message RemoveResponse {
|
|||
```
|
||||
|
||||
* 请求: RemoveRequest,包含需要删除对象的key
|
||||
* 响应: RemoveResponse,包含状态码 status_code
|
||||
* 响应: RemoveResponse,包含状态码 status_code
|
||||
|
||||
用于删除指定 key 对应的对象及其所有副本。Master Service 将对应对象的所有副本状态标记为删除。
|
||||
|
||||
|
|
@ -620,7 +695,7 @@ mooncake提供了DFS可用空间的配置,用户可以在启动master时指定
|
|||
**注意** 当前还没有提供DFS上文件驱逐的能力
|
||||
|
||||
#### 数据访问机制
|
||||
持久化功能同样遵循了mooncake store中控制流和数据流分离的设计。kvcache object的读\写操作在client端完成,kvcache object的查询和管理功能在master端完成。在文件系统中key -> kvcache object的索引信息是由固定的索引机制维护,每个文件对应一个kvcache object(文件名即为对应的key名称)。
|
||||
持久化功能同样遵循了mooncake store中控制流和数据流分离的设计。kvcache object的读\写操作在client端完成,kvcache object的查询和管理功能在master端完成。在文件系统中key -> kvcache object的索引信息是由固定的索引机制维护,每个文件对应一个kvcache object(文件名即为对应的key名称)。
|
||||
|
||||
启用持久化功能后,对于每次 `Put`或`BatchPut` 操作,都会发起一次同步的memory pool写入操作和一次异步的DFS持久化操作。之后执行 `Get`或 `BatchGet` 时,如果在memory pool中没有找到对应的kvcache,则会尝试从DFS中读取该文件数据,并返回给用户。
|
||||
|
||||
|
|
@ -647,6 +722,8 @@ HTTP 元数据服务器可通过以下参数进行配置:
|
|||
- **MC_STORE_MEMCPY**: 控制是否启用本地 memcpy 优化, 1/true 启用, 0/false 禁用
|
||||
- **MC_STORE_CLIENT_METRIC**: 启用客户端指标上报, 默认启用;设为 0/false 禁用
|
||||
- **MC_STORE_CLIENT_METRIC_INTERVAL**: 指标上报间隔(秒), 默认 0(仅收集不上报)
|
||||
- **MC_STORE_USE_HUGEPAGE**: 启用 hugepage 优化, 默认禁用, 设置为 1/true 启用
|
||||
- **MC_STORE_HUGEPAGE_SIZE**: hugepage 页大小, 默认 2M
|
||||
|
||||
#### 使用示例
|
||||
|
||||
|
|
@ -702,6 +779,8 @@ Max threads: 4
|
|||
Master service listening on 0.0.0.0:50051
|
||||
```
|
||||
|
||||
如果 Master 运行在容器中,而容器 IP 可能动态变化,建议使用 `--rpc-interface=<网卡名>`(例如 `--rpc-interface=eth0`)而不是写死 `--rpc-address`。Master 会在启动时解析该网卡当前的 IPv4 地址,并将其作为最终的 `rpc_address` 使用。
|
||||
|
||||
**高可用模式**:
|
||||
|
||||
高可用模式依赖于 etcd 服务进行协调。如果 Transfer Engine 也使用 etcd 作为其元数据服务,那么 Mooncake Store 使用的 etcd 集群可以与 Transfer Engine 使用的集群共用,也可以是独立的。
|
||||
|
|
@ -711,6 +790,7 @@ Master service listening on 0.0.0.0:50051
|
|||
--enable-ha:启用高可用模式
|
||||
--etcd-endpoints:指定 etcd 服务的多个入口,使用分号 ';' 分隔
|
||||
--rpc-address:该实例的 RPC 地址。注意,这里填写的地址应当是客户端可访问的地址。
|
||||
--rpc-interface:按网卡名解析当前实例的 IPv4 地址。设置后会覆盖 --rpc-address,适合容器 IP 会变化的场景。
|
||||
```
|
||||
|
||||
例如:
|
||||
|
|
@ -721,6 +801,14 @@ Master service listening on 0.0.0.0:50051
|
|||
--rpc-address=10.0.0.1
|
||||
```
|
||||
|
||||
容器部署示例:
|
||||
```
|
||||
./build/mooncake-store/src/mooncake_master \
|
||||
--enable-ha=true \
|
||||
--etcd-endpoints="0.0.0.0:2379;0.0.0.0:2479;0.0.0.0:2579" \
|
||||
--rpc-interface=eth0
|
||||
```
|
||||
|
||||
### 启动验证程序
|
||||
Mooncake Store 提供了多种验证程序,包括基于 C++ 和 Python 等接口形态。下面以 `stress_cluster_benchmark` 为例介绍一下如何运行。
|
||||
|
||||
|
|
@ -796,6 +884,8 @@ retcode = store.setup(
|
|||
|
||||
- **`host`**: (字符串, 默认: "0.0.0.0"): client 的主机名。
|
||||
|
||||
- **`port`**: (整型, 默认: 50052): client 监听的端口。
|
||||
|
||||
- **`global_segment_size`**: (字符串, 默认: "4GB"): client 向集群中挂载的 Segment 大小。
|
||||
|
||||
- **`master_server_address`**: (字符串, 默认: "localhost:50051"): Master 服务的地址。
|
||||
|
|
|
|||
|
|
@ -115,7 +115,7 @@ get_first_buffer_address(segment_name)
|
|||
- `segment_name` (str): 段的名称
|
||||
|
||||
**返回值:**
|
||||
- `int`: 段中第一个缓冲区的内存地址
|
||||
- `int`: 段中第一个缓冲区的内存地址,如果段不存在或没有已注册的缓冲区则返回 0
|
||||
|
||||
### 数据传输操作
|
||||
|
||||
|
|
@ -200,7 +200,7 @@ transfer_check_status(batch_id)
|
|||
- `batch_id` (int): 从transfer_submit_write()返回的批次ID
|
||||
|
||||
**返回值:**
|
||||
- `int`:
|
||||
- `int`:
|
||||
- 1: 传输成功完成
|
||||
- 0: 传输仍在进行中
|
||||
- -1: 传输失败
|
||||
|
|
@ -388,11 +388,11 @@ else:
|
|||
# Use the buffer
|
||||
test_data = b"Test data for managed buffer"
|
||||
engine.write_bytes_to_buffer(buffer_addr, test_data, len(test_data))
|
||||
|
||||
|
||||
# Read back
|
||||
read_data = engine.read_bytes_from_buffer(buffer_addr, len(test_data))
|
||||
print(f"Read data: {read_data}")
|
||||
|
||||
|
||||
# Free the buffer when done
|
||||
engine.free_managed_buffer(buffer_addr, buffer_size)
|
||||
```
|
||||
|
|
@ -459,4 +459,4 @@ Transfer Engine Python API在大多数操作中都是线程安全的。但是,
|
|||
2. **错误检查**: 检查所有API调用的返回值
|
||||
3. **配置优化**: 根据硬件环境调整环境变量
|
||||
4. **监控**: 启用指标报告以监控传输性能
|
||||
5. **测试**: 在生产环境中使用前进行充分的测试
|
||||
5. **测试**: 在生产环境中使用前进行充分的测试
|
||||
|
|
|
|||
|
|
@ -118,7 +118,7 @@ Transfer Engine 使用SIEVE算法来管理端点的逐出。如果由于链路
|
|||
```
|
||||
各个参数的含义如下(其余同前):
|
||||
- `--segment_id` 可以简单理解为目标节点对应的段名称,需要和启动目标节点时 `--local_server_name` 传入的值(如果有)保持一致。
|
||||
|
||||
|
||||
正常情况下,发起节点将开始进行传输操作,等待 10s 后回显“Test completed”信息,表明测试完成。
|
||||
|
||||
发起节点还可以配置下列测试参数:`--operation`(可为 `"read"` 或 `"write"`)、`batch_size`、`block_size`、`duration`、`threads` 等。
|
||||
|
|
@ -265,7 +265,7 @@ SegmentHandle openSegment(const std::string& segment_name);
|
|||
```
|
||||
- `segment_name`:segment 的唯一标志符。对于 RAM Segment,这需要与对端进程初始化 TransferEngine 对象时填写的 `server_name` 保持一致。
|
||||
- 返回值:若成功,返回对应的 SegmentHandle;否则返回负数值。
|
||||
|
||||
|
||||
```cpp
|
||||
int closeSegment(SegmentHandle segment_id);
|
||||
```
|
||||
|
|
@ -285,7 +285,7 @@ Value = {
|
|||
'rpc_port': 12345
|
||||
}
|
||||
|
||||
// 对于 segment,采用 mooncake/[proto]/[segment_name] 的 key 命名方式,segment name 可以采用 Server Name。
|
||||
// 对于 segment,采用 mooncake/[proto]/[segment_name] 的 key 命名方式,segment name 可以采用 Server Name。
|
||||
// Segment 对应机器,buffer 对应机器内的不同段内存或者不同的文件或者不同的盘。同一个 segment 的不同 buffer 处于同一个故障域。
|
||||
|
||||
// RAM Segment,用于 RDMA Transport 获取传输信息。
|
||||
|
|
@ -320,7 +320,7 @@ Key = mooncake/nvmeof/[segment_name]
|
|||
Value = {
|
||||
'server_name': server_name,
|
||||
'protocol': nvmeof,
|
||||
'buffers':[
|
||||
'buffers':[
|
||||
{
|
||||
'length': 1073741824,
|
||||
'file_path': "/mnt/nvme0" // 本机器上的文件路径
|
||||
|
|
@ -331,7 +331,7 @@ Value = {
|
|||
},
|
||||
{
|
||||
'length': 1073741824,
|
||||
'file_path': "/mnt/nvme1",
|
||||
'file_path': "/mnt/nvme1",
|
||||
'local_path_map': {
|
||||
"node02": "/mnt/transfer_engine/node02/nvme1",
|
||||
....
|
||||
|
|
@ -408,6 +408,7 @@ int init(const std::string &metadata_conn_string,
|
|||
- `MC_RETRY_CNT` Transfer Engine 中最大重试次数
|
||||
- `MC_LOG_LEVEL` 该选项可以设置成`TRACE`/`INFO`/`WARNING`/`ERROR`(详情见 [glog doc](https://github.com/google/glog/blob/master/docs/logging.md)),则在运行时会输出更详细的日志
|
||||
- `MC_HANDSHAKE_LISTEN_BACKLOG` 监听握手连接的 backlog 大小, 默认值 128
|
||||
- `MC_HANDSHAKE_MAX_LENGTH` P2P 模式下握手消息的最大长度(字节)。有效范围:1MB 到 128MB。默认值为 1MB (1048576 字节)。当单个 RDMA 实例注册大量内存缓冲区(>10,000)时,需要增大此值以避免握手失败。示例:设置为 10485760 表示 10MB
|
||||
- `MC_LOG_DIR` 该选项指定存放日志重定向文件的目录路径。如果路径无效,glog将回退到向标准错误[stderr]输出日志。
|
||||
- `MC_REDIS_PASSWORD` Redis 存储插件的密码,仅在指定 Redis 作为 metadata server 时生效。如果未设置,将不会尝试进行密码认证登录 Redis。
|
||||
- `MC_REDIS_DB_INDEX` Redis 存储插件的数据库索引,必须为 0 到 255 之间的整数。仅在指定 Redis 作为 metadata server 时生效。如果未设置或无效,默认值为 0。
|
||||
|
|
@ -415,8 +416,9 @@ int init(const std::string &metadata_conn_string,
|
|||
- `MC_ENABLE_DEST_DEVICE_AFFINITY` 启用设备亲和性以优化 RDMA 性能。启用后,Transfer Engine 将优先选择和本地网卡同名的远端网卡进行通信,以减少 QP 数量并改善 Rail-optimized 拓扑中的网络性能。默认值为 false
|
||||
- `MC_FORCE_HCA` 强制使用RDMA作为主要传输方式,如果没有探测到有效的RDMA网卡,返回失败
|
||||
- `MC_FORCE_MNNVL` 强制使用 Multi-Node NVLink 作为主要传输方式,无论是否安装了有效的 RDMA 网卡
|
||||
- `MC_INTRA_NVLINK` 指定使用Intra-Node NVLink 作为主要传输方式,同时注意该设置不能与MC_FORCE_MNNVL一起使用
|
||||
- `MC_FORCE_TCP` 强制使用 TCP 作为主要传输方式,无论是否安装了有效的 RDMA 网卡
|
||||
- `MC_MIN_PRC_PORT` 指定 RPC 服务使用的最小端口号。默认值为 15000。
|
||||
- `MC_MAX_PRC_PORT` 指定 RPC 服务使用的最大端口号。默认值为 17000。
|
||||
- `MC_PATH_ROUNDROBIN` 指定 RDMA 路径选择使用 Round Robin 模式,这对于传输大块数据可能有利。
|
||||
- `MC_ENDPOINT_STORE_TYPE` 选择 FIFO Endpoint Store (`FIFO`) 或者 Sieve Endpoint Store (`SIEVE`),模式是 `SIEVE`。
|
||||
- `MC_ENDPOINT_STORE_TYPE` 选择 FIFO Endpoint Store (`FIFO`) 或者 Sieve Endpoint Store (`SIEVE`),模式是 `SIEVE`。
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
Subproject commit 73dea196d23ad8fcd4914c6ef1238f390b9a1c48
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 32 KiB |
|
|
@ -0,0 +1,81 @@
|
|||
# Monitoring Mooncake with Prometheus and Grafana
|
||||
|
||||
This document provides instructions on how to set up and run Prometheus and Grafana to monitor the `mooncake_master` service.
|
||||
|
||||
## 1. Running Monitoring Services
|
||||
|
||||
The Prometheus and Grafana services are configured to run inside Docker containers.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
- Docker
|
||||
- Docker Compose
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Navigate to the monitoring directory:**
|
||||
|
||||
```bash
|
||||
cd monitoring
|
||||
```
|
||||
|
||||
2. **Start the containers:**
|
||||
|
||||
```bash
|
||||
docker-compose up -d
|
||||
```
|
||||
|
||||
This command will pull the Prometheus and Grafana images and start the containers in detached mode.
|
||||
|
||||
3. **Access the Prometheus UI:**
|
||||
|
||||
You can access the Prometheus web interface by navigating to `http://localhost:9090` in your web browser.
|
||||
|
||||
4. **Access the Grafana UI:**
|
||||
|
||||
You can access the Grafana web interface by navigating to `http://localhost:3000` in your web browser.
|
||||
|
||||
- **Default credentials:** `admin` / `admin`
|
||||
|
||||
Grafana is pre-configured with the Prometheus datasource and a simple dashboard for the `mooncake_master`.
|
||||
|
||||
### Configuration
|
||||
|
||||
- **`docker-compose.yml`**: Defines the Prometheus and Grafana services.
|
||||
- **`prometheus/prometheus.yml`**: The Prometheus configuration file. By default, it's configured to scrape metrics from the `mooncake_master`.
|
||||
- **`grafana/`**: Contains Grafana provisioning files for the datasource and a sample dashboard.
|
||||
|
||||
The `prometheus.yml` is configured to scrape metrics from `host.docker.internal:9003`. The `host.docker.internal` hostname allows the Prometheus container to communicate with services running on the host machine. If you are running on Linux and `host.docker.internal` is not available, you may need to add `extra_hosts: ["host.docker.internal:host-gateway"]` to the prometheus service definition in `docker-compose.yml`.
|
||||
|
||||
## 2. Running Mooncake Master with Metrics
|
||||
|
||||
To be monitored by Prometheus, the `mooncake_master` must be started with metric reporting enabled.
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Start the `mooncake_master`:**
|
||||
|
||||
Run the `mooncake_master` executable with the following flags to enable the metrics endpoint. For a complete list of flags, see the [Mooncake Store Deployment Guide](../docs/source/deployment/mooncake-store-deployment-guide.md).
|
||||
|
||||
```bash
|
||||
# Make sure you are in the root directory of the project
|
||||
# The path to mooncake_master might vary depending on your build output directory
|
||||
./build/mooncake_master \
|
||||
--metrics_port=9003 \
|
||||
--enable_metric_reporting=true
|
||||
```
|
||||
|
||||
Ensure that `--metrics_port` is set to `9003`, as this is the port Prometheus is configured to scrape.
|
||||
|
||||
2. **Verify Metrics Endpoint:**
|
||||
|
||||
You can verify that the metrics endpoint is active by running:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:9003/metrics
|
||||
```
|
||||
|
||||
This should return a list of Prometheus-style metrics.
|
||||
|
||||
Once both Prometheus and `mooncake_master` are running, you can go to the Prometheus UI (`http://localhost:9090`), navigate to **Status -> Targets**, and you should see the `mooncake-master` job with a state of "UP". You can then explore the pre-built dashboard in Grafana.
|
||||
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
version: '3.8'
|
||||
|
||||
services:
|
||||
prometheus:
|
||||
image: prom/prometheus:v2.54.0
|
||||
container_name: prometheus
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "9090:9090"
|
||||
volumes:
|
||||
- ./prometheus/prometheus.yml:/etc/prometheus/prometheus.yml
|
||||
command:
|
||||
- '--config.file=/etc/prometheus/prometheus.yml'
|
||||
- '--storage.tsdb.retention.time=15d'
|
||||
- '--web.console.libraries=/usr/share/prometheus/console_libraries'
|
||||
- '--web.console.templates=/usr/share/prometheus/consoles'
|
||||
extra_hosts:
|
||||
- "host.docker.internal:host-gateway"
|
||||
networks:
|
||||
- monitoring
|
||||
|
||||
grafana:
|
||||
image: grafana/grafana-oss:11.0.0
|
||||
container_name: grafana
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "3000:3000"
|
||||
environment:
|
||||
- GF_SECURITY_ADMIN_USER=admin
|
||||
- GF_SECURITY_ADMIN_PASSWORD=admin
|
||||
- GF_USERS_ALLOW_SIGN_UP=false
|
||||
volumes:
|
||||
- grafana-data:/var/lib/grafana
|
||||
- ./grafana/provisioning:/etc/grafana/provisioning
|
||||
- ./grafana/dashboards:/etc/grafana/dashboards
|
||||
depends_on:
|
||||
- prometheus
|
||||
networks:
|
||||
- monitoring
|
||||
|
||||
volumes:
|
||||
grafana-data:
|
||||
|
||||
networks:
|
||||
monitoring:
|
||||
driver: bridge
|
||||
|
|
@ -0,0 +1,67 @@
|
|||
{
|
||||
"annotations": {
|
||||
"list": [
|
||||
{
|
||||
"builtIn": 1,
|
||||
"datasource": "-- Grafana --",
|
||||
"enable": true,
|
||||
"hide": true,
|
||||
"iconColor": "rgba(0, 211, 255, 1)",
|
||||
"name": "Annotations & Alerts",
|
||||
"type": "dashboard"
|
||||
}
|
||||
]
|
||||
},
|
||||
"editable": true,
|
||||
"gnetId": null,
|
||||
"graphTooltip": 0,
|
||||
"links": [],
|
||||
"panels": [
|
||||
{
|
||||
"title": "RPC Requests",
|
||||
"type": "graph",
|
||||
"datasource": "Prometheus",
|
||||
"gridPos": {
|
||||
"h": 9,
|
||||
"w": 12,
|
||||
"x": 0,
|
||||
"y": 0
|
||||
},
|
||||
"id": 2,
|
||||
"targets": [
|
||||
{
|
||||
"expr": "rate(mooncake_master_rpc_requests_total[5m])",
|
||||
"legendFormat": "{{method}}",
|
||||
"refId": "A"
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"schemaVersion": 22,
|
||||
"style": "dark",
|
||||
"tags": [],
|
||||
"templating": {
|
||||
"list": []
|
||||
},
|
||||
"title": "Mooncake Master",
|
||||
"time": {
|
||||
"from": "now-6h",
|
||||
"to": "now"
|
||||
},
|
||||
"timepicker": {
|
||||
"refresh_intervals": [
|
||||
"5s",
|
||||
"10s",
|
||||
"30s",
|
||||
"1m",
|
||||
"5m",
|
||||
"15m",
|
||||
"30m",
|
||||
"1h",
|
||||
"2h",
|
||||
"1d"
|
||||
]
|
||||
},
|
||||
"timezone": "browser",
|
||||
"version": 1
|
||||
}
|
||||
|
|
@ -0,0 +1,11 @@
|
|||
apiVersion: 1
|
||||
|
||||
providers:
|
||||
- name: 'default'
|
||||
orgId: 1
|
||||
folder: ''
|
||||
type: file
|
||||
disableDeletion: false
|
||||
editable: true
|
||||
options:
|
||||
path: /etc/grafana/dashboards
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
apiVersion: 1
|
||||
|
||||
datasources:
|
||||
- name: Prometheus
|
||||
type: prometheus
|
||||
access: proxy
|
||||
url: http://prometheus:9090
|
||||
isDefault: true
|
||||
|
|
@ -0,0 +1,13 @@
|
|||
global:
|
||||
scrape_interval: 15s
|
||||
evaluation_interval: 15s
|
||||
|
||||
scrape_configs:
|
||||
- job_name: 'prometheus'
|
||||
static_configs:
|
||||
- targets: ['localhost:9090']
|
||||
|
||||
- job_name: 'mooncake-master'
|
||||
scrape_interval: 5s
|
||||
static_configs:
|
||||
- targets: ['host.docker.internal:9003']
|
||||
|
|
@ -0,0 +1,48 @@
|
|||
#!/bin/bash
|
||||
|
||||
# Ensure we are currently in the script's directory to make relative paths work
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
SCRIPT_DIR=$(pwd)
|
||||
echo "Working directory: $SCRIPT_DIR"
|
||||
|
||||
# Create network if it doesn't exist
|
||||
docker network create monitoring 2>/dev/null
|
||||
|
||||
# Create volume if it doesn't exist
|
||||
docker volume create grafana-data 2>/dev/null
|
||||
|
||||
# Remove existing containers to avoid conflicts
|
||||
docker rm -f prometheus grafana 2>/dev/null
|
||||
|
||||
echo "Starting Prometheus..."
|
||||
docker run -d \
|
||||
--name prometheus \
|
||||
--restart unless-stopped \
|
||||
--network monitoring \
|
||||
-p 9090:9090 \
|
||||
-v "$SCRIPT_DIR/prometheus/prometheus.yml:/etc/prometheus/prometheus.yml" \
|
||||
--add-host host.docker.internal:host-gateway \
|
||||
prom/prometheus:v2.54.0 \
|
||||
--config.file=/etc/prometheus/prometheus.yml \
|
||||
--storage.tsdb.retention.time=15d \
|
||||
--web.console.libraries=/usr/share/prometheus/console_libraries \
|
||||
--web.console.templates=/usr/share/prometheus/consoles
|
||||
|
||||
echo "Starting Grafana..."
|
||||
docker run -d \
|
||||
--name grafana \
|
||||
--restart unless-stopped \
|
||||
--network monitoring \
|
||||
-p 3000:3000 \
|
||||
-e GF_SECURITY_ADMIN_USER=admin \
|
||||
-e GF_SECURITY_ADMIN_PASSWORD=admin \
|
||||
-e GF_USERS_ALLOW_SIGN_UP=false \
|
||||
-v grafana-data:/var/lib/grafana \
|
||||
-v "$SCRIPT_DIR/grafana/provisioning:/etc/grafana/provisioning" \
|
||||
-v "$SCRIPT_DIR/grafana/dashboards:/etc/grafana/dashboards" \
|
||||
grafana/grafana-oss:11.0.0
|
||||
|
||||
echo "Monitoring stack started!"
|
||||
echo "Prometheus: http://localhost:9090"
|
||||
echo "Grafana: http://localhost:3000"
|
||||
|
|
@ -2,6 +2,10 @@ if ((USE_ETCD AND NOT USE_ETCD_LEGACY) OR STORE_USE_ETCD)
|
|||
add_subdirectory(etcd)
|
||||
endif()
|
||||
|
||||
if (STORE_USE_K8S_LEASE)
|
||||
add_subdirectory(k8s-lease)
|
||||
endif()
|
||||
|
||||
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
|
||||
add_subdirectory(src)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,20 @@
|
|||
include(FetchContent)
|
||||
|
||||
# UMDK 头文件库
|
||||
FetchContent_Declare(
|
||||
urma
|
||||
GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git
|
||||
GIT_TAG v25.12.0
|
||||
)
|
||||
|
||||
FetchContent_MakeAvailable(urma)
|
||||
|
||||
# 输出实际路径,确认位置
|
||||
message(STATUS "URMA source dir: ${urma_SOURCE_DIR}")
|
||||
message(STATUS "URMA binary dir: ${urma_BINARY_DIR}")
|
||||
|
||||
# 假设 UMDK 头文件在其 include 目录下
|
||||
set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include)
|
||||
|
||||
# 添加到需要的目标
|
||||
message(STATUS "urma_INCLUDE_DIR: ${urma_INCLUDE_DIR}")
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
# SetupPyTorchEnv.cmake
|
||||
#
|
||||
# This file provides helper functions for building Mooncake Pytorch extensions
|
||||
# and is meant to be included by BuildEpExt.cmake and BuildPgExt.cmake.
|
||||
|
||||
# Ensure we have the correct Python interpreter (respects active virtualenvs)
|
||||
find_package(Python3 REQUIRED COMPONENTS Interpreter)
|
||||
|
||||
# Install PyTorch for a specific version with proper CUDA compatibility handling.
|
||||
#
|
||||
# Usage:
|
||||
# install_pytorch_wheel("<VERSION>" <CUDA_MAJOR> <CUDA_MINOR> "<MODULE_PREFIX>")
|
||||
#
|
||||
# Example:
|
||||
# install_pytorch_wheel("2.11.0" 12 8 "[EP]")
|
||||
function(install_pytorch_wheel _version _cuda_major _cuda_minor _module_prefix)
|
||||
message(STATUS "${_module_prefix} Installing PyTorch ${_version} via pip...")
|
||||
|
||||
set(_cu_tag "")
|
||||
|
||||
# Determine the specific CUDA tag for PyTorch wheels
|
||||
if(_cuda_major GREATER_EQUAL 13)
|
||||
# TODO: Fix when we need to support more CUDA 13 versions or when the CI env is fixed.
|
||||
set(_cu_tag "cu130")
|
||||
|
||||
elseif(_cuda_major EQUAL 12 AND _version VERSION_GREATER_EQUAL "2.11.0")
|
||||
# PyTorch 2.11.0+ defaults to CUDA 13.
|
||||
# We must explicitly point to CUDA 12 wheels for these newer versions.
|
||||
if(_cuda_minor GREATER_EQUAL 8)
|
||||
set(_cu_tag "cu128")
|
||||
elseif(_cuda_minor GREATER_EQUAL 6)
|
||||
set(_cu_tag "cu126")
|
||||
else()
|
||||
message(FATAL_ERROR
|
||||
"${_module_prefix} Can't find a matching PyTorch wheel for version ${_version} "
|
||||
"with CUDA ${_cuda_major}.${_cuda_minor}"
|
||||
)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# Construct pip command using the absolute path to the Python executable
|
||||
set(_pip_cmd ${Python3_EXECUTABLE} -m pip install "torch==${_version}")
|
||||
|
||||
if(_cu_tag)
|
||||
set(_index_url "https://download.pytorch.org/whl/${_cu_tag}")
|
||||
message(STATUS "${_module_prefix} Using specific CUDA wheel: ${_index_url}")
|
||||
list(APPEND _pip_cmd --index-url "${_index_url}")
|
||||
else()
|
||||
message(STATUS "${_module_prefix} Using default PyPI wheels for PyTorch ${_version}")
|
||||
endif()
|
||||
|
||||
# Execute pip install
|
||||
execute_process(
|
||||
COMMAND ${_pip_cmd}
|
||||
RESULT_VARIABLE _ret
|
||||
)
|
||||
|
||||
if(NOT _ret EQUAL 0)
|
||||
message(FATAL_ERROR "${_module_prefix} Failed to install PyTorch ${_version}."
|
||||
" Command run: '${_pip_cmd}'")
|
||||
endif()
|
||||
|
||||
message(STATUS "${_module_prefix} PyTorch ${_version} is ready.")
|
||||
endfunction()
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
# SetupPython.cmake — resolve the Python interpreter for execute_process() calls.
|
||||
#
|
||||
# Honour -DPython3_EXECUTABLE=... when provided (e.g. Docker builds that
|
||||
# install a non-system Python via deadsnakes), otherwise fall back to the
|
||||
# default "python3" on PATH. Sets PYTHON_EXECUTABLE for legacy callers.
|
||||
|
||||
if(NOT Python3_EXECUTABLE)
|
||||
set(Python3_EXECUTABLE "python3")
|
||||
endif()
|
||||
set(PYTHON_EXECUTABLE "${Python3_EXECUTABLE}")
|
||||
|
|
@ -40,6 +40,9 @@ add_definitions(-DCONFIG_ERDMA)
|
|||
|
||||
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
|
||||
|
||||
# Memory-aware build parallelism (compile vs. link job pools)
|
||||
include(${CMAKE_CURRENT_LIST_DIR}/limit_jobs.cmake)
|
||||
|
||||
option(ENABLE_SCCACHE "Whether to open sccache" OFF)
|
||||
if (ENABLE_SCCACHE)
|
||||
find_program(SCCACHE sccache REQUIRED)
|
||||
|
|
@ -57,16 +60,49 @@ option(BUILD_EXAMPLES "Build examples" ON)
|
|||
|
||||
option(BUILD_UNIT_TESTS "Build unit tests" ON)
|
||||
option(USE_CUDA "option for enabling gpu features for NVIDIA GPU" OFF)
|
||||
option(USE_MLU "option for enabling Cambricon MLU features" OFF)
|
||||
option(USE_MUSA "option for enabling gpu features for MTHREADS GPU" OFF)
|
||||
option(USE_MACA "option for enabling gpu features for MUXI GPU with MACA" OFF)
|
||||
option(USE_HIP "option for enabling gpu features for AMD GPU" OFF)
|
||||
option(USE_NVMEOF "option for using NVMe over Fabric" OFF)
|
||||
option(USE_TCP "option for using TCP transport" ON)
|
||||
option(USE_BAREX "option for using accl-barex transport" OFF)
|
||||
option(USE_ASCEND "option for using npu with HCCL" OFF)
|
||||
option(USE_ASCEND_DIRECT "option for using ascend npu with adxl engine" OFF)
|
||||
option(USE_UBSHMEM "option for using ascend npu with shmem" OFF)
|
||||
option(USE_ASCEND_HETEROGENEOUS "option for transferring between ascend npu and gpu" OFF)
|
||||
option(USE_MNNVL "option for using Multi-Node NVLink transport" OFF)
|
||||
option(USE_CXL "option for using CXL protocol" OFF)
|
||||
option(USE_EFA "option for using AWS EFA transport" OFF)
|
||||
option(USE_UB "option for using UB protocol transport" OFF)
|
||||
|
||||
if (USE_UB)
|
||||
add_compile_definitions(USE_UB)
|
||||
message(STATUS "ub transport is enabled")
|
||||
include(${CMAKE_CURRENT_LIST_DIR}/FindUrma.cmake)
|
||||
endif()
|
||||
|
||||
if (USE_EFA)
|
||||
# Find libfabric headers and library; default to AWS EFA installer path
|
||||
find_path(LIBFABRIC_INCLUDE_DIR rdma/fabric.h
|
||||
HINTS /opt/amazon/efa/include
|
||||
PATH_SUFFIXES include)
|
||||
find_library(LIBFABRIC_LIBRARY fabric
|
||||
HINTS /opt/amazon/efa/lib
|
||||
PATH_SUFFIXES lib lib64)
|
||||
|
||||
if (NOT LIBFABRIC_INCLUDE_DIR OR NOT LIBFABRIC_LIBRARY)
|
||||
message(FATAL_ERROR "libfabric not found. Install AWS EFA or set LIBFABRIC_INCLUDE_DIR/LIBFABRIC_LIBRARY.")
|
||||
endif()
|
||||
|
||||
get_filename_component(LIBFABRIC_LIB_DIR ${LIBFABRIC_LIBRARY} DIRECTORY)
|
||||
include_directories(${LIBFABRIC_INCLUDE_DIR})
|
||||
link_directories(${LIBFABRIC_LIB_DIR})
|
||||
add_compile_definitions(USE_EFA)
|
||||
message(STATUS "AWS EFA (libfabric) transport is enabled")
|
||||
message(STATUS " libfabric include: ${LIBFABRIC_INCLUDE_DIR}")
|
||||
message(STATUS " libfabric library: ${LIBFABRIC_LIBRARY}")
|
||||
endif()
|
||||
option(USE_ETCD "option for enable etcd as metadata server" OFF)
|
||||
option(USE_ETCD_LEGACY "option for enable etcd based on etcd-cpp-api-v3" OFF)
|
||||
option(USE_REDIS "option for enable redis as metadata server" OFF)
|
||||
|
|
@ -78,8 +114,13 @@ option(WITH_NVIDIA_PEERMEM "disable to support RDMA without nvidia-peermem. If W
|
|||
option(USE_EVENT_DRIVEN_COMPLETION "option for using event-driven completion (store & transfer engine)" OFF)
|
||||
|
||||
option(USE_TENT "option for building Mooncake TENT" OFF)
|
||||
|
||||
option(ENABLE_MULTI_PROTOCOL "option for enabling multi-protocol support in transfer engine" OFF)
|
||||
if (ENABLE_MULTI_PROTOCOL)
|
||||
add_compile_definitions(ENABLE_MULTI_PROTOCOL)
|
||||
message(STATUS "Multi-protocol support is enabled")
|
||||
endif()
|
||||
option(USE_LRU_MASTER "option for using LRU in master service" OFF)
|
||||
option(USE_INTRA_NVLINK "option for using IntraNode nvlink transport" OFF)
|
||||
set(LRU_MAX_CAPACITY 1000)
|
||||
|
||||
if (USE_LRU_MASTER)
|
||||
|
|
@ -101,7 +142,7 @@ if (USE_NVMEOF)
|
|||
endif()
|
||||
|
||||
if (USE_MNNVL)
|
||||
if (NOT USE_HIP AND NOT USE_MUSA)
|
||||
if (NOT USE_HIP AND NOT USE_MUSA AND NOT USE_MACA)
|
||||
set(USE_CUDA ON)
|
||||
endif()
|
||||
add_compile_definitions(USE_MNNVL)
|
||||
|
|
@ -118,6 +159,60 @@ if (USE_CUDA)
|
|||
)
|
||||
endif()
|
||||
|
||||
if (NOT DEFINED NEUWARE_ROOT OR NEUWARE_ROOT STREQUAL "")
|
||||
if (DEFINED ENV{NEUWARE_HOME} AND NOT "$ENV{NEUWARE_HOME}" STREQUAL "")
|
||||
set(NEUWARE_ROOT "$ENV{NEUWARE_HOME}" CACHE PATH "Path to Cambricon Neuware SDK" FORCE)
|
||||
else()
|
||||
set(NEUWARE_ROOT "/usr/local/neuware" CACHE PATH "Path to Cambricon Neuware SDK" FORCE)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (NOT DEFINED MLU_INCLUDE_DIR OR MLU_INCLUDE_DIR STREQUAL "")
|
||||
set(MLU_INCLUDE_DIR "${NEUWARE_ROOT}/include")
|
||||
endif()
|
||||
|
||||
if (NOT DEFINED MLU_LIB_DIR OR MLU_LIB_DIR STREQUAL "")
|
||||
set(MLU_LIB_DIR "${NEUWARE_ROOT}/lib64")
|
||||
endif()
|
||||
|
||||
if (NOT DEFINED MACA_ROOT OR MACA_ROOT STREQUAL "")
|
||||
if (DEFINED ENV{MACA_HOME} AND NOT "$ENV{MACA_HOME}" STREQUAL "")
|
||||
set(MACA_ROOT "$ENV{MACA_HOME}" CACHE PATH "Path to MACA SDK" FORCE)
|
||||
else()
|
||||
set(MACA_ROOT "/opt/maca" CACHE PATH "Path to MACA SDK" FORCE)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (NOT DEFINED MACA_INCLUDE_DIR OR MACA_INCLUDE_DIR STREQUAL "")
|
||||
set(MACA_INCLUDE_DIR "${MACA_ROOT}/include")
|
||||
endif()
|
||||
|
||||
if (NOT DEFINED MACA_LIB_DIR OR MACA_LIB_DIR STREQUAL "")
|
||||
if (EXISTS "${MACA_ROOT}/lib64")
|
||||
set(MACA_LIB_DIR "${MACA_ROOT}/lib64")
|
||||
else()
|
||||
set(MACA_LIB_DIR "${MACA_ROOT}/lib")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (USE_MLU)
|
||||
add_compile_definitions(USE_MLU)
|
||||
message(STATUS "MLU support is enabled")
|
||||
include_directories(${MLU_INCLUDE_DIR})
|
||||
if (EXISTS "${MLU_LIB_DIR}")
|
||||
link_directories(${MLU_LIB_DIR})
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (USE_MACA)
|
||||
add_compile_definitions(USE_MACA)
|
||||
message(STATUS "MACA support is enabled")
|
||||
include_directories(${MACA_INCLUDE_DIR})
|
||||
if (EXISTS "${MACA_LIB_DIR}")
|
||||
link_directories(${MACA_LIB_DIR})
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (USE_MUSA)
|
||||
add_compile_definitions(USE_MUSA)
|
||||
message(STATUS "MUSA support is enabled")
|
||||
|
|
@ -181,7 +276,7 @@ if (USE_BAREX)
|
|||
add_compile_definitions(USE_BAREX)
|
||||
endif()
|
||||
|
||||
if (USE_ASCEND OR USE_ASCEND_DIRECT)
|
||||
if (USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM)
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DOPEN_BUILD_PROJECT ")
|
||||
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -DOPEN_BUILD_PROJECT ")
|
||||
string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" CURRENT_CPU)
|
||||
|
|
@ -217,6 +312,11 @@ if (USE_ASCEND_DIRECT)
|
|||
add_compile_definitions(USE_ASCEND_DIRECT)
|
||||
endif()
|
||||
|
||||
if (USE_UBSHMEM)
|
||||
set(BUILD_SHARED_LIBS ON)
|
||||
add_compile_definitions(USE_UBSHMEM)
|
||||
endif()
|
||||
|
||||
if (USE_ASCEND_HETEROGENEOUS)
|
||||
file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
|
||||
set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64")
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ add_custom_command(
|
|||
COMMAND bash -c "go mod tidy" && bash -c "go build -buildmode=c-shared -o ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.so etcd_wrapper.go" && cp ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.h ${CMAKE_CURRENT_SOURCE_DIR}
|
||||
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
|
||||
COMMENT "Building Go shared library"
|
||||
DEPENDS etcd_wrapper.go
|
||||
DEPENDS etcd_wrapper.go go.mod go.sum build.sh
|
||||
)
|
||||
|
||||
set(ETCD_WRAPPER_INCLUDE ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.h)
|
||||
|
|
@ -17,4 +17,4 @@ add_custom_target(
|
|||
install(
|
||||
FILES ${ETCD_WRAPPER_LIB}
|
||||
DESTINATION lib
|
||||
)
|
||||
)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -1,27 +1,28 @@
|
|||
module github.com/kvcache-ai/Mooncake/mooncake-common/etcd
|
||||
|
||||
go 1.23.0
|
||||
go 1.25.0
|
||||
|
||||
toolchain go1.23.7
|
||||
toolchain go1.25.9
|
||||
|
||||
require go.etcd.io/etcd/client/v3 v3.5.21
|
||||
require (
|
||||
go.etcd.io/etcd/api/v3 v3.5.21
|
||||
go.etcd.io/etcd/client/v3 v3.5.21
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/coreos/go-semver v0.3.0 // indirect
|
||||
github.com/coreos/go-systemd/v22 v22.3.2 // indirect
|
||||
github.com/gogo/protobuf v1.3.2 // indirect
|
||||
github.com/golang/protobuf v1.5.4 // indirect
|
||||
go.etcd.io/etcd/api/v3 v3.5.21 // indirect
|
||||
go.etcd.io/etcd/client/pkg/v3 v3.5.21 // indirect
|
||||
go.uber.org/atomic v1.7.0 // indirect
|
||||
go.uber.org/multierr v1.6.0 // indirect
|
||||
go.uber.org/zap v1.17.0 // indirect
|
||||
golang.org/x/net v0.38.0 // indirect
|
||||
golang.org/x/sys v0.31.0 // indirect
|
||||
golang.org/x/text v0.23.0 // indirect
|
||||
google.golang.org/genproto v0.0.0-20230822172742-b8732ec3820d // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20230822172742-b8732ec3820d // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
|
||||
google.golang.org/grpc v1.59.0 // indirect
|
||||
google.golang.org/protobuf v1.33.0 // indirect
|
||||
golang.org/x/net v0.48.0 // indirect
|
||||
golang.org/x/sys v0.39.0 // indirect
|
||||
golang.org/x/text v0.32.0 // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect
|
||||
google.golang.org/grpc v1.79.3 // indirect
|
||||
google.golang.org/protobuf v1.36.10 // indirect
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,108 @@
|
|||
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/coreos/go-semver v0.3.0 h1:wkHLiw0WNATZnSG7epLsujiMCgPAc9xhjJ4tgnAxmfM=
|
||||
github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
|
||||
github.com/coreos/go-systemd/v22 v22.3.2 h1:D9/bQk5vlXQFZ6Kwuu6zaiXJ9oTPe68++AzAJc1DzSI=
|
||||
github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
|
||||
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
|
||||
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
|
||||
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
|
||||
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
|
||||
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
|
||||
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
|
||||
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
|
||||
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
|
||||
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
|
||||
github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
|
||||
github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
|
||||
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
|
||||
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
|
||||
go.etcd.io/etcd/api/v3 v3.5.21 h1:A6O2/JDb3tvHhiIz3xf9nJ7REHvtEFJJ3veW3FbCnS8=
|
||||
go.etcd.io/etcd/api/v3 v3.5.21/go.mod h1:c3aH5wcvXv/9dqIw2Y810LDXJfhSYdHQ0vxmP3CCHVY=
|
||||
go.etcd.io/etcd/client/pkg/v3 v3.5.21 h1:lPBu71Y7osQmzlflM9OfeIV2JlmpBjqBNlLtcoBqUTc=
|
||||
go.etcd.io/etcd/client/pkg/v3 v3.5.21/go.mod h1:BgqT/IXPjK9NkeSDjbzwsHySX3yIle2+ndz28nVsjUs=
|
||||
go.etcd.io/etcd/client/v3 v3.5.21 h1:T6b1Ow6fNjOLOtM0xSoKNQt1ASPCLWrF9XMHcH9pEyY=
|
||||
go.etcd.io/etcd/client/v3 v3.5.21/go.mod h1:mFYy67IOqmbRf/kRUvsHixzo3iG+1OF2W2+jVIQRAnU=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
|
||||
go.opentelemetry.io/otel v1.39.0 h1:8yPrr/S0ND9QEfTfdP9V+SiwT4E0G7Y5MO7p85nis48=
|
||||
go.opentelemetry.io/otel v1.39.0/go.mod h1:kLlFTywNWrFyEdH0oj2xK0bFYZtHRYUdv1NklR/tgc8=
|
||||
go.opentelemetry.io/otel/metric v1.39.0 h1:d1UzonvEZriVfpNKEVmHXbdf909uGTOQjA0HF0Ls5Q0=
|
||||
go.opentelemetry.io/otel/metric v1.39.0/go.mod h1:jrZSWL33sD7bBxg1xjrqyDjnuzTUB0x1nBERXd7Ftcs=
|
||||
go.opentelemetry.io/otel/sdk v1.39.0 h1:nMLYcjVsvdui1B/4FRkwjzoRVsMK8uL/cj0OyhKzt18=
|
||||
go.opentelemetry.io/otel/sdk v1.39.0/go.mod h1:vDojkC4/jsTJsE+kh+LXYQlbL8CgrEcwmt1ENZszdJE=
|
||||
go.opentelemetry.io/otel/sdk/metric v1.39.0 h1:cXMVVFVgsIf2YL6QkRF4Urbr/aMInf+2WKg+sEJTtB8=
|
||||
go.opentelemetry.io/otel/sdk/metric v1.39.0/go.mod h1:xq9HEVH7qeX69/JnwEfp6fVq5wosJsY1mt4lLfYdVew=
|
||||
go.opentelemetry.io/otel/trace v1.39.0 h1:2d2vfpEDmCJ5zVYz7ijaJdOF59xLomrvj7bjt6/qCJI=
|
||||
go.opentelemetry.io/otel/trace v1.39.0/go.mod h1:88w4/PnZSazkGzz/w84VHpQafiU4EtqqlVdxWy+rNOA=
|
||||
go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
|
||||
go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
|
||||
go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
|
||||
go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
|
||||
go.uber.org/zap v1.17.0 h1:MTjgFu6ZLKvY6Pvaqk97GlxNBuMpV4Hy/3P6tRGlI2U=
|
||||
go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo=
|
||||
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
|
||||
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
|
||||
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
|
||||
golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
|
||||
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
|
||||
golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU=
|
||||
golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY=
|
||||
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
|
||||
golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
|
||||
golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
|
||||
golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
|
||||
golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
|
||||
gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 h1:fCvbg86sFXwdrl5LgVcTEvNC+2txB5mgROGmRL5mrls=
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:+rXWjjaukWZun3mLfjmVnQi18E1AsFbDN9QdJ5YXLto=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 h1:gRkg/vSppuSQoDjxyiGfN4Upv/h/DQmIR10ZU8dh4Ww=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
|
||||
google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE=
|
||||
google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
|
||||
google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
|
||||
google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
|
||||
gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
|
||||
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
|
|
@ -70,6 +70,18 @@ class DefaultConfig {
|
|||
void GetUInt64(const std::string& key, uint64_t* val,
|
||||
uint64_t default_value = 0) const;
|
||||
|
||||
/**
|
||||
* @brief GetDurationMs retrieves a duration value from the configuration
|
||||
* and converts it to milliseconds.
|
||||
* @param key The key to look up in the configuration
|
||||
* @param val Pointer to store the retrieved value in milliseconds
|
||||
* @param default_value Default value to return if the key is not found
|
||||
* @note Duration strings may use ms, s, m, or h as suffixes. Bare numbers
|
||||
* are interpreted as milliseconds.
|
||||
*/
|
||||
void GetDurationMs(const std::string& key, uint64_t* val,
|
||||
uint64_t default_value = 0) const;
|
||||
|
||||
/**
|
||||
* @brief GetDouble retrieves a double value from the configuration
|
||||
* @param key The key to look up in the configuration
|
||||
|
|
|
|||
|
|
@ -0,0 +1,95 @@
|
|||
#pragma once
|
||||
|
||||
#include <cctype>
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
inline std::string_view TrimAsciiWhitespace(std::string_view value) {
|
||||
while (!value.empty() &&
|
||||
std::isspace(static_cast<unsigned char>(value.front()))) {
|
||||
value.remove_prefix(1);
|
||||
}
|
||||
while (!value.empty() &&
|
||||
std::isspace(static_cast<unsigned char>(value.back()))) {
|
||||
value.remove_suffix(1);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
inline bool ParseDurationMs(std::string_view value, uint64_t* result,
|
||||
std::string* error = nullptr) {
|
||||
auto set_error = [&](std::string message) {
|
||||
if (error != nullptr) {
|
||||
*error = std::move(message);
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
if (result == nullptr) {
|
||||
return set_error("duration output pointer is null");
|
||||
}
|
||||
|
||||
std::string_view trimmed = TrimAsciiWhitespace(value);
|
||||
if (trimmed.empty()) {
|
||||
return set_error(
|
||||
"duration is empty; expected a non-negative integer optionally "
|
||||
"followed by ms, s, m, or h");
|
||||
}
|
||||
|
||||
size_t number_end = 0;
|
||||
while (number_end < trimmed.size() &&
|
||||
std::isdigit(static_cast<unsigned char>(trimmed[number_end]))) {
|
||||
++number_end;
|
||||
}
|
||||
|
||||
if (number_end == 0) {
|
||||
return set_error(
|
||||
"duration must start with a non-negative integer and may use ms, "
|
||||
"s, m, or h as the unit suffix");
|
||||
}
|
||||
|
||||
uint64_t numeric_value = 0;
|
||||
for (size_t i = 0; i < number_end; ++i) {
|
||||
const uint64_t digit = static_cast<uint64_t>(trimmed[i] - '0');
|
||||
if (numeric_value >
|
||||
(std::numeric_limits<uint64_t>::max() - digit) / 10) {
|
||||
return set_error("duration value is too large");
|
||||
}
|
||||
numeric_value = numeric_value * 10 + digit;
|
||||
}
|
||||
|
||||
std::string_view suffix = TrimAsciiWhitespace(trimmed.substr(number_end));
|
||||
std::string normalized_suffix;
|
||||
normalized_suffix.reserve(suffix.size());
|
||||
for (char ch : suffix) {
|
||||
normalized_suffix.push_back(
|
||||
static_cast<char>(std::tolower(static_cast<unsigned char>(ch))));
|
||||
}
|
||||
|
||||
uint64_t multiplier = 1;
|
||||
if (normalized_suffix.empty() || normalized_suffix == "ms") {
|
||||
multiplier = 1;
|
||||
} else if (normalized_suffix == "s") {
|
||||
multiplier = 1000;
|
||||
} else if (normalized_suffix == "m") {
|
||||
multiplier = 60 * 1000;
|
||||
} else if (normalized_suffix == "h") {
|
||||
multiplier = 60 * 60 * 1000;
|
||||
} else {
|
||||
return set_error("unsupported duration unit '" + normalized_suffix +
|
||||
"'; supported units are ms, s, m, and h");
|
||||
}
|
||||
|
||||
if (numeric_value > std::numeric_limits<uint64_t>::max() / multiplier) {
|
||||
return set_error("duration value is too large after unit conversion");
|
||||
}
|
||||
|
||||
*result = numeric_value * multiplier;
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,104 @@
|
|||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
class Environ {
|
||||
public:
|
||||
// Singleton access
|
||||
static Environ& Get();
|
||||
|
||||
// Getters for Environment Variables
|
||||
int GetNumCqPerCtx() const { return num_cq_per_ctx_; }
|
||||
int GetNumCompChannelsPerCtx() const { return num_comp_channels_per_ctx_; }
|
||||
int GetIbPort() const { return ib_port_; }
|
||||
int GetIbTc() const { return ib_tc_; }
|
||||
int GetIbPciRelaxedOrdering() const { return ib_pci_relaxed_ordering_; }
|
||||
int GetGidIndex() const { return gid_index_; }
|
||||
int GetMaxCqePerCtx() const { return max_cqe_per_ctx_; }
|
||||
int GetMaxEpPerCtx() const { return max_ep_per_ctx_; }
|
||||
int GetNumQpPerEp() const { return num_qp_per_ep_; }
|
||||
int GetMaxSge() const { return max_sge_; }
|
||||
int GetMaxWr() const { return max_wr_; }
|
||||
int GetMaxInline() const { return max_inline_; }
|
||||
int GetMtu() const { return mtu_; }
|
||||
int GetWorkersPerCtx() const { return workers_per_ctx_; }
|
||||
size_t GetSliceSize() const { return slice_size_; }
|
||||
int GetRetryCnt() const { return retry_cnt_; }
|
||||
std::string GetLogLevel() const { return log_level_; }
|
||||
bool GetDisableMetacache() const { return disable_metacache_; }
|
||||
int GetHandshakeListenBacklog() const { return handshake_listen_backlog_; }
|
||||
int GetHandshakeMaxLength() const { return handshake_max_length_; }
|
||||
std::string GetLogDir() const { return log_dir_; }
|
||||
std::string GetRedisPassword() const { return redis_password_; }
|
||||
int GetRedisDbIndex() const { return redis_db_index_; }
|
||||
int GetFragmentRatio() const { return fragment_ratio_; }
|
||||
bool GetEnableDestDeviceAffinity() const {
|
||||
return enable_dest_device_affinity_;
|
||||
}
|
||||
bool GetUseIpv6() const { return use_ipv6_; }
|
||||
int GetMinPrcPort() const { return min_prc_port_; }
|
||||
int GetMaxPrcPort() const { return max_prc_port_; }
|
||||
int GetEnableParallelRegMr() const { return enable_parallel_reg_mr_; }
|
||||
std::string GetEndpointStoreType() const { return endpoint_store_type_; }
|
||||
bool GetForceTcp() const { return force_tcp_; }
|
||||
bool GetForceHca() const { return force_hca_; }
|
||||
bool GetForceMnnvl() const { return force_mnnvl_; }
|
||||
bool GetIntraNvlink() const { return intra_nvlink_; }
|
||||
bool GetPathRoundrobin() const { return path_roundrobin_; }
|
||||
|
||||
private:
|
||||
Environ();
|
||||
|
||||
// Helper method to get int from env
|
||||
static int GetInt(const char* name, int default_value);
|
||||
// Helper method to get size_t from env
|
||||
static size_t GetSizeT(const char* name, size_t default_value);
|
||||
// Helper method to get bool from env (checks for "1", "true", "TRUE")
|
||||
static bool GetBool(const char* name, bool default_value);
|
||||
// Helper method to get string from env
|
||||
static std::string GetString(const char* name,
|
||||
const std::string& default_value);
|
||||
|
||||
// Member variables
|
||||
int num_cq_per_ctx_;
|
||||
int num_comp_channels_per_ctx_;
|
||||
int ib_port_;
|
||||
int ib_tc_;
|
||||
int ib_pci_relaxed_ordering_;
|
||||
int gid_index_;
|
||||
int max_cqe_per_ctx_;
|
||||
int max_ep_per_ctx_;
|
||||
int num_qp_per_ep_;
|
||||
int max_sge_;
|
||||
int max_wr_;
|
||||
int max_inline_;
|
||||
int mtu_;
|
||||
int workers_per_ctx_;
|
||||
size_t slice_size_;
|
||||
int retry_cnt_;
|
||||
std::string log_level_;
|
||||
bool disable_metacache_;
|
||||
int handshake_listen_backlog_;
|
||||
int handshake_max_length_;
|
||||
std::string log_dir_;
|
||||
std::string redis_password_;
|
||||
int redis_db_index_;
|
||||
int fragment_ratio_;
|
||||
bool enable_dest_device_affinity_;
|
||||
bool use_ipv6_;
|
||||
int min_prc_port_;
|
||||
int max_prc_port_;
|
||||
int enable_parallel_reg_mr_;
|
||||
std::string endpoint_store_type_;
|
||||
bool force_tcp_;
|
||||
bool force_hca_;
|
||||
bool force_mnnvl_;
|
||||
bool intra_nvlink_;
|
||||
bool path_roundrobin_;
|
||||
};
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
add_custom_command(
|
||||
OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so
|
||||
COMMAND bash -c "go mod tidy" && bash -c "go build -buildmode=c-shared -o ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so k8s_lease_wrapper.go" && cp ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h ${CMAKE_CURRENT_SOURCE_DIR}
|
||||
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
|
||||
COMMENT "Building K8s Lease Go shared library"
|
||||
DEPENDS k8s_lease_wrapper.go
|
||||
)
|
||||
|
||||
set(K8S_LEASE_WRAPPER_INCLUDE ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h)
|
||||
set(K8S_LEASE_WRAPPER_LIB ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so)
|
||||
|
||||
add_custom_target(
|
||||
build_k8s_lease_wrapper
|
||||
DEPENDS ${K8S_LEASE_WRAPPER_LIB}
|
||||
)
|
||||
|
||||
install(
|
||||
FILES ${K8S_LEASE_WRAPPER_LIB}
|
||||
DESTINATION lib
|
||||
)
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
// envtest-server starts a real kube-apiserver + etcd via envtest, writes the
|
||||
// KUBECONFIG path to stdout, and blocks until SIGTERM or SIGINT. This lets
|
||||
// C++ tests launch it as a subprocess and talk to a real K8s API without a
|
||||
// full cluster.
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"syscall"
|
||||
|
||||
"k8s.io/client-go/tools/clientcmd"
|
||||
clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
|
||||
"sigs.k8s.io/controller-runtime/pkg/envtest"
|
||||
)
|
||||
|
||||
func main() {
|
||||
env := &envtest.Environment{}
|
||||
|
||||
cfg, err := env.Start()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "envtest start failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Write a KUBECONFIG file that points at the envtest kube-apiserver.
|
||||
kubeconfigPath := filepath.Join(os.TempDir(), fmt.Sprintf("envtest-kubeconfig-%d", os.Getpid()))
|
||||
kubeconfig := clientcmdapi.NewConfig()
|
||||
kubeconfig.Clusters["envtest"] = &clientcmdapi.Cluster{
|
||||
Server: cfg.Host,
|
||||
CertificateAuthorityData: cfg.CAData,
|
||||
}
|
||||
kubeconfig.AuthInfos["envtest"] = &clientcmdapi.AuthInfo{
|
||||
ClientCertificateData: cfg.CertData,
|
||||
ClientKeyData: cfg.KeyData,
|
||||
}
|
||||
kubeconfig.Contexts["envtest"] = &clientcmdapi.Context{
|
||||
Cluster: "envtest",
|
||||
AuthInfo: "envtest",
|
||||
}
|
||||
kubeconfig.CurrentContext = "envtest"
|
||||
|
||||
if err := clientcmd.WriteToFile(*kubeconfig, kubeconfigPath); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to write kubeconfig: %v\n", err)
|
||||
env.Stop()
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Print the kubeconfig path — the parent process reads this from stdout.
|
||||
fmt.Println(kubeconfigPath)
|
||||
|
||||
// Block until SIGTERM or SIGINT.
|
||||
sigCh := make(chan os.Signal, 1)
|
||||
signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
|
||||
<-sigCh
|
||||
|
||||
os.Remove(kubeconfigPath)
|
||||
env.Stop()
|
||||
}
|
||||
|
|
@ -0,0 +1,60 @@
|
|||
module github.com/kvcache-ai/Mooncake/mooncake-common/k8s-lease
|
||||
|
||||
go 1.24.0
|
||||
|
||||
require (
|
||||
k8s.io/api v0.34.3
|
||||
k8s.io/apimachinery v0.34.3
|
||||
k8s.io/client-go v0.34.3
|
||||
k8s.io/utils v0.0.0-20251002143259-bc988d571ff4
|
||||
sigs.k8s.io/controller-runtime v0.22.5
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/beorn7/perks v1.0.1 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
|
||||
github.com/emicklei/go-restful/v3 v3.12.2 // indirect
|
||||
github.com/evanphx/json-patch/v5 v5.9.11 // indirect
|
||||
github.com/fxamacker/cbor/v2 v2.9.0 // indirect
|
||||
github.com/go-logr/logr v1.4.3 // indirect
|
||||
github.com/go-openapi/jsonpointer v0.21.0 // indirect
|
||||
github.com/go-openapi/jsonreference v0.20.2 // indirect
|
||||
github.com/go-openapi/swag v0.23.0 // indirect
|
||||
github.com/gogo/protobuf v1.3.2 // indirect
|
||||
github.com/google/gnostic-models v0.7.0 // indirect
|
||||
github.com/google/go-cmp v0.7.0 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/josharian/intern v1.0.0 // indirect
|
||||
github.com/json-iterator/go v1.1.12 // indirect
|
||||
github.com/mailru/easyjson v0.7.7 // indirect
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
|
||||
github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/prometheus/client_golang v1.23.2 // indirect
|
||||
github.com/prometheus/client_model v0.6.2 // indirect
|
||||
github.com/prometheus/common v0.66.1 // indirect
|
||||
github.com/prometheus/procfs v0.16.1 // indirect
|
||||
github.com/spf13/pflag v1.0.9 // indirect
|
||||
github.com/x448/float16 v0.8.4 // indirect
|
||||
go.yaml.in/yaml/v2 v2.4.3 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.4 // indirect
|
||||
golang.org/x/net v0.47.0 // indirect
|
||||
golang.org/x/oauth2 v0.30.0 // indirect
|
||||
golang.org/x/sys v0.38.0 // indirect
|
||||
golang.org/x/term v0.37.0 // indirect
|
||||
golang.org/x/text v0.31.0 // indirect
|
||||
golang.org/x/time v0.9.0 // indirect
|
||||
google.golang.org/protobuf v1.36.8 // indirect
|
||||
gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect
|
||||
gopkg.in/inf.v0 v0.9.1 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
k8s.io/apiextensions-apiserver v0.34.3 // indirect
|
||||
k8s.io/klog/v2 v2.130.1 // indirect
|
||||
k8s.io/kube-openapi v0.0.0-20250910181357-589584f1c912 // indirect
|
||||
sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect
|
||||
sigs.k8s.io/randfill v1.0.0 // indirect
|
||||
sigs.k8s.io/structured-merge-diff/v6 v6.3.2-0.20260122202528-d9cc6641c482 // indirect
|
||||
sigs.k8s.io/yaml v1.6.0 // indirect
|
||||
)
|
||||
|
|
@ -0,0 +1,571 @@
|
|||
//go:build integration
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
"k8s.io/client-go/kubernetes"
|
||||
"k8s.io/client-go/rest"
|
||||
"k8s.io/client-go/tools/leaderelection"
|
||||
"k8s.io/client-go/tools/leaderelection/resourcelock"
|
||||
|
||||
"sigs.k8s.io/controller-runtime/pkg/envtest"
|
||||
)
|
||||
|
||||
var (
|
||||
testEnv *envtest.Environment
|
||||
testConfig *rest.Config
|
||||
)
|
||||
|
||||
type electionStateNoRelease struct {
|
||||
cancel context.CancelFunc
|
||||
elected chan struct{}
|
||||
lost chan struct{}
|
||||
}
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
testEnv = &envtest.Environment{}
|
||||
|
||||
var err error
|
||||
testConfig, err = testEnv.Start()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to start envtest: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Set up global client for the wrapper
|
||||
client, err := kubernetes.NewForConfig(testConfig)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to create clientset: %v\n", err)
|
||||
testEnv.Stop()
|
||||
os.Exit(1)
|
||||
}
|
||||
clientMutex.Lock()
|
||||
globalClient = client
|
||||
clientMutex.Unlock()
|
||||
|
||||
code := m.Run()
|
||||
|
||||
testEnv.Stop()
|
||||
os.Exit(code)
|
||||
}
|
||||
|
||||
func runElectionWithoutRelease(namespace, leaseName, identity string,
|
||||
leaseDurationSec, renewDeadlineSec, retryPeriodSec int) (*electionStateNoRelease, error) {
|
||||
if err := ensureClientInitialized(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
state := &electionStateNoRelease{
|
||||
cancel: cancel,
|
||||
elected: make(chan struct{}),
|
||||
lost: make(chan struct{}),
|
||||
}
|
||||
|
||||
lock := &resourcelock.LeaseLock{
|
||||
LeaseMeta: metav1.ObjectMeta{
|
||||
Name: leaseName,
|
||||
Namespace: namespace,
|
||||
},
|
||||
Client: globalClient.CoordinationV1(),
|
||||
LockConfig: resourcelock.ResourceLockConfig{
|
||||
Identity: identity,
|
||||
},
|
||||
}
|
||||
|
||||
le, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
|
||||
Lock: lock,
|
||||
LeaseDuration: time.Duration(leaseDurationSec) * time.Second,
|
||||
RenewDeadline: time.Duration(renewDeadlineSec) * time.Second,
|
||||
RetryPeriod: time.Duration(retryPeriodSec) * time.Second,
|
||||
ReleaseOnCancel: false,
|
||||
Callbacks: leaderelection.LeaderCallbacks{
|
||||
OnStartedLeading: func(ctx context.Context) {
|
||||
close(state.elected)
|
||||
<-ctx.Done()
|
||||
},
|
||||
OnStoppedLeading: func() {
|
||||
close(state.lost)
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, fmt.Errorf("failed to create leader elector: %w", err)
|
||||
}
|
||||
|
||||
go le.Run(ctx)
|
||||
return state, nil
|
||||
}
|
||||
|
||||
// TestSingleLeaderElection verifies a single candidate becomes leader.
|
||||
func TestSingleLeaderElection(t *testing.T) {
|
||||
ns := "default"
|
||||
lease := "single-election-test"
|
||||
identity := "node-1:8080"
|
||||
|
||||
err := runElection(ns, lease, identity, 5, 4, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("runElection failed: %v", err)
|
||||
}
|
||||
|
||||
// Wait for elected
|
||||
key := electionKey(ns, lease)
|
||||
electionMutex.Lock()
|
||||
state := elections[key]
|
||||
electionMutex.Unlock()
|
||||
|
||||
select {
|
||||
case <-state.elected:
|
||||
// success
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out waiting for election")
|
||||
}
|
||||
|
||||
// Verify holder via getHolder
|
||||
holder, transitions, err := getHolder(ns, lease)
|
||||
if err != nil {
|
||||
t.Fatalf("getHolder failed: %v", err)
|
||||
}
|
||||
if holder != identity {
|
||||
t.Errorf("expected holder %q, got %q", identity, holder)
|
||||
}
|
||||
// First election — transitions should be 0 or 1
|
||||
if transitions < 0 {
|
||||
t.Errorf("expected non-negative transitions, got %d", transitions)
|
||||
}
|
||||
|
||||
// Cancel the election
|
||||
electionMutex.Lock()
|
||||
state = elections[key]
|
||||
electionMutex.Unlock()
|
||||
state.cancel()
|
||||
|
||||
select {
|
||||
case <-state.lost:
|
||||
// success
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timed out waiting for election loss after cancel")
|
||||
}
|
||||
}
|
||||
|
||||
// TestLeaderEpoch verifies leaseTransitions increments across elections.
|
||||
func TestLeaderEpoch(t *testing.T) {
|
||||
ns := "default"
|
||||
lease := "epoch-test"
|
||||
|
||||
// First election
|
||||
err := runElection(ns, lease, "node-epoch-1:8080", 5, 4, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("first runElection failed: %v", err)
|
||||
}
|
||||
|
||||
key := electionKey(ns, lease)
|
||||
electionMutex.Lock()
|
||||
state1 := elections[key]
|
||||
electionMutex.Unlock()
|
||||
|
||||
select {
|
||||
case <-state1.elected:
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out on first election")
|
||||
}
|
||||
|
||||
_, trans1, _ := getHolder(ns, lease)
|
||||
|
||||
// Cancel first election and wait for loss
|
||||
state1.cancel()
|
||||
select {
|
||||
case <-state1.lost:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timed out waiting for first election loss")
|
||||
}
|
||||
|
||||
// Wait for lease to expire / be released
|
||||
time.Sleep(2 * time.Second)
|
||||
|
||||
// Second election
|
||||
err = runElection(ns, lease, "node-epoch-2:8080", 5, 4, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("second runElection failed: %v", err)
|
||||
}
|
||||
|
||||
electionMutex.Lock()
|
||||
state2 := elections[key]
|
||||
electionMutex.Unlock()
|
||||
|
||||
select {
|
||||
case <-state2.elected:
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out on second election")
|
||||
}
|
||||
|
||||
_, trans2, _ := getHolder(ns, lease)
|
||||
if trans2 <= trans1 {
|
||||
t.Errorf("expected transitions to increment: first=%d, second=%d", trans1, trans2)
|
||||
}
|
||||
|
||||
state2.cancel()
|
||||
select {
|
||||
case <-state2.lost:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timed out waiting for second election loss")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSequentialLeadershipHandoff tests that a second candidate can acquire
|
||||
// leadership after the first one releases it.
|
||||
func TestSequentialLeadershipHandoff(t *testing.T) {
|
||||
ns := "default"
|
||||
lease := "two-candidate-test"
|
||||
|
||||
err1 := runElection(ns, lease, "candidate-a:8080", 5, 4, 1)
|
||||
if err1 != nil {
|
||||
t.Fatalf("first runElection failed: %v", err1)
|
||||
}
|
||||
|
||||
key := electionKey(ns, lease)
|
||||
electionMutex.Lock()
|
||||
stateA := elections[key]
|
||||
electionMutex.Unlock()
|
||||
|
||||
// Wait for first candidate to win
|
||||
select {
|
||||
case <-stateA.elected:
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out waiting for first candidate")
|
||||
}
|
||||
|
||||
// Verify holder is candidate-a
|
||||
holder, _, err := getHolder(ns, lease)
|
||||
if err != nil {
|
||||
t.Fatalf("getHolder failed: %v", err)
|
||||
}
|
||||
if holder != "candidate-a:8080" {
|
||||
t.Errorf("expected candidate-a, got %q", holder)
|
||||
}
|
||||
|
||||
// Cancel candidate-a
|
||||
stateA.cancel()
|
||||
select {
|
||||
case <-stateA.lost:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timed out waiting for candidate-a loss")
|
||||
}
|
||||
|
||||
// Wait for lease to expire
|
||||
time.Sleep(2 * time.Second)
|
||||
|
||||
// Start candidate-b
|
||||
err2 := runElection(ns, lease, "candidate-b:8080", 5, 4, 1)
|
||||
if err2 != nil {
|
||||
t.Fatalf("second runElection failed: %v", err2)
|
||||
}
|
||||
|
||||
electionMutex.Lock()
|
||||
stateB := elections[key]
|
||||
electionMutex.Unlock()
|
||||
|
||||
select {
|
||||
case <-stateB.elected:
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out waiting for candidate-b")
|
||||
}
|
||||
|
||||
holder, _, err = getHolder(ns, lease)
|
||||
if err != nil {
|
||||
t.Fatalf("getHolder after takeover failed: %v", err)
|
||||
}
|
||||
if holder != "candidate-b:8080" {
|
||||
t.Errorf("expected candidate-b, got %q", holder)
|
||||
}
|
||||
|
||||
stateB.cancel()
|
||||
select {
|
||||
case <-stateB.lost:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timed out waiting for candidate-b loss")
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentCandidateElection starts two candidates simultaneously and
|
||||
// verifies that exactly one wins leadership.
|
||||
func TestConcurrentCandidateElection(t *testing.T) {
|
||||
ns := "default"
|
||||
lease := "concurrent-election-test"
|
||||
|
||||
type result struct {
|
||||
identity string
|
||||
elected bool
|
||||
}
|
||||
|
||||
candidates := []string{"candidate-a:8080", "candidate-b:8080"}
|
||||
results := make(chan result, len(candidates))
|
||||
|
||||
lock := func(identity string) *resourcelock.LeaseLock {
|
||||
return &resourcelock.LeaseLock{
|
||||
LeaseMeta: metav1.ObjectMeta{
|
||||
Name: lease,
|
||||
Namespace: ns,
|
||||
},
|
||||
Client: globalClient.CoordinationV1(),
|
||||
LockConfig: resourcelock.ResourceLockConfig{
|
||||
Identity: identity,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for _, id := range candidates {
|
||||
wg.Add(1)
|
||||
go func(identity string) {
|
||||
defer wg.Done()
|
||||
|
||||
// Short timeout: enough for one to acquire, but the loser
|
||||
// times out before the winner's lease could expire.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
|
||||
defer cancel()
|
||||
|
||||
elected := make(chan struct{})
|
||||
le, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
|
||||
Lock: lock(identity),
|
||||
LeaseDuration: 5 * time.Second,
|
||||
RenewDeadline: 3 * time.Second,
|
||||
RetryPeriod: 1 * time.Second,
|
||||
ReleaseOnCancel: true,
|
||||
Callbacks: leaderelection.LeaderCallbacks{
|
||||
OnStartedLeading: func(ctx context.Context) {
|
||||
close(elected)
|
||||
<-ctx.Done()
|
||||
},
|
||||
OnStoppedLeading: func() {},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Errorf("NewLeaderElector(%s): %v", identity, err)
|
||||
return
|
||||
}
|
||||
|
||||
go le.Run(ctx)
|
||||
|
||||
select {
|
||||
case <-elected:
|
||||
results <- result{identity, true}
|
||||
// Keep holding until context expires (8s total).
|
||||
// Winner does NOT release early, so loser cannot
|
||||
// re-acquire within its own 8s window.
|
||||
<-ctx.Done()
|
||||
case <-ctx.Done():
|
||||
results <- result{identity, false}
|
||||
}
|
||||
}(id)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
close(results)
|
||||
|
||||
winners := 0
|
||||
for r := range results {
|
||||
if r.elected {
|
||||
winners++
|
||||
t.Logf("winner: %s", r.identity)
|
||||
}
|
||||
}
|
||||
|
||||
if winners != 1 {
|
||||
t.Fatalf("expected exactly 1 winner, got %d", winners)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCancelElection tests that cancelling an election makes WaitLost return.
|
||||
func TestCancelElection(t *testing.T) {
|
||||
ns := "default"
|
||||
lease := "cancel-test"
|
||||
|
||||
err := runElection(ns, lease, "cancel-node:8080", 5, 4, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("runElection failed: %v", err)
|
||||
}
|
||||
|
||||
key := electionKey(ns, lease)
|
||||
electionMutex.Lock()
|
||||
state := elections[key]
|
||||
electionMutex.Unlock()
|
||||
|
||||
// Wait for elected
|
||||
select {
|
||||
case <-state.elected:
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out waiting for election")
|
||||
}
|
||||
|
||||
// Cancel
|
||||
state.cancel()
|
||||
|
||||
// WaitLost should return promptly
|
||||
select {
|
||||
case <-state.lost:
|
||||
// success
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("WaitLost did not return after cancel")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetHolderDuringElection verifies getHolder works while election is active.
|
||||
func TestGetHolderDuringElection(t *testing.T) {
|
||||
ns := "default"
|
||||
lease := "active-get-holder-test"
|
||||
identity := "active-node:8080"
|
||||
|
||||
err := runElection(ns, lease, identity, 5, 4, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("runElection failed: %v", err)
|
||||
}
|
||||
|
||||
key := electionKey(ns, lease)
|
||||
electionMutex.Lock()
|
||||
state := elections[key]
|
||||
electionMutex.Unlock()
|
||||
|
||||
select {
|
||||
case <-state.elected:
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out waiting for election")
|
||||
}
|
||||
|
||||
// Concurrent getHolder calls during active election
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 5; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
holder, _, err := getHolder(ns, lease)
|
||||
if err != nil {
|
||||
t.Errorf("getHolder during election failed: %v", err)
|
||||
return
|
||||
}
|
||||
if holder != identity {
|
||||
t.Errorf("expected %q, got %q", identity, holder)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
state.cancel()
|
||||
<-state.lost
|
||||
}
|
||||
|
||||
// TestGetHolderReturnsEmptyAfterLeaderDeath verifies that after a leader stops
|
||||
// renewing its lease without releasing it, getHolder returns an empty holder
|
||||
// once the lease expires. This is the integration-level counterpart to the
|
||||
// unit test TestGetHolderReturnsEmptyForExpiredLease.
|
||||
func TestGetHolderReturnsEmptyAfterLeaderDeath(t *testing.T) {
|
||||
ns := "default"
|
||||
lease := "expired-leader-test"
|
||||
identity := "doomed-leader:8080"
|
||||
|
||||
// Acquire leadership without ReleaseOnCancel so canceling simulates a dead
|
||||
// leader that stops renewing and leaves the old holder until expiry.
|
||||
state, err := runElectionWithoutRelease(ns, lease, identity, 5, 4, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("runElection failed: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-state.elected:
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out waiting for election")
|
||||
}
|
||||
|
||||
// Verify holder while active.
|
||||
holder, _, err := getHolder(ns, lease)
|
||||
if err != nil {
|
||||
t.Fatalf("getHolder (active) failed: %v", err)
|
||||
}
|
||||
if holder != identity {
|
||||
t.Fatalf("expected active holder %q, got %q", identity, holder)
|
||||
}
|
||||
|
||||
// Simulate leader death: stop renewing without explicitly releasing.
|
||||
state.cancel()
|
||||
select {
|
||||
case <-state.lost:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timed out waiting for loss")
|
||||
}
|
||||
|
||||
// Wait for the lease to expire (leaseDuration=5s, add margin).
|
||||
time.Sleep(7 * time.Second)
|
||||
|
||||
// After expiry, getHolder must return empty holder so that the
|
||||
// supervisor will attempt acquisition.
|
||||
holder, _, err = getHolder(ns, lease)
|
||||
if err != nil {
|
||||
t.Fatalf("getHolder (expired) failed: %v", err)
|
||||
}
|
||||
if holder != "" {
|
||||
t.Errorf("expected empty holder after lease expiry, got %q", holder)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFailoverAfterLeaderDeath verifies that a new candidate can acquire
|
||||
// leadership after the previous leader dies and its lease expires.
|
||||
func TestFailoverAfterLeaderDeath(t *testing.T) {
|
||||
ns := "default"
|
||||
lease := "failover-test"
|
||||
|
||||
// First leader acquires without ReleaseOnCancel so canceling leaves the
|
||||
// old holder in place until the lease naturally expires.
|
||||
state1, err := runElectionWithoutRelease(ns, lease, "leader-1:8080", 5, 4, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("first runElection failed: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-state1.elected:
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("timed out waiting for first election")
|
||||
}
|
||||
|
||||
// Simulate crash: cancel without release, wait for expiry.
|
||||
state1.cancel()
|
||||
<-state1.lost
|
||||
time.Sleep(7 * time.Second)
|
||||
|
||||
// Second candidate should be able to acquire.
|
||||
err = runElection(ns, lease, "leader-2:8080", 5, 4, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("second runElection failed: %v", err)
|
||||
}
|
||||
|
||||
key := electionKey(ns, lease)
|
||||
electionMutex.Lock()
|
||||
state2 := elections[key]
|
||||
electionMutex.Unlock()
|
||||
|
||||
select {
|
||||
case <-state2.elected:
|
||||
// success — failover worked
|
||||
case <-time.After(15 * time.Second):
|
||||
t.Fatal("second candidate failed to acquire after leader death")
|
||||
}
|
||||
|
||||
holder, _, err := getHolder(ns, lease)
|
||||
if err != nil {
|
||||
t.Fatalf("getHolder after failover failed: %v", err)
|
||||
}
|
||||
if holder != "leader-2:8080" {
|
||||
t.Errorf("expected new leader %q, got %q", "leader-2:8080", holder)
|
||||
}
|
||||
|
||||
state2.cancel()
|
||||
<-state2.lost
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue