MCPcopy Create free account
hub / github.com/OpenNMT/CTranslate2 / ReplicaPoolHelper

Method ReplicaPoolHelper

python/cpp/replica_pool.h:41–70  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

39 class ReplicaPoolHelper {
40 public:
41 ReplicaPoolHelper(const std::string& model_path,
42 const std::string& device,
43 const std::variant<int, std::vector<int>>& device_index,
44 const StringOrMap& compute_type,
45 size_t inter_threads,
46 size_t intra_threads,
47 long max_queued_batches,
48 bool flash_attention,
49 bool tensor_parallel,
50 py::object files)
51 : _model_loader(create_model_reader(model_path, files))
52 , _device(str_to_device(device))
53 , _num_replicas_per_device(inter_threads)
54 {
55 pybind11::gil_scoped_release nogil;
56
57 _model_loader.device = str_to_device(device);
58 _model_loader.device_indices = std::visit(DeviceIndexResolver(), device_index);
59 _model_loader.compute_type = std::visit(ComputeTypeResolver(device), compute_type);
60 _model_loader.num_replicas_per_device = inter_threads;
61 _model_loader.use_flash_attention = flash_attention;
62 _model_loader.tensor_parallel = tensor_parallel;
63
64 _pool_config.num_threads_per_replica = intra_threads;
65 _pool_config.max_queued_batches = max_queued_batches;
66
67 _pool = std::make_unique<T>(_model_loader, _pool_config);
68 _device_index = _model_loader.device_indices;
69 _model_is_loaded = true;
70 }
71
72 ~ReplicaPoolHelper() {
73 pybind11::gil_scoped_release nogil;

Callers

nothing calls this directly

Calls 4

create_model_readerFunction · 0.85
str_to_deviceFunction · 0.85
DeviceIndexResolverClass · 0.85
ComputeTypeResolverClass · 0.85

Tested by

no test coverage detected