Skip to content

[https://nvbugs/6104831][fix] disagg buffer-pool RAII and NIXL agent lifetime fixes #14739

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
chienchunhung wants to merge 1 commit into from
Closed

[https://nvbugs/6104831][fix] disagg buffer-pool RAII and NIXL agent lifetime fixes #14739

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
186 changes: 178 additions & 8 deletions cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,83 @@
#include "tensorrt_llm/common/logger.h"
#include "tensorrt_llm/common/opUtils.h"

#include <exception>
#include <mutex>

namespace tensorrt_llm::batch_manager
{

namespace
{

char const* bufferKindName(BufferKind kind)
{
switch (kind)
{
case BufferKind::kKV: return "kv";
case BufferKind::kKV_INDEXER: return "kv_indexer";
case BufferKind::kRNN: return "rnn";
}
return "unknown";
}

} // namespace

void BufferIndexHolder::release() noexcept
{
// Happy-path release: frees the slot and disarms the holder in one
// noexcept call. Used in place of an older detach() + explicit
// freeBufferIndex*() sequence so a throw between the two calls cannot
// leave the holder in a partially-released state.
if (!mHeld || mMgr == nullptr)
{
return;
}
try
{
if (mIsRecv)
{
mMgr->freeBufferIndexForRecv(mIndex);
}
else
{
mMgr->freeBufferIndexForSend(mIndex);
}
}
catch (...)
{
// Swallow; the destructor must be noexcept and any exit path that
// failed to release explicitly relies on this fallback to free the
// slot.
}
mHeld = false;
}

void BufferIndexHolder::poison() noexcept
{
if (!mHeld || mMgr == nullptr)
{
return;
}
try
{
if (mIsRecv)
{
mMgr->poisonBufferIndexForRecv(mIndex);
}
else
{
mMgr->poisonBufferIndexForSend(mIndex);
}
}
catch (...)
{
// poisonBufferIndex is noexcept; keep this as belt-and-suspenders so
// fail-closed cleanup cannot throw from an exception path.
}
mHeld = false;
}

BaseTransBufferManager::BaseTransBufferManager(
size_t transferBufferSize, nvinfer1::DataType dataType, std::optional<size_t> maxNumTokens)
: mDataType{dataType}
Expand Down Expand Up @@ -54,26 +126,40 @@ BaseTransBufferManager::BaseTransBufferManager(
allocateBuffer();
}

std::optional<int> BaseTransBufferManager::assignBufferIndexForSend()
std::optional<int> BaseTransBufferManager::assignBufferIndexForSend(
std::atomic<bool> const* perRequestCancel, int64_t waitSliceMs, std::optional<uint64_t> requestIdForLog)
{
return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer);
return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer, perRequestCancel,
waitSliceMs, requestIdForLog);
}

void BaseTransBufferManager::freeBufferIndexForSend(std::optional<int> bufferId)
{
freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer);
}

std::optional<int> BaseTransBufferManager::assignBufferIndexForRecv()
void BaseTransBufferManager::poisonBufferIndexForSend(std::optional<int> bufferId) noexcept
{
poisonBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer, "send");
}

std::optional<int> BaseTransBufferManager::assignBufferIndexForRecv(
std::atomic<bool> const* perRequestCancel, int64_t waitSliceMs, std::optional<uint64_t> requestIdForLog)
{
return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer);
return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer, perRequestCancel,
waitSliceMs, requestIdForLog);
}

void BaseTransBufferManager::freeBufferIndexForRecv(std::optional<int> bufferId)
{
freeBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer);
}

void BaseTransBufferManager::poisonBufferIndexForRecv(std::optional<int> bufferId) noexcept
{
poisonBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer, "recv");
}

std::tuple<std::vector<runtime::ITensor::SharedPtr>, size_t, bool> BaseTransBufferManager::getOrAllocateSendBuffers(
std::optional<int> bufferId, int targetNum, std::vector<size_t> const& requestedNumberOfElements,
runtime::BufferManager const& bufferManagerToUse)
Expand Down Expand Up @@ -225,16 +311,43 @@ void BaseTransBufferManager::allocateBuffer()
}
}

