diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-033/README.md b/.scratch/distributed-gguf-runtime/evidence/DGR-033/README.md index 8b8408a..6d4875e 100644 --- a/.scratch/distributed-gguf-runtime/evidence/DGR-033/README.md +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-033/README.md @@ -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. diff --git a/.scratch/distributed-gguf-runtime/prd.json b/.scratch/distributed-gguf-runtime/prd.json index 7f86fae..90425d5 100644 --- a/.scratch/distributed-gguf-runtime/prd.json +++ b/.scratch/distributed-gguf-runtime/prd.json @@ -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", diff --git a/packages/node/native/worker/fake_engine.h b/packages/node/native/worker/fake_engine.h index 6da4b38..bdbbc3a 100644 --- a/packages/node/native/worker/fake_engine.h +++ b/packages/node/native/worker/fake_engine.h @@ -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 diff --git a/packages/node/native/worker/shard_service.cpp b/packages/node/native/worker/shard_service.cpp index d72c295..c3875f3 100644 --- a/packages/node/native/worker/shard_service.cpp +++ b/packages/node/native/worker/shard_service.cpp @@ -1,5 +1,6 @@ #include "shard_service.h" +#include #include #include @@ -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::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(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(std::min(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(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 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 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(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 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; - } + 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 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; - } + 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 lk(sessions_mu_); + sessions_.erase(route_session_id); + } sp::SessionResponse response; sp::ShardStatus* status = response.mutable_status(); status->set_work_id(release.work_id()); diff --git a/packages/node/native/worker/shard_service.h b/packages/node/native/worker/shard_service.h index a53cfad..b9fe470 100644 --- a/packages/node/native/worker/shard_service.h +++ b/packages/node/native/worker/shard_service.h @@ -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 seen_steps; std::set 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_; diff --git a/tests/test_native_shard_worker.py b/tests/test_native_shard_worker.py index bad615a..6ce60db 100644 --- a/tests/test_native_shard_worker.py +++ b/tests/test_native_shard_worker.py @@ -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: +_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", +) + + +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=pb.SCHEMA_VERSION_1, + schema_version=schema_version, route_session_id=route_session_id, route_epoch=route_epoch, - fingerprint=pb.Fingerprint( - 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), + fingerprint=fp, + shard_range=sr, proposed_flow_control=pb.FlowControl( credits_granted=credits_granted, - max_inflight_chunks=16, - max_chunk_bytes=4 * 1024 * 1024, + 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"