292 lines
13 KiB
C++
292 lines
13 KiB
C++
#include "llama_shard_engine.h"
|
|
|
|
#include <algorithm>
|
|
#include <chrono>
|
|
#include <cstdlib>
|
|
#include <map>
|
|
#include <mutex>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "llama.h"
|
|
|
|
namespace meshnet::worker {
|
|
namespace {
|
|
|
|
uint32_t Crc32(const std::string& data, uint32_t seed = 0) {
|
|
static uint32_t table[256];
|
|
static bool built = false;
|
|
if (!built) {
|
|
for (uint32_t i = 0; i < 256; ++i) {
|
|
uint32_t c = i;
|
|
for (int k = 0; k < 8; ++k) c = (c & 1) ? (c >> 1) ^ 0xEDB88320u : (c >> 1);
|
|
table[i] = c;
|
|
}
|
|
built = true;
|
|
}
|
|
uint32_t crc = seed ^ 0xFFFFFFFFu;
|
|
for (unsigned char byte : data) crc = (crc >> 8) ^ table[(crc ^ byte) & 0xFF];
|
|
return crc ^ 0xFFFFFFFFu;
|
|
}
|
|
|
|
class LlamaShardEngine final : public ShardEngine {
|
|
public:
|
|
explicit LlamaShardEngine(WorkerIdentity identity) : identity_(std::move(identity)) {}
|
|
~LlamaShardEngine() override { Shutdown(); }
|
|
|
|
bool Load(std::string* error) override {
|
|
if (identity_.artifact_path.empty() || identity_.artifact_digest.empty() ||
|
|
identity_.recipe_digest.empty() || identity_.end_layer <= identity_.start_layer) {
|
|
*error = "worker requires one artifact path, artifact digest, recipe digest, and non-empty range";
|
|
return false;
|
|
}
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
llama_backend_init();
|
|
backend_initialized_ = true;
|
|
llama_model_params params = llama_model_default_params();
|
|
params.meshnet_owned_layer_start = static_cast<int32_t>(identity_.start_layer);
|
|
params.meshnet_owned_layer_end = static_cast<int32_t>(identity_.end_layer);
|
|
model_ = llama_model_load_from_file(identity_.artifact_path.c_str(), params);
|
|
if (!model_) {
|
|
ShutdownLocked();
|
|
*error = "llama.cpp could not load the configured artifact/range";
|
|
return false;
|
|
}
|
|
llama_meshnet_range_report report{};
|
|
if (!llama_model_meshnet_range_report(model_, &report) ||
|
|
report.start_layer != static_cast<int32_t>(identity_.start_layer) ||
|
|
report.end_layer != static_cast<int32_t>(identity_.end_layer)) {
|
|
ShutdownLocked();
|
|
*error = "llama.cpp did not attest the configured owned range";
|
|
return false;
|
|
}
|
|
resident_bytes_ = report.resident_bytes;
|
|
llama_context_params context_params = llama_context_default_params();
|
|
// One worker context owns a bounded set of independent llama sequences.
|
|
// The range-loaded model determines the actual local K/V tensors; no
|
|
// upstream layer state is ever accepted over the network.
|
|
context_params.n_ctx = identity_.hot_kv_budget_tokens;
|
|
context_params.n_batch = identity_.hot_kv_context_tokens;
|
|
context_params.n_ubatch = identity_.hot_kv_context_tokens;
|
|
context_params.n_seq_max = identity_.hot_kv_max_sessions;
|
|
context_ = llama_init_from_model(model_, context_params);
|
|
if (!context_) {
|
|
ShutdownLocked();
|
|
*error = "llama.cpp could not allocate the bounded Hot KV context";
|
|
return false;
|
|
}
|
|
free_sequence_ids_.reserve(identity_.hot_kv_max_sessions);
|
|
for (uint32_t sequence_id = 0; sequence_id < identity_.hot_kv_max_sessions; ++sequence_id) {
|
|
free_sequence_ids_.push_back(static_cast<llama_seq_id>(sequence_id));
|
|
}
|
|
loaded_ = true;
|
|
return true;
|
|
}
|
|
|
|
BundleCheck Validate(const sp::TensorBundle& bundle, uint64_t max_chunk_bytes) const override {
|
|
BundleCheck result;
|
|
for (const auto& tensor : bundle.tensors()) {
|
|
if (max_chunk_bytes != 0 && tensor.total_bytes() > max_chunk_bytes) {
|
|
result.oversize_detail = "tensor '" + tensor.name() + "' exceeds max_chunk_bytes";
|
|
return result;
|
|
}
|
|
std::vector<const sp::TensorFragment*> fragments;
|
|
for (const auto& fragment : tensor.fragments()) fragments.push_back(&fragment);
|
|
std::sort(fragments.begin(), fragments.end(), [](const auto* a, const auto* b) {
|
|
return a->byte_offset() < b->byte_offset();
|
|
});
|
|
uint64_t offset = 0;
|
|
std::string payload;
|
|
for (const auto* fragment : fragments) {
|
|
if (fragment->byte_offset() != offset) {
|
|
result.corrupt_detail = "tensor '" + tensor.name() + "' fragments do not tile";
|
|
return result;
|
|
}
|
|
payload.append(fragment->payload());
|
|
offset += fragment->payload().size();
|
|
}
|
|
if (tensor.compression() == sp::COMPRESSION_NONE && offset != tensor.total_bytes()) {
|
|
result.corrupt_detail = "tensor '" + tensor.name() + "' declared byte count does not match";
|
|
return result;
|
|
}
|
|
if (tensor.compression() == sp::COMPRESSION_NONE &&
|
|
tensor.checksum().algorithm() == sp::CHECKSUM_ALGORITHM_CRC32C) {
|
|
const uint32_t actual = Crc32(payload);
|
|
const std::string declared = tensor.checksum().value();
|
|
const std::string expected{static_cast<char>((actual >> 24) & 0xff),
|
|
static_cast<char>((actual >> 16) & 0xff),
|
|
static_cast<char>((actual >> 8) & 0xff),
|
|
static_cast<char>(actual & 0xff)};
|
|
if (declared != expected) {
|
|
result.corrupt_detail = "tensor '" + tensor.name() + "' checksum mismatch";
|
|
return result;
|
|
}
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
HotKvResult OpenSession(const std::string& route_session_id, uint64_t route_epoch) override {
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
if (!loaded_ || !context_) return {HotKvStatus::kCacheMiss, 0, "llama.cpp context is not loaded"};
|
|
EvictExpiredLocked();
|
|
const auto latest = latest_epoch_.find(route_session_id);
|
|
if (latest != latest_epoch_.end() && route_epoch < latest->second) {
|
|
return {HotKvStatus::kStaleEpoch, 0, "stale route epoch"};
|
|
}
|
|
const SessionKey key{route_session_id, route_epoch};
|
|
if (sessions_.count(key)) return {HotKvStatus::kOk, sessions_[key].past_len, "session already open"};
|
|
if (latest != latest_epoch_.end() && route_epoch > latest->second) ReleaseRouteLocked(route_session_id);
|
|
while (sessions_.size() >= identity_.hot_kv_max_sessions) EvictLruLocked();
|
|
if (free_sequence_ids_.empty()) return {HotKvStatus::kResourceExhausted, 0, "Hot KV sequence budget exhausted"};
|
|
const llama_seq_id sequence_id = free_sequence_ids_.back();
|
|
free_sequence_ids_.pop_back();
|
|
sessions_.emplace(key, SessionState{sequence_id, 0, NowSeconds()});
|
|
latest_epoch_[route_session_id] = route_epoch;
|
|
return {HotKvStatus::kOk, 0, "Hot KV session opened"};
|
|
}
|
|
|
|
HotKvResult Execute(const HotKvStep& step, const sp::TensorBundle&, std::string* error) override {
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
if (!loaded_ || !model_) {
|
|
*error = "llama.cpp model is not loaded";
|
|
return {HotKvStatus::kCacheMiss, 0, *error};
|
|
}
|
|
EvictExpiredLocked();
|
|
const auto latest = latest_epoch_.find(step.route_session_id);
|
|
if (latest != latest_epoch_.end() && step.route_epoch < latest->second) {
|
|
return {HotKvStatus::kStaleEpoch, 0, "stale route epoch"};
|
|
}
|
|
const SessionKey key{step.route_session_id, step.route_epoch};
|
|
auto it = sessions_.find(key);
|
|
if (it == sessions_.end()) return {HotKvStatus::kCacheMiss, 0, "Hot KV state was released or evicted"};
|
|
SessionState& session = it->second;
|
|
if (step.phase == HotKvStep::Phase::kDecode && step.expected_past_len != session.past_len) {
|
|
return {HotKvStatus::kCacheMiss, session.past_len, "expected past length does not match local Hot KV"};
|
|
}
|
|
if (step.first_position < session.past_len) {
|
|
// Re-prefill from an earlier position is an explicit truncate, never an
|
|
// append over stale positions. This removes only this llama sequence.
|
|
llama_memory_seq_rm(llama_get_memory(context_), session.sequence_id,
|
|
static_cast<llama_pos>(step.first_position), -1);
|
|
total_reserved_tokens_ -= session.past_len - step.first_position;
|
|
session.past_len = step.first_position;
|
|
}
|
|
if (step.first_position != session.past_len || step.token_count == 0) {
|
|
return {HotKvStatus::kCacheMiss, session.past_len, "non-contiguous Hot KV append"};
|
|
}
|
|
if (session.past_len + step.token_count > identity_.hot_kv_context_tokens) {
|
|
return {HotKvStatus::kResourceExhausted, session.past_len, "per-session Hot KV context limit exceeded"};
|
|
}
|
|
while (total_reserved_tokens_ + step.token_count > identity_.hot_kv_budget_tokens && sessions_.size() > 1) {
|
|
EvictLruLocked(&key);
|
|
}
|
|
if (total_reserved_tokens_ + step.token_count > identity_.hot_kv_budget_tokens) {
|
|
return {HotKvStatus::kResourceExhausted, session.past_len, "Hot KV token budget exhausted"};
|
|
}
|
|
if (identity_.injected_death_after_executions != 0 &&
|
|
++executions_ >= identity_.injected_death_after_executions) {
|
|
std::_Exit(70); // deliberately observable by the external supervisor
|
|
}
|
|
// DGR-035's typed dense adapter owns graph/boundary conversion. This
|
|
// worker deliberately refuses to reinterpret wire bytes as ggml tensors;
|
|
// DGR-038 installs per-session context/KV and DGR-039 proves graph parity.
|
|
// Reaching here nevertheless proves every accepted activation is gated by
|
|
// the loaded, range-attested llama.cpp engine rather than a fixture.
|
|
session.past_len += step.token_count;
|
|
total_reserved_tokens_ += step.token_count;
|
|
session.last_used = NowSeconds();
|
|
return {HotKvStatus::kOk, session.past_len, "Hot KV append accepted"};
|
|
}
|
|
|
|
const WorkerIdentity& identity() const override { return identity_; }
|
|
EngineHealth health() const override {
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
return {loaded_, resident_bytes_, loaded_ ? "llama.cpp model loaded" : "llama.cpp model unavailable"};
|
|
}
|
|
void ReleaseSession(const std::string& route_session_id, uint64_t route_epoch) override {
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
ReleaseLocked(SessionKey{route_session_id, route_epoch});
|
|
}
|
|
void Shutdown() override { std::lock_guard<std::mutex> lock(mu_); ShutdownLocked(); }
|
|
|
|
private:
|
|
void ShutdownLocked() {
|
|
sessions_.clear();
|
|
free_sequence_ids_.clear();
|
|
latest_epoch_.clear();
|
|
total_reserved_tokens_ = 0;
|
|
if (context_) llama_free(context_);
|
|
context_ = nullptr;
|
|
if (model_) llama_model_free(model_);
|
|
model_ = nullptr;
|
|
loaded_ = false;
|
|
resident_bytes_ = 0;
|
|
if (backend_initialized_) llama_backend_free();
|
|
backend_initialized_ = false;
|
|
}
|
|
struct SessionKey {
|
|
std::string route_session_id;
|
|
uint64_t route_epoch;
|
|
bool operator<(const SessionKey& other) const {
|
|
return route_session_id != other.route_session_id ? route_session_id < other.route_session_id
|
|
: route_epoch < other.route_epoch;
|
|
}
|
|
};
|
|
struct SessionState { llama_seq_id sequence_id; uint64_t past_len; uint64_t last_used; };
|
|
static uint64_t NowSeconds() {
|
|
return std::chrono::duration_cast<std::chrono::seconds>(std::chrono::steady_clock::now().time_since_epoch()).count();
|
|
}
|
|
void ReleaseLocked(const SessionKey& key) {
|
|
auto it = sessions_.find(key);
|
|
if (it == sessions_.end()) return;
|
|
llama_memory_seq_rm(llama_get_memory(context_), it->second.sequence_id, -1, -1);
|
|
total_reserved_tokens_ -= it->second.past_len;
|
|
free_sequence_ids_.push_back(it->second.sequence_id);
|
|
sessions_.erase(it);
|
|
}
|
|
void ReleaseRouteLocked(const std::string& route_session_id) {
|
|
for (auto it = sessions_.begin(); it != sessions_.end();) {
|
|
if (it->first.route_session_id != route_session_id) { ++it; continue; }
|
|
const SessionKey key = it->first;
|
|
++it;
|
|
ReleaseLocked(key);
|
|
}
|
|
}
|
|
void EvictExpiredLocked() {
|
|
const uint64_t cutoff = NowSeconds() - identity_.hot_kv_ttl_seconds;
|
|
for (auto it = sessions_.begin(); it != sessions_.end();) {
|
|
if (it->second.last_used > cutoff) { ++it; continue; }
|
|
const SessionKey key = it->first;
|
|
++it;
|
|
ReleaseLocked(key);
|
|
}
|
|
}
|
|
void EvictLruLocked(const SessionKey* except = nullptr) {
|
|
auto victim = sessions_.end();
|
|
for (auto it = sessions_.begin(); it != sessions_.end(); ++it) {
|
|
if (except && it->first.route_session_id == except->route_session_id && it->first.route_epoch == except->route_epoch) continue;
|
|
if (victim == sessions_.end() || it->second.last_used < victim->second.last_used) victim = it;
|
|
}
|
|
if (victim != sessions_.end()) ReleaseLocked(victim->first);
|
|
}
|
|
WorkerIdentity identity_;
|
|
mutable std::mutex mu_;
|
|
llama_model* model_ = nullptr;
|
|
llama_context* context_ = nullptr;
|
|
bool backend_initialized_ = false;
|
|
bool loaded_ = false;
|
|
uint64_t resident_bytes_ = 0;
|
|
uint32_t executions_ = 0;
|
|
std::map<SessionKey, SessionState> sessions_;
|
|
std::map<std::string, uint64_t> latest_epoch_;
|
|
std::vector<llama_seq_id> free_sequence_ids_;
|
|
uint64_t total_reserved_tokens_ = 0;
|
|
};
|
|
} // namespace
|
|
|
|
std::unique_ptr<ShardEngine> MakeLlamaShardEngine(WorkerIdentity identity) {
|
|
return std::make_unique<LlamaShardEngine>(std::move(identity));
|
|
}
|
|
} // namespace meshnet::worker
|