mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[CCF Archive] Store object type eviction policy submission #3

Closed
kancel wants to merge 382 commits from kancel:ccf-archive-pr2746 into main
pull from: kancel:ccf-archive-pr2746
merge into: mooncake-track:main
Conversation 0 Commits 382 Files Changed 654 +259 -16
7 changed files with 259 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit b0ae4a727f - Show all commits

15
mooncake-integration/store/store_py.cpp
View File

@ -1574,6 +1574,20 @@ class MooncakeHostMemAllocatorPyWrapper {
};
PYBIND11_MODULE(store, m) {
// Object data type classification
py::enum_<ObjectDataType>(m, "ObjectDataType")
.value("UNKNOWN", ObjectDataType::UNKNOWN)
.value("KVCACHE", ObjectDataType::KVCACHE)
.value("TENSOR", ObjectDataType::TENSOR)
.value("WEIGHT", ObjectDataType::WEIGHT)
.value("SAMPLE", ObjectDataType::SAMPLE)
.value("ACTIVATION", ObjectDataType::ACTIVATION)
.value("GRADIENT", ObjectDataType::GRADIENT)
.value("OPTIMIZER_STATE", ObjectDataType::OPTIMIZER_STATE)
.value("METADATA", ObjectDataType::METADATA)
.value("GENERAL", ObjectDataType::GENERAL)
.export_values();
// Define the ReplicateConfig class
py::class_<ReplicateConfig>(m, "ReplicateConfig")
.def(py::init<>())
@ -1585,6 +1599,7 @@ PYBIND11_MODULE(store, m) {
.def_readwrite("preferred_segment", &ReplicateConfig::preferred_segment)
.def_readwrite("prefer_alloc_in_same_node",
&ReplicateConfig::prefer_alloc_in_same_node)
.def_readwrite("data_type", &ReplicateConfig::data_type)
.def("__str__", [](const ReplicateConfig &config) {
std::ostringstream oss;
oss << config;

10
mooncake-store/include/master_service.h
View File

@ -598,10 +598,12 @@ class MasterService {
const UUID& client_id_,
const std::chrono::system_clock::time_point put_start_time_,
size_t value_length, std::vector<Replica>&& reps,
bool enable_soft_pin, bool enable_hard_pin = false)
bool enable_soft_pin, bool enable_hard_pin = false,
ObjectDataType data_type_ = ObjectDataType::UNKNOWN)
: client_id(client_id_),
put_start_time(put_start_time_),
size(value_length),
data_type(data_type_),
lease_timeout(),
soft_pin_timeout(std::nullopt),
hard_pinned(enable_hard_pin),
@ -624,6 +626,7 @@ class MasterService {
// Updated by UpsertStart (Case B) to reset the discard timeout.
std::chrono::system_clock::time_point put_start_time;
const size_t size;
const ObjectDataType data_type{ObjectDataType::UNKNOWN};
mutable SpinLock lock;
// Default constructor, creates a time_point representing
@ -1040,7 +1043,8 @@ class MasterService {
void Create(const UUID& client_id, uint64_t total_length,
std::vector<Replica> replicas, bool enable_soft_pin,
bool enable_hard_pin = false) {
bool enable_hard_pin = false,
ObjectDataType data_type = ObjectDataType::UNKNOWN) {
if (Exists()) {
throw std::logic_error("Already exists");
}
@ -1049,7 +1053,7 @@ class MasterService {
std::piecewise_construct, std::forward_as_tuple(key_),
std::forward_as_tuple(client_id, now, total_length,
std::move(replicas), enable_soft_pin,
enable_hard_pin));
enable_hard_pin, data_type));
it_ = result.first;
}

4
mooncake-store/include/replica.h
View File

@ -93,6 +93,7 @@ struct ReplicateConfig {
std::string preferred_segment{}; // Deprecated: Single preferred segment
// for backward compatibility
bool prefer_alloc_in_same_node{false};
ObjectDataType data_type{ObjectDataType::UNKNOWN};
friend std::ostream& operator<<(std::ostream& os,
const ReplicateConfig& config) noexcept {
@ -110,7 +111,8 @@ struct ReplicateConfig {
<< config.preferred_segment;
}
os << ", prefer_alloc_in_same_node: "
<< config.prefer_alloc_in_same_node << " }";
<< config.prefer_alloc_in_same_node
<< ", data_type: " << config.data_type << " }";
return os;
}
};

41
mooncake-store/include/types.h
View File

@ -5,6 +5,7 @@
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <limits>
#include <unordered_map>
#include <utility>
@ -123,6 +124,46 @@ static constexpr uint64_t DEFAULT_PROCESSING_TASK_TIMEOUT_SEC =
300; // 0 to be no timeout
static constexpr uint32_t DEFAULT_MAX_RETRY_ATTEMPTS = 10;
/**
* @brief Data type classification for objects stored in Mooncake Store.
*
* This allows the store to track what kind of data each object holds,
* enabling future type-aware policies (eviction priority, replication
* strategies, etc.). Defaults to UNKNOWN for backward compatibility.
*/
enum class ObjectDataType : uint8_t {
UNKNOWN = 0,
KVCACHE = 1,
TENSOR = 2,
WEIGHT = 3,
SAMPLE = 4,
ACTIVATION = 5,
GRADIENT = 6,
OPTIMIZER_STATE = 7,
METADATA = 8,
GENERAL = 9,
// 10-255 reserved for future types
};
inline std::ostream& operator<<(std::ostream& os,
const ObjectDataType& type) noexcept {
static const std::unordered_map<ObjectDataType, std::string_view>
type_strings{{ObjectDataType::UNKNOWN, "UNKNOWN"},
{ObjectDataType::KVCACHE, "KVCACHE"},
{ObjectDataType::TENSOR, "TENSOR"},
{ObjectDataType::WEIGHT, "WEIGHT"},
{ObjectDataType::SAMPLE, "SAMPLE"},
{ObjectDataType::ACTIVATION, "ACTIVATION"},
{ObjectDataType::GRADIENT, "GRADIENT"},
{ObjectDataType::OPTIMIZER_STATE, "OPTIMIZER_STATE"},
{ObjectDataType::METADATA, "METADATA"},
{ObjectDataType::GENERAL, "GENERAL"}};
auto it = type_strings.find(type);
os << (it != type_strings.end() ? it->second : "UNKNOWN");
return os;
}
// Forward declarations
class BufferAllocatorBase;
class CachelibBufferAllocator;

47
mooncake-store/src/master_service.cpp
View File

@ -845,7 +845,8 @@ auto MasterService::AllocateAndInsertMetadata(
shard->metadata.emplace(
std::piecewise_construct, std::forward_as_tuple(key),
std::forward_as_tuple(client_id, now, value_length, std::move(replicas),
config.with_soft_pin, config.with_hard_pin));
config.with_soft_pin, config.with_hard_pin,
config.data_type));
shard->processing_keys.insert(key);
return replica_list;
@ -4322,7 +4323,7 @@ MasterService::MetadataSerializer::DeserializeShard(const msgpack::object& obj,
metadata_ptr->client_id, metadata_ptr->put_start_time,
metadata_ptr->size, metadata_ptr->PopReplicas(),
metadata_ptr->soft_pin_timeout.has_value(),
metadata_ptr->IsHardPinned()));
metadata_ptr->IsHardPinned(), metadata_ptr->data_type));
it->second.lease_timeout = metadata_ptr->lease_timeout;
it->second.soft_pin_timeout = metadata_ptr->soft_pin_timeout;
@ -4337,12 +4338,12 @@ MasterService::MetadataSerializer::SerializeMetadata(
MsgpackPacker& packer) const {
// Pack ObjectMetadata using array structure for efficiency
// Format: [client_id, put_start_time, size, lease_timeout,
// has_soft_pin_timeout, soft_pin_timeout, replicas_count, replicas...,
// hard_pinned]
// has_soft_pin_timeout, soft_pin_timeout, replicas_count, data_type,
// replicas..., hard_pinned]
size_t array_size = 8; // client_id, put_start_time, size, lease_timeout,
size_t array_size = 9; // client_id, put_start_time, size, lease_timeout,
// has_soft_pin_timeout, soft_pin_timeout,
// replicas_count + hard_pinned
// replicas_count, data_type, hard_pinned
array_size += metadata.CountReplicas(); // One element per replica
packer.pack_array(array_size);
@ -4382,6 +4383,9 @@ MasterService::MetadataSerializer::SerializeMetadata(
// Serialize replicas count
packer.pack(static_cast<uint32_t>(metadata.CountReplicas()));
// Serialize data_type
packer.pack(static_cast<uint8_t>(metadata.data_type));
// Serialize replicas
for (const auto& replica : metadata.GetAllReplicas()) {
auto result = Serializer<Replica>::serialize(
@ -4408,7 +4412,6 @@ MasterService::MetadataSerializer::DeserializeMetadata(
// Need at least 7 elements: client_id, put_start_time, size, lease_timeout,
// has_soft_pin_timeout, soft_pin_timeout, replicas_count
// (8th element = hard_pinned is optional for backward compat)
if (obj.via.array.size < 7) {
return tl::unexpected(SerializationError(
ErrorCode::DESERIALIZE_FAIL,
@ -4441,15 +4444,35 @@ MasterService::MetadataSerializer::DeserializeMetadata(
// Deserialize replicas count
uint32_t replicas_count = array[index++].as<uint32_t>();
// Array size: 7 + replicas_count (old format) or 8 + replicas_count (new
// format with hard_pinned)
if (obj.via.array.size != 7 + replicas_count &&
obj.via.array.size != 8 + replicas_count) {
// Format detection:
// v1: 7 + replicas_count, no data_type or hard_pinned
// v2: 8 + replicas_count, either data_type or trailing hard_pinned
// v3: 9 + replicas_count, data_type plus trailing hard_pinned
constexpr uint32_t kOldFieldCount = 7;
constexpr uint32_t kOneExtraFieldCount = 8;
constexpr uint32_t kCurrentFieldCount = 9;
const uint32_t total_elements = obj.via.array.size;
const bool is_old_format =
(total_elements == kOldFieldCount + replicas_count);
const bool is_one_extra_format =
(total_elements == kOneExtraFieldCount + replicas_count);
const bool is_current_format =
(total_elements == kCurrentFieldCount + replicas_count);
if (!is_current_format && !is_one_extra_format && !is_old_format) {
return tl::unexpected(SerializationError(
ErrorCode::DESERIALIZE_FAIL,
"deserialize ObjectMetadata array size mismatch"));
}
ObjectDataType data_type = ObjectDataType::UNKNOWN;
if (is_current_format) {
data_type = static_cast<ObjectDataType>(array[index++].as<uint8_t>());
} else if (is_one_extra_format &&
array[index].type == msgpack::type::POSITIVE_INTEGER) {
data_type = static_cast<ObjectDataType>(array[index++].as<uint8_t>());
}
// Deserialize replicas
std::vector<Replica> replicas;
replicas.reserve(replicas_count);
@ -4475,7 +4498,7 @@ MasterService::MetadataSerializer::DeserializeMetadata(
client_id,
std::chrono::system_clock::time_point(
std::chrono::milliseconds(put_start_time_timestamp)),
size, std::move(replicas), enable_soft_pin, is_hard_pinned);
size, std::move(replicas), enable_soft_pin, is_hard_pinned, data_type);
metadata->lease_timeout = std::chrono::system_clock::time_point(
std::chrono::milliseconds(lease_timestamp));

1
mooncake-store/tests/CMakeLists.txt
View File

@ -85,6 +85,7 @@ add_store_test(dummy_client_get_buffer_test dummy_client_get_buffer_test.cpp)
add_store_test(health_check_test health_check_test.cpp)
add_store_test(mmap_arena_test mmap_arena_test.cpp)
add_store_test(mmap_arena_fallback_test mmap_arena_fallback_test.cpp)
add_store_test(object_data_type_test object_data_type_test.cpp)
add_subdirectory(e2e)
add_executable(high_availability_test ha/leadership/high_availability_test.cpp)

157
mooncake-store/tests/object_data_type_test.cpp Normal file
View File

@ -0,0 +1,157 @@
#include "types.h"
#include "replica.h"
#include "master_service.h"
#include <glog/logging.h>
#include <gtest/gtest.h>
#include <sstream>
#include <vector>
namespace mooncake::test {
class ObjectDataTypeTest : public ::testing::Test {
protected:
void SetUp() override {
google::InitGoogleLogging("ObjectDataTypeTest");
FLAGS_logtostderr = true;
}
void TearDown() override { google::ShutdownGoogleLogging(); }
static constexpr size_t kDefaultSegmentBase = 0x300000000;
static constexpr size_t kDefaultSegmentSize = 1024 * 1024 * 16;
Segment MakeSegment(std::string name = "test_segment",
size_t base = kDefaultSegmentBase,
size_t size = kDefaultSegmentSize) const {
Segment segment;
segment.id = generate_uuid();
segment.name = std::move(name);
segment.base = base;
segment.size = size;
segment.te_endpoint = segment.name;
return segment;
}
};
// Verify enum values match the RFC spec
TEST_F(ObjectDataTypeTest, EnumValues) {
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::UNKNOWN), 0);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::KVCACHE), 1);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::TENSOR), 2);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::WEIGHT), 3);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::SAMPLE), 4);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::ACTIVATION), 5);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::GRADIENT), 6);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::OPTIMIZER_STATE), 7);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::METADATA), 8);
EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::GENERAL), 9);
}
// Verify stream operator produces readable output
TEST_F(ObjectDataTypeTest, StreamOperator) {
std::ostringstream oss;
oss << ObjectDataType::KVCACHE;
EXPECT_EQ(oss.str(), "KVCACHE");
oss.str("");
oss << ObjectDataType::UNKNOWN;
EXPECT_EQ(oss.str(), "UNKNOWN");
oss.str("");
oss << ObjectDataType::OPTIMIZER_STATE;
EXPECT_EQ(oss.str(), "OPTIMIZER_STATE");
oss.str("");
oss << ObjectDataType::GENERAL;
EXPECT_EQ(oss.str(), "GENERAL");
// Out-of-range value should print "UNKNOWN"
oss.str("");
oss << static_cast<ObjectDataType>(200);
EXPECT_EQ(oss.str(), "UNKNOWN");
}
// ReplicateConfig defaults to UNKNOWN
TEST_F(ObjectDataTypeTest, ReplicateConfigDefaultDataType) {
ReplicateConfig config;
EXPECT_EQ(config.data_type, ObjectDataType::UNKNOWN);
}
// ReplicateConfig can be set to other types
TEST_F(ObjectDataTypeTest, ReplicateConfigSetDataType) {
ReplicateConfig config;
config.data_type = ObjectDataType::WEIGHT;
EXPECT_EQ(config.data_type, ObjectDataType::WEIGHT);
}
// ReplicateConfig stream output includes data_type
TEST_F(ObjectDataTypeTest, ReplicateConfigStreamIncludesDataType) {
ReplicateConfig config;
config.data_type = ObjectDataType::TENSOR;
std::ostringstream oss;
oss << config;
EXPECT_NE(oss.str().find("data_type: TENSOR"), std::string::npos);
}
// PutStart with data_type propagates to ObjectMetadata
TEST_F(ObjectDataTypeTest, PutStartWithDataType) {
std::unique_ptr<MasterService> service(new MasterService());
Segment segment = MakeSegment();
UUID client_id = generate_uuid();
auto mount_result = service->MountSegment(segment, client_id);
ASSERT_TRUE(mount_result.has_value());
UUID put_client = generate_uuid();
// Put with WEIGHT type
ReplicateConfig config;
config.replica_num = 1;
config.data_type = ObjectDataType::WEIGHT;
auto result = service->PutStart(put_client, "key_weight", 1024, config);
ASSERT_TRUE(result.has_value());
EXPECT_FALSE(result.value().empty());
auto end_result =
service->PutEnd(put_client, "key_weight", ReplicaType::MEMORY);
EXPECT_TRUE(end_result.has_value());
}
// PutStart with default UNKNOWN data_type still works (backward compat)
TEST_F(ObjectDataTypeTest, PutStartDefaultDataType) {
std::unique_ptr<MasterService> service(new MasterService());
Segment segment = MakeSegment();
UUID client_id = generate_uuid();
auto mount_result = service->MountSegment(segment, client_id);
ASSERT_TRUE(mount_result.has_value());
UUID put_client = generate_uuid();
ReplicateConfig config;
config.replica_num = 1;
// data_type left as default (UNKNOWN)
auto result = service->PutStart(put_client, "key_default", 1024, config);
ASSERT_TRUE(result.has_value());
EXPECT_FALSE(result.value().empty());
}
// Verify all enum values can roundtrip through uint8_t cast
TEST_F(ObjectDataTypeTest, EnumRoundtrip) {
std::vector<ObjectDataType> all_types = {
ObjectDataType::UNKNOWN, ObjectDataType::KVCACHE,
ObjectDataType::TENSOR, ObjectDataType::WEIGHT,
ObjectDataType::SAMPLE, ObjectDataType::ACTIVATION,
ObjectDataType::GRADIENT, ObjectDataType::OPTIMIZER_STATE,
ObjectDataType::METADATA, ObjectDataType::GENERAL,
};
for (auto type : all_types) {
uint8_t raw = static_cast<uint8_t>(type);
auto recovered = static_cast<ObjectDataType>(raw);
EXPECT_EQ(type, recovered);
}
}
} // namespace mooncake::test