cont : fixes + add test [no ci]

ggerganov · ggerganov · commit c1845a9512da · 2024-09-11T13:33:49.000+03:00
diff --git a/Makefile b/Makefile
@@ -54,6 +54,7 @@ TEST_TARGETS = \
 	tests/test-grammar-parser \
 	tests/test-json-schema-to-grammar \
 	tests/test-llama-grammar \
+	tests/test-log \
 	tests/test-model-load-cancel \
 	tests/test-opt \
 	tests/test-quantize-fns \
@@ -1528,6 +1529,11 @@ tests/test-llama-grammar: tests/test-llama-grammar.cpp \
 	$(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS)
 
+tests/test-log: tests/test-log.cpp \
+	$(OBJ_ALL)
+	$(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<)
+	$(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS)
+
 tests/test-grammar-parser: tests/test-grammar-parser.cpp \
 	$(OBJ_ALL)
 	$(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<)
diff --git a/common/log.cpp b/common/log.cpp
@@ -5,8 +5,6 @@
 #include <cstdio>
 #include <condition_variable>
 
-#define LOG_MAX_MESSAGE_SIZE (256)
-
 #define LOG_COLORS // TMP
 
 #ifdef LOG_COLORS
@@ -41,11 +39,7 @@ struct gpt_log_entry {
     int verbosity;
     int64_t timestamp;
 
-    // static-sized message
-    char msg_stck[LOG_MAX_MESSAGE_SIZE];
-
-    // if it doesn't fit in the stack, it goes here
-    std::vector<char> msg_heap;
+    std::vector<char> msg;
 
     // signals the worker thread to stop
     bool is_end;
@@ -79,11 +73,7 @@ struct gpt_log_entry {
             }
         }
 
-        if (msg_heap.empty()) {
-            fprintf(file, "%s", msg_stck);
-        } else {
-            fprintf(file, "%s", msg_heap.data());
-        }
+        fprintf(file, "%s", msg.data());
 
         fflush(file);
     }
@@ -98,7 +88,6 @@ struct gpt_log {
         entries.resize(capacity);
         head = 0;
         tail = 0;
-        buffer.resize(1024);
 
         resume();
     }
@@ -111,9 +100,7 @@ struct gpt_log {
     }
 
 private:
-    std::mutex mtx_inp;
-    std::mutex mtx_wrk;
-
+    std::mutex mtx;
     std::thread thrd;
     std::condition_variable cv;
 
