97 lines
3.7 KiB
C++
97 lines
3.7 KiB
C++
// The native Shard worker's ShardRuntime service (DGR-033).
|
|
//
|
|
// A faithful C++ port of `ShardRuntimeServicer` in `shard_runtime_server.py`:
|
|
// the same per-`route_session_id` identity/credit/dedup state, the same
|
|
// fail-closed negative paths (stale epoch, expired deadline, corrupt/oversize
|
|
// payload, exhausted flow-control credit, duplicate idempotency step, in-band
|
|
// and out-of-band cancellation), and the same lifecycle (open/prefill/decode/
|
|
// flow-control/release/cancel). The only compute it does is the fake engine's
|
|
// bounded forward — there is no llama.cpp linkage and no arbitrary-graph RPC.
|
|
|
|
#ifndef MESHNET_NATIVE_WORKER_SHARD_SERVICE_H_
|
|
#define MESHNET_NATIVE_WORKER_SHARD_SERVICE_H_
|
|
|
|
#include <cstdint>
|
|
#include <map>
|
|
#include <mutex>
|
|
#include <set>
|
|
#include <string>
|
|
|
|
#include <grpcpp/grpcpp.h>
|
|
|
|
#include "llama_shard_engine.h"
|
|
#include "shard_runtime.grpc.pb.h"
|
|
#include "shard_runtime.pb.h"
|
|
|
|
namespace meshnet::worker {
|
|
|
|
namespace sp = ::meshnet::shard::v1;
|
|
|
|
struct FlowLimits {
|
|
uint32_t credits_granted = 16;
|
|
uint32_t max_inflight_chunks = 16;
|
|
uint64_t max_chunk_bytes = 4u * 1024u * 1024u;
|
|
uint32_t max_prefill_chunk_tokens = 512;
|
|
};
|
|
|
|
// Per-route-session identity/credit/dedup state, kept on the servicer instance
|
|
// (guarded by a lock) so an out-of-band unary Cancel from a different handler
|
|
// thread can reach a session a concurrent Session stream is still iterating.
|
|
struct SessionState {
|
|
uint64_t epoch = 0;
|
|
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:
|
|
ShardRuntimeServiceImpl(FlowLimits limits, ShardEngine& engine) : limits_(limits), engine_(engine) {}
|
|
|
|
grpc::Status GetCapability(grpc::ServerContext* context,
|
|
const sp::CapabilityRequest* request,
|
|
sp::CapabilityReport* response) override;
|
|
|
|
grpc::Status Health(grpc::ServerContext* context, const sp::HealthRequest* request,
|
|
sp::HealthReport* response) override;
|
|
|
|
grpc::Status Session(
|
|
grpc::ServerContext* context,
|
|
grpc::ServerReaderWriter<sp::SessionResponse, sp::SessionRequest>* stream) override;
|
|
|
|
grpc::Status Release(grpc::ServerContext* context, const sp::ReleaseRequest* request,
|
|
sp::ReleaseResponse* response) override;
|
|
|
|
grpc::Status Cancel(grpc::ServerContext* context, const sp::CancelRequest* request,
|
|
sp::CancelResponse* response) override;
|
|
|
|
private:
|
|
// Returns the number of items newly marked cancelled, creating session state
|
|
// 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_;
|
|
ShardEngine& engine_;
|
|
std::mutex sessions_mu_;
|
|
std::map<std::string, SessionState> sessions_;
|
|
};
|
|
|
|
} // namespace meshnet::worker
|
|
|
|
#endif // MESHNET_NATIVE_WORKER_SHARD_SERVICE_H_
|