wehub-resource-sync
418 lines
16 KiB
C++
418 lines
16 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/threaded_engine.cc
|
|
* \brief The implementation for threaded serving engine in MLC LLM.
|
|
*/
|
|
#include "threaded_engine.h"
|
|
|
|
#include <tvm/ffi/extra/module.h>
|
|
#include <tvm/ffi/function.h>
|
|
#include <tvm/ffi/reflection/registry.h>
|
|
|
|
#include <atomic>
|
|
#include <condition_variable>
|
|
#include <mutex>
|
|
#include <optional>
|
|
|
|
#include "../support/json_parser.h"
|
|
#include "../support/module_vtable.h"
|
|
#include "../support/result.h"
|
|
#include "engine.h"
|
|
#include "request.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
using tvm::Device;
|
|
using namespace tvm::runtime;
|
|
|
|
/*! \brief The threaded engine instruction kind. */
|
|
enum class InstructionKind : int {
|
|
kAddRequest = 0,
|
|
kAbortRequest = 1,
|
|
kUnloadEngine = 2,
|
|
kReloadEngine = 3,
|
|
kResetEngine = 4,
|
|
kDebugCallFuncOnAllAllWorker = 5,
|
|
};
|
|
|
|
/*! \brief The implementation of ThreadedEngine. */
|
|
class ThreadedEngineImpl : public ThreadedEngine {
|
|
public:
|
|
void InitThreadedEngine(Device device, Optional<Function> request_stream_callback,
|
|
Optional<EventTraceRecorder> trace_recorder) final {
|
|
device_ = device;
|
|
TVM_FFI_ICHECK(request_stream_callback.has_value())
|
|
<< "ThreadedEngine requires request stream callback function, but it is not given.";
|
|
request_stream_callback_ = request_stream_callback.value();
|
|
trace_recorder_ = trace_recorder;
|
|
}
|
|
|
|
void Reload(String engine_config_json_str) final {
|
|
// NOTE: important to set this before, we send out
|
|
// reload instruction to the other threads
|
|
// otherwise there can be deadlocks
|
|
reload_finished_ = false;
|
|
bool need_notify = false;
|
|
{
|
|
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
|
instruction_queue_.emplace_back(InstructionKind::kReloadEngine,
|
|
std::move(engine_config_json_str));
|
|
++pending_request_operation_cnt_;
|
|
need_notify = engine_waiting_;
|
|
}
|
|
if (need_notify) {
|
|
background_loop_cv_.notify_one();
|
|
}
|
|
{
|
|
std::unique_lock<std::mutex> lock(reload_unload_mutex_);
|
|
reload_unload_cv_.wait(lock, [this] { return reload_finished_; });
|
|
}
|
|
}
|
|
|
|
void Unload() final {
|
|
// NOTE: important to set this before, we send out
|
|
// reload instruction to the other threads
|
|
// otherwise there can be deadlocks
|
|
// e.g. the other thread finish unload job and set the flag to true
|
|
// then we set it back to false
|
|
unload_finished_ = false;
|
|
bool need_notify = false;
|
|
{
|
|
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
|
instruction_queue_.emplace_back(InstructionKind::kUnloadEngine, ObjectRef(nullptr));
|
|
++pending_request_operation_cnt_;
|
|
need_notify = engine_waiting_;
|
|
}
|
|
if (need_notify) {
|
|
background_loop_cv_.notify_one();
|
|
}
|
|
{
|
|
std::unique_lock<std::mutex> lock(reload_unload_mutex_);
|
|
reload_unload_cv_.wait(lock, [this] { return unload_finished_; });
|
|
}
|
|
}
|
|
|
|
void Reset() final {
|
|
bool need_notify = false;
|
|
{
|
|
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
|
instruction_queue_.emplace_back(InstructionKind::kResetEngine, ObjectRef(nullptr));
|
|
++pending_request_operation_cnt_;
|
|
need_notify = engine_waiting_;
|
|
}
|
|
if (need_notify) {
|
|
background_loop_cv_.notify_one();
|
|
}
|
|
}
|
|
|
|
void AddRequest(Request request) final {
|
|
bool need_notify = false;
|
|
{
|
|
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
|
instruction_queue_.emplace_back(InstructionKind::kAddRequest, request);
|
|
++pending_request_operation_cnt_;
|
|
need_notify = engine_waiting_;
|
|
}
|
|
if (need_notify) {
|
|
background_loop_cv_.notify_one();
|
|
}
|
|
}
|
|
|
|
void AbortRequest(const String& request_id) final {
|
|
bool need_notify = false;
|
|
{
|
|
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
|
instruction_queue_.emplace_back(InstructionKind::kAbortRequest, request_id);
|
|
++pending_request_operation_cnt_;
|
|
need_notify = engine_waiting_;
|
|
}
|
|
if (need_notify) {
|
|
background_loop_cv_.notify_one();
|
|
}
|
|
}
|
|
|
|
void RunBackgroundLoop() final {
|
|
// The local vectors that load the requests from critical regions.
|
|
std::vector<std::pair<InstructionKind, Any>> local_instruction_queue;
|
|
|
|
while (!exit_now_.load(std::memory_order_relaxed)) {
|
|
{
|
|
std::unique_lock<std::mutex> lock(background_loop_mutex_);
|
|
engine_waiting_ = true;
|
|
background_loop_cv_.wait(lock, [this] {
|
|
return (background_engine_ != nullptr && !background_engine_->Empty()) ||
|
|
pending_request_operation_cnt_.load() > 0 ||
|
|
exit_now_.load(std::memory_order_relaxed);
|
|
});
|
|
engine_waiting_ = false;
|
|
local_instruction_queue = instruction_queue_;
|
|
instruction_queue_.clear();
|
|
pending_request_operation_cnt_ = 0;
|
|
}
|
|
for (const auto& [kind, arg] : local_instruction_queue) {
|
|
if (kind == InstructionKind::kAddRequest) {
|
|
TVM_FFI_ICHECK(background_engine_ != nullptr) << "Background engine is not loaded.";
|
|
background_engine_->AddRequest(arg.as_or_throw<Request>());
|
|
} else if (kind == InstructionKind::kAbortRequest) {
|
|
// in a rare case, abort request can happen after unloading
|
|
// aka background engine is nullptr
|
|
// this happens when the on going generation was interrupted
|
|
// the engine get unloaded, and then abort was called.
|
|
// it is safe to ignore these abort in such case
|
|
if (background_engine_ != nullptr) {
|
|
background_engine_->AbortRequest(arg.as_or_throw<String>());
|
|
}
|
|
} else if (kind == InstructionKind::kUnloadEngine) {
|
|
EngineUnloadImpl();
|
|
} else if (kind == InstructionKind::kReloadEngine) {
|
|
EngineUnloadImpl();
|
|
EngineReloadImpl(arg.as_or_throw<String>());
|
|
} else if (kind == InstructionKind::kResetEngine) {
|
|
if (background_engine_ != nullptr) {
|
|
background_engine_->Reset();
|
|
}
|
|
} else if (kind == InstructionKind::kDebugCallFuncOnAllAllWorker) {
|
|
TVM_FFI_ICHECK(background_engine_ != nullptr) << "Background engine is not loaded.";
|
|
Array<Any> packed_args = arg.as_or_throw<Array<Any>>();
|
|
background_engine_->DebugCallFuncOnAllAllWorker(
|
|
packed_args[0].as_or_throw<String>(), packed_args[1].as_or_throw<Optional<String>>());
|
|
} else {
|
|
LOG(FATAL) << "Cannot reach here";
|
|
}
|
|
}
|
|
if (background_engine_ != nullptr) {
|
|
background_engine_->Step();
|
|
}
|
|
}
|
|
}
|
|
|
|
void RunBackgroundStreamBackLoop() final {
|
|
// The local vectors that load the request stream callback inputs from critical regions.
|
|
std::vector<Array<RequestStreamOutput>> local_request_stream_callback_inputs;
|
|
std::vector<RequestStreamOutput> flattened_callback_inputs;
|
|
|
|
while (!exit_now_.load(std::memory_order_relaxed)) {
|
|
{
|
|
std::unique_lock<std::mutex> lock(request_stream_callback_mutex_);
|
|
stream_callback_waiting_ = true;
|
|
request_stream_callback_cv_.wait(lock, [this] {
|
|
return pending_request_stream_callback_cnt_.load() > 0 ||
|
|
exit_now_.load(std::memory_order_relaxed);
|
|
});
|
|
stream_callback_waiting_ = false;
|
|
|
|
local_request_stream_callback_inputs = request_stream_callback_inputs_;
|
|
request_stream_callback_inputs_.clear();
|
|
pending_request_stream_callback_cnt_ = 0;
|
|
}
|
|
for (const Array<RequestStreamOutput>& callback_inputs :
|
|
local_request_stream_callback_inputs) {
|
|
for (const RequestStreamOutput& callback_input : callback_inputs) {
|
|
flattened_callback_inputs.push_back(callback_input);
|
|
}
|
|
}
|
|
if (!flattened_callback_inputs.empty()) {
|
|
request_stream_callback_(Array<RequestStreamOutput>(flattened_callback_inputs));
|
|
}
|
|
flattened_callback_inputs.clear();
|
|
}
|
|
}
|
|
|
|
void ExitBackgroundLoop() final {
|
|
{
|
|
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
|
exit_now_.store(true);
|
|
}
|
|
background_loop_cv_.notify_one();
|
|
request_stream_callback_cv_.notify_one();
|
|
}
|
|
|
|
/************** Query/Profile/Debug **************/
|
|
|
|
GenerationConfig GetDefaultGenerationConfig() const final {
|
|
TVM_FFI_ICHECK(default_generation_config_.has_value())
|
|
<< "The default generation config has not been set.";
|
|
return default_generation_config_.value();
|
|
}
|
|
|
|
Request CreateRequest(String id, Array<Data> inputs, String generation_cfg_json_str) const {
|
|
json::Object config = json::ParseToJSONObject(generation_cfg_json_str);
|
|
auto gen_config = GenerationConfig::FromJSON(config, GetDefaultGenerationConfig());
|
|
TVM_FFI_ICHECK(gen_config.IsOk()) << gen_config.UnwrapErr();
|
|
return Request(std::move(id), std::move(inputs), gen_config.Unwrap());
|
|
}
|
|
|
|
EngineConfig GetCompleteEngineConfig() const final {
|
|
TVM_FFI_ICHECK(complete_engine_config_.has_value()) << "The engine config has not been set.";
|
|
return complete_engine_config_.value();
|
|
}
|
|
|
|
String GetCompleteEngineConfigJSONString() const {
|
|
return GetCompleteEngineConfig()->AsJSONString();
|
|
}
|
|
|
|
void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) final {
|
|
bool need_notify = false;
|
|
{
|
|
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
|
instruction_queue_.emplace_back(InstructionKind::kDebugCallFuncOnAllAllWorker,
|
|
Array<Any>{func_name, func_args});
|
|
++pending_request_operation_cnt_;
|
|
need_notify = engine_waiting_;
|
|
}
|
|
if (need_notify) {
|
|
background_loop_cv_.notify_one();
|
|
}
|
|
}
|
|
|
|
private:
|
|
void EngineReloadImpl(const std::string& engine_config_json_str) {
|
|
auto frequest_stream_callback_wrapper = [this](Array<RequestStreamOutput> delta_outputs) {
|
|
bool need_notify = false;
|
|
{
|
|
std::lock_guard<std::mutex> lock(request_stream_callback_mutex_);
|
|
request_stream_callback_inputs_.push_back(std::move(delta_outputs));
|
|
++pending_request_stream_callback_cnt_;
|
|
need_notify = stream_callback_waiting_;
|
|
}
|
|
if (need_notify) {
|
|
request_stream_callback_cv_.notify_one();
|
|
}
|
|
};
|
|
|
|
FRequestStreamCallback request_stream_callback(frequest_stream_callback_wrapper);
|
|
Result<EngineCreationOutput> output_res =
|
|
Engine::Create(engine_config_json_str, device_, request_stream_callback, trace_recorder_);
|
|
TVM_FFI_ICHECK(output_res.IsOk()) << output_res.UnwrapErr();
|
|
EngineCreationOutput output = output_res.Unwrap();
|
|
background_engine_ = std::move(output.reloaded_engine);
|
|
default_generation_config_ = output.default_generation_cfg;
|
|
complete_engine_config_ = output.completed_engine_config;
|
|
{
|
|
// Wake up the thread waiting for reload finish.
|
|
std::lock_guard<std::mutex> lock(reload_unload_mutex_);
|
|
reload_finished_ = true;
|
|
}
|
|
reload_unload_cv_.notify_one();
|
|
}
|
|
|
|
void EngineUnloadImpl() {
|
|
if (background_engine_ != nullptr) {
|
|
background_engine_->AbortAllRequests();
|
|
background_engine_ = nullptr;
|
|
// Clear the allocated memory in cached memory pool.
|
|
static Function fclear_memory_manager =
|
|
Function::GetGlobalRequired("vm.builtin.memory_manager.clear");
|
|
fclear_memory_manager();
|
|
default_generation_config_ = std::nullopt;
|
|
complete_engine_config_ = std::nullopt;
|
|
}
|
|
{
|
|
// Wake up the thread waiting for unload finish.
|
|
std::lock_guard<std::mutex> lock(reload_unload_mutex_);
|
|
unload_finished_ = true;
|
|
}
|
|
reload_unload_cv_.notify_one();
|
|
}
|
|
|
|
/*! \brief The device to run models on. */
|
|
Device device_;
|
|
/*! \brief The background normal engine for request processing. */
|
|
std::unique_ptr<Engine> background_engine_;
|
|
/*! \brief The request stream callback. */
|
|
Function request_stream_callback_;
|
|
/*! \brief Event trace recorder. */
|
|
Optional<EventTraceRecorder> trace_recorder_;
|
|
|
|
/*! \brief complete engine config. */
|
|
Optional<EngineConfig> complete_engine_config_;
|
|
/*! \brief The default generation config. */
|
|
Optional<GenerationConfig> default_generation_config_;
|
|
|
|
/*! \brief The mutex ensuring only one thread can access critical regions. */
|
|
std::mutex background_loop_mutex_;
|
|
std::mutex request_stream_callback_mutex_;
|
|
std::mutex reload_unload_mutex_;
|
|
/*! \brief The condition variable preventing threaded engine from spinning. */
|
|
std::condition_variable background_loop_cv_;
|
|
std::condition_variable request_stream_callback_cv_;
|
|
std::condition_variable reload_unload_cv_;
|
|
/*! \brief A boolean flag denoting if the engine needs to exit background loop. */
|
|
std::atomic<bool> exit_now_ = false;
|
|
|
|
/************** Critical Regions **************/
|
|
/*!
|
|
* \brief The instruction queue for the threaded engine.
|
|
* The instructions include:
|
|
* - requests to add into the background engine,
|
|
* - requests to abort from the background engine,
|
|
* - engine unload/reload,
|
|
* - and other debugging instructions.
|
|
* Elements are sended from other threads and consumed by
|
|
* the threaded engine in the background loop.
|
|
*/
|
|
std::vector<std::pair<InstructionKind, Any>> instruction_queue_;
|
|
/*!
|
|
* \brief The delta outputs to pass through callback.
|
|
* Elements are sended from the background loop thread and
|
|
* consumed by the foreground thread.
|
|
*/
|
|
std::vector<Array<RequestStreamOutput>> request_stream_callback_inputs_;
|
|
/*!
|
|
* \brief Number of pending request operations, should be the size of
|
|
* `requests_to_add_` and `requests_to_abort_`.
|
|
*/
|
|
std::atomic<int> pending_request_operation_cnt_ = 0;
|
|
/*!
|
|
* \brief Number of pending request stream callback invocations.
|
|
* It should be the size of `request_stream_callback_inputs_`.
|
|
*/
|
|
std::atomic<int> pending_request_stream_callback_cnt_ = 0;
|
|
/*! \brief A boolean flag indicating if the engine is waiting for new requests/aborts. */
|
|
bool engine_waiting_ = false;
|
|
/*! \brief A boolean flag indicating if the stream callback loop is waiting. */
|
|
bool stream_callback_waiting_ = false;
|
|
/*! \brief A boolean indicating if the engine reload has finished. */
|
|
bool reload_finished_ = false;
|
|
/*! \brief A boolean indicating if the engine unload has finished. */
|
|
bool unload_finished_ = false;
|
|
};
|
|
|
|
/*! \brief The implementation of ThreadedEngine. */
|
|
class ThreadedEngineModule : public ThreadedEngineImpl, public ffi::ModuleObj {
|
|
public:
|
|
TVM_MODULE_VTABLE_BEGIN("mlc.serve.async_threaded_engine");
|
|
TVM_MODULE_VTABLE_ENTRY("init_threaded_engine", &ThreadedEngineImpl::InitThreadedEngine);
|
|
TVM_MODULE_VTABLE_ENTRY("reload", &ThreadedEngineImpl::Reload);
|
|
TVM_MODULE_VTABLE_ENTRY("add_request", &ThreadedEngineImpl::AddRequest);
|
|
TVM_MODULE_VTABLE_ENTRY("create_request", &ThreadedEngineImpl::CreateRequest);
|
|
TVM_MODULE_VTABLE_ENTRY("abort_request", &ThreadedEngineImpl::AbortRequest);
|
|
TVM_MODULE_VTABLE_ENTRY("run_background_loop", &ThreadedEngineImpl::RunBackgroundLoop);
|
|
TVM_MODULE_VTABLE_ENTRY("run_background_stream_back_loop",
|
|
&ThreadedEngineImpl::RunBackgroundStreamBackLoop);
|
|
TVM_MODULE_VTABLE_ENTRY("exit_background_loop", &ThreadedEngineImpl::ExitBackgroundLoop);
|
|
TVM_MODULE_VTABLE_ENTRY("get_complete_engine_config",
|
|
&ThreadedEngineImpl::GetCompleteEngineConfigJSONString);
|
|
TVM_MODULE_VTABLE_ENTRY("reset", &ThreadedEngineImpl::Reset);
|
|
TVM_MODULE_VTABLE_ENTRY("debug_call_func_on_all_worker",
|
|
&ThreadedEngineImpl::DebugCallFuncOnAllAllWorker);
|
|
TVM_MODULE_VTABLE_END();
|
|
};
|
|
|
|
TVM_FFI_STATIC_INIT_BLOCK() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::GlobalDef().def("mlc.serve.create_threaded_engine",
|
|
[]() { return Module(tvm::ffi::make_object<ThreadedEngineModule>()); });
|
|
}
|
|
|
|
std::unique_ptr<ThreadedEngine> ThreadedEngine::Create() {
|
|
std::unique_ptr<ThreadedEngineImpl> threaded_engine = std::make_unique<ThreadedEngineImpl>();
|
|
return std::move(threaded_engine);
|
|
}
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|