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

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;