secondary_model_loader.cpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#include <entropic/inference/secondary_model_loader.h>
#include <entropic/types/logging.h>
 
#include "llama_cpp_backend.h"
 
#include <algorithm>
 
namespace entropic {
 
namespace {
auto logger = entropic::log::get("inference.secondary_loader");
} // anonymous namespace
 
bool SecondaryModelLoader::ensure_loaded(
    const std::string& role, const ModelConfig& config) {
    std::lock_guard<std::mutex> lock(slots_mutex_);
 
    const std::string new_path = config.path.string();
    auto path_it = slot_paths_.find(role);
    if (path_it != slot_paths_.end() && path_it->second == new_path) {
        auto it = slots_.find(role);
        if (it != slots_.end() && it->second->is_loaded()) {
            return true;
        }
    }
 
    auto backend = std::make_shared<LlamaCppBackend>();
    if (!backend->load_and_activate(config)) {
        logger->error("Failed to activate role '{}' from path: {}",
                      role, new_path);
        return false;
    }
 
    slots_[role] = backend;
    slot_paths_[role] = new_path;
    logger->info("Activated secondary role '{}' from {}", role, new_path);
    return true;
}
 
InferenceBackend* SecondaryModelLoader::get(const std::string& role) const {
    std::lock_guard<std::mutex> lock(slots_mutex_);
    auto it = slots_.find(role);
    return (it == slots_.end()) ? nullptr : it->second.get();
}
 
std::shared_ptr<InferenceBackend> SecondaryModelLoader::get_shared(
    const std::string& role) const {
    std::lock_guard<std::mutex> lock(slots_mutex_);
    auto it = slots_.find(role);
    return (it == slots_.end()) ? std::shared_ptr<InferenceBackend>{}
                                : it->second;
}
 
bool SecondaryModelLoader::release_role(const std::string& role) {
    std::lock_guard<std::mutex> lock(slots_mutex_);
    auto it = slots_.find(role);
    if (it == slots_.end()) {
        return false;
    }
    if (it->second->is_loaded()) {
        it->second->unload();
    }
    slots_.erase(it);
    slot_paths_.erase(role);
    logger->info("Released secondary role '{}'", role);
    return true;
}
 
bool SecondaryModelLoader::is_loaded(const std::string& role) const {
    std::lock_guard<std::mutex> lock(slots_mutex_);
    auto it = slots_.find(role);
    return it != slots_.end() && it->second->is_loaded();
}
 
std::vector<std::string> SecondaryModelLoader::loaded_roles() const {
    std::lock_guard<std::mutex> lock(slots_mutex_);
    std::vector<std::string> out;
    out.reserve(slots_.size());
    for (const auto& [role, backend] : slots_) {
        if (backend->is_loaded()) {
            out.push_back(role);
        }
    }
    std::sort(out.begin(), out.end());
    return out;
}
 
void SecondaryModelLoader::clear_all_prompt_caches() {
    std::lock_guard<std::mutex> lock(slots_mutex_);
    for (auto& [role, backend] : slots_) {
        backend->clear_prompt_cache();
    }
}
 
void SecondaryModelLoader::shutdown() {
    std::lock_guard<std::mutex> lock(slots_mutex_);
    for (auto& [role, backend] : slots_) {
        if (backend->is_loaded()) {
            backend->unload();
        }
    }
    slots_.clear();
    slot_paths_.clear();
}
 
} // namespace entropic