rate limit

Author: Lijiachen1018

ucm/integration/vllm/ucm_connector.py

@@ -354,7 +354,7 @@ def get_num_new_matched_tokens(
         except RuntimeError as e:
             external_hit_blocks = 0
             logger.error(f"request {request.request_id} look up error. {e}")
-        logger.info(
+        logger.info_once(
             f"request_id: {request.request_id}, "
             f"total_blocks_num: {len(ucm_block_ids)}, "
             f"hit hbm: {hbm_hit_block_num}, "

ucm/logger.py

@@ -105,7 +105,7 @@ def format_log_msg(msg, *args) -> str:
             return msg % args
         return msg
 
-    def log(self, levelno, message, *args, exc_info=None, scope=None):
+    def log(self, levelno, message, *args, exc_info=None, scope=None, rate_limit=False):
         level = LevelMap[levelno]
         frame = inspect.currentframe()
         caller_frame = frame.f_back.f_back
@@ -116,6 +116,9 @@ def log(self, levelno, message, *args, exc_info=None, scope=None):
         if exc_info:
             exc_text = self.format_exception(exc_info)
             msg = msg + "\n" + exc_text
+        if rate_limit:
+            ucmlogger.log_rate_limit(level, file, func, line, msg, message)
+            return
         ucmlogger.log(level, file, func, line, msg)
 
     @staticmethod
@@ -148,6 +151,15 @@ def debug_once(self, message: str, *args: Hashable, **kwargs: Hashable):
     def exception(self, message: str, *args: Hashable, **kwargs: Hashable):
         self.log(logging.ERROR, message, *args, **kwargs, exc_info=True)
 
+    def info_limit(self, message: str, *args, **kwargs):
+        self.log(logging.INFO, message, *args, **kwargs, rate_limit=True)
+
+    def warning_limit(self, message: str, *args, **kwargs):
+        self.log(logging.WARNING, message, *args, **kwargs, rate_limit=True)
+
+    def debug_limit(self, message: str, *args, **kwargs):
+        self.log(logging.DEBUG, message, *args, **kwargs, rate_limit=True)
+
 
 def init_logger(name: str = "UC") -> Logger:
     return Logger(name)

ucm/shared/infra/logger/cc/spdlog_logger.cc

@@ -45,6 +45,39 @@ void Logger::Log(Level&& lv, SourceLocation&& loc, std::string&& msg)
     this->logger_->log(spdlog::source_loc{loc.file, loc.line, loc.func}, level, std::move(msg));
 }
 
+bool Logger::Filter(std::string&& msg)
+{
+    using Clock = std::chrono::steady_clock;
+    const auto now = Clock::now();
+
+    const std::size_t key = std::hash<std::string>{}(msg);
+
+    std::lock_guard<std::mutex> lg(this->mutex_);
+
+    auto it = cache_.find(key);
+    if (it != cache_.end()) {
+        auto& entry = it->second;
+        auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(now - entry.last_time);
+        if (elapsed < rate_limit_window_) {
+            lru_list_.splice(lru_list_.begin(), lru_list_, entry.lru_it);
+            return false;
+        }
+        entry.last_time = now;
+        lru_list_.splice(lru_list_.begin(), lru_list_, entry.lru_it);
+        return true;
+    }
+
+    if (cache_.size() >= cache_capacity_ && !lru_list_.empty()) {
+        auto lru_it = std::prev(lru_list_.end());
+        cache_.erase(*lru_it);
+        lru_list_.pop_back();
+    }
+
+    lru_list_.push_front(key);
+    cache_.emplace(lru_list_.front(), CacheEntry{now, lru_list_.begin()});
+    return true;
+}
+
 std::shared_ptr<spdlog::logger> Logger::Make()
 {
     if (this->logger_) { return this->logger_; }

ucm/shared/infra/logger/cc/spdlog_logger.h

@@ -23,9 +23,14 @@
  * */
 #ifndef UNIFIEDCACHE_INFRA_LOGGER_SPDLOG_LOGGER_H
 #define UNIFIEDCACHE_INFRA_LOGGER_SPDLOG_LOGGER_H
+#include <chrono>
 #include <csignal>
 #include <cstdlib>
+#include <list>
+#include <mutex>
 #include <spdlog/spdlog.h>
+#include <string>
+#include <unordered_map>
 namespace UC::Logger {
 
 enum class Level { DEBUG, INFO, WARN, ERROR, CRITICAL };
@@ -68,7 +73,20 @@ class Logger {
 
     bool IsEnabledFor(Level lv);
 
+    bool Filter(std::string&& msg);
+
 private:
+    using TimePoint = std::chrono::steady_clock::time_point;
+    struct CacheEntry {
+        TimePoint last_time;
+        std::list<std::size_t>::iterator lru_it;
+    };
+
+    std::list<std::size_t> lru_list_;
+    std::unordered_map<std::size_t, CacheEntry> cache_;
+    std::size_t cache_capacity_{128};
+    std::chrono::seconds rate_limit_window_{std::chrono::seconds(60)};
+
     std::shared_ptr<spdlog::logger> Make();
     std::string path_{"log"};
     int max_files_{3};

ucm/shared/infra/logger/cpy/spdlog_logger.py.cc

@@ -32,11 +32,19 @@ void LogWrapper(Level lv, std::string file, std::string func, int line, std::str
     Log(std::move(lv), std::move(file), std::move(func), line, std::move(msg));
 }
 
+void RateLimitLogWrapper(Level lv, std::string file, std::string func, int line, std::string msg,
+                         std::string ori_msg)
+{
+    LogRateLimit(std::move(lv), std::move(file), std::move(func), line, std::move(msg),
+                 std::move(ori_msg));
+}
+
 PYBIND11_MODULE(ucmlogger, m)
 {
     m.def("setup", &Setup);
     m.def("flush", &Flush);
     m.def("log", &LogWrapper);
+    m.def("log_rate_limit", &RateLimitLogWrapper);
     m.def("isEnabledFor", &isEnabledFor);
     py::enum_<Level>(m, "Level")
         .value("DEBUG", Level::DEBUG)

ucm/shared/infra/logger/logger.cc

@@ -33,6 +33,16 @@ void Log(Level lv, std::string file, std::string func, int line, std::string msg
                               std::move(msg));
 }
 
+void LogRateLimit(Level lv, std::string file, std::string func, int line, std::string msg,
+                  std::string ori_msg)
+{
+    std::string full_msg = ori_msg + file + func + std::to_string(line);
+    if (Logger::GetInstance().Filter(std::move(full_msg))) {
+        Logger::GetInstance().Log(std::move(lv), SourceLocation{file.c_str(), func.c_str(), line},
+                                  std::move(msg));
+    }
+}
+
 void Setup(const std::string& path, int max_files, int max_size)
 {
     Logger::GetInstance().Setup(path, max_files, max_size);

ucm/shared/infra/logger/logger.h

@@ -32,6 +32,8 @@
 namespace UC::Logger {
 
 void Log(Level lv, std::string file, std::string func, int line, std::string msg);
+void LogRateLimit(Level lv, std::string file, std::string func, int line, std::string msg,
+                  std::string ori_msg);
 
 template <typename... Args>
 void Log(Level lv, const SourceLocation& loc, fmt::format_string<Args...> fmt, Args&&... args)
@@ -46,9 +48,14 @@ bool isEnabledFor(Level lv);
 }  // namespace UC::Logger
 #define UC_SOURCE_LOCATION {__FILE__, __FUNCTION__, __LINE__}
 #define UC_LOG(lv, fmt, ...) UC::Logger::Log(lv, UC_SOURCE_LOCATION, FMT_STRING(fmt), ##__VA_ARGS__)
+#define UC_LOG_LIMIT(lv, fmt, ...) \
+    UC::Logger::LogRateLimit(lv, UC_SOURCE_LOCATION, FMT_STRING(fmt), ##__VA_ARGS__)
 #define UC_DEBUG(fmt, ...) UC_LOG(UC::Logger::Level::DEBUG, fmt, ##__VA_ARGS__)
 #define UC_INFO(fmt, ...) UC_LOG(UC::Logger::Level::INFO, fmt, ##__VA_ARGS__)
 #define UC_WARN(fmt, ...) UC_LOG(UC::Logger::Level::WARN, fmt, ##__VA_ARGS__)
 #define UC_ERROR(fmt, ...) UC_LOG(UC::Logger::Level::ERROR, fmt, ##__VA_ARGS__)
-
+#define UC_DEBUG_LIMIT(fmt, ...) UC_LOG_LIMIT(UC::Logger::Level::DEBUG, fmt, ##__VA_ARGS__)
+#define UC_INFO_LIMIT(fmt, ...) UC_LOG_LIMIT(UC::Logger::Level::INFO, fmt, ##__VA_ARGS__)
+#define UC_WARN_LIMIT(fmt, ...) UC_LOG_LIMIT(UC::Logger::Level::WARN, fmt, ##__VA_ARGS__)
+#define UC_ERROR_LIMIT(fmt, ...) UC_LOG_LIMIT(UC::Logger::Level::ERROR, fmt, ##__VA_ARGS__)
 #endif