EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

support multi npu partially

Browse Source
This commit is contained in:
starkwj
2026-01-08 06:54:33 +00:00
parent fa0fb46853
commit 2a571d8bc8
12 changed files with 331 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
csrc/camem_allocator.cpp
View File

@@ -20,6 +20,7 @@
#include <atomic>
#include "idle_offload/shm_worker.h"
#include "idle_offload/npu_helper.h"
extern "C" {
@@ -474,8 +475,10 @@ static PyObject* python_create_and_map(PyObject* self, PyObject* args) {
static PyObject* py_init_module_offload(PyObject* self, PyObject* args) {
PyObject* malloc_callback = nullptr;
PyObject* free_callback = nullptr;
unsigned long long device = 0;
if (!PyArg_ParseTuple(args, "OO", &malloc_callback, &free_callback)) {
if (!PyArg_ParseTuple(args, "OOK", &malloc_callback, &free_callback,
&device)) {
return nullptr;
}
@@ -497,7 +500,13 @@ static PyObject* py_init_module_offload(PyObject* self, PyObject* args) {
}
g_initialized.store(true);
shm_worker = new ShmWorker();
std::vector<int> gpu_ids = get_npu_ids();
if (device >= gpu_ids.size()) {
throw std::runtime_error("Invalid device id: " + std::to_string(device) +
" " + __FILE__ + ":" + std::to_string(__LINE__));
}
int gpu_id = gpu_ids[device];
// get pid
aclError error_code;
int32_t pid;
@@ -508,11 +517,12 @@ static PyObject* py_init_module_offload(PyObject* self, PyObject* args) {
std::to_string(error_code) + " " + __FILE__ + ":" +
std::to_string(__LINE__));
}
shm_worker = new ShmWorker();
uint64_t shareable_handle;
shm_worker->register_worker(pid, &shareable_handle, &g_size);
shm_worker->register_worker(pid, gpu_id, &shareable_handle, &g_size);
// import shareable handle
uint32_t device = 0;
aclrtDrvMemHandle memHandle;
error_code =
aclrtMemImportFromShareableHandle(shareable_handle, device, &memHandle);
@@ -570,9 +580,16 @@ static PyObject* python_get_mem_info_offload(PyObject* self, PyObject* args) {
return tuple;
}
static PyObject* python_lock_gpu_offload(PyObject* self, PyObject* args) {
bool prev_is_self = shm_worker->lock_gpu();
return PyBool_FromLong(prev_is_self);
static PyObject* python_try_lock_gpu_offload(PyObject* self, PyObject* args) {
bool prev_is_self = false;
bool success = shm_worker->try_lock_gpu(prev_is_self);
PyObject* tuple = PyTuple_New(2);
if (!tuple) {
return nullptr;
}
PyTuple_SetItem(tuple, 0, PyBool_FromLong(success));
PyTuple_SetItem(tuple, 1, PyBool_FromLong(prev_is_self));
return tuple;
}
static PyObject* python_unlock_gpu_offload(PyObject* self, PyObject* args) {
@@ -597,7 +614,7 @@ static PyMethodDef module_methods[] = {
"Unmap and release memory on the device."},
{"python_get_mem_info_offload", (PyCFunction)python_get_mem_info_offload,
METH_NOARGS, "Get mem info in the reserved pool."},
{"python_lock_gpu_offload", (PyCFunction)python_lock_gpu_offload,
{"python_try_lock_gpu_offload", (PyCFunction)python_try_lock_gpu_offload,
METH_NOARGS, "Lock GPU."},
{"python_unlock_gpu_offload", (PyCFunction)python_unlock_gpu_offload,
METH_NOARGS, "Unlock GPU."},

86
csrc/idle_offload/npu_helper.h Normal file
View File

@@ -0,0 +1,86 @@
#include <vector>
#include <string>
#include <stdint.h>
#include <filesystem>
#include <algorithm>
#include "spdlog/spdlog.h"
#include "acl/acl.h"
std::vector<int> get_available_devices() {
namespace fs = std::filesystem;
std::vector<int> devices;
const std::string dev_path = "/dev";
const std::string prefix = "davinci";
if (!fs::exists(dev_path)) {
return devices;
}
try {
for (const auto &entry : fs::directory_iterator(dev_path)) {
std::string filename = entry.path().filename().string();
if (filename.rfind(prefix, 0) == 0) {
std::string suffix = filename.substr(prefix.length());
// filter not digit suffix
if (!suffix.empty() &&
std::all_of(suffix.begin(), suffix.end(),
[](unsigned char c) { return std::isdigit(c); })) {
try {
int id = std::stoi(suffix);
devices.push_back(id);
} catch (...) {
}
}
}
}
} catch (const fs::filesystem_error &e) {
spdlog::error("Error accessing /dev: {}", e.what());
}
std::sort(devices.begin(), devices.end());
return devices;
}
std::vector<int> get_npu_ids() {
std::vector<int> available_devices = get_available_devices();
std::vector<int> npu_ids;
uint32_t device_count = 0;
aclError error_code = aclrtGetDeviceCount(&device_count);
if (error_code != 0) {
spdlog::error("Failed to get NPU device count, error code: {}", error_code);
throw std::runtime_error("Failed to get NPU device count");
}
if (device_count > available_devices.size()) {
spdlog::error("The number of available NPU devices ({}) is less than the "
"number of devices reported by ACL ({}).",
available_devices.size(), device_count);
throw std::runtime_error("Inconsistent NPU device count");
}
const char *env_available_npu = getenv("ASCEND_RT_VISIBLE_DEVICES");
if (env_available_npu) {
std::string npu_str(env_available_npu);
size_t start = 0;
while (start < npu_str.size()) {
size_t next = npu_str.find(',', start);
if (next == std::string::npos) {
next = npu_str.size();
}
int device_id = std::stoi(npu_str.substr(start, next - start));
npu_ids.push_back(available_devices[device_id]);
start = next + 1;
if (npu_ids.size() >= device_count) {
break;
}
}
} else {
npu_ids.insert(npu_ids.end(), available_devices.begin(),
available_devices.begin() + device_count);
}
return npu_ids;
}

135
csrc/idle_offload/offload_daemon.cpp
View File

@@ -13,6 +13,7 @@
#include "acl/acl.h"
#include "shm_manager.h"
#include "npu_helper.h"
#include "spdlog/spdlog.h"
@@ -49,8 +50,6 @@ void ensure_context(unsigned long long device) {
void init_acl() {
int32_t deviceId=0;
// aclrtStream stream;
bool g_isDevice;
aclError ret = aclrtSetDevice(deviceId);
if (ret != ACL_ERROR_NONE) {
@@ -59,7 +58,8 @@ void init_acl() {
}
}
void reset_pids(const std::vector<int32_t> &pids, uint64_t shareable_handle) {
void reset_pids(const std::vector<int32_t> &pids,
const std::vector<uint64_t> &shareable_handles) {
int cnt = pids.size();
if (cnt <= 0) {
return;
@@ -68,21 +68,21 @@ void reset_pids(const std::vector<int32_t> &pids, uint64_t shareable_handle) {
int32_t pids_data[cnt];
memcpy(pids_data, pids.data(), cnt * sizeof(int32_t));
aclError error_code =
aclrtMemSetPidToShareableHandle(shareable_handle, pids_data, cnt);
if (error_code != 0) {
spdlog::error("aclrtMemSetPidToShareableHandle failed, error_code: {}",
error_code);
throw std::runtime_error("aclrtMemSetPidToShareableHandle failed");
} else {
spdlog::info("aclrtMemSetPidToShareableHandle succeeded, num_pids: {}",
cnt);
for (int i = 0; i < shareable_handles.size(); ++i) {
uint64_t shareable_handle = shareable_handles[i];
aclError error_code =
aclrtMemSetPidToShareableHandle(shareable_handle, pids_data, cnt);
if (error_code != 0) {
spdlog::error("aclrtMemSetPidToShareableHandle failed, error_code: {}",
error_code);
throw std::runtime_error("aclrtMemSetPidToShareableHandle failed");
}
}
spdlog::info("aclrtMemSetPidToShareableHandle succeeded, num_pids: {}", cnt);
}
void start_daemon() {
init_acl();
void alloc_physical(uint32_t device_id, aclrtDrvMemHandle &out_mem_handle,
size_t &out_g_size) {
aclError error_code;
size_t free_mem = 0, total = 0;
error_code = aclrtGetMemInfo(ACL_HBM_MEM, &free_mem, &total);
@@ -94,12 +94,11 @@ void start_daemon() {
total);
}
uint32_t device = 0;
aclrtPhysicalMemProp prop = {};
prop.handleType = ACL_MEM_HANDLE_TYPE_NONE;
prop.allocationType = ACL_MEM_ALLOCATION_TYPE_PINNED;
prop.memAttr = ACL_HBM_MEM_HUGE;
prop.location.id = device;
prop.location.id = device_id;
prop.location.type = ACL_MEM_LOCATION_TYPE_DEVICE;
prop.reserve = 0;
@@ -107,7 +106,8 @@ void start_daemon() {
error_code = aclrtMemGetAllocationGranularity(
&prop, ACL_RT_MEM_ALLOC_GRANULARITY_MINIMUM, &granularity);
if (error_code != 0) {
spdlog::error("aclrtMemGetAllocationGranularity failed, error_code: {}", error_code);
spdlog::error("aclrtMemGetAllocationGranularity failed, error_code: {}",
error_code);
throw std::runtime_error("aclrtMemGetAllocationGranularity failed");
} else {
spdlog::info("aclrtMemGetAllocationGranularity succeeded, granularity: {}",
@@ -118,59 +118,68 @@ void start_daemon() {
reserved_mem_size, free_mem);
throw std::runtime_error("Not enough free memory to reserve");
}
size_t g_size = free_mem - reserved_mem_size;
g_size = (g_size / granularity) * granularity;
out_g_size = free_mem - reserved_mem_size;
out_g_size = (out_g_size / granularity) * granularity;
// allocate physical memory
aclrtDrvMemHandle mem_handle;
error_code = aclrtMallocPhysical(&mem_handle, g_size, &prop, 0);
error_code = aclrtMallocPhysical(&out_mem_handle, out_g_size, &prop, 0);
if (error_code != 0) {
spdlog::error("aclrtMallocPhysical failed, error_code: {}", error_code);
throw std::runtime_error("aclrtMallocPhysical failed");
} else {
spdlog::info("aclrtMallocPhysical succeeded, size: {}", g_size);
spdlog::info("device {} aclrtMallocPhysical succeeded, size: {}", device_id,
out_g_size);
}
}
// // reserve address
// void *vmem_addr = nullptr;
// error_code = aclrtReserveMemAddress(&vmem_addr, g_size, 0, nullptr, 0);
// if (error_code != 0) {
// spdlog::error("aclrtReserveMemAddress failed, error_code: {}", error_code);
// throw std::runtime_error("aclrtReserveMemAddress failed");
// } else {
// spdlog::info("aclrtReserveMemAddress succeeded, vmem_addr: {}", vmem_addr);
// }
// // map
// error_code = aclrtMapMem(vmem_addr, g_size, 0, mem_handle, 0);
// if (error_code != 0) {
// spdlog::error("aclrtMapMem failed, error_code: {}", error_code);
// throw std::runtime_error("aclrtMapMem failed");
// } else {
// spdlog::info("aclrtMapMem succeeded, vmem_addr: {}", vmem_addr);
// }
// export
uint64_t shareable_handle;
error_code = aclrtMemExportToShareableHandle(
mem_handle, ACL_MEM_HANDLE_TYPE_NONE, ACL_RT_VMM_EXPORT_FLAG_DEFAULT,
&shareable_handle);
if (error_code != 0) {
spdlog::error("aclrtMemExportToShareableHandle failed, error_code: {}",
error_code);
throw std::runtime_error("aclrtMemExportToShareableHandle failed");
} else {
spdlog::info(
"aclrtMemExportToShareableHandle succeeded, shareable_handle: {}",
shareable_handle);
}
void start_daemon() {
init_acl();
std::vector<int> npu_ids = get_npu_ids();
std::vector<aclrtDrvMemHandle> mem_handles;
std::vector<uint64_t> shareable_handles;
// shm
shm_manager = new ShmManager();
shm_manager->set_gpu_info(g_size, shareable_handle);
for (int i = 0; i < npu_ids.size(); ++i) {
uint32_t device_id = i;
int npu_id = npu_ids[i];
spdlog::info("Setting up device id {} - npu id {}", device_id, npu_id);
aclError error_code = aclrtSetDevice(device_id);
if (error_code != ACL_ERROR_NONE) {
throw std::runtime_error("aclrtSetDevice failed with acl error code: " +
std::to_string(error_code) + " " + __FILE__ +
":" + std::to_string(__LINE__));
}
// alloc physical
aclrtDrvMemHandle mem_handle;
size_t g_size;
alloc_physical(device_id, mem_handle, g_size);
mem_handles.push_back(mem_handle);
// export
uint64_t shareable_handle;
error_code = aclrtMemExportToShareableHandle(
mem_handle, ACL_MEM_HANDLE_TYPE_NONE, ACL_RT_VMM_EXPORT_FLAG_DEFAULT,
&shareable_handle);
if (error_code != 0) {
spdlog::error("aclrtMemExportToShareableHandle failed, error_code: {}",
error_code);
throw std::runtime_error("aclrtMemExportToShareableHandle failed");
} else {
spdlog::info(
"aclrtMemExportToShareableHandle succeeded, shareable_handle: {}",
shareable_handle);
}
shm_manager->set_gpu_info(npu_id, g_size, shareable_handle);
shareable_handles.push_back(shareable_handle);
}
shm_manager->register_callback_on_worker_change(
[&](const std::vector<int32_t> &pids) {
reset_pids(pids, shareable_handle);
reset_pids(pids, shareable_handles);
});
// start busy loop
@@ -181,10 +190,12 @@ void start_daemon() {
shm_manager = nullptr;
// free physical memory
error_code = aclrtFreePhysical(mem_handle);
if (error_code != 0) {
spdlog::error("aclrtFreePhysical failed, error_code: {}", error_code);
throw std::runtime_error("aclrtFreePhysical failed");
for (auto mem_handle : mem_handles) {
aclError error_code = aclrtFreePhysical(mem_handle);
if (error_code != 0) {
spdlog::error("aclrtFreePhysical failed, error_code: {}", error_code);
throw std::runtime_error("aclrtFreePhysical failed");
}
}
}

24
csrc/idle_offload/shm_helper.h
View File

@@ -19,6 +19,7 @@
#define MAX_WORKERS 60
#define MAX_DEVICES 16
// static constexpr const char *SHM_NAME = "/vllm_acl_vnpu_offload_shm";
static inline std::string get_shm_name() {
const char *env_shm_name = getenv("VLLM_IDLE_OFFLOAD_SHM_NAME");
@@ -69,11 +70,14 @@ static inline uint64_t pack_unlocked_tgid(int32_t tgid) {
// mmap usually page-aligned
struct alignas(64) ShmHelper {
struct VramInfo {
uint64_t total_vmem_size;
uint64_t shareable_handle;
};
VramInfo vram_info[MAX_DEVICES]; // support max 16 NPUs
// GPU lock flag
std::atomic<uint64_t> gpu_flag;
uint64_t total_vmem_size;
uint64_t shareable_handle;
uint8_t _padding[64 - sizeof(std::atomic<uint64_t>) - 2 * sizeof(uint64_t)];
std::atomic<uint64_t> gpu_flag[MAX_DEVICES];
// uint8_t _padding1[64 - sizeof(std::atomic<uint64_t>)];
// request
enum RequestType: uint32_t {
@@ -105,14 +109,16 @@ struct alignas(64) ShmHelper {
WorkerHeartBeat heart_beats[MAX_WORKERS];
void init() {
gpu_flag.store(0, std::memory_order_release);
memset(vram_info, 0, sizeof(vram_info));
for (size_t i = 0; i < MAX_DEVICES; ++i) {
gpu_flag[i].store(0, std::memory_order_release);
}
req_ready.store(READY_STATE_NO_REQUEST, std::memory_order_release);
}
void set_gpu_info(uint64_t vmem_size, uint64_t shared_handle) {
total_vmem_size = vmem_size;
shareable_handle = shared_handle;
init();
void set_gpu_info(int gpu_id, uint64_t vmem_size, uint64_t shared_handle) {
vram_info[gpu_id].total_vmem_size = vmem_size;
vram_info[gpu_id].shareable_handle = shared_handle;
}
};

27
csrc/idle_offload/shm_manager.cpp
View File

@@ -1,7 +1,6 @@
#include "shm_manager.h"
#include <algorithm>
// === ShmManager ===
ShmManager::ShmManager() {
std::string shm_name = get_shm_name();
@@ -37,11 +36,12 @@ ShmManager::~ShmManager() {
shm_unlink(shm_name.c_str());
}
void ShmManager::set_gpu_info(uint64_t vmem_size, uint64_t shared_handle) {
shm_helper->set_gpu_info(vmem_size, shared_handle);
void ShmManager::set_gpu_info(int gpu_id, uint64_t vmem_size,
uint64_t shared_handle) {
shm_helper->set_gpu_info(gpu_id, vmem_size, shared_handle);
this->valid_gpu_ids.push_back(gpu_id);
}
void ShmManager::run_busy_loop() {
if (!cb_on_worker_change) {
spdlog::error("cb_on_worker_change is not set");
@@ -155,14 +155,17 @@ void ShmManager::check_heart_beats() {
shm_helper->heart_beats[i].timestamp.store(0,
std::memory_order_release);
// check dead lock
uint64_t gpu_flag =
shm_helper->gpu_flag.load(std::memory_order_acquire);
if (unpack_lock_field(gpu_flag) == 1 &&
unpack_tgid_field(gpu_flag) == tgid) {
// release lock held by dead worker
spdlog::warn("Releasing GPU lock held by dead worker TGID {}", tgid);
shm_helper->gpu_flag.store(pack_unlocked_tgid(tgid),
std::memory_order_release);
for (int gpu_id : valid_gpu_ids) {
uint64_t gpu_flag =
shm_helper->gpu_flag[gpu_id].load(std::memory_order_acquire);
if (unpack_lock_field(gpu_flag) == 1 &&
unpack_tgid_field(gpu_flag) == tgid) {
// release lock held by dead worker
spdlog::warn("Releasing GPU {} lock held by dead worker TGID {}",
gpu_id, tgid);
shm_helper->gpu_flag[gpu_id].store(pack_unlocked_tgid(tgid),
std::memory_order_release);
}
}
local_worker_tgids[i] = 0;
alive_worker_tgids.erase(std::remove(alive_worker_tgids.begin(),

3
csrc/idle_offload/shm_manager.h
View File

@@ -7,7 +7,7 @@ class ShmManager {
ShmManager();
~ShmManager();
void set_gpu_info(uint64_t vmem_size, uint64_t shared_handle);
void set_gpu_info(int gpu_id, uint64_t vmem_size, uint64_t shared_handle);
// request
void process_requests();
@@ -29,6 +29,7 @@ class ShmManager {
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;
};

90
csrc/idle_offload/shm_worker.cpp
View File

@@ -1,6 +1,5 @@
#include "shm_worker.h"
// === ShmWorker ===
ShmWorker::ShmWorker() {
std::string shm_name = get_shm_name();
@@ -29,16 +28,22 @@ ShmWorker::~ShmWorker() {
munmap(shm_helper, SHM_SIZE);
}
bool ShmWorker::register_worker(int32_t tgid, uint64_t *out_shareable_handle,
bool ShmWorker::register_worker(int32_t tgid, int gpu_id,
uint64_t *out_shareable_handle,
uint64_t *out_vmem_size) {
if (gpu_id < 0 || gpu_id >= MAX_DEVICES) {
spdlog::error("Invalid GPU ID {}", gpu_id);
throw std::runtime_error("Invalid GPU ID");
}
this->tgid = tgid;
this->gpu_id = gpu_id;
int slot = register_worker_shm();
if (slot == -1) {
return false;
}
*out_shareable_handle = shm_helper->shareable_handle;
*out_vmem_size = shm_helper->total_vmem_size;
*out_shareable_handle = shm_helper->vram_info[gpu_id].shareable_handle;
*out_vmem_size = shm_helper->vram_info[gpu_id].total_vmem_size;
stop_heart_beat.store(false, std::memory_order_release);
heart_beat_thread = std::thread(&ShmWorker::heart_beat_loop, this, slot);
@@ -68,47 +73,62 @@ void ShmWorker::heart_beat_loop(int slot) {
}
}
bool ShmWorker::lock_gpu() {
int retry_cnt = 0;
uint64_t old_flag = shm_helper->gpu_flag.load(std::memory_order_acquire);
bool ShmWorker::try_lock_gpu(bool &out_self_hold) {
static int retry_cnt = 0;
uint64_t old_flag =
shm_helper->gpu_flag[gpu_id].load(std::memory_order_acquire);
if (unpack_lock_field(old_flag) == 0) { // free
uint64_t new_flag = pack_locked_tgid(tgid);
if (shm_helper->gpu_flag[gpu_id].compare_exchange_weak(
old_flag, new_flag, std::memory_order_acq_rel,
std::memory_order_acquire)) {
spdlog::info("TGID {} acquired GPU {} lock", tgid, gpu_id);
int32_t prev_tgid = unpack_tgid_field(old_flag);
out_self_hold = prev_tgid == tgid;
retry_cnt = 0;
return true;
}
} else { // locked
if (unpack_tgid_field(old_flag) == tgid) {
spdlog::info("TGID {} already holds the GPU {} lock", tgid, gpu_id);
out_self_hold = true;
retry_cnt = 0;
return true;
}
}
// failed
if (++retry_cnt % 2000 == 0) {
spdlog::info(
"TGID {} trying to acquire GPU {} lock, current lock holder TGID {}",
tgid, gpu_id, unpack_tgid_field(old_flag));
}
out_self_hold = false;
return false;
}
bool ShmWorker::lock_gpu(bool &out_self_hold) {
while (true) {
if (unpack_lock_field(old_flag) == 0) {
uint64_t new_flag = pack_locked_tgid(tgid);
if (shm_helper->gpu_flag.compare_exchange_weak(old_flag, new_flag,
std::memory_order_acq_rel,
std::memory_order_acquire)) {
spdlog::info("TGID {} acquired GPU lock", tgid);
int32_t old_tgid = unpack_tgid_field(old_flag);
return old_tgid == tgid;
}
} else {
if (unpack_tgid_field(old_flag) == tgid) {
spdlog::info("TGID {} already holds the GPU lock", tgid);
return true;
}
if (try_lock_gpu(out_self_hold)) {
return true;
}
// failed
++retry_cnt;
if (retry_cnt % 1000 == 0) {
spdlog::info(
"TGID {} waiting for GPU lock, current lock holder TGID {}", tgid,
unpack_tgid_field(old_flag));
}
usleep(1000);
old_flag = shm_helper->gpu_flag.load(std::memory_order_acquire);
}
}
void ShmWorker::unlock_gpu() {
uint64_t old_flag = shm_helper->gpu_flag.load(std::memory_order_acquire);
uint64_t old_flag =
shm_helper->gpu_flag[gpu_id].load(std::memory_order_acquire);
if (unpack_tgid_field(old_flag) != tgid) {
spdlog::warn("previous gpu flag {} does not match expected locked flag for "
"TGID {}. This may be a bug, unless during startup.",
old_flag, tgid);
// spdlog::warn("previous gpu flag {} does not match expected locked flag for "
// "TGID {}. This may be a bug, unless during startup.",
// old_flag, tgid);
spdlog::info("TGID {} does not hold GPU {} lock", tgid, gpu_id);
} else {
uint64_t new_flag = pack_unlocked_tgid(tgid);
shm_helper->gpu_flag.store(new_flag, std::memory_order_release);
spdlog::info("TGID {} released GPU lock", tgid);
shm_helper->gpu_flag[gpu_id].store(new_flag, std::memory_order_release);
spdlog::info("TGID {} released GPU {} lock", tgid, gpu_id);
}
}
@@ -142,7 +162,7 @@ uint64_t ShmWorker::make_request(uint32_t type, uint64_t parameter) {
uint64_t response = shm_helper->request.response;
// set ready to 0
shm_helper->req_ready.store(ShmHelper::READY_STATE_NO_REQUEST,
std::memory_order_release);
std::memory_order_release);
return response;
}

12
csrc/idle_offload/shm_worker.h
View File

@@ -1,21 +1,27 @@
#pragma once
#include <vector>
#include <atomic>
#include <thread>
#include "shm_helper.h"
class ShmWorker {
public:
ShmWorker();
~ShmWorker();
bool register_worker(int32_t tgid, uint64_t *out_shareable_handle,
bool register_worker(int32_t tgid, int gpu_id, uint64_t *out_shareable_handle,
uint64_t *out_vmem_size);
bool lock_gpu();
bool try_lock_gpu(bool &out_self_hold);
bool lock_gpu(bool &out_self_hold);
void unlock_gpu();
private:
int32_t tgid;
int gpu_id;
ShmHelper *shm_helper;
std::thread heart_beat_thread;
std::atomic<bool> stop_heart_beat;