adapter_manager.cpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#include <entropic/inference/adapter_manager.h>
#include <entropic/types/logging.h>
 
#include <llama.h>
#include <nlohmann/json.hpp>
 
#include <cstring>
 
namespace entropic {
 
namespace {
auto logger = entropic::log::get("inference.adapter_manager");
using entropic::log::now;
using entropic::log::elapsed_ms;
 
void apply_adapter(llama_context* ctx,
                   llama_adapter_lora* adapter, float scale)
{
    llama_set_adapters_lora(ctx, &adapter, 1, &scale);
}
 
void clear_adapters(llama_context* ctx) {
    llama_set_adapters_lora(ctx, nullptr, 0, nullptr);
}
 
} // anonymous namespace
 
// ── Load ────────────────────────────────────────────────────
 
bool AdapterManager::load(
    const std::string& name,
    const std::filesystem::path& adapter_path,
    llama_model* model,
    float scale)
{
    std::lock_guard<std::mutex> lock(adapter_mutex_);
 
    bool precondition_failed =
        adapters_.find(name) != adapters_.end() || !model;
    if (precondition_failed) {
        logger->error("Cannot load adapter '{}': {}",
                     name, !model ? "base model is null" : "duplicate name");
        return false;
    }
 
    auto t_start = now();
    auto* lora = llama_adapter_lora_init(model, adapter_path.c_str());
    if (!lora) {
        logger->error("llama_adapter_lora_init failed for '{}' at {}",
                     name, adapter_path.string());
        return false;
    }
 
    AdapterEntry entry;
    entry.name = name;
    entry.path = adapter_path;
    entry.handle = lora;
    entry.model = model;
    entry.scale = scale;
    entry.state = AdapterState::WARM;
    adapters_.emplace(name, std::move(entry));
 
    logger->info("Loaded adapter '{}' from {} in {:.1f}ms (scale={:.2f})",
                name, adapter_path.string(),
                elapsed_ms(t_start, now()), scale);
    return true;
}
 
// ── Unload ──────────────────────────────────────────────────
 
void AdapterManager::unload(const std::string& name, llama_context* ctx) {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
 
    auto it = adapters_.find(name);
    if (it == adapters_.end()) {
        return;
    }
 
    auto& entry = it->second;
    if (entry.state == AdapterState::HOT && ctx) {
        clear_adapters(ctx);
        active_name_.clear();
    }
 
    if (entry.handle) {
        llama_adapter_lora_free(entry.handle);
        entry.handle = nullptr;
    }
 
    entry.state = AdapterState::COLD;
    logger->info("Unloaded adapter '{}'", name);
 
    adapters_.erase(it);
}
 
// ── Activate ────────────────────────────────────────────────
 
bool AdapterManager::activate(const std::string& name, llama_context* ctx) {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
 
    auto it = adapters_.find(name);
    bool cannot_activate =
        it == adapters_.end() || it->second.state == AdapterState::COLD;
    if (it != adapters_.end() && it->second.state == AdapterState::HOT) {
        return true;  // Already active — no-op
    }
    if (cannot_activate) {
        logger->error("Cannot activate adapter '{}': {}",
                     name, it == adapters_.end() ? "not found" : "state is COLD");
        return false;
    }
 
    auto& entry = it->second;
 
    // Mark previous HOT adapter as WARM
    if (!active_name_.empty() && active_name_ != name) {
        auto active_it = adapters_.find(active_name_);
        if (active_it != adapters_.end()) {
            active_it->second.state = AdapterState::WARM;
        }
    }
 
    apply_adapter(ctx, entry.handle, entry.scale);
    entry.state = AdapterState::HOT;
    active_name_ = name;
 
    logger->info("Activated adapter '{}' (scale={:.2f})", name, entry.scale);
    return true;
}
 
// ── Deactivate ──────────────────────────────────────────────
 
void AdapterManager::deactivate(llama_context* ctx) {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
 
    if (active_name_.empty()) {
        return;
    }
 
    auto it = adapters_.find(active_name_);
    if (it != adapters_.end()) {
        if (ctx) {
            clear_adapters(ctx);
        }
        it->second.state = AdapterState::WARM;
    }
 
    logger->info("Deactivated adapter '{}'", active_name_);
    active_name_.clear();
}
 
// ── Swap ────────────────────────────────────────────────────
 
bool AdapterManager::swap(const std::string& name, llama_context* ctx) {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
 
    auto it = adapters_.find(name);
    bool cannot_swap = it == adapters_.end()
                    || it->second.state == AdapterState::COLD;
    if (active_name_ == name && !cannot_swap) {
        return true;  // Already active — no-op
    }
    if (cannot_swap || !fire_swap_hook(active_name_, name, it->second.path)) {
        logger->error("Cannot swap to adapter '{}': {}",
                     name, cannot_swap ? "not found or COLD" : "cancelled by hook");
        return false;
    }
 
    auto t_start = now();
 
    // Mark current HOT as WARM
    if (!active_name_.empty()) {
        auto active_it = adapters_.find(active_name_);
        if (active_it != adapters_.end()) {
            active_it->second.state = AdapterState::WARM;
        }
    }
 
    // Apply target
    auto& target = it->second;
    apply_adapter(ctx, target.handle, target.scale);
    target.state = AdapterState::HOT;
 
    std::string previous = active_name_;
    active_name_ = name;
 
    logger->info("Swapped adapter '{}' -> '{}' in {:.1f}ms",
                previous, name, elapsed_ms(t_start, now()));
    return true;
}
 
// ── Unload All ──────────────────────────────────────────────
 
void AdapterManager::unload_all_for_model(
    llama_model* model, llama_context* ctx)
{
    std::lock_guard<std::mutex> lock(adapter_mutex_);
 
    bool cleared_context = false;
    std::vector<std::string> to_remove;
 
    for (auto& [name, entry] : adapters_) {
        if (entry.model != model) {
            continue;
        }
 
        if (entry.state == AdapterState::HOT
            && ctx && !cleared_context)
        {
            clear_adapters(ctx);
            cleared_context = true;
        }
 
        if (entry.handle) {
            llama_adapter_lora_free(entry.handle);
            entry.handle = nullptr;
        }
 
        entry.state = AdapterState::COLD;
        to_remove.push_back(name);
    }
 
    for (const auto& name : to_remove) {
        if (name == active_name_) {
            active_name_.clear();
        }
        adapters_.erase(name);
    }
 
    logger->info("Unloaded {} adapter(s) for model", to_remove.size());
}
 
void AdapterManager::unload_all() {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
    size_t freed = 0;
    for (auto& [name, entry] : adapters_) {
        if (entry.handle) {
            llama_adapter_lora_free(entry.handle);
            entry.handle = nullptr;
            ++freed;
        }
        entry.state = AdapterState::COLD;
    }
    adapters_.clear();
    active_name_.clear();
    if (freed > 0) {
        logger->info("Unloaded all {} adapter(s) on shutdown", freed);
    }
}
 
// ── Queries ─────────────────────────────────────────────────
 
AdapterState AdapterManager::state(const std::string& name) const {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
    auto it = adapters_.find(name);
    if (it == adapters_.end()) {
        return AdapterState::COLD;
    }
    return it->second.state;
}
 
AdapterInfo AdapterManager::info(const std::string& name) const {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
    auto it = adapters_.find(name);
    if (it == adapters_.end()) {
        return {};
    }
    return make_info(it->second);
}
 
std::vector<AdapterInfo> AdapterManager::list_adapters() const {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
    std::vector<AdapterInfo> result;
    result.reserve(adapters_.size());
    for (const auto& [name, entry] : adapters_) {
        result.push_back(make_info(entry));
    }
    return result;
}
 
std::string AdapterManager::active_adapter() const {
    std::lock_guard<std::mutex> lock(adapter_mutex_);
    return active_name_;
}
 
void AdapterManager::set_hook_interface(const HookInterface& hooks) {
    hooks_ = hooks;
}
 
// ── Private ─────────────────────────────────────────────────
 
AdapterInfo AdapterManager::make_info(const AdapterEntry& entry) {
    AdapterInfo info;
    info.name = entry.name;
    info.path = entry.path;
    info.state = entry.state;
    info.scale = entry.scale;
    info.tier_name = entry.tier_name;
    info.ram_bytes = entry.ram_bytes;
    info.metadata = entry.metadata;
    return info;
}
 
bool AdapterManager::fire_swap_hook(
    const std::string& current,
    const std::string& target,
    const std::filesystem::path& target_path)
{
    if (!hooks_.fire_pre || !hooks_.registry) {
        return true;  // No hook registered — proceed
    }
 
    nlohmann::json ctx;
    ctx["current_adapter"] = current;
    ctx["target_adapter"] = target;
    ctx["adapter_path"] = target_path.string();
    std::string ctx_str = ctx.dump();
 
    char* modified = nullptr;
    int rc = hooks_.fire_pre(
        hooks_.registry,
        ENTROPIC_HOOK_ON_ADAPTER_SWAP,
        ctx_str.c_str(),
        &modified);
 
    if (modified) {
        free(modified);
    }
 
    return rc == 0;
}
 
} // namespace entropic