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

fix: DGR-033 repair native worker protocol per cross-review BLOCK

Browse Source 
Address the Codex GPT-5.5 review of the standalone fake C++ gRPC Shard
worker. Four root protocol defects fixed:

- Fail closed before SessionOpen: a per-session `opened` flag gates
  chunk/decode so no activation bypasses lifecycle, cancellation, epoch
  or flow-control state (terminal ERROR_CODE_INTERNAL), even when an
  out-of-band Cancel created placeholder state.
- Strict flow-control negotiation: NegotiateFlow takes the strictest of
  peer-vs-worker bounds (mirrors codec.negotiate_flow_control) and the
  negotiated per-session max_chunk_bytes is enforced on every bundle
  instead of trusting the peer proposal.
- In-stream ReleaseSignal now erases session state immediately.
- SessionOpen rejects incompatible schema, fingerprint, and shard-range
  identity and reports the worker's own served fingerprint rather than
  echoing the caller.

Adds 9 regression tests (worker suite 18 -> 27). Real gates on the
rebuilt pinned-gRPC binary: cmake build exit 0; ctest 2/2; worker
pytest 27 passed; harness+protocol 63 passed; compileall 0; diff --check
clean; ldd/nm show 0 llama/ggml linkage. DGR-033 passes -> true.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-07-26 22:57:03 +03:00
parent c073826374
commit 7473bb7e44
6 changed files with 399 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
.scratch/distributed-gguf-runtime/evidence/DGR-033/README.md
View File

@@ -1,6 +1,7 @@
# DGR-033 evidence — standalone fake C++ gRPC Shard worker
**Completed:** 2026-07-25
**Completed:** 2026-07-25 (initial); **repaired:** 2026-07-26 after Codex
GPT-5.5 cross-review BLOCK (see "Cross-review repair" below).
**Branch:** `ralph/distributed-gguf-opus`
**Authority:** `.scratch/distributed-gguf-runtime/prd.json`
**Dependencies:** DGR-022 (lifecycle/status contract), DGR-024 (real generated
@@ -206,3 +207,75 @@ DGR-029/030); the Python client uses that venv's `grpcio==1.82.1`,
supervision primitives — a readiness line for start detection, `SIGTERM`
graceful drain with a clean-exit line, and a `--selftest` liveness probe.
A supervisor can start/monitor/restart the process around these.
## Cross-review repair (2026-07-26)
An independent Codex GPT-5.5 review BLOCKED the initial implementation. Four
root protocol defects in the native worker were fixed in this worktree
(`.claude/worktrees/distributed-gguf-opus`); the fake-engine echo semantics and
supervision shape are unchanged.
### Defects fixed
1. **Activation before SessionOpen bypassed all state.** A chunk/decode whose
`route_session_id` had no opened session fell through every `if (state && ...)`
guard and was echoed — bypassing lifecycle, cancellation, epoch and
flow-control. `SessionState` now carries an `opened` flag set only by a valid
`SessionOpen`; chunk and decode fail closed with a terminal
`ERROR_CODE_INTERNAL` and end the stream when it is false. A placeholder state
created by an out-of-band `Cancel` that races `Open` has `opened == false`, so
it can never admit work either.
2. **Flow control blindly trusted the peer proposal.** `SessionOpen` copied the
proposed `credits/max_inflight/max_chunk_bytes` verbatim into session state and
the accepted reply. New `ShardRuntimeServiceImpl::NegotiateFlow` takes the
strictest bound of peer-vs-worker for every field (mirroring
`negotiate_flow_control` in `native_protocol/codec.py`), stores the negotiated
ceilings on the session, and enforces the negotiated per-session
`max_chunk_bytes` on every bundle (`FakeShardEngine::Validate` now takes the
ceiling as an argument instead of a fixed construction-time value).
3. **In-stream `ReleaseSignal` leaked session state.** The stream `release` arm
wrote a terminal status but never dropped the session. It now erases the
session under the lock before responding, so KV/credits/dedup are freed
immediately (the out-of-band `Release` RPC already erased).
4. **`SessionOpen` echoed caller identity instead of validating it.** The handshake
now rejects an incompatible `schema_version` (`SCHEMA_UNSUPPORTED`), a
mismatched model/recipe `Fingerprint` (`FINGERPRINT_MISMATCH`), and a
`ShardRange` outside the worker's served range (`SHARD_RANGE_MISMATCH`), each
terminal; `SessionAccepted` now reports the worker's own served fingerprint
rather than a copy of the caller's.
### Changed files (repair)
- `packages/node/native/worker/shard_service.h` — `opened` +
`max_prefill_chunk_tokens` on `SessionState`; `NegotiateFlow` decl; engine now
default-constructed.
- `packages/node/native/worker/shard_service.cpp` — worker-identity constants +
fill helpers; `NegotiateFlow`; `SessionOpen` validation/negotiation; fail-closed
chunk/decode; per-session `max_chunk_bytes`; in-stream release erase.
- `packages/node/native/worker/fake_engine.h` — `Validate(bundle, max_chunk_bytes)`.
- `tests/test_native_shard_worker.py` — extended `_open` (schema/fingerprint/range/
flow overrides); fixed `test_release_rpc_is_idempotent` for the new erase
semantics; added 9 regression tests (chunk/decode before open, flow-control
clamp, negotiated-ceiling cap, in-stream release erase, schema/fingerprint/range
rejection, worker-fingerprint-not-caller).
### Re-run gates (real, rebuilt binary)
Build driven through the pinned `cmake` (Unix Makefiles + `gmake`, gRPC 1.82.1):
```text
cmake --build build/native --parallel 8 -> BUILD_EXIT 0
ctest --test-dir build/native --output-on-failure -> 100% (2/2) passed
shard_worker_selftest ....... Passed
shard_protocol_conformance .. Passed
python -m pytest -q tests/test_native_shard_worker.py -> 27 passed
python -m pytest -q tests/test_shard_runtime_harness.py \
tests/test_native_shard_protocol.py -> 63 passed
python -m compileall -q packages tests -> exit 0
git diff --check -> clean
ldd build/native/shard_worker | grep -iE 'llama|ggml' -> NONE
nm -C build/native/shard_worker | grep -cE 'llama_|ggml_' -> 0
```
The worker integration suite grew from 18 to 27 tests; all pass against the
freshly compiled binary. No `.ralph-lane` runtime artifacts were touched.

