mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[CCF Archive] Store object type eviction policy submission #3

Closed
kancel wants to merge 382 commits from kancel:ccf-archive-pr2746 into main
pull from: kancel:ccf-archive-pr2746
merge into: mooncake-track:main
Conversation 0 Commits 382 Files Changed 654 +41 -33
1 changed files with 41 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit c7b7526a5a - Show all commits

74
mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp
View File

@ -62,8 +62,18 @@ static bool isCudaMemory(void* addr) {
cudaPointerAttributes attributes;
auto status = cudaPointerGetAttributes(&attributes, addr);
if (status != cudaSuccess) return false;
if (attributes.type == cudaMemoryTypeDevice) return true;
return false;
return attributes.type == cudaMemoryTypeDevice;
}
// Returns the CUDA device ordinal if addr is device memory, or -1 otherwise.
// Callers must call cudaSetDevice before any cudaMemcpy to avoid implicit
// GPU 0 context creation.
static int getCudaDeviceId(void* addr) {
cudaPointerAttributes attributes;
auto status = cudaPointerGetAttributes(&attributes, addr);
if (status != cudaSuccess) return -1;
if (attributes.type == cudaMemoryTypeDevice) return attributes.device;
return -1;
}
#endif
@ -143,12 +153,15 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
}
char* dram_buffer = addr + total_transferred_bytes_;
int cuda_device = -1;
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
if (isCudaMemory(addr)) {
cuda_device = getCudaDeviceId(addr);
if (cuda_device >= 0) {
dram_buffer = new char[buffer_size];
cudaSetDevice(cuda_device);
cudaError_t cuda_status =
cudaMemcpy(dram_buffer, addr + total_transferred_bytes_,
buffer_size, cudaMemcpyDefault);
@ -165,12 +178,12 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
asio::async_write(
*socket_, asio::buffer(dram_buffer, buffer_size),
[this, addr, dram_buffer, self](const asio::error_code& ec,
std::size_t transferred_bytes) {
[this, addr, dram_buffer, cuda_device, self](
const asio::error_code& ec, std::size_t transferred_bytes) {
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
if (isCudaMemory(addr)) {
if (cuda_device >= 0) {
delete[] dram_buffer;
}
#endif
@ -204,21 +217,20 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
}
char* dram_buffer = addr + total_transferred_bytes_;
int cuda_device = -1;
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
bool is_cuda_memory = isCudaMemory(addr);
if (is_cuda_memory) {
cuda_device = getCudaDeviceId(addr);
if (cuda_device >= 0) {
dram_buffer = new char[buffer_size];
}
#else
bool is_cuda_memory = false;
#endif
asio::async_read(
*socket_, asio::buffer(dram_buffer, buffer_size),
[this, addr, dram_buffer, is_cuda_memory, self](
[this, addr, dram_buffer, cuda_device, self](
const asio::error_code& ec, std::size_t transferred_bytes) {
if (ec) {
// If client closed connection (EOF), this is normal - don't
@ -233,18 +245,15 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
<< " (value: " << ec.value() << ")";
}
session_mutex_.unlock();
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
if (is_cuda_memory) delete[] dram_buffer;
#endif
if (cuda_device >= 0) delete[] dram_buffer;
return; // Connection will be closed
}
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
if (is_cuda_memory) {
if (cuda_device >= 0) {
cudaSetDevice(cuda_device);
cudaError_t cuda_status =
cudaMemcpy(addr + total_transferred_bytes_, dram_buffer,
transferred_bytes, cudaMemcpyDefault);
@ -340,21 +349,20 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
}
char* dram_buffer = addr + total_transferred_bytes_;
int cuda_device = -1;
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
bool is_cuda_memory = isCudaMemory(addr);
if (is_cuda_memory) {
cuda_device = getCudaDeviceId(addr);
if (cuda_device >= 0) {
dram_buffer = new char[buffer_size];
}
#else
bool is_cuda_memory = false;
#endif
asio::async_read(
*socket_, asio::buffer(dram_buffer, buffer_size),
[this, addr, dram_buffer, is_cuda_memory, self](
[this, addr, dram_buffer, cuda_device, self](
const asio::error_code& ec, std::size_t transferred_bytes) {
if (ec) {
LOG(ERROR)
@ -366,7 +374,7 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
// Post entire cleanup to ensure it runs after callback
// returns
asio::post(socket_->get_executor(),
[this, self, dram_buffer, is_cuda_memory,
[this, self, dram_buffer, cuda_device,
on_finalize = std::move(on_finalize_),
on_complete = std::move(on_complete_)]() {
if (on_finalize)
@ -374,7 +382,7 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
if (is_cuda_memory) delete[] dram_buffer;
if (cuda_device >= 0) delete[] dram_buffer;
#endif
session_mutex_.unlock();
if (on_complete) on_complete();
@ -385,7 +393,8 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
if (is_cuda_memory) {
if (cuda_device >= 0) {
cudaSetDevice(cuda_device);
cudaError_t cuda_status =
cudaMemcpy(addr + total_transferred_bytes_, dram_buffer,
transferred_bytes, cudaMemcpyDefault);
@ -438,12 +447,15 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
}
char* dram_buffer = addr + total_transferred_bytes_;
int cuda_device = -1;
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
if (isCudaMemory(addr)) {
cuda_device = getCudaDeviceId(addr);
if (cuda_device >= 0) {
dram_buffer = new char[buffer_size];
cudaSetDevice(cuda_device);
cudaError_t cuda_status =
cudaMemcpy(dram_buffer, addr + total_transferred_bytes_,
buffer_size, cudaMemcpyDefault);
@ -469,15 +481,11 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
asio::async_write(
*socket_, asio::buffer(dram_buffer, buffer_size),
[this, addr, dram_buffer, self](const asio::error_code& ec,
std::size_t transferred_bytes) {
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
defined(USE_COREX)
if (isCudaMemory(addr)) {
[this, addr, dram_buffer, cuda_device, self](
const asio::error_code& ec, std::size_t transferred_bytes) {
if (cuda_device >= 0) {
delete[] dram_buffer;
}
#endif
if (ec) {
LOG(ERROR)
<< "ClientSession::writeBody failed. "