From d5cff81c1efcc82fcc10eb9920ad93300898d944 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Thu, 5 Aug 2021 18:13:50 +0800 Subject: [PATCH] fix(mge): limit task queue size GitOrigin-RevId: a25c3390fc9742b48b6617f2d8887b3769aa4160 --- imperative/src/impl/interpreter/interpreter_impl.h | 13 +++++++++++-- src/core/include/megbrain/utils/thread_impl_1.h | 7 +++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/imperative/src/impl/interpreter/interpreter_impl.h b/imperative/src/impl/interpreter/interpreter_impl.h index bea3cfeb7..28dd963db 100644 --- a/imperative/src/impl/interpreter/interpreter_impl.h +++ b/imperative/src/impl/interpreter/interpreter_impl.h @@ -157,10 +157,19 @@ private: // set max_spin=0 to prevent Queue fetch task in busy wait manner. // this won't affect throughput when python interpreter is sending enough task, // but will significantly save CPU time when waiting for task, e.g. wait for data input - // limit pending tasks to 1000000 + // limit pending tasks to 10000 WorkQueue(ChannelImpl* owner) - : AsyncQueueSC(0, 1000000), m_owner(owner) { + : AsyncQueueSC(0, 10000), m_owner(owner) { sys::set_thread_name("interpreter"); + if (const char* env_val = MGB_GETENV("MEGENGINE_ASYNC_QUEUE_SIZE")) { + int len = strlen(env_val); + for (int i = 0; i < len; i ++) { + mgb_assert(env_val[i] >= '0' && env_val[i] <= '9', "async queue size should be an integer"); + } + size_t val; + sscanf(env_val, "%zu", &val); + update_max_items(val); + } } void process_one_task(IdentifiedCommand& icmd) { m_owner->process_one_task(icmd); diff --git a/src/core/include/megbrain/utils/thread_impl_1.h b/src/core/include/megbrain/utils/thread_impl_1.h index ed59ef26a..965220119 100644 --- a/src/core/include/megbrain/utils/thread_impl_1.h +++ b/src/core/include/megbrain/utils/thread_impl_1.h @@ -264,6 +264,13 @@ namespace mgb { return m_synchronizer.check_finished(); } + void update_max_items(ptrdiff_t max_items) { + if (max_items >= 0) { + // -1 / 2 == 0 + m_block_quota = (max_items - 1) / BLOCK_SIZE + 1; + } + } + protected: ~AsyncQueueSC() noexcept = default; -- GitLab