Skip to content

[Store] feat: add workload-aware KV cache hints #2830

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

Open
Lin-z-w wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Store] feat: add workload-aware KV cache hints #2830

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 7 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 @@ -1826,6 +1826,12 @@ PYBIND11_MODULE(store, m) {
.value("GENERAL", ObjectDataType::GENERAL)
.export_values();

py::class_<WorkloadHints>(m, "WorkloadHints")
.def(py::init<>())
.def_readwrite("session_id", &WorkloadHints::session_id)
.def_readwrite("retention_priority",
&WorkloadHints::retention_priority);

// Define the ReplicateConfig class
py::class_<ReplicateConfig>(m, "ReplicateConfig")
.def(py::init<>())
Expand All @@ -1842,6 +1848,7 @@ PYBIND11_MODULE(store, m) {
&ReplicateConfig::prefer_alloc_in_same_node)
.def_readwrite("data_type", &ReplicateConfig::data_type)
.def_readwrite("group_ids", &ReplicateConfig::group_ids)
.def_readwrite("workload_hints", &ReplicateConfig::workload_hints)
.def("__str__", [](const ReplicateConfig &config) {
std::ostringstream oss;
oss << config;
Expand Down
5 changes: 1 addition & 4 deletions mooncake-integration/store/store_py_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -861,10 +861,7 @@ bool parallelism_specs_equal_by_kind(const TensorParallelismSpec &lhs,
}

bool is_default_replicate_config(const ReplicateConfig &config) {
return config.replica_num == 1 && !config.with_soft_pin &&
!config.with_hard_pin && config.preferred_segments.empty() &&
config.preferred_segment.empty() &&
!config.prefer_alloc_in_same_node && !config.group_ids.has_value();
return config.IsDefault();
}

std::optional<ParallelAxisSpec> parse_parallel_axis_spec(
Expand Down
6 changes: 5 additions & 1 deletion mooncake-store/include/master_service.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ struct MetadataStoragePlugin;
namespace test {
class MasterServiceSnapshotTestBase;
class SnapshotChildProcessTest;
class MasterServiceTest;
// Friended so the promotion-on-hit tests can drive a serialize/reset/
// deserialize cycle directly via the otherwise-private
// MetadataSerializer, and inspect private clamp fields. This avoids
Expand Down Expand Up @@ -85,6 +86,7 @@ class MasterService {
// Test friend class for snapshot/restore testing
friend class test::MasterServiceSnapshotTestBase;
friend class test::SnapshotChildProcessTest;
friend class test::MasterServiceTest;
friend class test::PromotionOnHitTest;
friend class benchmarks::BatchEvictBench;
friend class test::MasterServiceTenantQuotaTest;
Expand Down Expand Up @@ -894,14 +896,15 @@ class MasterService {
bool enable_soft_pin, bool enable_hard_pin = false,
ObjectDataType data_type_ = ObjectDataType::UNKNOWN,
std::string group_id_ = "", std::string tenant_id_ = "default",
std::string user_key_ = {})
std::string user_key_ = {}, WorkloadHints workload_hints_ = {})
: client_id(client_id_),
put_start_time(put_start_time_),
size(value_length),
data_type(data_type_),
group_id(std::move(group_id_)),
tenant_id(std::move(tenant_id_)),
user_key(std::move(user_key_)),
workload_hints(std::move(workload_hints_)),
lease_timeout(),
soft_pin_timeout(std::nullopt),
hard_pinned(enable_hard_pin),
Expand All @@ -928,6 +931,7 @@ class MasterService {
const std::string group_id;
const std::string tenant_id;
const std::string user_key;
const WorkloadHints workload_hints;

mutable SpinLock lock;
// Default constructor, creates a time_point representing
Expand Down
23 changes: 23 additions & 0 deletions mooncake-store/include/replica.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,23 @@ inline std::ostream& operator<<(std::ostream& os,
return os;
}

/**
* @brief Workload-level hints associated with an object.
*/
struct WorkloadHints {
std::string session_id{}; // Empty means unset.
int32_t retention_priority{0}; // Zero means unset.

bool operator==(const WorkloadHints&) const = default;

friend std::ostream& operator<<(std::ostream& os,
const WorkloadHints& hints) noexcept {
Comment on lines +87 to +88

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

medium

Marking operator<< as noexcept is unsafe because stream insertion operations (such as writing to os or copying/formatting hints.session_id) can throw exceptions (e.g., if the stream has its exception mask configured to throw on failure, or due to allocation failures). If an exception is thrown inside a noexcept function, std::terminate will be called immediately. It is standard practice to omit noexcept from stream insertion operators.

Suggested change
friend std::ostream& operator<<(std::ostream& os,
const WorkloadHints& hints) noexcept {
friend std::ostream& operator<<(std::ostream& os,
const WorkloadHints& hints) {

return os << "{ session_id: " << hints.session_id
<< ", retention_priority: " << hints.retention_priority
<< " }";
}
};

/**
* @brief Configuration for replica management
*/
Expand All @@ -96,6 +113,11 @@ struct ReplicateConfig {
// ungrouped. Grouped keys share metadata routing, coalesced lease refresh,
// and memory eviction behavior.
std::optional<std::vector<std::string>> group_ids{};
WorkloadHints workload_hints{};

bool operator==(const ReplicateConfig&) const = default;

bool IsDefault() const { return *this == ReplicateConfig{}; }

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

medium

To avoid constructing a temporary ReplicateConfig object (which may involve initializing strings, vectors, and optionals) on every call to IsDefault(), you can compare against a static const instance.

    bool IsDefault() const {
        static const ReplicateConfig kDefault{};
        return *this == kDefault;
    }


ReplicateConfig ForSingleKey(size_t key_index) const {
ReplicateConfig key_config = *this;
Expand Down Expand Up @@ -142,6 +164,7 @@ struct ReplicateConfig {
}
os << "]";
}
os << ", workload_hints: " << config.workload_hints;
os << " }";
return os;
}
Expand Down
6 changes: 6 additions & 0 deletions mooncake-store/include/rpc_service.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,15 @@

namespace mooncake {

namespace test {
class MasterServiceTest;
}

// Forward declaration
class HttpMetadataServer;
class WrappedMasterService {
friend class test::MasterServiceTest;

public:
// Constructor with optional metadata-cleanup-on-timeout configuration.
// - http_metadata_server: in-process pointer used when the HTTP metadata
Expand Down
6 changes: 4 additions & 2 deletions mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,17 +135,19 @@ DeserializeStandbyObjectMetadata(
// compatibility with MasterService::MetadataSerializer, which appends
// them over time:
// data_type (positive int) appears before the replicas;
// hard_pinned (bool) and group_id (str) trail them.
// hard_pinned (bool), group_id (str), and workload_hints (array)
// trail them.
// v1: 7 + replica_count, no optional fields
// v2: 8 + replica_count, either data_type or hard_pinned
// v3: 9 + replica_count, data_type + hard_pinned or
// hard_pinned + group_id
// v4: 10 + replica_count, data_type + hard_pinned + group_id
// v5: 11 + replica_count, v4 + workload_hints
// 64-bit arithmetic keeps an attacker-controlled near-UINT32_MAX
// replica_count from wrapping the bounds and slipping an out-of-bounds
// index through.
constexpr uint64_t kBaseFieldCount = 7;
constexpr uint64_t kMaxOptionalFieldCount = 3;
constexpr uint64_t kMaxOptionalFieldCount = 4;
const uint64_t total_elements = object.via.array.size;
const uint64_t min_elements = kBaseFieldCount + replica_count;
if (total_elements < min_elements ||
Expand Down
55 changes: 47 additions & 8 deletions mooncake-store/src/master_service.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3012,7 +3012,8 @@ auto MasterService::AllocateAndInsertMetadata(
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.data_type, group_id, tenant_id, key));
config.data_type, group_id, tenant_id, key,
config.workload_hints));
if (!inserted) {
LOG(INFO) << "key=" << key << ", info=object_already_exists";
abort_reserved_quota();
Expand Down Expand Up @@ -3713,6 +3714,7 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key,
merged_config.with_hard_pin || metadata.IsHardPinned();
merged_config.with_soft_pin =
merged_config.with_soft_pin || metadata.IsSoftPinned();
merged_config.workload_hints = metadata.workload_hints;

const std::string existing_group_id = metadata.group_id;
const uint64_t old_quota_charge =
Expand Down Expand Up @@ -8557,7 +8559,8 @@ MasterService::MetadataSerializer::DeserializeShard(const msgpack::object& obj,
metadata_ptr->size, metadata_ptr->PopReplicas(),
metadata_ptr->soft_pin_timeout.has_value(),
metadata_ptr->IsHardPinned(), metadata_ptr->data_type,
metadata_ptr->group_id, tenant_id, user_key));
metadata_ptr->group_id, tenant_id, user_key,
metadata_ptr->workload_hints));

it->second.lease_timeout = metadata_ptr->lease_timeout;
it->second.soft_pin_timeout = metadata_ptr->soft_pin_timeout;
Expand All @@ -8580,11 +8583,13 @@ MasterService::MetadataSerializer::SerializeMetadata(
// 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, data_type,
// replicas..., hard_pinned, group_id]
// replicas..., hard_pinned, group_id,
// [workload_session_id, workload_retention_priority]]

size_t array_size = 10; // client_id, put_start_time, size, lease_timeout,
size_t array_size = 11; // client_id, put_start_time, size, lease_timeout,
// has_soft_pin_timeout, soft_pin_timeout,
// replicas_count, data_type, hard_pinned, group_id
// replicas_count, data_type, hard_pinned,
// group_id, workload_hints
array_size += metadata.CountReplicas(); // One element per replica
packer.pack_array(array_size);

Expand Down Expand Up @@ -8638,6 +8643,9 @@ MasterService::MetadataSerializer::SerializeMetadata(

packer.pack(metadata.IsHardPinned());
packer.pack(metadata.group_id);
packer.pack_array(2);
packer.pack(metadata.workload_hints.session_id);
packer.pack(metadata.workload_hints.retention_priority);

return {};
}
Expand Down Expand Up @@ -8690,12 +8698,14 @@ MasterService::MetadataSerializer::DeserializeMetadata(
// v1: 7 + replicas_count, no optional fields
// v2: 8 + replicas_count, either data_type or hard_pinned
// v3: 9 + replicas_count, data_type + hard_pinned or hard_pinned +
// group_id v4: 10 + replicas_count, data_type + hard_pinned + group_id
// group_id
// v4: 10 + replicas_count, data_type + hard_pinned + group_id
// v5: 11 + replicas_count, v4 + nested workload_hints
// 64-bit arithmetic keeps an attacker-controlled near-UINT32_MAX
// replicas_count from wrapping the bounds and slipping an out-of-bounds
// index past the size check.
constexpr uint64_t kBaseFieldCount = 7;
constexpr uint64_t kMaxOptionalFieldCount = 3;
constexpr uint64_t kMaxOptionalFieldCount = 4;
const uint64_t total_elements = obj.via.array.size;
const uint64_t min_elements = kBaseFieldCount + replicas_count;
if (total_elements < min_elements ||
Expand Down Expand Up @@ -8745,14 +8755,43 @@ MasterService::MetadataSerializer::DeserializeMetadata(
group_id = array[index++].as<std::string>();
}

WorkloadHints workload_hints;
if (index < total_elements) {
const auto& hints_object = array[index++];
if (hints_object.type != msgpack::type::ARRAY ||
hints_object.via.array.size != 2) {
Comment on lines +8759 to +8762

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

medium

To ensure robust forward-compatibility, we should check if the next element is actually an array before consuming it as workload_hints. If a future version of the snapshot format omits workload_hints but appends other optional fields of different types, unconditionally advancing index and expecting an array will cause deserialization to fail. Checking the type first allows the deserializer to safely skip workload_hints and ignore unknown trailing fields.

    if (index < total_elements && array[index].type == msgpack::type::ARRAY) {
        const auto& hints_object = array[index++];
        if (hints_object.via.array.size != 2) {

return tl::unexpected(SerializationError(
ErrorCode::DESERIALIZE_FAIL,
"deserialize workload_hints must be a two-element array"));
}

const auto* hints_array = hints_object.via.array.ptr;
if (hints_array[0].type != msgpack::type::STR ||
(hints_array[1].type != msgpack::type::POSITIVE_INTEGER &&
hints_array[1].type != msgpack::type::NEGATIVE_INTEGER)) {
return tl::unexpected(SerializationError(
ErrorCode::DESERIALIZE_FAIL,
"deserialize workload_hints has invalid field types"));
}

try {
workload_hints.session_id = hints_array[0].as<std::string>();
workload_hints.retention_priority = hints_array[1].as<int32_t>();
} catch (const std::exception& e) {
return tl::unexpected(SerializationError(
ErrorCode::DESERIALIZE_FAIL,
"deserialize workload_hints failed: " + std::string(e.what())));
}
}

// Create ObjectMetadata instance
bool enable_soft_pin = has_soft_pin_timeout;
auto metadata = std::make_unique<ObjectMetadata>(
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, data_type,
group_id);
group_id, "default", std::string{}, std::move(workload_hints));
metadata->lease_timeout = std::chrono::system_clock::time_point(
std::chrono::milliseconds(lease_timestamp));

Expand Down
12 changes: 9 additions & 3 deletions mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,13 +177,19 @@ TEST_P(CatalogBackedSnapshotProviderTest,
}

TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotWithGroupId) {
// 10 + replica_count: current writer format (data_type + hard_pinned +
// trailing group_id). Regression test for the live snapshot restore
// failure against the latest metadata layout.
// 10 + replica_count: data_type + hard_pinned + trailing group_id.
PublishSnapshotPayload(SnapshotMetadataFormat::kWithGroupId);
ExpectLoadsDefaultObject();
}

TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotWithWorkloadHints) {
// 11 + replica_count: current writer format, with workload hints encoded
// as one trailing [session_id, retention_priority] array. The standby does
// not use the hints, but it must accept snapshots that contain them.
PublishSnapshotPayload(SnapshotMetadataFormat::kWithWorkloadHints);
ExpectLoadsDefaultObject();
}

TEST_P(CatalogBackedSnapshotProviderTest, RejectsOverflowingReplicaCount) {
// A near-UINT32_MAX replica_count must not wrap the format-detection
// arithmetic into a valid-looking total and slip an out-of-bounds index
Expand Down
Loading
Loading