story: DGR-033 Build a standalone fake C++ gRPC Shard worker
This commit is contained in:
Dobromir Popov
302
packages/node/native/worker/shard_worker_main.cpp
Normal file
302
packages/node/native/worker/shard_worker_main.cpp
Normal file
@@ -0,0 +1,302 @@
|
||||
// Standalone native Shard worker executable (DGR-033).
|
||||
//
|
||||
// Serves the complete ShardRuntime lifecycle/stream contract over real
|
||||
// gRPC/HTTP2 using the model-free FakeShardEngine. It links neither llama.cpp
|
||||
// nor any graph-execution entry point: the only surface it exposes is the
|
||||
// ShardRuntime service defined in shard_runtime.proto.
|
||||
//
|
||||
// Usage:
|
||||
// shard_worker [listen_addr] serve until SIGTERM/SIGINT (graceful drain)
|
||||
// shard_worker --selftest bind an ephemeral port, self-drive the
|
||||
// lifecycle over a real loopback channel, exit
|
||||
//
|
||||
// Environment:
|
||||
// MESHNET_SHARD_LISTEN_ADDR host:port to bind (default localhost:50051)
|
||||
// MESHNET_MAX_CHUNK_BYTES per-chunk byte ceiling the worker enforces
|
||||
//
|
||||
// On a normal run it prints one readiness line — "ShardRuntime worker listening
|
||||
// on <addr>" — once the socket is bound, so a supervisor/harness has a real
|
||||
// readiness signal instead of a sleep.
|
||||
|
||||
#include <atomic>
|
||||
#include <cerrno>
|
||||
#include <csignal>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <unistd.h>
|
||||
|
||||
#include <grpcpp/grpcpp.h>
|
||||
|
||||
#include "shard_service.h"
|
||||
#include "shard_runtime.grpc.pb.h"
|
||||
|
||||
namespace {
|
||||
|
||||
namespace sp = ::meshnet::shard::v1;
|
||||
|
||||
// Self-pipe: the signal handler must stay async-signal-safe, so it only writes
|
||||
// one byte; a helper thread reads it and performs the (non-signal-safe) server
|
||||
// Shutdown(). Set once in main() before installing the handler.
|
||||
volatile std::sig_atomic_t g_signal_pipe_write_fd = -1;
|
||||
|
||||
extern "C" void HandleTermination(int /*signum*/) {
|
||||
if (g_signal_pipe_write_fd >= 0) {
|
||||
const char byte = 1;
|
||||
ssize_t rc = ::write(g_signal_pipe_write_fd, &byte, 1);
|
||||
(void)rc; // best-effort; nothing safe to do on failure inside a handler
|
||||
}
|
||||
}
|
||||
|
||||
meshnet::worker::FlowLimits LimitsFromEnv() {
|
||||
meshnet::worker::FlowLimits limits;
|
||||
if (const char* raw = std::getenv("MESHNET_MAX_CHUNK_BYTES")) {
|
||||
char* end = nullptr;
|
||||
const unsigned long long value = std::strtoull(raw, &end, 10);
|
||||
if (end != raw && value > 0) {
|
||||
limits.max_chunk_bytes = static_cast<uint64_t>(value);
|
||||
}
|
||||
}
|
||||
return limits;
|
||||
}
|
||||
|
||||
int RunSelfTest() {
|
||||
meshnet::worker::ShardRuntimeServiceImpl service(LimitsFromEnv());
|
||||
int selected_port = 0;
|
||||
grpc::ServerBuilder builder;
|
||||
builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &selected_port);
|
||||
builder.RegisterService(&service);
|
||||
std::unique_ptr<grpc::Server> server(builder.BuildAndStart());
|
||||
if (!server || selected_port == 0) {
|
||||
std::cerr << "selftest: failed to bind ephemeral port\n";
|
||||
return 1;
|
||||
}
|
||||
const std::string target = "127.0.0.1:" + std::to_string(selected_port);
|
||||
auto channel = grpc::CreateChannel(target, grpc::InsecureChannelCredentials());
|
||||
auto stub = sp::ShardRuntime::NewStub(channel);
|
||||
|
||||
int failures = 0;
|
||||
auto check = [&](bool cond, const char* what) {
|
||||
if (!cond) {
|
||||
std::cerr << "selftest FAIL: " << what << "\n";
|
||||
++failures;
|
||||
}
|
||||
};
|
||||
|
||||
// Capability + health.
|
||||
{
|
||||
grpc::ClientContext ctx;
|
||||
sp::CapabilityRequest req;
|
||||
req.set_schema_version(sp::SCHEMA_VERSION_1);
|
||||
sp::CapabilityReport rep;
|
||||
grpc::Status status = stub->GetCapability(&ctx, req, &rep);
|
||||
check(status.ok(), "GetCapability RPC");
|
||||
check(rep.validated(), "capability validated");
|
||||
check(rep.schema_version() == sp::SCHEMA_VERSION_1, "capability schema version");
|
||||
}
|
||||
{
|
||||
grpc::ClientContext ctx;
|
||||
sp::HealthRequest req;
|
||||
req.set_schema_version(sp::SCHEMA_VERSION_1);
|
||||
sp::HealthReport rep;
|
||||
grpc::Status status = stub->Health(&ctx, req, &rep);
|
||||
check(status.ok(), "Health RPC");
|
||||
check(rep.state() == sp::SERVING_STATE_SERVING, "health serving");
|
||||
}
|
||||
|
||||
// A minimal session: open -> fragmented prefill -> decode -> release.
|
||||
{
|
||||
grpc::ClientContext ctx;
|
||||
auto stream = stub->Session(&ctx);
|
||||
|
||||
sp::SessionRequest open;
|
||||
sp::SessionOpen* o = open.mutable_open();
|
||||
o->set_schema_version(sp::SCHEMA_VERSION_1);
|
||||
o->set_route_session_id("selftest");
|
||||
o->set_route_epoch(1);
|
||||
sp::FlowControl* fc = o->mutable_proposed_flow_control();
|
||||
fc->set_credits_granted(16);
|
||||
fc->set_max_inflight_chunks(16);
|
||||
fc->set_max_chunk_bytes(4u * 1024u * 1024u);
|
||||
check(stream->Write(open), "write open");
|
||||
|
||||
sp::SessionResponse accepted;
|
||||
check(stream->Read(&accepted), "read accepted");
|
||||
check(accepted.kind_case() == sp::SessionResponse::kAccepted, "accepted kind");
|
||||
|
||||
// Fragmented prefill: two fragments tiling a 6-byte payload.
|
||||
const std::string payload = "ABCDEF";
|
||||
sp::SessionRequest chunk;
|
||||
sp::ActivationChunk* ac = chunk.mutable_chunk();
|
||||
sp::Envelope* env = ac->mutable_envelope();
|
||||
env->set_schema_version(sp::SCHEMA_VERSION_1);
|
||||
env->set_work_id("w1");
|
||||
env->set_route_session_id("selftest");
|
||||
env->set_route_epoch(1);
|
||||
env->set_idempotency_step(1);
|
||||
env->set_phase(sp::PHASE_PREFILL);
|
||||
sp::TensorBundle* bundle = ac->mutable_bundle();
|
||||
bundle->set_bundle_version(1);
|
||||
sp::NamedTensor* tensor = bundle->add_tensors();
|
||||
tensor->set_name("hidden_states");
|
||||
tensor->set_dtype(sp::DTYPE_BFLOAT16);
|
||||
tensor->set_byte_order(sp::BYTE_ORDER_LITTLE_ENDIAN);
|
||||
tensor->set_total_bytes(payload.size());
|
||||
tensor->set_compression(sp::COMPRESSION_NONE);
|
||||
sp::Checksum* cksum = tensor->mutable_checksum();
|
||||
cksum->set_algorithm(sp::CHECKSUM_ALGORITHM_CRC32C);
|
||||
const uint32_t crc = meshnet::worker::Crc32(payload);
|
||||
std::string crc_be(4, '\0');
|
||||
crc_be[0] = static_cast<char>((crc >> 24) & 0xFF);
|
||||
crc_be[1] = static_cast<char>((crc >> 16) & 0xFF);
|
||||
crc_be[2] = static_cast<char>((crc >> 8) & 0xFF);
|
||||
crc_be[3] = static_cast<char>(crc & 0xFF);
|
||||
cksum->set_value(crc_be);
|
||||
sp::TensorFragment* f0 = tensor->add_fragments();
|
||||
f0->set_fragment_index(0);
|
||||
f0->set_fragment_count(2);
|
||||
f0->set_byte_offset(0);
|
||||
f0->set_payload(payload.substr(0, 3));
|
||||
sp::TensorFragment* f1 = tensor->add_fragments();
|
||||
f1->set_fragment_index(1);
|
||||
f1->set_fragment_count(2);
|
||||
f1->set_byte_offset(3);
|
||||
f1->set_payload(payload.substr(3));
|
||||
check(stream->Write(chunk), "write chunk");
|
||||
|
||||
sp::SessionResponse echoed;
|
||||
check(stream->Read(&echoed), "read chunk echo");
|
||||
check(echoed.kind_case() == sp::SessionResponse::kChunk, "chunk echo kind");
|
||||
|
||||
sp::SessionRequest decode;
|
||||
sp::DecodeStep* ds = decode.mutable_decode();
|
||||
ds->set_idempotency_step(2);
|
||||
ds->set_position(1);
|
||||
ds->set_work_id("w2");
|
||||
sp::TensorBundle* dbundle = ds->mutable_bundle();
|
||||
dbundle->set_bundle_version(1);
|
||||
sp::NamedTensor* dt = dbundle->add_tensors();
|
||||
dt->set_name("hidden_states");
|
||||
dt->set_dtype(sp::DTYPE_BFLOAT16);
|
||||
dt->set_byte_order(sp::BYTE_ORDER_LITTLE_ENDIAN);
|
||||
dt->set_total_bytes(payload.size());
|
||||
dt->set_compression(sp::COMPRESSION_NONE);
|
||||
sp::Checksum* dck = dt->mutable_checksum();
|
||||
dck->set_algorithm(sp::CHECKSUM_ALGORITHM_CRC32C);
|
||||
dck->set_value(crc_be);
|
||||
sp::TensorFragment* df = dt->add_fragments();
|
||||
df->set_fragment_index(0);
|
||||
df->set_fragment_count(1);
|
||||
df->set_byte_offset(0);
|
||||
df->set_payload(payload);
|
||||
check(stream->Write(decode), "write decode");
|
||||
|
||||
sp::SessionResponse decode_echo;
|
||||
check(stream->Read(&decode_echo), "read decode echo");
|
||||
check(decode_echo.kind_case() == sp::SessionResponse::kChunk, "decode echo kind");
|
||||
|
||||
sp::SessionRequest release;
|
||||
sp::ReleaseSignal* rs = release.mutable_release();
|
||||
rs->set_route_session_id("selftest");
|
||||
rs->set_work_id("w-final");
|
||||
check(stream->Write(release), "write release");
|
||||
stream->WritesDone();
|
||||
|
||||
sp::SessionResponse terminal;
|
||||
check(stream->Read(&terminal), "read terminal");
|
||||
check(terminal.kind_case() == sp::SessionResponse::kStatus && terminal.status().terminal(),
|
||||
"terminal status");
|
||||
|
||||
grpc::Status status = stream->Finish();
|
||||
check(status.ok(), "stream finish");
|
||||
}
|
||||
|
||||
server->Shutdown();
|
||||
server->Wait();
|
||||
|
||||
if (failures == 0) {
|
||||
std::cout << "selftest: all lifecycle checks passed\n";
|
||||
return 0;
|
||||
}
|
||||
std::cerr << "selftest: " << failures << " check(s) failed\n";
|
||||
return 1;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
GOOGLE_PROTOBUF_VERIFY_VERSION;
|
||||
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
if (std::strcmp(argv[i], "--selftest") == 0) {
|
||||
return RunSelfTest();
|
||||
}
|
||||
}
|
||||
|
||||
std::string listen_addr = "localhost:50051";
|
||||
if (const char* env = std::getenv("MESHNET_SHARD_LISTEN_ADDR")) {
|
||||
listen_addr = env;
|
||||
}
|
||||
if (argc > 1 && argv[1][0] != '-') {
|
||||
listen_addr = argv[1];
|
||||
}
|
||||
|
||||
meshnet::worker::FlowLimits limits = LimitsFromEnv();
|
||||
meshnet::worker::ShardRuntimeServiceImpl service(limits);
|
||||
|
||||
grpc::ServerBuilder builder;
|
||||
int selected_port = 0;
|
||||
builder.AddListeningPort(listen_addr, grpc::InsecureServerCredentials(), &selected_port);
|
||||
// Bounded messages, two layers: a hard transport receive ceiling (never below
|
||||
// 4 MiB so the handshake and normal chunks always fit) plus the finer
|
||||
// app-level per-tensor RESOURCE_EXHAUSTED check the service enforces against
|
||||
// the negotiated max_chunk_bytes. Neither path lets an unbounded frame in.
|
||||
constexpr int kTransportFloor = 4 * 1024 * 1024;
|
||||
const int transport_max = limits.max_chunk_bytes > static_cast<uint64_t>(kTransportFloor)
|
||||
? static_cast<int>(limits.max_chunk_bytes)
|
||||
: kTransportFloor;
|
||||
builder.SetMaxReceiveMessageSize(transport_max);
|
||||
builder.RegisterService(&service);
|
||||
std::unique_ptr<grpc::Server> server(builder.BuildAndStart());
|
||||
if (!server || selected_port == 0) {
|
||||
std::cerr << "failed to bind " << listen_addr << "\n";
|
||||
return 1;
|
||||
}
|
||||
|
||||
int pipe_fds[2];
|
||||
if (::pipe(pipe_fds) != 0) {
|
||||
std::cerr << "failed to create shutdown pipe\n";
|
||||
return 1;
|
||||
}
|
||||
g_signal_pipe_write_fd = pipe_fds[1];
|
||||
|
||||
struct sigaction sa;
|
||||
std::memset(&sa, 0, sizeof(sa));
|
||||
sa.sa_handler = HandleTermination;
|
||||
::sigaction(SIGTERM, &sa, nullptr);
|
||||
::sigaction(SIGINT, &sa, nullptr);
|
||||
|
||||
// Drain thread: wakes on the first termination signal and shuts the server
|
||||
// down gracefully so in-flight sessions finish rather than being severed.
|
||||
std::thread drain([&server, read_fd = pipe_fds[0]]() {
|
||||
char byte = 0;
|
||||
ssize_t rc = 0;
|
||||
do {
|
||||
rc = ::read(read_fd, &byte, 1);
|
||||
} while (rc < 0 && errno == EINTR);
|
||||
server->Shutdown();
|
||||
});
|
||||
|
||||
std::cout << "ShardRuntime worker listening on " << listen_addr << std::endl;
|
||||
|
||||
server->Wait();
|
||||
drain.join();
|
||||
::close(pipe_fds[0]);
|
||||
::close(pipe_fds[1]);
|
||||
std::cout << "ShardRuntime worker shut down cleanly" << std::endl;
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user