forked from mooncake-track/Mooncake
992 lines
40 KiB
Python
992 lines
40 KiB
Python
import ctypes
|
|
import os
|
|
import sys
|
|
import json
|
|
import time
|
|
import argparse
|
|
import unittest
|
|
import torch
|
|
import numpy as np
|
|
from dataclasses import dataclass
|
|
from mooncake.store import MooncakeDistributedStore
|
|
|
|
from mooncake.mooncake_config import MooncakeConfig
|
|
|
|
import concurrent.futures
|
|
|
|
# ==========================================
|
|
# Global Variables & Configuration
|
|
# ==========================================
|
|
|
|
# Global Store instance to ensure only one connection is established during the entire test session
|
|
GLOBAL_STORE = None
|
|
GLOBAL_CONFIG = None
|
|
|
|
DEFAULT_MOONCAKE_CONFIG_PATH_ENV = "MOONCAKE_CONFIG_PATH"
|
|
DEFAULT_GLOBAL_SEGMENT_SIZE = 16 * 1024 * 1024 * 1024 # 16 GiB
|
|
DEFAULT_LOCAL_BUFFER_SIZE = 8 * 1024 * 1024 * 1024 # 8 GB
|
|
DEFAULT_MASTER_METRICS_PORT = 9003
|
|
DEFAULT_CHECK_SERVER = False
|
|
|
|
def verify_tensor_equality(original, received, rtol=0, atol=0, verbose=True):
|
|
"""
|
|
compare two tensors.
|
|
"""
|
|
def to_numpy(x):
|
|
if isinstance(x, torch.Tensor):
|
|
if x.is_cuda:
|
|
x = x.cpu()
|
|
return x.detach().numpy()
|
|
elif isinstance(x, np.ndarray):
|
|
return x
|
|
else:
|
|
raise TypeError(f"Unsupported tensor type: {type(x)}")
|
|
|
|
try:
|
|
orig_np = to_numpy(original)
|
|
recv_np = to_numpy(received)
|
|
except Exception as e:
|
|
if verbose:
|
|
print(f"❌ Error converting tensors: {e}")
|
|
return False
|
|
|
|
if orig_np.shape != recv_np.shape:
|
|
if verbose:
|
|
print(f"❌ Shape mismatch: original {orig_np.shape} vs received {recv_np.shape}")
|
|
return False
|
|
|
|
if orig_np.dtype != recv_np.dtype:
|
|
if verbose:
|
|
print(f"❌ Dtype mismatch: original {orig_np.dtype} vs received {recv_np.dtype}")
|
|
return False
|
|
|
|
if np.array_equal(orig_np, recv_np):
|
|
# if verbose:
|
|
# print("✅ Tensors are identical!")
|
|
return True
|
|
else:
|
|
diff_mask = orig_np != recv_np
|
|
diff_indices = np.where(diff_mask)
|
|
if len(diff_indices[0]) > 0:
|
|
first_diff_idx = tuple(idx[0] for idx in diff_indices)
|
|
orig_val = orig_np[first_diff_idx]
|
|
recv_val = recv_np[first_diff_idx]
|
|
if verbose:
|
|
print(f"❌ Tensors differ at index {first_diff_idx}")
|
|
print(f" Original: {orig_val}")
|
|
print(f" Received: {recv_val}")
|
|
print(f" Difference: {abs(orig_val - recv_val)}")
|
|
return False
|
|
|
|
def parse_global_segment_size(value) -> int:
|
|
"""Parse human-readable size strings (e.g., '4GB') into bytes."""
|
|
if isinstance(value, int): return value
|
|
if isinstance(value, str):
|
|
s = value.strip().lower()
|
|
if s.endswith("gb"):
|
|
return int(s[:-2].strip()) * 1024**3
|
|
return int(s)
|
|
return int(value)
|
|
|
|
def create_store_connection():
|
|
"""Create and connect to the Store (called only once by setUpModule)."""
|
|
store = MooncakeDistributedStore()
|
|
config = MooncakeConfig.load_from_env()
|
|
print(f"[{os.getpid()}] Connecting to Mooncake Master at {config.master_server_address} using {config.protocol}...")
|
|
rc = store.setup(
|
|
config.local_hostname,
|
|
config.metadata_server,
|
|
config.global_segment_size,
|
|
config.local_buffer_size,
|
|
config.protocol,
|
|
config.device_name,
|
|
config.master_server_address,
|
|
)
|
|
if rc != 0:
|
|
raise RuntimeError(f"Failed to setup mooncake store, error code: {rc}")
|
|
return store, config
|
|
|
|
def generate_tensors(num_tensors, size_mb):
|
|
"""Generate random Tensors suitable for TP testing."""
|
|
size_bytes = int(size_mb * 1024 * 1024)
|
|
element_size = 4 # float32
|
|
num_elements = size_bytes // element_size
|
|
dim = int(np.sqrt(num_elements))
|
|
dim = (dim // 8) * 8 # Adjust dimension to be divisible by common TP sizes (2, 4, 8)
|
|
|
|
# Use random data and ensure the tensor is contiguous in memory
|
|
tensors = [torch.randn(dim, dim, dtype=torch.float32).contiguous() for _ in range(num_tensors)]
|
|
# Use timestamp to prevent key collision in rare edge cases (though we remove_all anyway)
|
|
keys = [f"test_tensor_{i}_{int(time.time()*1000)}" for i in range(num_tensors)]
|
|
return keys, tensors
|
|
|
|
# ==========================================
|
|
# Module Level Setup/Teardown (Hooks)
|
|
# ==========================================
|
|
|
|
def setUpModule():
|
|
"""Executed once before all tests in this file: establishes the global connection."""
|
|
global GLOBAL_STORE, GLOBAL_CONFIG
|
|
try:
|
|
GLOBAL_STORE, GLOBAL_CONFIG = create_store_connection()
|
|
print("✅ Global Store connection established.")
|
|
except Exception as e:
|
|
print(f"❌ Failed to establish global store connection: {e}")
|
|
sys.exit(1)
|
|
|
|
def tearDownModule():
|
|
"""Executed once after all tests in this file: closes the global connection."""
|
|
global GLOBAL_STORE
|
|
if GLOBAL_STORE:
|
|
print("\nClosing global store connection...")
|
|
GLOBAL_STORE.close()
|
|
GLOBAL_STORE = None
|
|
|
|
# ==========================================
|
|
# Base Test Class
|
|
# ==========================================
|
|
|
|
class MooncakeTestBase(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Executed before each test method (test_xxx)."""
|
|
# 1. Access the global connection
|
|
if GLOBAL_STORE is None:
|
|
self.skipTest("Store not initialized")
|
|
|
|
self.store = GLOBAL_STORE
|
|
self.config = GLOBAL_CONFIG
|
|
|
|
# 2. [Critical] Clean environment before the test starts
|
|
# This ensures no stale data from previous tests affects the current one
|
|
self.store.remove_all()
|
|
|
|
# ==========================================
|
|
# Functional Tests
|
|
# ==========================================
|
|
|
|
class TestMooncakeFunctional(MooncakeTestBase):
|
|
def test_01_basic_put_get(self):
|
|
"""Verify basic put and get functionality."""
|
|
key = "func_test_single"
|
|
tensor = torch.randn(1024, 1024, dtype=torch.float32)
|
|
|
|
# Perform Put
|
|
rc = self.store.put_tensor(key, tensor)
|
|
self.assertEqual(rc, 0, f"put_tensor failed with rc={rc}")
|
|
self.assertTrue(self.store.is_exist(key), "Key not found after put")
|
|
|
|
# Perform Get
|
|
retrieved = self.store.get_tensor(key)
|
|
self.assertIsNotNone(retrieved, "Get returned None")
|
|
self.assertTrue(torch.equal(tensor, retrieved), "Data mismatch between original and retrieved tensor")
|
|
|
|
def test_02_tp_single_tensor(self):
|
|
"""Verify TP (Tensor Parallelism) splitting and reconstruction for a single Tensor."""
|
|
tp_size = 4
|
|
split_dim = 1
|
|
key = "func_test_tp_single"
|
|
|
|
# Create a small tensor (e.g., 16MB)
|
|
_, tensors = generate_tensors(1, 16)
|
|
target_tensor = tensors[0]
|
|
|
|
# 1. Put with TP
|
|
rc = self.store.put_tensor_with_tp(key, target_tensor, tp_size=tp_size, split_dim=split_dim)
|
|
self.assertEqual(rc, 0, "put_tensor_with_tp failed")
|
|
|
|
# 2. Verify existence of shards (White-box check: key_tp_0, key_tp_1...)
|
|
for rank in range(tp_size):
|
|
shard_key = f"{key}_tp_{rank}"
|
|
self.assertTrue(self.store.is_exist(shard_key), f"Shard key {shard_key} is missing in store")
|
|
|
|
# 3. Get shards and Reconstruct
|
|
slices = []
|
|
expected_chunks = target_tensor.chunk(tp_size, split_dim)
|
|
for rank in range(tp_size):
|
|
t_slice = self.store.get_tensor_with_tp(key, tp_rank=rank, tp_size=tp_size, split_dim=split_dim)
|
|
self.assertIsNotNone(t_slice, f"Slice for rank {rank} is None")
|
|
self.assertTrue(torch.equal(t_slice, expected_chunks[rank]), f"Data mismatch for rank {rank}")
|
|
slices.append(t_slice)
|
|
|
|
reconstructed = torch.cat(slices, dim=split_dim)
|
|
self.assertTrue(torch.equal(reconstructed, target_tensor), "Reconstructed tensor does not match original")
|
|
|
|
def test_03_tp_batch(self):
|
|
"""Verify TP splitting and reconstruction for a Batch of Tensors."""
|
|
tp_size = 2
|
|
split_dim = 0
|
|
num_tensors = 4
|
|
keys, tensors = generate_tensors(num_tensors, 8) # Small size for functional testing
|
|
|
|
# 1. Batch Put with TP
|
|
results = self.store.batch_put_tensor_with_tp(keys, tensors, tp_size=tp_size, split_dim=split_dim)
|
|
self.assertTrue(all(r == 0 for r in results), f"Batch put failed. Results: {results}")
|
|
|
|
# 2. Batch Get per Rank
|
|
all_shards = [] # List of lists: [ [shards_rank0...], [shards_rank1...] ]
|
|
for rank in range(tp_size):
|
|
shards = self.store.batch_get_tensor_with_tp(keys, tp_rank=rank, tp_size=tp_size)
|
|
self.assertEqual(len(shards), num_tensors)
|
|
all_shards.append(shards)
|
|
|
|
# 3. Verify & Reconstruct
|
|
for i in range(num_tensors):
|
|
original = tensors[i]
|
|
expected_chunks = original.chunk(tp_size, split_dim)
|
|
reconstruction_parts = []
|
|
|
|
for rank in range(tp_size):
|
|
shard = all_shards[rank][i]
|
|
self.assertTrue(torch.equal(shard, expected_chunks[rank]),
|
|
f"Tensor {i} Rank {rank} data mismatch")
|
|
reconstruction_parts.append(shard)
|
|
|
|
recon = torch.cat(reconstruction_parts, dim=split_dim)
|
|
self.assertTrue(torch.equal(recon, original), f"Tensor {i} final reconstruction mismatch")
|
|
|
|
def test_04_tp_consistency(self):
|
|
input_tensor = torch.arange(12).view(3, 4)
|
|
tp_size=2
|
|
self.store.batch_put_tensor_with_tp(['key'], [input_tensor], tp_size=tp_size, split_dim=1)
|
|
chunked_tensors = input_tensor.chunk(chunks=2, dim=1)
|
|
tmp_tensor_0 = self.store.batch_get_tensor_with_tp(['key'], tp_rank=0, tp_size=tp_size)[0]
|
|
tmp_tensor_1 = self.store.batch_get_tensor_with_tp(['key'], tp_rank=1, tp_size=tp_size)[0]
|
|
self.assertTrue(tmp_tensor_0.sum() == chunked_tensors[0].sum())
|
|
self.assertTrue(tmp_tensor_1.sum() == chunked_tensors[1].sum())
|
|
|
|
buffer_spacing = 1 * 1024 * 1024
|
|
buffer_2 = (ctypes.c_ubyte * buffer_spacing)()
|
|
buffer_3 = (ctypes.c_ubyte * buffer_spacing)()
|
|
buffer_ptr_2 = ctypes.addressof(buffer_2)
|
|
buffer_ptr_3 = ctypes.addressof(buffer_3)
|
|
res = self.store.register_buffer(buffer_ptr_2, buffer_spacing)
|
|
self.assertEqual(res, 0, f"Buffer registration failed for buffer at {buffer_ptr_2}")
|
|
res = self.store.register_buffer(buffer_ptr_3, buffer_spacing)
|
|
self.assertEqual(res, 0, f"Buffer registration failed for buffer at {buffer_ptr_3}")
|
|
tmp_tensor_2 = self.store.batch_get_tensor_with_tp_into(['key'], [buffer_ptr_2], [buffer_spacing], tp_rank=0, tp_size=tp_size)[0]
|
|
tmp_tensor_3 = self.store.batch_get_tensor_with_tp_into(['key'], [buffer_ptr_3], [buffer_spacing], tp_rank=1, tp_size=tp_size)[0]
|
|
self.assertTrue(tmp_tensor_2.sum() == chunked_tensors[0].sum())
|
|
self.assertTrue(tmp_tensor_3.sum() == chunked_tensors[1].sum())
|
|
res = self.store.unregister_buffer(buffer_ptr_2)
|
|
self.assertEqual(res, 0, f"Buffer unregistration failed for buffer at {buffer_ptr_2}")
|
|
res = self.store.unregister_buffer(buffer_ptr_3)
|
|
self.assertEqual(res, 0, f"Buffer unregistration failed for buffer at {buffer_ptr_3}")
|
|
|
|
def test_05_put_get_into(self):
|
|
"""Verify basic put and get into functionality."""
|
|
key = "get_into_test"
|
|
tensor = torch.randn(1024, 1024, dtype=torch.float32)
|
|
buffer_spacing = 64 * 1024 * 1024
|
|
total_buffer_size = buffer_spacing
|
|
|
|
# Allocate large contiguous buffer
|
|
large_buffer = (ctypes.c_ubyte * total_buffer_size)()
|
|
large_buffer_ptr = ctypes.addressof(large_buffer)
|
|
|
|
res = self.store.register_buffer(large_buffer_ptr, total_buffer_size)
|
|
self.assertEqual(res, 0, "Buffer registration failed for buffer at {large_buffer_ptr}")
|
|
|
|
# Perform Put
|
|
rc = self.store.put_tensor(key, tensor)
|
|
self.assertEqual(rc, 0, f"put_tensor failed with rc={rc}")
|
|
self.assertTrue(self.store.is_exist(key), "Key not found after put")
|
|
|
|
# Perform Get
|
|
retrieved = self.store.get_tensor_into(key, large_buffer_ptr, total_buffer_size)
|
|
self.assertIsNotNone(retrieved, "Get returned None")
|
|
self.assertTrue(torch.equal(tensor, retrieved), f"Data mismatch between original and retrieved tensor, tensor: {tensor}, retrieved: {retrieved}")
|
|
# Unregister buffer
|
|
res = self.store.unregister_buffer(large_buffer_ptr),
|
|
for r in res:
|
|
self.assertEqual(r, 0, f"Buffer unregistration failed for buffer at {large_buffer_ptr}")
|
|
|
|
def test_06_batch_put_get_into(self):
|
|
"""Zero copy Batch Get."""
|
|
num_tensors = 4
|
|
keys, tensors = generate_tensors(num_tensors, 8)
|
|
buffer_spacing = 64 * 1024 * 1024 # 64MB per tensor slot
|
|
batch_size = len(keys)
|
|
total_buffer_size = buffer_spacing * batch_size
|
|
|
|
# Allocate large contiguous buffer
|
|
large_buffer = (ctypes.c_ubyte * total_buffer_size)()
|
|
large_buffer_ptr = ctypes.addressof(large_buffer)
|
|
|
|
# Prepare pointers (only addresses, no ctypes array objects)
|
|
buffer_ptrs = []
|
|
buffer_sizes = []
|
|
for i in range(batch_size):
|
|
offset = i * buffer_spacing
|
|
ptr = large_buffer_ptr + offset
|
|
buffer_ptrs.append(ptr)
|
|
buffer_sizes.append(buffer_spacing)
|
|
|
|
# Register the entire buffer with the store
|
|
res = self.store.register_buffer(large_buffer_ptr, total_buffer_size)
|
|
self.assertEqual(res, 0, "Buffer registration failed for buffer at {large_buffer_ptr}")
|
|
|
|
results = self.store.batch_put_tensor(keys, tensors)
|
|
self.assertTrue(all(r == 0 for r in results), f"Batch put failed. Results: {results}")
|
|
|
|
res = self.store.batch_get_tensor_into(keys, buffer_ptrs, buffer_sizes)
|
|
self.assertEqual(len(res), len(tensors))
|
|
for j in range(batch_size):
|
|
self.assertTrue(
|
|
verify_tensor_equality(tensors[j], res[j]),
|
|
f"Tensor {j} content mismatch, tensor: {tensors[j]}, res: {res[j]}"
|
|
)
|
|
# Unregister buffer
|
|
res = self.store.unregister_buffer(large_buffer_ptr),
|
|
for r in res:
|
|
self.assertEqual(r, 0, f"Buffer unregistration failed for buffer at {large_buffer_ptr}")
|
|
|
|
def test_07_put_get_into_with_tp(self):
|
|
"""Zero copy Batch Get with TP — each rank has its own buffer."""
|
|
tp_size = 4
|
|
split_dim = 0
|
|
key = "get_into_with_tp_test"
|
|
tensor = torch.randn(1024, 1024, dtype=torch.float32)
|
|
|
|
# Step 1: Put full tensors with TP splitting
|
|
result = self.store.put_tensor_with_tp(
|
|
key, tensor, tp_size=tp_size, split_dim=split_dim
|
|
)
|
|
self.assertEqual(result, 0, f"Put failed. Result: {result}")
|
|
|
|
# Step 2: For each TP rank, allocate its own buffer and get shards
|
|
all_shards = []
|
|
registered_buffers = [] # Keep track of (ptr, size) for cleanup
|
|
|
|
for rank in range(tp_size):
|
|
buffer_spacing = 64 * 1024 * 1024 # 64MB per tensor slot
|
|
total_buffer_size = buffer_spacing
|
|
|
|
# Allocate buffer for this rank
|
|
large_buffer = (ctypes.c_ubyte * total_buffer_size)()
|
|
large_buffer_ptr = ctypes.addressof(large_buffer)
|
|
|
|
# Register buffer for this rank
|
|
res = self.store.register_buffer(large_buffer_ptr, total_buffer_size)
|
|
self.assertEqual(res, 0, f"Buffer registration failed for rank {rank}")
|
|
|
|
# Keep buffer alive and record for cleanup
|
|
# We store (large_buffer, large_buffer_ptr, total_buffer_size)
|
|
# large_buffer must be kept alive until unregister!
|
|
registered_buffers.append((large_buffer, large_buffer_ptr, total_buffer_size))
|
|
|
|
# Get shard for this rank
|
|
shard = self.store.get_tensor_with_tp_into(
|
|
key, large_buffer_ptr, total_buffer_size,
|
|
tp_rank=rank, tp_size=tp_size
|
|
)
|
|
all_shards.append(shard)
|
|
|
|
# Step 3: Validate reconstruction
|
|
original = tensor
|
|
expected_chunks = original.chunk(tp_size, split_dim)
|
|
reconstruction_parts = []
|
|
|
|
for rank in range(tp_size):
|
|
shard = all_shards[rank]
|
|
self.assertTrue(
|
|
torch.equal(shard, expected_chunks[rank]),
|
|
f"Tensor Rank {rank} data mismatch"
|
|
)
|
|
reconstruction_parts.append(shard)
|
|
|
|
recon = torch.cat(reconstruction_parts, dim=split_dim)
|
|
self.assertTrue(
|
|
torch.equal(recon, original),
|
|
f"Tensor final reconstruction mismatch"
|
|
)
|
|
|
|
# Step 4: Unregister all buffers
|
|
for large_buffer, ptr, size in registered_buffers:
|
|
res = self.store.unregister_buffer(ptr)
|
|
self.assertEqual(res, 0, f"Buffer unregistration failed for buffer at {ptr}")
|
|
# large_buffer will be GC'd after this; no need to explicitly delete
|
|
|
|
def test_08_batch_put_get_into_with_tp(self):
|
|
"""Zero copy Batch Get with TP — each rank has its own buffer."""
|
|
tp_size = 4
|
|
split_dim = 0
|
|
num_tensors = 4
|
|
keys, tensors = generate_tensors(num_tensors, 8)
|
|
|
|
# Step 1: Put full tensors with TP splitting
|
|
results = self.store.batch_put_tensor_with_tp(
|
|
keys, tensors, tp_size=tp_size, split_dim=split_dim
|
|
)
|
|
self.assertTrue(all(r == 0 for r in results), f"Batch put failed. Results: {results}")
|
|
|
|
# Step 2: For each TP rank, allocate its own buffer and get shards
|
|
all_shards = []
|
|
registered_buffers = [] # Keep track of (ptr, size) for cleanup
|
|
|
|
for rank in range(tp_size):
|
|
batch_size = len(keys)
|
|
buffer_spacing = 64 * 1024 * 1024 # 64MB per tensor slot
|
|
total_buffer_size = buffer_spacing * batch_size
|
|
|
|
# Allocate buffer for this rank
|
|
large_buffer = (ctypes.c_ubyte * total_buffer_size)()
|
|
large_buffer_ptr = ctypes.addressof(large_buffer)
|
|
|
|
# Prepare buffer pointers for this rank
|
|
buffer_ptrs = []
|
|
buffer_sizes = []
|
|
for i in range(batch_size):
|
|
offset = i * buffer_spacing
|
|
ptr = large_buffer_ptr + offset
|
|
buffer_ptrs.append(ptr)
|
|
buffer_sizes.append(buffer_spacing)
|
|
|
|
# Register buffer for this rank
|
|
res = self.store.register_buffer(large_buffer_ptr, total_buffer_size)
|
|
self.assertEqual(res, 0, f"Buffer registration failed for rank {rank}, at {large_buffer_ptr}")
|
|
|
|
# Keep buffer alive and record for cleanup
|
|
# We store (large_buffer, large_buffer_ptr, total_buffer_size)
|
|
# large_buffer must be kept alive until unregister!
|
|
registered_buffers.append((large_buffer, large_buffer_ptr, total_buffer_size))
|
|
|
|
# Get shards for this rank
|
|
shards = self.store.batch_get_tensor_with_tp_into(
|
|
keys, buffer_ptrs, buffer_sizes,
|
|
tp_rank=rank, tp_size=tp_size
|
|
)
|
|
self.assertEqual(len(shards), num_tensors)
|
|
all_shards.append(shards)
|
|
|
|
# Step 3: Validate reconstruction
|
|
for i in range(num_tensors):
|
|
original = tensors[i]
|
|
expected_chunks = original.chunk(tp_size, split_dim)
|
|
reconstruction_parts = []
|
|
|
|
for rank in range(tp_size):
|
|
shard = all_shards[rank][i]
|
|
self.assertTrue(
|
|
torch.equal(shard, expected_chunks[rank]),
|
|
f"Tensor {i} Rank {rank} data mismatch"
|
|
)
|
|
reconstruction_parts.append(shard)
|
|
|
|
recon = torch.cat(reconstruction_parts, dim=split_dim)
|
|
self.assertTrue(torch.equal(recon, original), f"Tensor {i} final reconstruction mismatch")
|
|
|
|
# Step 4: Unregister all buffers
|
|
for large_buffer, ptr, size in registered_buffers:
|
|
res = self.store.unregister_buffer(ptr)
|
|
self.assertEqual(res, 0, f"Buffer unregistration failed for buffer at {ptr}")
|
|
# large_buffer will be GC'd after this; no need to explicitly delete
|
|
|
|
# ==========================================
|
|
# Performance/Benchmark Tests
|
|
# ==========================================
|
|
|
|
class TestMooncakeBenchmark(MooncakeTestBase):
|
|
# Benchmark Settings
|
|
BENCH_ITERATIONS = 5
|
|
TENSOR_SIZE_MB = 16
|
|
TOTAL_SIZE_MB = 256
|
|
|
|
def setUp(self):
|
|
"""Benchmark-specific setUp."""
|
|
# 1. Call parent setUp to clean the store (remove_all)
|
|
super().setUp()
|
|
|
|
# 2. Generate test data
|
|
total_bytes = int(self.TOTAL_SIZE_MB * 1024**2)
|
|
tensor_bytes = self.TENSOR_SIZE_MB * 1024**2
|
|
self.num_tensors = max(1, total_bytes // tensor_bytes)
|
|
|
|
print(f"\n[Gen] Generating {self.num_tensors} tensors (~{self.TENSOR_SIZE_MB}MB each)...")
|
|
self.keys, self.tensors = generate_tensors(self.num_tensors, self.TENSOR_SIZE_MB)
|
|
self.total_bits = (tensor_bytes * self.num_tensors) * 8
|
|
|
|
def _print_perf(self, name, times):
|
|
avg_time = np.mean(times)
|
|
avg_gbps = (self.total_bits / 1e9) / avg_time
|
|
print(f"👉 [Result] {name:30} | Avg Time: {avg_time:.4f}s | Throughput: {avg_gbps:.2f} Gbps")
|
|
|
|
def test_benchmark_01_batch_put_get(self):
|
|
"""Benchmark: Standard Batch Put/Get."""
|
|
put_times = []
|
|
get_times = []
|
|
|
|
print(f"--- Running Standard Batch Benchmark ({self.BENCH_ITERATIONS} iters) ---")
|
|
for i in range(self.BENCH_ITERATIONS):
|
|
# Clean store before each iteration for "cold" writes
|
|
self.store.remove_all()
|
|
|
|
# Measure Put
|
|
t0 = time.perf_counter()
|
|
self.store.batch_put_tensor(self.keys, self.tensors)
|
|
put_times.append(time.perf_counter() - t0)
|
|
|
|
# Measure Get
|
|
t0 = time.perf_counter()
|
|
res = self.store.batch_get_tensor(self.keys)
|
|
get_times.append(time.perf_counter() - t0)
|
|
self.assertEqual(len(res), len(self.tensors))
|
|
|
|
self._print_perf("Standard Batch Put", put_times)
|
|
self._print_perf("Standard Batch Get", get_times)
|
|
|
|
def test_benchmark_02_tp_batch(self):
|
|
"""Benchmark: TP Batch Put/Get."""
|
|
tp_size = 4
|
|
split_dim = 0
|
|
put_times = []
|
|
get_times = []
|
|
|
|
print(f"--- Running TP Batch Benchmark (TP={tp_size}) ---")
|
|
for i in range(self.BENCH_ITERATIONS):
|
|
self.store.remove_all()
|
|
|
|
# Measure TP Put (Auto-chunking)
|
|
t0 = time.perf_counter()
|
|
self.store.batch_put_tensor_with_tp(self.keys, self.tensors, tp_size=tp_size, split_dim=split_dim)
|
|
put_times.append(time.perf_counter() - t0)
|
|
|
|
# Measure TP Get (Simulating gathering all ranks)
|
|
t_get_start = time.perf_counter()
|
|
for rank in range(tp_size):
|
|
res = self.store.batch_get_tensor_with_tp(self.keys, tp_rank=rank, tp_size=tp_size)
|
|
self.assertEqual(len(res), len(self.tensors))
|
|
get_times.append(time.perf_counter() - t_get_start)
|
|
|
|
self._print_perf(f"TP Batch Put (TP={tp_size})", put_times)
|
|
self._print_perf(f"TP Batch Get (TP={tp_size})", get_times)
|
|
|
|
def test_benchmark_03_batch_put_get_into(self):
|
|
"""Benchmark: Zero copy Batch Get."""
|
|
self.store.remove_all()
|
|
buffer_spacing = 300 * 1024 * 1024 # 300MB per tensor slot
|
|
batch_size = len(self.keys)
|
|
total_buffer_size = buffer_spacing * batch_size
|
|
|
|
# Allocate large contiguous buffer
|
|
large_buffer = (ctypes.c_ubyte * total_buffer_size)()
|
|
large_buffer_ptr = ctypes.addressof(large_buffer)
|
|
|
|
# Prepare pointers (only addresses, no ctypes array objects)
|
|
buffer_ptrs = []
|
|
buffer_sizes = []
|
|
for i in range(batch_size):
|
|
offset = i * buffer_spacing
|
|
ptr = large_buffer_ptr + offset
|
|
buffer_ptrs.append(ptr)
|
|
buffer_sizes.append(buffer_spacing)
|
|
|
|
# Register the entire buffer with the store
|
|
res = self.store.register_buffer(large_buffer_ptr, total_buffer_size)
|
|
self.assertEqual(res, 0, "Buffer registration should succeed")
|
|
|
|
print(f"--- Running zero copy Batch Benchmark ({self.BENCH_ITERATIONS} iters) ---")
|
|
put_times = []
|
|
get_times = []
|
|
|
|
for i in range(self.BENCH_ITERATIONS):
|
|
# Clean store before each iteration for "cold" writes
|
|
self.store.remove_all()
|
|
|
|
# Measure Put
|
|
t0 = time.perf_counter()
|
|
self.store.batch_put_tensor(self.keys, self.tensors)
|
|
put_times.append(time.perf_counter() - t0)
|
|
|
|
# Measure Get
|
|
t0 = time.perf_counter()
|
|
res = self.store.batch_get_tensor_into(self.keys, buffer_ptrs, buffer_sizes)
|
|
get_times.append(time.perf_counter() - t0)
|
|
self.assertEqual(len(res), len(self.tensors))
|
|
for j in range(batch_size):
|
|
self.assertTrue(
|
|
verify_tensor_equality(self.tensors[j], res[j]),
|
|
f"Tensor {j} content mismatch"
|
|
)
|
|
|
|
self._print_perf("Standard Batch Put", put_times)
|
|
self._print_perf("Zero copy Batch Get", get_times)
|
|
|
|
# Unregister buffer
|
|
res = self.store.unregister_buffer(large_buffer_ptr),
|
|
for r in res:
|
|
self.assertEqual(r, 0, f"Buffer unregistration failed for buffer at {large_buffer_ptr}")
|
|
|
|
def test_benchmark_04_batch_put_get_into_with_tp(self):
|
|
"""Benchmark: Zero copy Batch Get with tp."""
|
|
tp_size = 4
|
|
split_dim = 0
|
|
batch_size = len(self.keys)
|
|
self.store.remove_all()
|
|
buffer_spacing = 64 * 1024 * 1024 # 64MB per tensor slot
|
|
|
|
# Allocate and register a separate buffer for each TP rank
|
|
rank_buffers = [] # Store metadata for cleanup
|
|
|
|
for rank in range(tp_size):
|
|
total_buffer_size = buffer_spacing * batch_size
|
|
large_buffer = (ctypes.c_ubyte * total_buffer_size)()
|
|
large_buffer_ptr = ctypes.addressof(large_buffer)
|
|
|
|
buffer_ptrs = []
|
|
buffer_sizes = []
|
|
for i in range(batch_size):
|
|
offset = i * buffer_spacing
|
|
ptr = large_buffer_ptr + offset
|
|
buffer_ptrs.append(ptr)
|
|
buffer_sizes.append(buffer_spacing)
|
|
|
|
# Register buffer for this rank
|
|
res = self.store.register_buffer(large_buffer_ptr, total_buffer_size)
|
|
self.assertEqual(res, 0, f"Buffer registration failed for rank {rank}")
|
|
|
|
rank_buffers.append({
|
|
'buffer_obj': large_buffer, # keep alive
|
|
'ptrs': buffer_ptrs,
|
|
'sizes': buffer_sizes,
|
|
'base_ptr': large_buffer_ptr,
|
|
'total_size': total_buffer_size
|
|
})
|
|
|
|
print(f"--- Running zero copy Batch Benchmark (TP={tp_size}, {self.BENCH_ITERATIONS} iters) ---")
|
|
put_times = []
|
|
get_times = []
|
|
|
|
for i in range(self.BENCH_ITERATIONS):
|
|
# Clean store before each iteration for "cold" writes
|
|
self.store.remove_all()
|
|
|
|
# Measure Put
|
|
t0 = time.perf_counter()
|
|
self.store.batch_put_tensor_with_tp(
|
|
self.keys, self.tensors, tp_size=tp_size, split_dim=split_dim
|
|
)
|
|
put_times.append(time.perf_counter() - t0)
|
|
|
|
# Measure Get: each rank uses its own buffer
|
|
t0 = time.perf_counter()
|
|
all_res = []
|
|
for rank in range(tp_size):
|
|
res = self.store.batch_get_tensor_with_tp_into(
|
|
self.keys,
|
|
rank_buffers[rank]['ptrs'],
|
|
rank_buffers[rank]['sizes'],
|
|
tp_rank=rank,
|
|
tp_size=tp_size
|
|
)
|
|
self.assertEqual(len(res), batch_size)
|
|
all_res.append(res)
|
|
get_times.append(time.perf_counter() - t0)
|
|
|
|
# Verify correctness using rank 0's result
|
|
for j in range(batch_size):
|
|
original = self.tensors[j]
|
|
expected_shard = original.chunk(tp_size, split_dim)[0] # rank 0 shard
|
|
actual = all_res[0][j]
|
|
self.assertTrue(
|
|
torch.equal(actual, expected_shard),
|
|
f"Tensor {j} content mismatch on rank 0"
|
|
)
|
|
|
|
self._print_perf(f"Standard Batch Put with tp (TP={tp_size})", put_times)
|
|
self._print_perf(f"Zero copy Batch Get with tp (TP={tp_size})", get_times)
|
|
|
|
# Unregister all buffers
|
|
for buf_info in rank_buffers:
|
|
res = self.store.unregister_buffer(buf_info['base_ptr'])
|
|
self.assertEqual(res, 0, f"Buffer unregistration failed for buffer at {buf_info['base_ptr']}")
|
|
|
|
# ==========================================
|
|
# Stress/Concurrency Tests
|
|
# ==========================================
|
|
class TestMooncakeStress(MooncakeTestBase):
|
|
"""
|
|
Stress tests with Fixed Operation Count and Pre-generated Data.
|
|
"""
|
|
# Default Config (Overridden by main)
|
|
NUM_THREADS = 8
|
|
TOTAL_ITEMS = 800 # Total number of items to process across all threads
|
|
TENSOR_SIZE_MB = 4 # Size per tensor
|
|
|
|
def _run_stress_worker(self, thread_id, items_per_thread):
|
|
"""
|
|
Worker function:
|
|
1. PRE-GENERATES data (to exclude generation time from benchmark).
|
|
2. Performs Put -> Get -> Verify loop.
|
|
"""
|
|
ops_count = 0
|
|
failure_msg = None
|
|
|
|
# Pre-calculate dimensions
|
|
element_size = 4 # float32
|
|
num_elements = (self.TENSOR_SIZE_MB * 1024 * 1024) // element_size
|
|
dim = int(np.sqrt(num_elements))
|
|
|
|
# --- Phase 1: Pre-generate Data ---
|
|
# "Don't keep generating random data" -> We generate a pool first.
|
|
# This ensures we measure store performance, not RNG performance.
|
|
print(f" [Thread {thread_id}] Pre-generating {items_per_thread} tensors...")
|
|
data_pool = []
|
|
for i in range(items_per_thread):
|
|
key = f"stress_fixed_t{thread_id}_{i}"
|
|
# Create random tensor
|
|
tensor = torch.randn(dim, dim, dtype=torch.float32)
|
|
data_pool.append((key, tensor))
|
|
|
|
# Barrier logic simulation: wait for main test to indicate start?
|
|
# In simple unittest, we just start processing.
|
|
|
|
# --- Phase 2: Execution (Timed) ---
|
|
t_start = time.perf_counter()
|
|
|
|
try:
|
|
for key, original_tensor in data_pool:
|
|
# 1. WRITE (Put)
|
|
rc = self.store.put_tensor(key, original_tensor)
|
|
if rc != 0:
|
|
raise RuntimeError(f"Put failed for {key}, rc={rc}")
|
|
|
|
# 2. READ (Get)
|
|
retrieved_tensor = self.store.get_tensor(key)
|
|
|
|
# 3. VALIDATE
|
|
if retrieved_tensor is None:
|
|
raise RuntimeError(f"Get returned None for key {key}")
|
|
|
|
if not torch.equal(original_tensor, retrieved_tensor):
|
|
raise RuntimeError(f"Data Mismatch for {key}!")
|
|
|
|
ops_count += 1
|
|
|
|
except Exception as e:
|
|
failure_msg = str(e)
|
|
|
|
t_duration = time.perf_counter() - t_start
|
|
return ops_count, t_duration, failure_msg
|
|
|
|
def test_stress_consistency_fixed(self):
|
|
"""
|
|
Run a fixed number of operations with data consistency checks.
|
|
"""
|
|
items_per_thread = self.TOTAL_ITEMS // self.NUM_THREADS
|
|
# Adjust for remainder if any
|
|
|
|
print(f"\n--- [Stress] Running Fixed Count Test ({self.TOTAL_ITEMS} items total) ---")
|
|
print(f"--- Config: {self.NUM_THREADS} Threads, ~{items_per_thread} items/thread, {self.TENSOR_SIZE_MB}MB each ---")
|
|
|
|
futures = []
|
|
total_ops = 0
|
|
errors = []
|
|
|
|
# We measure wall time from when threads are submitted until all are done
|
|
t0 = time.perf_counter()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_THREADS) as executor:
|
|
# Distribute work
|
|
for i in range(self.NUM_THREADS):
|
|
count = items_per_thread + (1 if i < (self.TOTAL_ITEMS % self.NUM_THREADS) else 0)
|
|
futures.append(executor.submit(self._run_stress_worker, i, count))
|
|
|
|
# Gather results
|
|
for future in concurrent.futures.as_completed(futures):
|
|
ops, duration, error = future.result()
|
|
total_ops += ops
|
|
if error:
|
|
errors.append(error)
|
|
|
|
elapsed = time.perf_counter() - t0
|
|
|
|
# Reporting
|
|
print(f"\n--- [Stress Report] ---")
|
|
if errors:
|
|
print(f"❌ FAILED with {len(errors)} errors.")
|
|
print(f"First Error: {errors[0]}")
|
|
self.fail(f"Stress test failed with {len(errors)} errors.")
|
|
else:
|
|
total_data_gb = (total_ops * self.TENSOR_SIZE_MB) / 1024
|
|
throughput_gbps = (total_data_gb * 8) / elapsed
|
|
|
|
print(f"✅ PASSED (No Consistency Errors)")
|
|
print(f"Total Items: {total_ops}")
|
|
print(f"Wall Time: {elapsed:.4f} s")
|
|
print(f"Avg QPS: {total_ops / elapsed:.2f} ops/s")
|
|
print(f"Avg Goodput: {throughput_gbps:.2f} Gbps")
|
|
|
|
|
|
# ==========================================
|
|
# Data Type & Precision Tests (Full Enum)
|
|
# ==========================================
|
|
|
|
class TestMooncakeDataTypes(MooncakeTestBase):
|
|
def _test_dtype_roundtrip(self, dtype, name, expected_enum_name=None):
|
|
"""
|
|
Generic test for put/get consistency.
|
|
Args:
|
|
dtype: The torch.dtype to test.
|
|
name: Readable name for logging.
|
|
expected_enum_name: (Optional) If we could inspect the C++ enum value, we would check this.
|
|
"""
|
|
key = f"dtype_check_{name}"
|
|
shape = (64, 64)
|
|
|
|
if dtype == torch.bool:
|
|
original = torch.randint(0, 2, shape).bool()
|
|
elif dtype.is_floating_point:
|
|
original = torch.randn(shape, dtype=torch.float32).to(dtype)
|
|
else:
|
|
if dtype == torch.int8:
|
|
original = torch.randint(-128, 127, shape, dtype=dtype)
|
|
elif dtype == torch.uint8:
|
|
original = torch.randint(0, 255, shape, dtype=dtype)
|
|
else:
|
|
original = torch.randint(-1000, 1000, shape, dtype=dtype)
|
|
|
|
# The C++ store will infer the Enum based on original.dtype
|
|
rc = self.store.put_tensor(key, original)
|
|
if rc != 0:
|
|
print(f" [Fail] {name:<15} Put failed with rc={rc}")
|
|
self.fail(f"Put failed for {name}")
|
|
|
|
retrieved = self.store.get_tensor(key)
|
|
if retrieved is None:
|
|
print(f" [Fail] {name:<15} Get returned None")
|
|
self.fail(f"Get returned None for {name}")
|
|
|
|
# We expect the retrieved tensor to have the same dtype as input
|
|
if original.dtype != retrieved.dtype:
|
|
msg = f"Dtype mismatch for {name}! Input: {original.dtype}, Output: {retrieved.dtype}"
|
|
print(f" [Fail] {name:<15} {msg}")
|
|
self.fail(msg)
|
|
|
|
# Use byte-view comparison for robustness (especially for FP8/BF16 on CPU)
|
|
try:
|
|
# Cast to untyped storage byte view (or uint8 view)
|
|
t1_bytes = original.view(torch.uint8) if original.element_size() > 0 else original
|
|
t2_bytes = retrieved.view(torch.uint8) if retrieved.element_size() > 0 else retrieved
|
|
is_equal = torch.equal(t1_bytes, t2_bytes)
|
|
except Exception:
|
|
# Fallback for types that might fail view() or equal()
|
|
is_equal = torch.equal(original.cpu(), retrieved.cpu())
|
|
|
|
if not is_equal:
|
|
print(f" [Fail] {name:<15} Data content mismatch")
|
|
self.fail(f"Data content mismatch for {name}")
|
|
|
|
buffer_spacing = 1 * 1024 * 1024
|
|
buffer = (ctypes.c_ubyte * buffer_spacing)()
|
|
buffer_ptr = ctypes.addressof(buffer)
|
|
res = self.store.register_buffer(buffer_ptr, buffer_spacing)
|
|
self.assertEqual(res, 0, f"Buffer registration failed for buffer at {buffer_ptr}")
|
|
retrieved = self.store.get_tensor_into(key, buffer_ptr, buffer_spacing)
|
|
if retrieved is None:
|
|
print(f" [Fail] {name:<15} Get returned None")
|
|
self.fail(f"Get returned None for {name}")
|
|
|
|
# We expect the retrieved tensor to have the same dtype as input
|
|
if original.dtype != retrieved.dtype:
|
|
msg = f"Dtype mismatch for {name}! Input: {original.dtype}, Output: {retrieved.dtype}"
|
|
print(f" [Fail] {name:<15} {msg}")
|
|
self.fail(msg)
|
|
|
|
# Use byte-view comparison for robustness (especially for FP8/BF16 on CPU)
|
|
try:
|
|
# Cast to untyped storage byte view (or uint8 view)
|
|
t1_bytes = original.view(torch.uint8) if original.element_size() > 0 else original
|
|
t2_bytes = retrieved.view(torch.uint8) if retrieved.element_size() > 0 else retrieved
|
|
is_equal = torch.equal(t1_bytes, t2_bytes)
|
|
except Exception:
|
|
# Fallback for types that might fail view() or equal()
|
|
is_equal = torch.equal(original.cpu(), retrieved.cpu())
|
|
|
|
if not is_equal:
|
|
print(f" [Fail] {name:<15} Data content mismatch")
|
|
self.fail(f"Data content mismatch for {name}")
|
|
res = self.store.unregister_buffer(buffer_ptr)
|
|
self.assertEqual(res, 0, f"Buffer unregistration failed for buffer at {buffer_ptr}")
|
|
|
|
print(f" [Pass] {name:<15} {str(dtype)}")
|
|
|
|
|
|
def test_all_dtypes(self):
|
|
print("\n--- Testing All Supported PyTorch Data Types ---")
|
|
|
|
test_cases = [
|
|
("FLOAT32", torch.float32),
|
|
("FLOAT64", torch.float64),
|
|
("INT8", torch.int8),
|
|
("UINT8", torch.uint8),
|
|
("INT16", torch.int16),
|
|
("INT32", torch.int32),
|
|
("INT64", torch.int64),
|
|
("BOOL", torch.bool),
|
|
("FLOAT16", torch.float16),
|
|
("BFLOAT16", torch.bfloat16),
|
|
]
|
|
|
|
for name, dtype in test_cases:
|
|
with self.subTest(dtype=name):
|
|
self._test_dtype_roundtrip(dtype, name)
|
|
|
|
def test_fp8_types(self):
|
|
print("\n--- Testing FP8 Types ---")
|
|
|
|
fp8_cases = []
|
|
# Check support dynamically
|
|
if hasattr(torch, 'float8_e4m3fn'):
|
|
fp8_cases.append(("FLOAT8_E4M3", torch.float8_e4m3fn)) # Enum 13
|
|
else:
|
|
print(" [Skip] FLOAT8_E4M3 (Not supported in this PyTorch version)")
|
|
|
|
if hasattr(torch, 'float8_e5m2'):
|
|
fp8_cases.append(("FLOAT8_E5M2", torch.float8_e5m2)) # Enum 14
|
|
else:
|
|
print(" [Skip] FLOAT8_E5M2 (Not supported in this PyTorch version)")
|
|
|
|
for name, dtype in fp8_cases:
|
|
with self.subTest(dtype=name):
|
|
self._test_dtype_roundtrip(dtype, name)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Mooncake Distributed Store Tests")
|
|
parser.add_argument("--mode", type=str, default="all", choices=["all", "func", "perf", "stress", "types"],
|
|
help="Run mode")
|
|
parser.add_argument("--threads", type=int, default=8, help="Number of threads")
|
|
parser.add_argument("--count", type=int, default=800, help="Total number of items to process")
|
|
parser.add_argument("--size_mb", type=float, default=0.5, help="Tensor size in MB")
|
|
|
|
args, unknown = parser.parse_known_args()
|
|
|
|
# Update Stress Test Config
|
|
TestMooncakeStress.NUM_THREADS = args.threads
|
|
TestMooncakeStress.TOTAL_ITEMS = args.count
|
|
TestMooncakeStress.TENSOR_SIZE_MB = args.size_mb
|
|
|
|
suite = unittest.TestSuite()
|
|
loader = unittest.TestLoader()
|
|
|
|
if args.mode in ["all", "func"]:
|
|
print(">> Loading Functional Tests...")
|
|
suite.addTests(loader.loadTestsFromTestCase(TestMooncakeFunctional))
|
|
|
|
if args.mode in ["all", "perf"]:
|
|
print(">> Loading Performance Benchmark Tests...")
|
|
suite.addTests(loader.loadTestsFromTestCase(TestMooncakeBenchmark))
|
|
|
|
if args.mode in ["all", "stress"]:
|
|
print(f">> Loading Stress Tests ({args.count} items, {args.threads} threads)...")
|
|
suite.addTests(loader.loadTestsFromTestCase(TestMooncakeStress))
|
|
|
|
if args.mode in ["all", "types", "func"]: # 'types' can be part of 'func' or standalone
|
|
print(">> Loading Data Type Tests...")
|
|
suite.addTests(loader.loadTestsFromTestCase(TestMooncakeDataTypes))
|
|
|
|
runner = unittest.TextTestRunner(verbosity=2)
|
|
result = runner.run(suite)
|
|
|
|
sys.exit(not result.wasSuccessful())
|