36 lines
797 B
C++
36 lines
797 B
C++
#pragma once
|
|
|
|
#include "shm_helper.h"
|
|
|
|
class ShmManager {
|
|
public:
|
|
ShmManager();
|
|
~ShmManager();
|
|
|
|
void set_gpu_info(int gpu_id, uint64_t vmem_size, uint64_t shared_handle);
|
|
|
|
// request
|
|
void process_requests();
|
|
// heart beat
|
|
void check_heart_beats();
|
|
|
|
void run_busy_loop();
|
|
|
|
void stop_busy_loop() {
|
|
stop_loop_flag.store(true, std::memory_order_release);
|
|
}
|
|
|
|
void register_callback_on_worker_change(
|
|
std::function<void(const std::vector<int32_t> &)> cb) {
|
|
cb_on_worker_change = std::move(cb);
|
|
}
|
|
|
|
private:
|
|
ShmHelper *shm_helper;
|
|
std::vector<int32_t> local_worker_tgids;
|
|
std::vector<int32_t> alive_worker_tgids;
|
|
std::vector<int> valid_gpu_ids;
|
|
std::atomic<bool> stop_loop_flag;
|
|
std::function<void(const std::vector<int32_t> &)> cb_on_worker_change;
|
|
};
|