#pragma once

#include <shared_mutex>

#include <ctranslate2/replica_pool.h>

#include "utils.h"

namespace ctranslate2 {

namespace python {

inline std::shared_ptr<models::ModelReader>

create_model_reader(const std::string& model, py::object files) {

if (files.is_none())

return std::make_shared<models::ModelFileReader>(model);

if (!py::isinstance<py::dict>(files))

throw pybind11::type_error("files argument must be a dictionary mapping file names "

"to the file contents");

auto reader = std::make_shared<models::ModelMemoryReader>(model);

for (const auto& pair : files.cast<py::dict>()) {

auto filename = pair.first;

auto content = pair.second;

auto read = py::getattr(content, "read", py::none());

if (!read.is_none())

content = read();

else if (!py::isinstance<py::bytes>(content))

throw pybind11::type_error("File content must be a file-like or bytes object");

reader->register_file(filename.cast<std::string>(), content.cast<std::string>());

}

return reader;

}

template <typename T>

class ReplicaPoolHelper {

public:

ReplicaPoolHelper(const std::string& model_path,

const std::string& device,

const std::variant<int, std::vector<int>>& device_index,

const StringOrMap& compute_type,

size_t inter_threads,

size_t intra_threads,

long max_queued_batches,

bool flash_attention,

bool tensor_parallel,

py::object files)

: _model_loader(create_model_reader(model_path, files))

, _device(str_to_device(device))

, _num_replicas_per_device(inter_threads)

{

pybind11::gil_scoped_release nogil;

_model_loader.device = str_to_device(device);

_model_loader.device_indices = std::visit(DeviceIndexResolver(), device_index);

_model_loader.compute_type = std::visit(ComputeTypeResolver(device), compute_type);

_model_loader.num_replicas_per_device = inter_threads;

_model_loader.use_flash_attention = flash_attention;

_model_loader.tensor_parallel = tensor_parallel;

_pool_config.num_threads_per_replica = intra_threads;

_pool_config.max_queued_batches = max_queued_batches;

_pool = std::make_unique<T>(_model_loader, _pool_config);

_device_index = _model_loader.device_indices;

_model_is_loaded = true;

}

~ReplicaPoolHelper() {

pybind11::gil_scoped_release nogil;

_pool.reset();

}

std::string device() const {

return device_to_str(_model_loader.device);

}

const std::vector<int>& device_index() const {

return _model_loader.device_indices;

}

std::string compute_type() const {

return compute_type_to_str(model()->effective_compute_type());

}

bool tensor_parallel() const {

return _model_loader.tensor_parallel;

}

size_t num_replicas() const {

return _pool->num_replicas();

}

size_t num_queued_batches() const {

return _pool->num_queued_batches();

}

size_t num_active_batches() const {

return _pool->num_active_batches();

}

bool model_is_loaded() {

std::shared_lock lock(_mutex);

return _model_is_loaded;

}

void unload_model(const bool to_cpu) {

if (to_cpu && _device == Device::CPU)

return;

// Do not unload the model if some batches are still being processed.

if (_pool->num_active_batches() > 0)

return;

// If the lock is not acquired immediately it means the model is being used

// in another thread and we can't unload it at this time.

std::unique_lock lock(_mutex, std::try_to_lock);

if (!lock)

return;

std::vector<std::shared_ptr<const models::Model>> loaded_models;

if (_model_is_loaded)

loaded_models = _pool->detach_models();

if (to_cpu && _cached_models.empty())

_cached_models = clone_models(Device::CPU, std::vector<int>(loaded_models.size(), 0), loaded_models);

else if (!to_cpu)

_cached_models.clear();

loaded_models.clear();

// We clear the CUDA allocator cache to further reduce the memory after unloading the model.

if (_device == Device::CUDA)

_pool->clear_cache();

_model_is_loaded = false;

}

void load_model(const bool keep_cache) {

std::unique_lock lock(_mutex);

if (_model_is_loaded)

return;

std::vector<std::shared_ptr<const models::Model>> loaded_models;

if (_cached_models.empty())

loaded_models = _model_loader.load();

else

loaded_models = clone_models(_device, _device_index, _cached_models, _num_replicas_per_device);

_pool->set_models(loaded_models);

if (!keep_cache)

_cached_models.clear();

_model_is_loaded = true;

}

protected:

std::unique_ptr<T> _pool;

models::ModelLoader _model_loader;

ReplicaPoolConfig _pool_config;

const Device _device;

const size_t _num_replicas_per_device;

std::vector<int> _device_index;

std::vector<std::shared_ptr<const models::Model>> _cached_models;

bool _model_is_loaded;

// Use a shared mutex to protect the model state (loaded/unloaded).

// Multiple threads can read the model at the same time, but a single thread can change

// the model state (e.g. load or unload the model).

std::shared_mutex _mutex;

std::vector<std::shared_ptr<const models::Model>> clone_models(Device device,

const std::vector<int>& device_index,

std::vector<std::shared_ptr<const models::Model>> cached_models,

size_t num_models_per_device = 1) {

std::vector<std::shared_ptr<const models::Model>> copied_models;

for (size_t i = 0; i < cached_models.size(); ++i) {

auto& model = const_cast<models::Model&>(*cached_models[i]);

auto copied_model = model.copy_to(device, device_index[i / num_models_per_device]);

copied_models.push_back(copied_model);

}

return copied_models;

}

const std::shared_ptr<const models::Model>& model() const {

return _pool->get_first_replica().model();

}

void assert_model_is_ready() const {

if (!_model_is_loaded)

throw std::runtime_error("The model for this translator was unloaded");

}

};

}

}