Codeprocastinator/optimized-tinyllama-covalent

1#include "log.h"2 3#include <chrono>4#include <condition_variable>5#include <cstdarg>6#include <cstdio>7#include <mutex>8#include <sstream>9#include <thread>10#include <vector>11 12int common_log_verbosity_thold = LOG_DEFAULT_LLAMA;13 14void common_log_set_verbosity_thold(int verbosity) {15    common_log_verbosity_thold = verbosity;16}17 18static int64_t t_us() {19    return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count();20}21 22// colors23enum common_log_col : int {24    COMMON_LOG_COL_DEFAULT = 0,25    COMMON_LOG_COL_BOLD,26    COMMON_LOG_COL_RED,27    COMMON_LOG_COL_GREEN,28    COMMON_LOG_COL_YELLOW,29    COMMON_LOG_COL_BLUE,30    COMMON_LOG_COL_MAGENTA,31    COMMON_LOG_COL_CYAN,32    COMMON_LOG_COL_WHITE,33};34 35// disable colors by default36static std::vector<const char *> g_col = {37    "",38    "",39    "",40    "",41    "",42    "",43    "",44    "",45    "",46};47 48struct common_log_entry {49    enum ggml_log_level level;50 51    bool prefix;52 53    int64_t timestamp;54 55    std::vector<char> msg;56 57    // signals the worker thread to stop58    bool is_end;59 60    void print(FILE * file = nullptr) const {61        FILE * fcur = file;62        if (!fcur) {63            // stderr displays DBG messages only when their verbosity level is not higher than the threshold64            // these messages will still be logged to a file65            if (level == GGML_LOG_LEVEL_DEBUG && common_log_verbosity_thold < LOG_DEFAULT_DEBUG) {66                return;67            }68 69            fcur = stdout;70 71            if (level != GGML_LOG_LEVEL_NONE) {72                fcur = stderr;73            }74        }75 76        if (level != GGML_LOG_LEVEL_NONE && level != GGML_LOG_LEVEL_CONT && prefix) {77            if (timestamp) {78                // [M.s.ms.us]79                fprintf(fcur, "%s%d.%02d.%03d.%03d%s ",80                        g_col[COMMON_LOG_COL_BLUE],81                        (int) (timestamp / 1000000 / 60),82                        (int) (timestamp / 1000000 % 60),83                        (int) (timestamp / 1000 % 1000),84                        (int) (timestamp % 1000),85                        g_col[COMMON_LOG_COL_DEFAULT]);86            }87 88            switch (level) {89                case GGML_LOG_LEVEL_INFO:  fprintf(fcur, "%sI %s", g_col[COMMON_LOG_COL_GREEN],   g_col[COMMON_LOG_COL_DEFAULT]); break;90                case GGML_LOG_LEVEL_WARN:  fprintf(fcur, "%sW %s", g_col[COMMON_LOG_COL_MAGENTA], ""                        ); break;91                case GGML_LOG_LEVEL_ERROR: fprintf(fcur, "%sE %s", g_col[COMMON_LOG_COL_RED],     ""                        ); break;92                case GGML_LOG_LEVEL_DEBUG: fprintf(fcur, "%sD %s", g_col[COMMON_LOG_COL_YELLOW],  ""                        ); break;93                default:94                    break;95            }96        }97 98        fprintf(fcur, "%s", msg.data());99 100        if (level == GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_ERROR || level == GGML_LOG_LEVEL_DEBUG) {101            fprintf(fcur, "%s", g_col[COMMON_LOG_COL_DEFAULT]);102        }103 104        fflush(fcur);105    }106};107 108struct common_log {109    // default capacity - will be expanded if needed110    common_log() : common_log(256) {}111 112    common_log(size_t capacity) {113        file = nullptr;114        prefix = false;115        timestamps = false;116        running = false;117        t_start = t_us();118 119        // initial message size - will be expanded if longer messages arrive120        entries.resize(capacity);121        for (auto & entry : entries) {122            entry.msg.resize(256);123        }124 125        head = 0;126        tail = 0;127 128        resume();129    }130 131    ~common_log() {132        pause();133        if (file) {134            fclose(file);135        }136    }137 138private:139    std::mutex mtx;140    std::thread thrd;141    std::condition_variable cv;142 143    FILE * file;144 145    bool prefix;146    bool timestamps;147    bool running;148 149    int64_t t_start;150 151    // ring buffer of entries152    std::vector<common_log_entry> entries;153    size_t head;154    size_t tail;155 156    // worker thread copies into this157    common_log_entry cur;158 159public:160    void add(enum ggml_log_level level, const char * fmt, va_list args) {161        std::lock_guard<std::mutex> lock(mtx);162 163        if (!running) {164            // discard messages while the worker thread is paused165            return;166        }167 168        auto & entry = entries[tail];169 170        {171            // cannot use args twice, so make a copy in case we need to expand the buffer172            va_list args_copy;173            va_copy(args_copy, args);174 175#if 1176            const size_t n = vsnprintf(entry.msg.data(), entry.msg.size(), fmt, args);177            if (n >= entry.msg.size()) {178                entry.msg.resize(n + 1);179                vsnprintf(entry.msg.data(), entry.msg.size(), fmt, args_copy);180            }181#else182            // hack for bolding arguments183 184            std::stringstream ss;185            for (int i = 0; fmt[i] != 0; i++) {186                if (fmt[i] == '%') {187                    ss << LOG_COL_BOLD;188                    while (fmt[i] != ' ' && fmt[i] != ')' && fmt[i] != ']' && fmt[i] != 0) ss << fmt[i++];189                    ss << LOG_COL_DEFAULT;190                    if (fmt[i] == 0) break;191                }192                ss << fmt[i];193            }194            const size_t n = vsnprintf(entry.msg.data(), entry.msg.size(), ss.str().c_str(), args);195            if (n >= entry.msg.size()) {196                entry.msg.resize(n + 1);197                vsnprintf(entry.msg.data(), entry.msg.size(), ss.str().c_str(), args_copy);198            }199#endif200            va_end(args_copy);201        }202 203        entry.level = level;204        entry.prefix = prefix;205        entry.timestamp = 0;206        if (timestamps) {207            entry.timestamp = t_us() - t_start;208        }209        entry.is_end = false;210 211        tail = (tail + 1) % entries.size();212        if (tail == head) {213            // expand the buffer214            std::vector<common_log_entry> new_entries(2*entries.size());215 216            size_t new_tail = 0;217 218            do {219                new_entries[new_tail] = std::move(entries[head]);220 221                head     = (head     + 1) % entries.size();222                new_tail = (new_tail + 1);223            } while (head != tail);224 225            head = 0;226            tail = new_tail;227 228            for (size_t i = tail; i < new_entries.size(); i++) {229                new_entries[i].msg.resize(256);230            }231 232            entries = std::move(new_entries);233        }234 235        cv.notify_one();236    }237 238    void resume() {239        std::lock_guard<std::mutex> lock(mtx);240 241        if (running) {242            return;243        }244 245        running = true;246 247        thrd = std::thread([this]() {248            while (true) {249                {250                    std::unique_lock<std::mutex> lock(mtx);251                    cv.wait(lock, [this]() { return head != tail; });252 253                    cur = entries[head];254 255                    head = (head + 1) % entries.size();256                }257 258                if (cur.is_end) {259                    break;260                }261 262                cur.print(); // stdout and stderr263 264                if (file) {265                    cur.print(file);266                }267            }268        });269    }270 271    void pause() {272        {273            std::lock_guard<std::mutex> lock(mtx);274 275            if (!running) {276                return;277            }278 279            running = false;280 281            // push an entry to signal the worker thread to stop282            {283                auto & entry = entries[tail];284                entry.is_end = true;285 286                tail = (tail + 1) % entries.size();287            }288 289            cv.notify_one();290        }291 292        thrd.join();293    }294 295    void set_file(const char * path) {296        pause();297 298        if (file) {299            fclose(file);300        }301 302        if (path) {303            file = fopen(path, "w");304        } else {305            file = nullptr;306        }307 308        resume();309    }310 311    void set_colors(bool colors) {312        pause();313 314        if (colors) {315            g_col[COMMON_LOG_COL_DEFAULT] = LOG_COL_DEFAULT;316            g_col[COMMON_LOG_COL_BOLD]    = LOG_COL_BOLD;317            g_col[COMMON_LOG_COL_RED]     = LOG_COL_RED;318            g_col[COMMON_LOG_COL_GREEN]   = LOG_COL_GREEN;319            g_col[COMMON_LOG_COL_YELLOW]  = LOG_COL_YELLOW;320            g_col[COMMON_LOG_COL_BLUE]    = LOG_COL_BLUE;321            g_col[COMMON_LOG_COL_MAGENTA] = LOG_COL_MAGENTA;322            g_col[COMMON_LOG_COL_CYAN]    = LOG_COL_CYAN;323            g_col[COMMON_LOG_COL_WHITE]   = LOG_COL_WHITE;324        } else {325            for (size_t i = 0; i < g_col.size(); i++) {326                g_col[i] = "";327            }328        }329 330        resume();331    }332 333    void set_prefix(bool prefix) {334        std::lock_guard<std::mutex> lock(mtx);335 336        this->prefix = prefix;337    }338 339    void set_timestamps(bool timestamps) {340        std::lock_guard<std::mutex> lock(mtx);341 342        this->timestamps = timestamps;343    }344};345 346//347// public API348//349 350struct common_log * common_log_init() {351    return new common_log;352}353 354struct common_log * common_log_main() {355    static struct common_log log;356 357    return &log;358}359 360void common_log_pause(struct common_log * log) {361    log->pause();362}363 364void common_log_resume(struct common_log * log) {365    log->resume();366}367 368void common_log_free(struct common_log * log) {369    delete log;370}371 372void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...) {373    va_list args;374    va_start(args, fmt);375    log->add(level, fmt, args);376    va_end(args);377}378 379void common_log_set_file(struct common_log * log, const char * file) {380    log->set_file(file);381}382 383void common_log_set_colors(struct common_log * log, bool colors) {384    log->set_colors(colors);385}386 387void common_log_set_prefix(struct common_log * log, bool prefix) {388    log->set_prefix(prefix);389}390 391void common_log_set_timestamps(struct common_log * log, bool timestamps) {392    log->set_timestamps(timestamps);393}394