4
.scratch/distributed-gguf-runtime/prd.json
View File

@@ -656,13 +656,13 @@
"The worker exposes neither llama.cpp RPC nor arbitrary graph execution.",
"Applicable shared quality gates in `prd.json` pass, and the evidence handoff records exact commands/results, changed files, limitations, and dependency handoff."
],
"passes": false,
"passes": true,
"notes": "Generated source issue: .scratch/distributed-gguf-runtime/issues/033-build-a-standalone-fake-c-grpc-shard-worker.md; prd.json is authoritative.",
"blocks": [
"DGR-036",
"DGR-040"
],
"completionNotes": "Cross-review BLOCKED the initial implementation on pre-open activation acceptance, non-strict flow-control negotiation, in-stream release state retention, and fail-closed identity/version/range validation. Repair required before completion credit."
"completionNotes": "Cross-review (Codex GPT-5.5) BLOCK repaired in worktree distributed-gguf-opus. Root protocol defects fixed in the native worker: (1) chunk/decode now fail closed before SessionOpen via a per-session opened flag (terminal ERROR_CODE_INTERNAL), so no activation bypasses lifecycle/cancellation/epoch/flow-control state even when an out-of-band Cancel created placeholder state; (2) flow control is negotiated with strict worker bounds (ShardRuntimeServiceImpl::NegotiateFlow mirrors native_protocol/codec.py negotiate_flow_control) and the negotiated per-session max_chunk_bytes is enforced on every bundle instead of trusting the peer proposal; (3) an in-stream ReleaseSignal now erases session state immediately; (4) SessionOpen rejects incompatible schema, artifact/recipe fingerprint, and shard-range identity and reports the worker own served fingerprint rather than echoing the caller. Nine regression tests added. Real gates on the rebuilt pinned-gRPC binary: cmake --build exit 0; ctest 2/2 passed (shard_worker_selftest, shard_protocol_conformance); tests/test_native_shard_worker.py 27 passed; DGR-024 harness + native protocol 63 passed; compileall exit 0; git diff --check clean; ldd/nm show 0 llama/ggml linkage. Evidence: .scratch/distributed-gguf-runtime/evidence/DGR-033/README.md."
},
{
"id": "DGR-034",

15
packages/node/native/worker/fake_engine.h
View File

@@ -78,18 +78,22 @@ class FakeShardEngine {
// (DGR-036) can assert this is a fixture, not a real engine.
static constexpr const char* kEvidenceClass = "fixture";
explicit FakeShardEngine(uint64_t max_chunk_bytes) : max_chunk_bytes_(max_chunk_bytes) {}
FakeShardEngine() = default;
BundleCheck Validate(const sp::TensorBundle& bundle) const {
// `max_chunk_bytes` is the per-session *negotiated* ceiling (the strictest of
// the worker's own limit and the peer's proposal), passed in on every call so
// the engine enforces exactly what the SessionOpen handshake settled — never a
// value the peer proposed unilaterally.
BundleCheck Validate(const sp::TensorBundle& bundle, uint64_t max_chunk_bytes) const {
BundleCheck result;
for (const auto& tensor : bundle.tensors()) {
// Bounded message: a declared payload larger than the ceiling is refused
// before any reassembly work is done.
if (max_chunk_bytes_ != 0 && tensor.total_bytes() > max_chunk_bytes_) {
if (max_chunk_bytes != 0 && tensor.total_bytes() > max_chunk_bytes) {
result.oversize_detail =
"tensor '" + tensor.name() + "': declared total_bytes " +
std::to_string(tensor.total_bytes()) + " exceeds max_chunk_bytes " +
std::to_string(max_chunk_bytes_);
std::to_string(max_chunk_bytes);
return result;
}
@@ -154,9 +158,6 @@ class FakeShardEngine {
}
return digest;
}
private:
uint64_t max_chunk_bytes_;
};
} // namespace meshnet::worker

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

@@ -1,5 +1,6 @@
#include "shard_service.h"
#include <algorithm>
#include <chrono>
#include <utility>
@@ -7,6 +8,42 @@ namespace meshnet::worker {
namespace {
// The exact identity this fixture worker serves. SessionOpen is validated
// against these — not echoed back from the caller — so an incompatible peer
// fails closed at open rather than being silently accepted with its own claimed
// identity. Kept in one place so GetCapability and the open handshake agree.
constexpr const char* kModelArtifactDigest = "sha256:native-test-artifact";
constexpr const char* kRuntimeRecipeDigest = "sha256:native-test-recipe";
constexpr const char* kRecipeId = "native-test";
constexpr const char* kRecipeVersion = "1";
constexpr const char* kCatalogueVersion = "1";
constexpr uint32_t kShardStartLayer = 0;
constexpr uint32_t kShardEndLayer = 32;
constexpr uint32_t kShardEffectiveStartLayer = 0;
void FillWorkerFingerprint(sp::Fingerprint* fp) {
fp->set_model_artifact_digest(kModelArtifactDigest);
fp->set_runtime_recipe_digest(kRuntimeRecipeDigest);
fp->set_recipe_id(kRecipeId);
fp->set_recipe_version(kRecipeVersion);
fp->set_catalogue_version(kCatalogueVersion);
}
void FillWorkerShardRange(sp::ShardRange* range) {
range->set_start_layer(kShardStartLayer);
range->set_end_layer(kShardEndLayer);
range->set_effective_start_layer(kShardEffectiveStartLayer);
}
// Strictest-of-both bound: the smallest positive of `a`/`b`, or `fallback` when
// neither is set. Mirrors the `_min` helper in `native_protocol/codec.py`.
uint64_t MinPositive(uint64_t a, uint64_t b, uint64_t fallback) {
if (a > 0 && b > 0) return std::min(a, b);
if (a > 0) return a;
if (b > 0) return b;
return fallback;
}
int64_t NowUnixNanos() {
return std::chrono::duration_cast<std::chrono::nanoseconds>(
std::chrono::system_clock::now().time_since_epoch())
@@ -52,16 +89,8 @@ grpc::Status ShardRuntimeServiceImpl::GetCapability(grpc::ServerContext*,
const sp::CapabilityRequest*,
sp::CapabilityReport* response) {
response->set_schema_version(sp::SCHEMA_VERSION_1);
sp::Fingerprint* fp = response->mutable_fingerprint();
fp->set_model_artifact_digest("sha256:native-test-artifact");
fp->set_runtime_recipe_digest("sha256:native-test-recipe");
fp->set_recipe_id("native-test");
fp->set_recipe_version("1");
fp->set_catalogue_version("1");
sp::ShardRange* range = response->mutable_shard_range();
range->set_start_layer(0);
range->set_end_layer(32);
range->set_effective_start_layer(0);
FillWorkerFingerprint(response->mutable_fingerprint());
FillWorkerShardRange(response->mutable_shard_range());
response->set_backend("grpc-native-cpp");
response->set_device("cpu");
response->set_validated(true);
@@ -88,6 +117,22 @@ grpc::Status ShardRuntimeServiceImpl::Health(grpc::ServerContext*, const sp::Hea
return grpc::Status::OK;
}
FlowLimits ShardRuntimeServiceImpl::NegotiateFlow(const sp::FlowControl& proposed) const {
FlowLimits out;
out.max_inflight_chunks = static_cast<uint32_t>(MinPositive(
proposed.max_inflight_chunks(), limits_.max_inflight_chunks, limits_.max_inflight_chunks));
const uint64_t credits = MinPositive(proposed.credits_granted(), limits_.credits_granted,
limits_.credits_granted);
out.credits_granted =
static_cast<uint32_t>(std::min<uint64_t>(credits, out.max_inflight_chunks));
out.max_chunk_bytes =
MinPositive(proposed.max_chunk_bytes(), limits_.max_chunk_bytes, limits_.max_chunk_bytes);
out.max_prefill_chunk_tokens = static_cast<uint32_t>(MinPositive(
proposed.max_prefill_chunk_tokens(), limits_.max_prefill_chunk_tokens,
limits_.max_prefill_chunk_tokens));
return out;
}
uint32_t ShardRuntimeServiceImpl::MarkCancelled(const std::string& route_session_id,
const std::string& work_id) {
std::lock_guard<std::mutex> lk(sessions_mu_);
@@ -119,20 +164,59 @@ grpc::Status ShardRuntimeServiceImpl::Session(
case sp::SessionRequest::kOpen: {
const sp::SessionOpen& open = request.open();
route_session_id = open.route_session_id();
// Reject an incompatible peer at open rather than mid-generation. The
// worker validates the caller's schema, artifact/recipe identity and
// requested layer range against its own — it never adopts the caller's
// claimed identity.
auto reject_open = [&](sp::ErrorCode code, const std::string& detail) {
stream->Write(MakeFail(route_session_id, /*work_id=*/"", /*step=*/0, code, detail,
/*terminal=*/true, /*retryable=*/false));
};
if (open.schema_version() != sp::SCHEMA_VERSION_1) {
reject_open(sp::ERROR_CODE_SCHEMA_UNSUPPORTED,
"worker serves schema version 1 only");
return grpc::Status::OK;
}
const sp::Fingerprint& fp = open.fingerprint();
if ((!fp.model_artifact_digest().empty() &&
fp.model_artifact_digest() != kModelArtifactDigest) ||
(!fp.runtime_recipe_digest().empty() &&
fp.runtime_recipe_digest() != kRuntimeRecipeDigest)) {
reject_open(sp::ERROR_CODE_FINGERPRINT_MISMATCH,
"model artifact or runtime recipe digest does not match this worker");
return grpc::Status::OK;
}
if (open.has_shard_range()) {
const sp::ShardRange& r = open.shard_range();
const bool within = r.start_layer() >= kShardStartLayer &&
r.end_layer() <= kShardEndLayer &&
r.start_layer() < r.end_layer() &&
r.effective_start_layer() >= r.start_layer() &&
r.effective_start_layer() < r.end_layer();
if (!within) {
reject_open(sp::ERROR_CODE_SHARD_RANGE_MISMATCH,
"requested layer range is not served by this worker");
return grpc::Status::OK;
}
}
// Settle the flow-control window with strict worker bounds, then keep
// the negotiated ceilings on the session so every later check enforces
// exactly what was agreed — not what the peer proposed.
const FlowLimits negotiated =
open.has_proposed_flow_control()
? NegotiateFlow(open.proposed_flow_control())
: limits_;
{
std::lock_guard<std::mutex> lk(sessions_mu_);
SessionState state;
state.epoch = open.route_epoch();
if (open.has_proposed_flow_control()) {
const sp::FlowControl& fc = open.proposed_flow_control();
state.credits = fc.credits_granted();
state.max_inflight = fc.max_inflight_chunks();
state.max_chunk_bytes = fc.max_chunk_bytes();
} else {
state.credits = limits_.credits_granted;
state.max_inflight = limits_.max_inflight_chunks;
state.max_chunk_bytes = limits_.max_chunk_bytes;
}
state.credits = negotiated.credits_granted;
state.max_inflight = negotiated.max_inflight_chunks;
state.max_chunk_bytes = negotiated.max_chunk_bytes;
state.max_prefill_chunk_tokens = negotiated.max_prefill_chunk_tokens;
state.opened = true;
auto it = sessions_.find(route_session_id);
if (it != sessions_.end()) {
// A prior out-of-band Cancel may have marked this session cancelled
@@ -147,11 +231,7 @@ grpc::Status ShardRuntimeServiceImpl::Session(
accepted->set_schema_version(sp::SCHEMA_VERSION_1);
accepted->set_route_session_id(open.route_session_id());
accepted->set_route_epoch(open.route_epoch());
if (open.has_proposed_flow_control()) {
*accepted->mutable_flow_control() = open.proposed_flow_control();
} else {
FillDefaultFlow(accepted->mutable_flow_control(), limits_);
}
FillDefaultFlow(accepted->mutable_flow_control(), negotiated);
if (open.accepted_compression_size() > 0) {
for (int c : open.accepted_compression()) {
accepted->add_accepted_compression(static_cast<sp::Compression>(c));
@@ -159,7 +239,9 @@ grpc::Status ShardRuntimeServiceImpl::Session(
} else {
accepted->add_accepted_compression(sp::COMPRESSION_NONE);
}
*accepted->mutable_fingerprint() = open.fingerprint();
// Report the fingerprint the worker actually serves, so a mismatch is
// visible at open — never a copy of the caller's claimed identity.
FillWorkerFingerprint(accepted->mutable_fingerprint());
stream->Write(response);
break;
}
@@ -174,28 +256,35 @@ grpc::Status ShardRuntimeServiceImpl::Session(
// holding the lock across a (possibly blocking) Write would deadlock an
// out-of-band Cancel RPC that needs the same lock.
sp::SessionResponse response;
bool terminate = false;
{
std::lock_guard<std::mutex> lk(sessions_mu_);
auto it = sessions_.find(route_session_id);
SessionState* state = it != sessions_.end() ? &it->second : nullptr;
if (state && (state->cancelled_session || state->cancelled_work.count(work_id))) {
if (state == nullptr || !state->opened) {
// Fail closed: an activation before a valid SessionOpen must never
// bypass lifecycle, cancellation, epoch or flow-control state.
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL,
"activation received before SessionOpen", true, false);
terminate = true;
} else if (state->cancelled_session || state->cancelled_work.count(work_id)) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_CANCELLED,
"work was cancelled", false, false);
} else if (state && envelope.route_epoch() < state->epoch) {
} else if (envelope.route_epoch() < state->epoch) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_EPOCH_STALE,
"stale route epoch", false, false);
} else if (envelope.deadline_unix_nanos() != 0 &&
NowUnixNanos() > envelope.deadline_unix_nanos()) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_DEADLINE_EXCEEDED,
"deadline already passed", false, false);
} else if (state && state->seen_steps.count(step)) {
} else if (state->seen_steps.count(step)) {
response = MakeAck(work_id, step, /*duplicate=*/true);
} else if (state && state->credits <= 0) {
} else if (state->credits <= 0) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_FLOW_CONTROL_VIOLATION,
"no flow-control credit remaining", false, true);
} else {
const BundleCheck check = engine_.Validate(chunk.bundle());
const BundleCheck check = engine_.Validate(chunk.bundle(), state->max_chunk_bytes);
if (check.oversize_detail) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_RESOURCE_EXHAUSTED,
*check.oversize_detail, false, false);
@@ -203,16 +292,17 @@ grpc::Status ShardRuntimeServiceImpl::Session(
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_PAYLOAD_CORRUPT,
*check.corrupt_detail, false, false);
} else {
if (state) {
state->seen_steps.insert(step);
state->credits -= 1;
}
engine_.BoundedForward(chunk.bundle()); // real bounded forward over wire bytes
*response.mutable_chunk() = chunk; // echo the exact bundle back
}
}
}
stream->Write(response);
if (terminate) {
return grpc::Status::OK;
}
break;
}
@@ -230,25 +320,32 @@ grpc::Status ShardRuntimeServiceImpl::Session(
}
sp::SessionResponse response;
bool terminate = false;
{
std::lock_guard<std::mutex> lk(sessions_mu_);
auto it = sessions_.find(route_session_id);
SessionState* state = it != sessions_.end() ? &it->second : nullptr;
if (state && (state->cancelled_session || state->cancelled_work.count(work_id))) {
if (state == nullptr || !state->opened) {
// Fail closed: a decode step before a valid SessionOpen must never
// bypass lifecycle, cancellation, epoch or flow-control state.
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL,
"activation received before SessionOpen", true, false);
terminate = true;
} else if (state->cancelled_session || state->cancelled_work.count(work_id)) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_CANCELLED,
"work was cancelled", false, false);
} else if (step_msg.deadline_unix_nanos() != 0 &&
NowUnixNanos() > step_msg.deadline_unix_nanos()) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_DEADLINE_EXCEEDED,
"deadline already passed", false, false);
} else if (state && state->seen_steps.count(step)) {
} else if (state->seen_steps.count(step)) {
response = MakeAck(work_id, step, /*duplicate=*/true);
} else if (state && state->credits <= 0) {
} else if (state->credits <= 0) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_FLOW_CONTROL_VIOLATION,
"no flow-control credit remaining", false, true);
} else {
const BundleCheck check = engine_.Validate(bundle);
const BundleCheck check = engine_.Validate(bundle, state->max_chunk_bytes);
if (check.oversize_detail) {
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_RESOURCE_EXHAUSTED,
*check.oversize_detail, false, false);
@@ -256,10 +353,8 @@ grpc::Status ShardRuntimeServiceImpl::Session(
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_PAYLOAD_CORRUPT,
*check.corrupt_detail, false, false);
} else {
if (state) {
state->seen_steps.insert(step);
state->credits -= 1;
}
engine_.BoundedForward(bundle);
// No decode response field exists; echo the step back as a
// chunk-bearing SessionResponse per the proto's relayed-frame design.
@@ -277,6 +372,9 @@ grpc::Status ShardRuntimeServiceImpl::Session(
}
}
stream->Write(response);
if (terminate) {
return grpc::Status::OK;
}
break;
}
@@ -308,6 +406,14 @@ grpc::Status ShardRuntimeServiceImpl::Session(
case sp::SessionRequest::kRelease: {
const sp::ReleaseSignal& release = request.release();
// An explicit release drops session state immediately (KV, credits,
// dedup) instead of holding it for the TTL — the whole point of the
// signal. Erase the session this stream opened so its resources are
// freed the moment the terminal status is sent.
{
std::lock_guard<std::mutex> lk(sessions_mu_);
sessions_.erase(route_session_id);
}
sp::SessionResponse response;
sp::ShardStatus* status = response.mutable_status();
status->set_work_id(release.work_id());

15
packages/node/native/worker/shard_service.h
View File

@@ -42,15 +42,20 @@ struct SessionState {
int64_t credits = 0;
uint32_t max_inflight = 0;
uint64_t max_chunk_bytes = 0;
uint32_t max_prefill_chunk_tokens = 0;
std::set<uint64_t> seen_steps;
std::set<std::string> cancelled_work;
bool cancelled_session = false;
// True only after a valid SessionOpen handshake completed for this
// route_session_id. An activation (chunk/decode) that arrives while this is
// false fails closed: no work may bypass the lifecycle handshake, even when a
// placeholder state already exists from an out-of-band Cancel that raced Open.
bool opened = false;
};
class ShardRuntimeServiceImpl final : public sp::ShardRuntime::Service {
public:
explicit ShardRuntimeServiceImpl(FlowLimits limits)
: limits_(limits), engine_(limits.max_chunk_bytes) {}
explicit ShardRuntimeServiceImpl(FlowLimits limits) : limits_(limits) {}
grpc::Status GetCapability(grpc::ServerContext* context,
const sp::CapabilityRequest* request,
@@ -74,6 +79,12 @@ class ShardRuntimeServiceImpl final : public sp::ShardRuntime::Service {
// if the Cancel raced ahead of SessionOpen.
uint32_t MarkCancelled(const std::string& route_session_id, const std::string& work_id);
// Settle a stream's flow-control window against this worker's own limits: the
// strictest bound of either peer wins for every field, so a peer can never
// raise the worker's ceilings by proposing a larger window. Mirrors
// `negotiate_flow_control` in `native_protocol/codec.py`.
FlowLimits NegotiateFlow(const sp::FlowControl& proposed) const;
FlowLimits limits_;
FakeShardEngine engine_;
std::mutex sessions_mu_;

166
tests/test_native_shard_worker.py
View File

@@ -139,24 +139,43 @@ def _crc32c(payload: bytes) -> bytes:
return zlib.crc32(payload).to_bytes(4, "big")
def _open(*, route_session_id="rs-1", route_epoch=7, credits_granted=16) -> pb.SessionRequest:
return pb.SessionRequest(
open=pb.SessionOpen(
schema_version=pb.SCHEMA_VERSION_1,
route_session_id=route_session_id,
route_epoch=route_epoch,
fingerprint=pb.Fingerprint(
_WORKER_FINGERPRINT = dict(
model_artifact_digest="sha256:native-test-artifact",
runtime_recipe_digest="sha256:native-test-recipe",
recipe_id="native-test",
recipe_version="1",
catalogue_version="1",
),
shard_range=pb.ShardRange(start_layer=0, end_layer=32, effective_start_layer=0),
proposed_flow_control=pb.FlowControl(
credits_granted=credits_granted,
)
def _open(
*,
route_session_id="rs-1",
route_epoch=7,
credits_granted=16,
max_inflight_chunks=16,
max_chunk_bytes=4 * 1024 * 1024,
schema_version=pb.SCHEMA_VERSION_1,
fingerprint=None,
shard_range=None,
) -> pb.SessionRequest:
fp = pb.Fingerprint(**_WORKER_FINGERPRINT) if fingerprint is None else fingerprint
sr = (
pb.ShardRange(start_layer=0, end_layer=32, effective_start_layer=0)
if shard_range is None
else shard_range
)
return pb.SessionRequest(
open=pb.SessionOpen(
schema_version=schema_version,
route_session_id=route_session_id,
route_epoch=route_epoch,
fingerprint=fp,
shard_range=sr,
proposed_flow_control=pb.FlowControl(
credits_granted=credits_granted,
max_inflight_chunks=max_inflight_chunks,
max_chunk_bytes=max_chunk_bytes,
max_prefill_chunk_tokens=512,
),
accepted_compression=[pb.COMPRESSION_NONE],
@@ -434,9 +453,9 @@ def test_out_of_band_cancel_rpc_races_ahead_of_open(worker):
def test_release_rpc_is_idempotent(worker):
stub = worker.stub()
# Open a session so state exists, then release it out of band twice.
worker.session([_open(route_session_id="rs-rel"), _release()])
# (release signal in-stream does not erase state; the unary Release RPC does)
# Open a session WITHOUT an in-stream release so state persists on the
# servicer, then drop it out of band twice.
worker.session([_open(route_session_id="rs-rel")])
first = stub.Release(pb.ReleaseRequest(schema_version=pb.SCHEMA_VERSION_1, route_session_id="rs-rel", route_epoch=7))
second = stub.Release(pb.ReleaseRequest(schema_version=pb.SCHEMA_VERSION_1, route_session_id="rs-rel", route_epoch=7))
assert first.released is True
@@ -496,3 +515,122 @@ def test_direct_and_opaque_relay_yield_identical_responses(worker):
assert len(direct_resp) == len(relay_resp) == 3
for i, (d, r) in enumerate(zip(direct_resp, relay_resp)):
assert d == r, f"response #{i} differs between direct and opaque relay"
# --- fail-closed before SessionOpen ----------------------------------------
def test_chunk_before_open_is_rejected(worker):
# An activation with no preceding SessionOpen must fail closed and end the
# stream: no work may bypass the lifecycle handshake.
responses = worker.session([_chunk("w-noopen", b"payload", step=1)])
assert len(responses) == 1
assert responses[0].WhichOneof("kind") == "status"
assert responses[0].status.error.code == pb.ERROR_CODE_INTERNAL
assert responses[0].status.terminal is True
assert "SessionOpen" in responses[0].status.error.detail
def test_decode_before_open_is_rejected(worker):
responses = worker.session([_decode("w-noopen", b"payload", step=1, position=0)])
assert len(responses) == 1
assert responses[0].WhichOneof("kind") == "status"
assert responses[0].status.error.code == pb.ERROR_CODE_INTERNAL
assert responses[0].status.terminal is True
# --- flow-control negotiation with strict worker bounds --------------------
def test_flow_control_proposal_is_clamped_to_worker_bounds(worker):
# A peer proposing a window far above the worker limits must be clamped to
# the worker own ceilings, never granted the inflated proposal.
responses = worker.session(
[_open(credits_granted=9999, max_inflight_chunks=9999, max_chunk_bytes=1073741824)]
)
fc = responses[0].accepted.flow_control
assert fc.max_inflight_chunks == 16
assert fc.credits_granted == 16
assert fc.max_chunk_bytes == 4 * 1024 * 1024
def test_negotiated_max_chunk_bytes_caps_peer_proposal():
# Worker ceiling is 64 bytes; the peer proposes 4 MiB. The negotiated per
# session ceiling is the stricter 64, so a 128-byte tensor is refused even
# though the peer allowed it — the worker never adopts the peer proposal.
w = _Worker(extra_env={"MESHNET_MAX_CHUNK_BYTES": "64"})
try:
big = b"x" * 128
responses = w.session(
[_open(max_chunk_bytes=4 * 1024 * 1024), _chunk("w-big", big, step=1, total_bytes=128)]
)
assert responses[0].accepted.flow_control.max_chunk_bytes == 64
assert responses[1].status.error.code == pb.ERROR_CODE_RESOURCE_EXHAUSTED
assert "max_chunk_bytes" in responses[1].status.error.detail
finally:
if w.proc.poll() is None:
w.close()
# --- in-stream release erases session state --------------------------------
def test_in_stream_release_erases_session_state(worker):
stub = worker.stub()
resp = worker.session(
[
_open(route_session_id="rs-erase"),
pb.SessionRequest(
release=pb.ReleaseSignal(route_session_id="rs-erase", route_epoch=7, work_id="w-final")
),
]
)
assert resp[-1].status.terminal is True
# The state is already gone: an out-of-band Release finds nothing to drop.
after = stub.Release(
pb.ReleaseRequest(schema_version=pb.SCHEMA_VERSION_1, route_session_id="rs-erase", route_epoch=7)
)
assert after.released is False
# --- SessionOpen identity validation ---------------------------------------
def test_incompatible_schema_is_rejected_at_open(worker):
responses = worker.session([_open(schema_version=pb.SCHEMA_VERSION_UNSPECIFIED)])
assert responses[0].WhichOneof("kind") == "status"
assert responses[0].status.error.code == pb.ERROR_CODE_SCHEMA_UNSUPPORTED
assert responses[0].status.terminal is True
def test_incompatible_fingerprint_is_rejected_at_open(worker):
bad_fp = pb.Fingerprint(
model_artifact_digest="sha256:some-other-model",
runtime_recipe_digest="sha256:native-test-recipe",
recipe_id="native-test",
recipe_version="1",
catalogue_version="1",
)
responses = worker.session([_open(fingerprint=bad_fp)])
assert responses[0].WhichOneof("kind") == "status"
assert responses[0].status.error.code == pb.ERROR_CODE_FINGERPRINT_MISMATCH
assert responses[0].status.terminal is True
def test_shard_range_mismatch_is_rejected_at_open(worker):
responses = worker.session(
[_open(shard_range=pb.ShardRange(start_layer=0, end_layer=64, effective_start_layer=0))]
)
assert responses[0].WhichOneof("kind") == "status"
assert responses[0].status.error.code == pb.ERROR_CODE_SHARD_RANGE_MISMATCH
assert responses[0].status.terminal is True
def test_session_accepted_reports_worker_fingerprint_not_caller(worker):
# The caller asserts no fingerprint; SessionAccepted must carry the worker
# OWN served identity, not a copy of the caller (empty) fingerprint.
responses = worker.session([_open(fingerprint=pb.Fingerprint())])
assert responses[0].WhichOneof("kind") == "accepted"
accepted = responses[0].accepted
assert accepted.fingerprint.model_artifact_digest == "sha256:native-test-artifact"
assert accepted.fingerprint.runtime_recipe_digest == "sha256:native-test-recipe"