popov/neuron-tai
1
0
Fork 0
Code Issues 42 Pull Requests Actions Packages Projects Releases Wiki Activity

story: DGR-038 Implement isolated shard-local Hot KV State

Browse Source
This commit is contained in:
Dobromir Popov
2026-08-01 01:32:46 +03:00
parent a1df87deb6
commit 49560b396f
6 changed files with 318 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

142
packages/node/native/worker/llama_shard_engine.cpp
View File

@@ -1,7 +1,9 @@
#include "llama_shard_engine.h"
#include <algorithm>
#include <chrono>
#include <cstdlib>
#include <map>
#include <mutex>
#include <utility>
#include <vector>
@@ -59,6 +61,24 @@ class LlamaShardEngine final : public ShardEngine {
return false;
}
resident_bytes_ = report.resident_bytes;
llama_context_params context_params = llama_context_default_params();
// One worker context owns a bounded set of independent llama sequences.
// The range-loaded model determines the actual local K/V tensors; no
// upstream layer state is ever accepted over the network.
context_params.n_ctx = identity_.hot_kv_budget_tokens;
context_params.n_batch = identity_.hot_kv_context_tokens;
context_params.n_ubatch = identity_.hot_kv_context_tokens;
context_params.n_seq_max = identity_.hot_kv_max_sessions;
context_ = llama_init_from_model(model_, context_params);
if (!context_) {
ShutdownLocked();
*error = "llama.cpp could not allocate the bounded Hot KV context";
return false;
}
free_sequence_ids_.reserve(identity_.hot_kv_max_sessions);
for (uint32_t sequence_id = 0; sequence_id < identity_.hot_kv_max_sessions; ++sequence_id) {
free_sequence_ids_.push_back(static_cast<llama_seq_id>(sequence_id));
}
loaded_ = true;
return true;
}
@@ -106,11 +126,63 @@ class LlamaShardEngine final : public ShardEngine {
return result;
}
bool Execute(const sp::TensorBundle&, std::string* error) override {
HotKvResult OpenSession(const std::string& route_session_id, uint64_t route_epoch) override {
std::lock_guard<std::mutex> lock(mu_);
if (!loaded_ || !context_) return {HotKvStatus::kCacheMiss, 0, "llama.cpp context is not loaded"};
EvictExpiredLocked();
const auto latest = latest_epoch_.find(route_session_id);
if (latest != latest_epoch_.end() && route_epoch < latest->second) {
return {HotKvStatus::kStaleEpoch, 0, "stale route epoch"};
}
const SessionKey key{route_session_id, route_epoch};
if (sessions_.count(key)) return {HotKvStatus::kOk, sessions_[key].past_len, "session already open"};
if (latest != latest_epoch_.end() && route_epoch > latest->second) ReleaseRouteLocked(route_session_id);
while (sessions_.size() >= identity_.hot_kv_max_sessions) EvictLruLocked();
if (free_sequence_ids_.empty()) return {HotKvStatus::kResourceExhausted, 0, "Hot KV sequence budget exhausted"};
const llama_seq_id sequence_id = free_sequence_ids_.back();
free_sequence_ids_.pop_back();
sessions_.emplace(key, SessionState{sequence_id, 0, NowSeconds()});
latest_epoch_[route_session_id] = route_epoch;
return {HotKvStatus::kOk, 0, "Hot KV session opened"};
}
HotKvResult Execute(const HotKvStep& step, const sp::TensorBundle&, std::string* error) override {
std::lock_guard<std::mutex> lock(mu_);
if (!loaded_ || !model_) {
*error = "llama.cpp model is not loaded";
return false;
return {HotKvStatus::kCacheMiss, 0, *error};
}
EvictExpiredLocked();
const auto latest = latest_epoch_.find(step.route_session_id);
if (latest != latest_epoch_.end() && step.route_epoch < latest->second) {
return {HotKvStatus::kStaleEpoch, 0, "stale route epoch"};
}
const SessionKey key{step.route_session_id, step.route_epoch};
auto it = sessions_.find(key);
if (it == sessions_.end()) return {HotKvStatus::kCacheMiss, 0, "Hot KV state was released or evicted"};
SessionState& session = it->second;
if (step.phase == HotKvStep::Phase::kDecode && step.expected_past_len != session.past_len) {
return {HotKvStatus::kCacheMiss, session.past_len, "expected past length does not match local Hot KV"};
}
if (step.first_position < session.past_len) {
// Re-prefill from an earlier position is an explicit truncate, never an
// append over stale positions. This removes only this llama sequence.
llama_memory_seq_rm(llama_get_memory(context_), session.sequence_id,
static_cast<llama_pos>(step.first_position), -1);
total_reserved_tokens_ -= session.past_len - step.first_position;
session.past_len = step.first_position;
}
if (step.first_position != session.past_len || step.token_count == 0) {
return {HotKvStatus::kCacheMiss, session.past_len, "non-contiguous Hot KV append"};
}
if (session.past_len + step.token_count > identity_.hot_kv_context_tokens) {
return {HotKvStatus::kResourceExhausted, session.past_len, "per-session Hot KV context limit exceeded"};
}
while (total_reserved_tokens_ + step.token_count > identity_.hot_kv_budget_tokens && sessions_.size() > 1) {
EvictLruLocked(&key);
}
if (total_reserved_tokens_ + step.token_count > identity_.hot_kv_budget_tokens) {
return {HotKvStatus::kResourceExhausted, session.past_len, "Hot KV token budget exhausted"};
}
if (identity_.injected_death_after_executions != 0 &&
++executions_ >= identity_.injected_death_after_executions) {
@@ -121,7 +193,10 @@ class LlamaShardEngine final : public ShardEngine {
// DGR-038 installs per-session context/KV and DGR-039 proves graph parity.
// Reaching here nevertheless proves every accepted activation is gated by
// the loaded, range-attested llama.cpp engine rather than a fixture.
return true;
session.past_len += step.token_count;
total_reserved_tokens_ += step.token_count;
session.last_used = NowSeconds();
return {HotKvStatus::kOk, session.past_len, "Hot KV append accepted"};
}
const WorkerIdentity& identity() const override { return identity_; }
@@ -129,11 +204,20 @@ class LlamaShardEngine final : public ShardEngine {
std::lock_guard<std::mutex> lock(mu_);
return {loaded_, resident_bytes_, loaded_ ? "llama.cpp model loaded" : "llama.cpp model unavailable"};
}
void ReleaseSession(const std::string&) override {}
void ReleaseSession(const std::string& route_session_id, uint64_t route_epoch) override {
std::lock_guard<std::mutex> lock(mu_);
ReleaseLocked(SessionKey{route_session_id, route_epoch});
}
void Shutdown() override { std::lock_guard<std::mutex> lock(mu_); ShutdownLocked(); }
private:
void ShutdownLocked() {
sessions_.clear();
free_sequence_ids_.clear();
latest_epoch_.clear();
total_reserved_tokens_ = 0;
if (context_) llama_free(context_);
context_ = nullptr;
if (model_) llama_model_free(model_);
model_ = nullptr;
loaded_ = false;
@@ -141,13 +225,63 @@ class LlamaShardEngine final : public ShardEngine {
if (backend_initialized_) llama_backend_free();
backend_initialized_ = false;
}
struct SessionKey {
std::string route_session_id;
uint64_t route_epoch;
bool operator<(const SessionKey& other) const {
return route_session_id != other.route_session_id ? route_session_id < other.route_session_id
: route_epoch < other.route_epoch;
}
};
struct SessionState { llama_seq_id sequence_id; uint64_t past_len; uint64_t last_used; };
static uint64_t NowSeconds() {
return std::chrono::duration_cast<std::chrono::seconds>(std::chrono::steady_clock::now().time_since_epoch()).count();
}
void ReleaseLocked(const SessionKey& key) {
auto it = sessions_.find(key);
if (it == sessions_.end()) return;
llama_memory_seq_rm(llama_get_memory(context_), it->second.sequence_id, -1, -1);
total_reserved_tokens_ -= it->second.past_len;
free_sequence_ids_.push_back(it->second.sequence_id);
sessions_.erase(it);
}
void ReleaseRouteLocked(const std::string& route_session_id) {
for (auto it = sessions_.begin(); it != sessions_.end();) {
if (it->first.route_session_id != route_session_id) { ++it; continue; }
const SessionKey key = it->first;
++it;
ReleaseLocked(key);
}
}
void EvictExpiredLocked() {
const uint64_t cutoff = NowSeconds() - identity_.hot_kv_ttl_seconds;
for (auto it = sessions_.begin(); it != sessions_.end();) {
if (it->second.last_used > cutoff) { ++it; continue; }
const SessionKey key = it->first;
++it;
ReleaseLocked(key);
}
}
void EvictLruLocked(const SessionKey* except = nullptr) {
auto victim = sessions_.end();
for (auto it = sessions_.begin(); it != sessions_.end(); ++it) {
if (except && it->first.route_session_id == except->route_session_id && it->first.route_epoch == except->route_epoch) continue;
if (victim == sessions_.end() || it->second.last_used < victim->second.last_used) victim = it;
}
if (victim != sessions_.end()) ReleaseLocked(victim->first);
}
WorkerIdentity identity_;
mutable std::mutex mu_;
llama_model* model_ = nullptr;
llama_context* context_ = nullptr;
bool backend_initialized_ = false;
bool loaded_ = false;
uint64_t resident_bytes_ = 0;
uint32_t executions_ = 0;
std::map<SessionKey, SessionState> sessions_;
std::map<std::string, uint64_t> latest_epoch_;
std::vector<llama_seq_id> free_sequence_ids_;
uint64_t total_reserved_tokens_ = 0;
};
} // namespace

31
packages/node/native/worker/llama_shard_engine.h
View File

@@ -25,6 +25,10 @@ struct WorkerIdentity {
uint32_t start_layer = 0;
uint32_t end_layer = 0; // half-open, as on the wire
uint32_t injected_death_after_executions = 0; // opt-in test hook; zero disables
uint32_t hot_kv_max_sessions = 8;
uint32_t hot_kv_context_tokens = 4096;
uint32_t hot_kv_budget_tokens = 32768;
uint32_t hot_kv_ttl_seconds = 300;
};
struct BundleCheck {
@@ -38,15 +42,38 @@ struct EngineHealth {
std::string detail;
};
// This is deliberately expressed in tokens, rather than guessed bytes: llama.cpp
// owns the actual K/V layout for the loaded range and backend. The worker uses
// the token reservation to keep its local KV arena bounded before a graph
// adapter materializes the typed boundary (DGR-039).
struct HotKvStep {
enum class Phase { kPrefill, kDecode };
std::string route_session_id;
uint64_t route_epoch = 0;
Phase phase = Phase::kPrefill;
uint64_t first_position = 0;
uint32_t token_count = 0;
uint64_t expected_past_len = 0;
};
enum class HotKvStatus { kOk, kCacheMiss, kStaleEpoch, kResourceExhausted, kCancelled };
struct HotKvResult {
HotKvStatus status = HotKvStatus::kOk;
uint64_t past_len = 0;
std::string detail;
};
class ShardEngine {
public:
virtual ~ShardEngine() = default;
virtual bool Load(std::string* error) = 0;
virtual BundleCheck Validate(const sp::TensorBundle&, uint64_t max_chunk_bytes) const = 0;
virtual bool Execute(const sp::TensorBundle&, std::string* error) = 0;
virtual HotKvResult OpenSession(const std::string& route_session_id, uint64_t route_epoch) = 0;
virtual HotKvResult Execute(const HotKvStep&, const sp::TensorBundle&, std::string* error) = 0;
virtual const WorkerIdentity& identity() const = 0;
virtual EngineHealth health() const = 0;
virtual void ReleaseSession(const std::string& route_session_id) = 0;
virtual void ReleaseSession(const std::string& route_session_id, uint64_t route_epoch) = 0;
virtual void Shutdown() = 0;
};

61
packages/node/native/worker/shard_service.cpp
View File

@@ -70,6 +70,23 @@ void FillDefaultFlow(sp::FlowControl* fc, const FlowLimits& limits) {
fc->set_max_prefill_chunk_tokens(limits.max_prefill_chunk_tokens);
}
sp::SessionResponse HotKvFailure(const std::string& route_session_id, const std::string& work_id,
uint64_t step, const HotKvResult& result) {
switch (result.status) {
case HotKvStatus::kStaleEpoch:
return MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_EPOCH_STALE, result.detail, false, false);
case HotKvStatus::kCacheMiss:
return MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_CACHE_MISS, result.detail, false, true);
case HotKvStatus::kResourceExhausted:
return MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_RESOURCE_EXHAUSTED, result.detail, false, true);
case HotKvStatus::kCancelled:
return MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_CANCELLED, result.detail, false, false);
case HotKvStatus::kOk:
break;
}
return MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL, "unexpected Hot KV result", false, true);
}
} // namespace
grpc::Status ShardRuntimeServiceImpl::GetCapability(grpc::ServerContext*,
@@ -198,6 +215,11 @@ grpc::Status ShardRuntimeServiceImpl::Session(
open.has_proposed_flow_control()
? NegotiateFlow(open.proposed_flow_control())
: limits_;
const HotKvResult hot_kv = engine_.OpenSession(route_session_id, open.route_epoch());
if (hot_kv.status != HotKvStatus::kOk) {
stream->Write(HotKvFailure(route_session_id, "", 0, hot_kv));
return grpc::Status::OK;
}
{
std::lock_guard<std::mutex> lk(sessions_mu_);
SessionState state;
@@ -282,13 +304,17 @@ grpc::Status ShardRuntimeServiceImpl::Session(
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_PAYLOAD_CORRUPT,
*check.corrupt_detail, false, false);
} else {
std::string execution_error;
const HotKvResult executed = engine_.Execute(
HotKvStep{route_session_id, envelope.route_epoch(), HotKvStep::Phase::kPrefill,
envelope.position().first_position(), envelope.position().token_count(),
envelope.cache_expectation().expected_past_len()},
chunk.bundle(), &execution_error);
if (executed.status != HotKvStatus::kOk) {
response = HotKvFailure(route_session_id, work_id, step, executed);
} else {
state->seen_steps.insert(step);
state->credits -= 1;
std::string execution_error;
if (!engine_.Execute(chunk.bundle(), &execution_error)) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL,
execution_error, false, true);
} else {
*response.mutable_chunk() = chunk; // echo the exact bundle back
}
}
@@ -348,13 +374,15 @@ grpc::Status ShardRuntimeServiceImpl::Session(
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_PAYLOAD_CORRUPT,
*check.corrupt_detail, false, false);
} else {
std::string execution_error;
const HotKvResult executed = engine_.Execute(
HotKvStep{route_session_id, state->epoch, HotKvStep::Phase::kDecode,
step_msg.position(), 1, step_msg.expected_past_len()}, bundle, &execution_error);
if (executed.status != HotKvStatus::kOk) {
response = HotKvFailure(route_session_id, work_id, step, executed);
} else {
state->seen_steps.insert(step);
state->credits -= 1;
std::string execution_error;
if (!engine_.Execute(bundle, &execution_error)) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL,
execution_error, false, true);
} else {
// No decode response field exists; echo the step back as a
// chunk-bearing SessionResponse per the proto's relayed-frame design.
sp::ActivationChunk* out = response.mutable_chunk();
@@ -412,8 +440,11 @@ grpc::Status ShardRuntimeServiceImpl::Session(
// freed the moment the terminal status is sent.
{
std::lock_guard<std::mutex> lk(sessions_mu_);
sessions_.erase(route_session_id);
engine_.ReleaseSession(route_session_id);
auto it = sessions_.find(release.route_session_id());
if (it != sessions_.end() && it->second.epoch == release.route_epoch()) {
sessions_.erase(it);
}
engine_.ReleaseSession(release.route_session_id(), release.route_epoch());
}
sp::SessionResponse response;
sp::ShardStatus* status = response.mutable_status();
@@ -454,8 +485,10 @@ grpc::Status ShardRuntimeServiceImpl::Release(grpc::ServerContext*,
bool existed;
{
std::lock_guard<std::mutex> lk(sessions_mu_);
existed = sessions_.erase(request->route_session_id()) != 0;
engine_.ReleaseSession(request->route_session_id());
auto it = sessions_.find(request->route_session_id());
existed = it != sessions_.end() && it->second.epoch == request->route_epoch();
if (existed) sessions_.erase(it);
engine_.ReleaseSession(request->route_session_id(), request->route_epoch());
}
response->set_released(existed);
return grpc::Status::OK;

21
packages/node/native/worker/shard_worker_main.cpp
View File

@@ -64,6 +64,19 @@ meshnet::worker::FlowLimits LimitsFromEnv() {
return limits;
}
bool PositiveEnv(const char* name, uint32_t* out, std::string* error) {
if (const char* value = std::getenv(name)) {
char* end = nullptr;
const unsigned long parsed = std::strtoul(value, &end, 10);
if (end == value || *end != '\0' || parsed == 0 || parsed > UINT32_MAX) {
*error = std::string("invalid ") + name;
return false;
}
*out = static_cast<uint32_t>(parsed);
}
return true;
}
bool IdentityFromEnv(meshnet::worker::WorkerIdentity* identity, std::string* error) {
const auto required = [&](const char* name, std::string* out) -> bool {
const char* value = std::getenv(name);
@@ -100,6 +113,14 @@ bool IdentityFromEnv(meshnet::worker::WorkerIdentity* identity, std::string* err
}
identity->injected_death_after_executions = static_cast<uint32_t>(parsed);
}
if (!PositiveEnv("MESHNET_HOT_KV_MAX_SESSIONS", &identity->hot_kv_max_sessions, error) ||
!PositiveEnv("MESHNET_HOT_KV_CONTEXT_TOKENS", &identity->hot_kv_context_tokens, error) ||
!PositiveEnv("MESHNET_HOT_KV_BUDGET_TOKENS", &identity->hot_kv_budget_tokens, error) ||
!PositiveEnv("MESHNET_HOT_KV_TTL_SECONDS", &identity->hot_kv_ttl_seconds, error)) return false;
if (identity->hot_kv_budget_tokens < identity->hot_kv_context_tokens) {
*error = "MESHNET_HOT_KV_BUDGET_TOKENS must cover one session context";
return false;
}
return true;
}