@@ -129,39 +116,27 @@ struct gpt_log {
     size_t head;
     size_t tail;
 
-    // print the message here before pushing
-    std::vector<char> buffer;
+    // worker thread copies into this
+    gpt_log_entry cur;
 
 public:
     void add(enum ggml_log_level level, int verbosity, const char * fmt, va_list args) {
-        std::unique_lock<std::mutex> lock_inp(mtx_inp);
+        std::lock_guard<std::mutex> lock(mtx);
 
         if (!running) {
             return;
         }
 
-        const size_t n = vsnprintf(buffer.data(), buffer.size(), fmt, args);
-        if (n >= buffer.size()) {
-            buffer.resize(n + 1);
-            vsnprintf(buffer.data(), buffer.size(), fmt, args);
-        }
-
-        std::lock_guard<std::mutex> lock_wrk(mtx_wrk);
-
         auto & entry = entries[tail];
 
-        if (n < LOG_MAX_MESSAGE_SIZE) {
-            memcpy(entry.msg_stck, buffer.data(), n);
-            entry.msg_stck[n] = '\0';
-            entry.msg_heap.clear();
-        } else {
-            entry.msg_heap.resize(n + 1);
-            memcpy(entry.msg_heap.data(), buffer.data(), n);
-            entry.msg_heap[n] = '\0';
+        {
+            const size_t n = vsnprintf(entry.msg.data(), entry.msg.size(), fmt, args);
+            if (n >= entry.msg.size()) {
+                entry.msg.resize(n + 1);
+                vsnprintf(entry.msg.data(), entry.msg.size(), fmt, args);
+            }
         }
 
-        lock_inp.unlock();
-
         entry.level = level;
         entry.verbosity = verbosity;
         entry.timestamp = 0;
@@ -173,20 +148,18 @@ struct gpt_log {
         tail = (tail + 1) % entries.size();
         if (tail == head) {
             // expand the buffer
-            size_t new_size = entries.size() * 2;
-            std::vector<gpt_log_entry> new_entries(new_size);
+            std::vector<gpt_log_entry> new_entries(2*entries.size());
 
-            size_t new_head = 0;
             size_t new_tail = 0;
 
-            while (head != tail) {
-                new_entries[new_tail] = entries[head];
+            do {
+                new_entries[new_tail] = std::move(entries[head]);
 
-                head = (head + 1) % entries.size();
-                new_tail = (new_tail + 1) % new_size;
-            }
+                head     = (head     + 1) % entries.size();
+                new_tail = (new_tail + 1);
+            } while (head != tail);
 
-            head = new_head;
+            head = 0;
             tail = new_tail;
 
             entries = std::move(new_entries);
@@ -196,7 +169,7 @@ struct gpt_log {
     }
 
     void resume() {
-        std::lock_guard<std::mutex> lock_inp(mtx_inp);
+        std::lock_guard<std::mutex> lock(mtx);
 
         if (running) {
             return;
@@ -206,35 +179,37 @@ struct gpt_log {
 
         thrd = std::thread([this]() {
             while (true) {
-                std::unique_lock<std::mutex> lock_wrk(mtx_wrk);
-                cv.wait(lock_wrk, [this]() { return head != tail; });
+                {
+                    std::unique_lock<std::mutex> lock(mtx);
+                    cv.wait(lock, [this]() { return head != tail; });
 
-                auto & entry = entries[head];
+                    cur = entries[head];
 
-                if (entry.is_end) {
+                    head = (head + 1) % entries.size();
+                }
+
+                if (cur.is_end) {
                     break;
                 }
 
-                entry.print(stdout);
+                cur.print(stdout);
 
                 if (file) {
-                    entry.print(file);
+                    cur.print(file);
                 }
-
-                head = (head + 1) % entries.size();
             }
         });
     }
 
     void pause() {
-        std::lock_guard<std::mutex> lock_inp(mtx_inp);
+        {
+            std::lock_guard<std::mutex> lock(mtx);
 
-        if (!running) {
-            return;
-        }
+            if (!running) {
+                return;
+            }
 
-        {
-            std::lock_guard<std::mutex> lock_wrk(mtx_wrk);
+            running = false;
 
             auto & entry = entries[tail];
 
@@ -246,13 +221,10 @@ struct gpt_log {
         }
 
         thrd.join();
-
-        running = false;
     }
 
     void set_file(const char * path) {
-        std::lock_guard<std::mutex> lock_inp(mtx_inp);
-        std::lock_guard<std::mutex> lock_wrk(mtx_wrk);
+        pause();
 
         if (file) {
             fclose(file);
@@ -263,11 +235,12 @@ struct gpt_log {
         } else {
             file = nullptr;
         }
+
+        resume();
     }
 
     void set_timestamps(bool timestamps) {
-        std::lock_guard<std::mutex> lock_inp(mtx_inp);
-        std::lock_guard<std::mutex> lock_wrk(mtx_wrk);
+        std::lock_guard<std::mutex> lock(mtx);
 
         this->timestamps = timestamps;
     }
diff --git a/common/log.h b/common/log.h
@@ -27,7 +27,7 @@ void             gpt_log_free  (struct gpt_log * log);
 LOG_ATTRIBUTE_FORMAT(4, 5)
 void gpt_log_add(struct gpt_log * log, enum ggml_log_level level, int verbosity, const char * fmt, ...);
 
-void gpt_log_set_file      (struct gpt_log * log, const char * file);
+void gpt_log_set_file      (struct gpt_log * log, const char * file); // not thread-safe
 void gpt_log_set_timestamps(struct gpt_log * log, bool timestamps);
 
 #define LOG_TMPL(level, verbosity, ...) \
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -108,6 +108,7 @@ llama_test(test-tokenizer-1-spm  NAME test-tokenizer-1-llama-spm ARGS ${CMAKE_CU
 #llama_test(test-tokenizer-1-spm  NAME test-tokenizer-1-baichuan  ARGS ${CMAKE_CURRENT_SOURCE_DIR}/../models/ggml-vocab-baichuan.gguf)
 
 # llama_target_and_test(test-double-float.cpp) # SLOW
+llama_target_and_test(test-log.cpp)
 llama_target_and_test(test-arg-parser.cpp)
 llama_target_and_test(test-quantize-fns.cpp)
 llama_target_and_test(test-quantize-perf.cpp)
diff --git a/tests/test-log.cpp b/tests/test-log.cpp
@@ -0,0 +1,38 @@
+#include "log.h"
+
+#include <thread>
+
+int main() {
+    const int n_thread = 8;
+    const int n_msg = 1000;
+
+    std::thread threads[n_thread];
+    for (int i = 0; i < n_thread; i++) {
+        threads[i] = std::thread([i]() {
+            for (int j = 0; j < n_msg; j++) {
+                const int log_type = std::rand() % 4;
+
+                switch (log_type) {
+                    case 0: LOG_INF("Thread %d: %d\n", i, j); break;
+                    case 1: LOG_WRN("Thread %d: %d\n", i, j); break;
+                    case 2: LOG_ERR("Thread %d: %d\n", i, j); break;
+                    case 3: LOG_DBG("Thread %d: %d\n", i, j); break;
+                    default:
+                        break;
+                }
+
+                if (rand () % 10 < 5) {
+                    gpt_log_set_timestamps(gpt_log_main(), rand() % 2);
+                }
+            }
+        });
+    }
+
+    for (int i = 0; i < n_thread; i++) {
+        threads[i].join();
+    }
+
+    gpt_log_pause(gpt_log_main());
+
+    return 0;
+}
