Skip to content

[Store] Add ObjectDataType enum for type-aware metadata #1719

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stmatengss merged 7 commits into from
May 13, 2026
Merged

[Store] Add ObjectDataType enum for type-aware metadata #1719

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
db0c825
feat: add ObjectDataType enum and metadata propagation (Phase 1 of #1…
he-yufeng Mar 22, 2026
85d9dad
style: fix clang-format violations
he-yufeng Mar 23, 2026
fd963c0
ci: retrigger CI after transient infra failures
he-yufeng Mar 25, 2026
093daa4
retrigger CI
he-yufeng Mar 26, 2026
dd8683f
ci: retrigger sglang integration test
he-yufeng Mar 26, 2026
10f5819
address review: add GENERAL type, clarify deserialization logic
he-yufeng Mar 28, 2026
1190d8f
style: fix clang-format issues
he-yufeng Mar 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions mooncake-integration/store/store_py.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1574,6 +1574,20 @@ class MooncakeHostMemAllocatorPyWrapper {
};

PYBIND11_MODULE(store, m) {
// Object data type classification
py::enum_<ObjectDataType>(m, "ObjectDataType")
.value("UNKNOWN", ObjectDataType::UNKNOWN)

@stmatengss stmatengss Mar 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we should add a general type?

.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<>())
Expand All @@ -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;
Expand Down
10 changes: 7 additions & 3 deletions mooncake-store/include/master_service.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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
Expand Down Expand Up @@ -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");
}
Expand All @@ -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;
}

Expand Down
4 changes: 3 additions & 1 deletion mooncake-store/include/replica.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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;
}
};
Expand Down
41 changes: 41 additions & 0 deletions mooncake-store/include/types.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <limits>
#include <unordered_map>
#include <utility>
Expand Down Expand Up @@ -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;
}
Comment on lines +148 to +165

Copilot AI Mar 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

types.h now uses std::string_view in the ObjectDataType stream operator, but the header does not include <string_view> (and also relies on indirect includes for std::ostream). This makes the header non-self-contained and can cause build breaks for translation units that include types.h directly. Add the missing standard includes in types.h.

Copilot uses AI. Check for mistakes.

// Forward declarations
class BufferAllocatorBase;
class CachelibBufferAllocator;
Expand Down
47 changes: 35 additions & 12 deletions mooncake-store/src/master_service.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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);

Expand Down Expand Up @@ -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));

Comment on lines 4339 to +4388

Copilot AI Mar 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Snapshot serialization now always writes the new array shape (8 + replicas_count) by inserting data_type. While the new deserializer can read both formats, this is only backward-compatible in the "new code reads old snapshots" direction. If any mixed-version deployment expects older binaries to load snapshots produced by newer binaries, they will likely fail on the extra field. Consider adding an explicit snapshot format/version marker or a compatibility mode to write the old format during rolling upgrades, and/or clarify the upgrade guarantees.

Copilot uses AI. Check for mistakes.
// Serialize replicas
for (const auto& replica : metadata.GetAllReplicas()) {
auto result = Serializer<Replica>::serialize(
Expand All @@ -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,
Expand Down Expand Up @@ -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);
Expand All @@ -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));

Expand Down
1 change: 1 addition & 0 deletions mooncake-store/tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,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)
Expand Down
Loading
Loading