365 lines
16 KiB
C++
365 lines
16 KiB
C++
#include "shard_service.h"
|
|
|
|
#include <chrono>
|
|
#include <utility>
|
|
|
|
namespace meshnet::worker {
|
|
|
|
namespace {
|
|
|
|
int64_t NowUnixNanos() {
|
|
return std::chrono::duration_cast<std::chrono::nanoseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count();
|
|
}
|
|
|
|
// Build the standard fail response (a terminal-or-not ShardStatus).
|
|
sp::SessionResponse MakeFail(const std::string& route_session_id, const std::string& work_id,
|
|
uint64_t step, sp::ErrorCode code, const std::string& detail,
|
|
bool terminal, bool retryable) {
|
|
sp::SessionResponse response;
|
|
sp::ShardStatus* status = response.mutable_status();
|
|
status->set_work_id(work_id);
|
|
status->set_route_session_id(route_session_id);
|
|
status->set_idempotency_step(step);
|
|
status->set_terminal(terminal);
|
|
sp::ShardError* error = status->mutable_error();
|
|
error->set_code(code);
|
|
error->set_detail(detail);
|
|
error->set_retryable(retryable);
|
|
return response;
|
|
}
|
|
|
|
sp::SessionResponse MakeAck(const std::string& work_id, uint64_t step, bool duplicate) {
|
|
sp::SessionResponse response;
|
|
sp::Ack* ack = response.mutable_ack();
|
|
ack->set_work_id(work_id);
|
|
ack->set_idempotency_step(step);
|
|
ack->set_duplicate(duplicate);
|
|
return response;
|
|
}
|
|
|
|
void FillDefaultFlow(sp::FlowControl* fc, const FlowLimits& limits) {
|
|
fc->set_credits_granted(limits.credits_granted);
|
|
fc->set_max_inflight_chunks(limits.max_inflight_chunks);
|
|
fc->set_max_chunk_bytes(limits.max_chunk_bytes);
|
|
fc->set_max_prefill_chunk_tokens(limits.max_prefill_chunk_tokens);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
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);
|
|
response->set_backend("grpc-native-cpp");
|
|
response->set_device("cpu");
|
|
response->set_validated(true);
|
|
response->set_detail("bounded real forward passed for fixture artifact");
|
|
response->set_max_concurrent_sessions(8);
|
|
response->set_max_context_tokens(131072);
|
|
FillDefaultFlow(response->mutable_flow_control(), limits_);
|
|
response->add_accepted_compression(sp::COMPRESSION_NONE);
|
|
response->add_supported_schema_versions(sp::SCHEMA_VERSION_1);
|
|
response->set_validated_at_unix_nanos(0);
|
|
return grpc::Status::OK;
|
|
}
|
|
|
|
grpc::Status ShardRuntimeServiceImpl::Health(grpc::ServerContext*, const sp::HealthRequest*,
|
|
sp::HealthReport* response) {
|
|
response->set_schema_version(sp::SCHEMA_VERSION_1);
|
|
response->set_state(sp::SERVING_STATE_SERVING);
|
|
response->set_active_sessions(1);
|
|
response->set_queued_chunks(0);
|
|
response->set_batch_occupancy(0);
|
|
response->set_kv_pressure(0.0f);
|
|
response->set_resident_bytes(0);
|
|
response->set_detail("native fixture worker serving");
|
|
return grpc::Status::OK;
|
|
}
|
|
|
|
uint32_t ShardRuntimeServiceImpl::MarkCancelled(const std::string& route_session_id,
|
|
const std::string& work_id) {
|
|
std::lock_guard<std::mutex> lk(sessions_mu_);
|
|
SessionState& state = sessions_[route_session_id]; // creates on first cancel-before-open
|
|
if (state.max_inflight == 0) {
|
|
// Freshly created placeholder for a Cancel that raced ahead of Open.
|
|
state.credits = limits_.credits_granted;
|
|
state.max_inflight = limits_.max_inflight_chunks;
|
|
state.max_chunk_bytes = limits_.max_chunk_bytes;
|
|
}
|
|
if (work_id.empty()) {
|
|
const bool already = state.cancelled_session;
|
|
state.cancelled_session = true;
|
|
return already ? 0 : 1;
|
|
}
|
|
const bool already = state.cancelled_work.count(work_id) != 0;
|
|
state.cancelled_work.insert(work_id);
|
|
return already ? 0 : 1;
|
|
}
|
|
|
|
grpc::Status ShardRuntimeServiceImpl::Session(
|
|
grpc::ServerContext*,
|
|
grpc::ServerReaderWriter<sp::SessionResponse, sp::SessionRequest>* stream) {
|
|
std::string route_session_id;
|
|
sp::SessionRequest request;
|
|
|
|
while (stream->Read(&request)) {
|
|
switch (request.kind_case()) {
|
|
case sp::SessionRequest::kOpen: {
|
|
const sp::SessionOpen& open = request.open();
|
|
route_session_id = open.route_session_id();
|
|
{
|
|
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;
|
|
}
|
|
auto it = sessions_.find(route_session_id);
|
|
if (it != sessions_.end()) {
|
|
// A prior out-of-band Cancel may have marked this session cancelled
|
|
// before Open arrived; preserve that so the work still fails closed.
|
|
state.cancelled_session = it->second.cancelled_session;
|
|
state.cancelled_work = it->second.cancelled_work;
|
|
}
|
|
sessions_[route_session_id] = std::move(state);
|
|
}
|
|
sp::SessionResponse response;
|
|
sp::SessionAccepted* accepted = response.mutable_accepted();
|
|
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_);
|
|
}
|
|
if (open.accepted_compression_size() > 0) {
|
|
for (int c : open.accepted_compression()) {
|
|
accepted->add_accepted_compression(static_cast<sp::Compression>(c));
|
|
}
|
|
} else {
|
|
accepted->add_accepted_compression(sp::COMPRESSION_NONE);
|
|
}
|
|
*accepted->mutable_fingerprint() = open.fingerprint();
|
|
stream->Write(response);
|
|
break;
|
|
}
|
|
|
|
case sp::SessionRequest::kChunk: {
|
|
const sp::ActivationChunk& chunk = request.chunk();
|
|
const sp::Envelope& envelope = chunk.envelope();
|
|
const std::string work_id = envelope.work_id();
|
|
const uint64_t step = envelope.idempotency_step();
|
|
|
|
// Compute the response under the lock, then write it *after* releasing —
|
|
// holding the lock across a (possibly blocking) Write would deadlock an
|
|
// out-of-band Cancel RPC that needs the same lock.
|
|
sp::SessionResponse response;
|
|
{
|
|
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))) {
|
|
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) {
|
|
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)) {
|
|
response = MakeAck(work_id, step, /*duplicate=*/true);
|
|
} else if (state && 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());
|
|
if (check.oversize_detail) {
|
|
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_RESOURCE_EXHAUSTED,
|
|
*check.oversize_detail, false, false);
|
|
} else if (check.corrupt_detail) {
|
|
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);
|
|
break;
|
|
}
|
|
|
|
case sp::SessionRequest::kDecode: {
|
|
const sp::DecodeStep& step_msg = request.decode();
|
|
const std::string work_id = step_msg.work_id();
|
|
const uint64_t step = step_msg.idempotency_step();
|
|
|
|
sp::TensorBundle bundle;
|
|
if (step_msg.bundle().tensors_size() > 0) {
|
|
bundle = step_msg.bundle();
|
|
} else {
|
|
bundle.set_bundle_version(1);
|
|
*bundle.add_tensors() = step_msg.tensor();
|
|
}
|
|
|
|
sp::SessionResponse response;
|
|
{
|
|
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))) {
|
|
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)) {
|
|
response = MakeAck(work_id, step, /*duplicate=*/true);
|
|
} else if (state && 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);
|
|
if (check.oversize_detail) {
|
|
response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_RESOURCE_EXHAUSTED,
|
|
*check.oversize_detail, false, false);
|
|
} else if (check.corrupt_detail) {
|
|
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.
|
|
sp::ActivationChunk* out = response.mutable_chunk();
|
|
sp::Envelope* out_env = out->mutable_envelope();
|
|
out_env->set_schema_version(sp::SCHEMA_VERSION_1);
|
|
out_env->set_work_id(work_id);
|
|
out_env->set_idempotency_step(step);
|
|
out_env->set_phase(sp::PHASE_DECODE);
|
|
sp::PositionSpan* pos = out_env->mutable_position();
|
|
pos->set_first_position(step_msg.position());
|
|
pos->set_token_count(1);
|
|
*out->mutable_bundle() = bundle;
|
|
}
|
|
}
|
|
}
|
|
stream->Write(response);
|
|
break;
|
|
}
|
|
|
|
case sp::SessionRequest::kFlowControl: {
|
|
const uint32_t topup = request.flow_control().credits_granted();
|
|
sp::SessionResponse response;
|
|
{
|
|
std::lock_guard<std::mutex> lk(sessions_mu_);
|
|
auto it = sessions_.find(route_session_id);
|
|
sp::FlowControl* fc = response.mutable_flow_control();
|
|
if (it != sessions_.end()) {
|
|
SessionState& state = it->second;
|
|
int64_t granted = std::min<int64_t>(state.credits + topup,
|
|
static_cast<int64_t>(state.max_inflight));
|
|
state.credits = granted;
|
|
fc->set_credits_granted(static_cast<uint32_t>(granted));
|
|
fc->set_max_inflight_chunks(state.max_inflight);
|
|
fc->set_max_chunk_bytes(state.max_chunk_bytes);
|
|
} else {
|
|
fc->set_credits_granted(topup != 0 ? topup : limits_.credits_granted);
|
|
fc->set_max_inflight_chunks(limits_.max_inflight_chunks);
|
|
fc->set_max_chunk_bytes(limits_.max_chunk_bytes);
|
|
}
|
|
fc->set_max_prefill_chunk_tokens(limits_.max_prefill_chunk_tokens);
|
|
}
|
|
stream->Write(response);
|
|
break;
|
|
}
|
|
|
|
case sp::SessionRequest::kRelease: {
|
|
const sp::ReleaseSignal& release = request.release();
|
|
sp::SessionResponse response;
|
|
sp::ShardStatus* status = response.mutable_status();
|
|
status->set_work_id(release.work_id());
|
|
status->set_route_session_id(release.route_session_id());
|
|
status->set_terminal(true);
|
|
stream->Write(response);
|
|
return grpc::Status::OK;
|
|
}
|
|
|
|
case sp::SessionRequest::kCancel: {
|
|
const sp::CancelSignal& signal = request.cancel();
|
|
MarkCancelled(route_session_id, signal.work_id());
|
|
const bool whole_session = signal.work_id().empty();
|
|
stream->Write(MakeFail(route_session_id, signal.work_id(), 0, sp::ERROR_CODE_CANCELLED,
|
|
signal.reason().empty() ? "cancelled" : signal.reason(),
|
|
whole_session, false));
|
|
if (whole_session) {
|
|
return grpc::Status::OK;
|
|
}
|
|
break;
|
|
}
|
|
|
|
default: {
|
|
sp::SessionResponse response;
|
|
response.mutable_status()->set_terminal(true);
|
|
stream->Write(response);
|
|
return grpc::Status::OK;
|
|
}
|
|
}
|
|
}
|
|
return grpc::Status::OK;
|
|
}
|
|
|
|
grpc::Status ShardRuntimeServiceImpl::Release(grpc::ServerContext*,
|
|
const sp::ReleaseRequest* request,
|
|
sp::ReleaseResponse* response) {
|
|
bool existed;
|
|
{
|
|
std::lock_guard<std::mutex> lk(sessions_mu_);
|
|
existed = sessions_.erase(request->route_session_id()) != 0;
|
|
}
|
|
response->set_released(existed);
|
|
return grpc::Status::OK;
|
|
}
|
|
|
|
grpc::Status ShardRuntimeServiceImpl::Cancel(grpc::ServerContext*,
|
|
const sp::CancelRequest* request,
|
|
sp::CancelResponse* response) {
|
|
const uint32_t newly = MarkCancelled(request->route_session_id(), request->work_id());
|
|
response->set_cancelled_work_items(newly);
|
|
return grpc::Status::OK;
|
|
}
|
|
|
|
} // namespace meshnet::worker
|