std::optional<int> BaseTransBufferManager::assignBufferIndex(
ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer)
std::optional<int> BaseTransBufferManager::assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount,
bool onlyUseDynamicBuffer, std::atomic<bool> const* perRequestCancel, int64_t waitSliceMs,
std::optional<uint64_t> requestIdForLog)
{
if (onlyUseDynamicBuffer)
{
TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed),
"Cannot assign dynamic cache transfer buffer kind=%s because a previous transfer left dynamic transfer "
"memory poisoned. The process must restart before these memory ranges can be safely reused.",
bufferKindName(getBufferKind()));
return std::nullopt;
}
// Bounded wait_for loop so a cancel fired on this request while parked
// here can interrupt the wait via the per-request cancel atomic, and so
// mTerminate (flipped between slices) keeps the drain worker responsive
// to shutdown.
std::unique_lock lk(resource.mBuffersMutex);
resource.mBuffersCV.wait(
lk, [&resource, bufferCount]() { return static_cast<size_t>(resource.mConcurrence) < bufferCount; });
auto const predicate = [&resource, bufferCount]()
{
return resource.mPoisoned.load(std::memory_order_relaxed)
|| static_cast<size_t>(resource.mConcurrence) < bufferCount;
};
auto const slice = std::chrono::milliseconds{waitSliceMs};
while (!predicate())
{
resource.mBuffersCV.wait_for(lk, slice);
if (perRequestCancel != nullptr && perRequestCancel->load(std::memory_order_relaxed))
{
auto const reqIdStr
= requestIdForLog.has_value() ? std::to_string(requestIdForLog.value()) : std::string{"?"};
TLLM_THROW("assignBufferIndex cancelled via perRequestCancel (reqId=%s)", reqIdStr.c_str());
}
}
TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed),
"Cannot assign cache transfer buffer kind=%s because a previous transfer left the buffer pool poisoned. "
"The process must restart before these memory ranges can be safely reused.",
bufferKindName(getBufferKind()));
int bufferId = -1;
for (size_t i = 0; i < bufferCount; i++)
{
Expand Down Expand Up @@ -264,13 +377,70 @@ void BaseTransBufferManager::freeBufferIndex(
TLLM_CHECK(static_cast<size_t>(bufferId.value()) < bufferCount);
{
std::scoped_lock lk(resource.mBuffersMutex);
if (resource.mBufferIndexFlag[bufferId.value()] == 2)
{
TLLM_LOG_ERROR("Refusing to free poisoned cache transfer buffer kind=%s index=%d",
bufferKindName(getBufferKind()), bufferId.value());
return;
}
resource.mBufferIndexFlag[bufferId.value()] = 0;
}
resource.mConcurrence--;
resource.mBuffersCV.notify_one();
}
}

void BaseTransBufferManager::poisonBufferIndex(ConcurrenceResource& resource, std::optional<int> bufferId,
size_t bufferCount, bool onlyUseDynamicBuffer, char const* direction) noexcept
{
resource.mPoisoned.store(true, std::memory_order_relaxed);

if (onlyUseDynamicBuffer)
{
TLLM_LOG_ERROR(
"Poisoned dynamic %s cache transfer buffer kind=%s. Dynamic transfer memory cannot be safely reused; "
"the process must restart.",
direction, bufferKindName(getBufferKind()));
resource.mBuffersCV.notify_all();
return;
}

if (!bufferId.has_value())
{
TLLM_LOG_ERROR("Poisoned unknown %s cache transfer buffer kind=%s. The process must restart.", direction,
bufferKindName(getBufferKind()));
resource.mBuffersCV.notify_all();
return;
}

try
{
TLLM_CHECK(static_cast<size_t>(bufferId.value()) < bufferCount);
{
std::scoped_lock lk(resource.mBuffersMutex);
if (resource.mBufferIndexFlag[bufferId.value()] == 1)
{
resource.mBufferIndexFlag[bufferId.value()] = 2;
}
}
TLLM_LOG_ERROR(
"Poisoned %s cache transfer buffer kind=%s index=%d. The slot will not be returned to the pool because "
"transport quiescence is unknown; restart the process before serving more KV transfers.",
direction, bufferKindName(getBufferKind()), bufferId.value());
}
catch (std::exception const& e)
{
TLLM_LOG_ERROR("Exception while poisoning %s cache transfer buffer kind=%s index=%d: %s", direction,
bufferKindName(getBufferKind()), bufferId.value_or(-1), e.what());
}
catch (...)
{
TLLM_LOG_ERROR("Unknown exception while poisoning %s cache transfer buffer kind=%s index=%d", direction,
bufferKindName(getBufferKind()), bufferId.value_or(-1));
}
resource.mBuffersCV.notify_all();
}

size_t BaseTransBufferManager::getRecvBufferCount()
{
return mRecvBufferCount;
Expand Down
Loading
Loading