feat: improved api reliability. (#268)

Author: b4rtaz

src/app.cpp

@@ -23,6 +23,7 @@ static ChatTemplateType parseChatTemplateType(char *val) {
 
 AppCliArgs AppCliArgs::parse(int argc, char* *argv, bool requireMode) {
     AppCliArgs args;
+    args.info = true;
     args.help = false;
     args.mode = nullptr;
     args.nBatches = 32;
@@ -236,7 +237,7 @@ void runInferenceApp(AppCliArgs *args, void (*handler)(AppInferenceContext *cont
         throw std::runtime_error("This version supports only Q40 weights with Q80 sync type");
 
     Tokenizer tokenizer(args->tokenizerPath);
-    if (tokenizer.vocabSize != header.vocabSize)
+    if (args->info && tokenizer.vocabSize != header.vocabSize)
         printf("Tokenizer vocab size (%d) does not match the model vocab size (%d)\n", tokenizer.vocabSize, header.vocabSize);
 
     Sampler sampler(tokenizer.vocabSize, args->temperature, args->topp, args->seed);
@@ -246,8 +247,11 @@ void runInferenceApp(AppCliArgs *args, void (*handler)(AppInferenceContext *cont
 
     NnNodeConfig *rootNodeConfig = &net.nodeConfigs[0];
 
-    printLlmHeader(&header);
-    printNodeRequiredMemory(&net.netConfig, rootNodeConfig);
+    if (args->info) {
+        tokenizer.printHeader();
+        printLlmHeader(&header);
+        printNodeRequiredMemory(&net.netConfig, rootNodeConfig);
+    }
 
     NnNetExecution execution(args->nThreads, &net.netConfig);
 
@@ -346,11 +350,11 @@ void runWorkerApp(AppCliArgs *args) {
                 }
                 executor.forward();
                 isFirstAttempt = true;
-            } catch (const NnReadNetworkException &e) {
-                printf("Read network exception: %s\n", e.message);
+            } catch (const NnTransferSocketException &e) {
+                printf("🚨 Network error: %s\n", e.what());
                 break;
-            } catch (const NnWriteNetworkException &e) {
-                printf("Write network exception: %s\n", e.message);
+            } catch (const NnExecutorException &e) {
+                printf("🚨 Inference error: %s\n", e.what());
                 break;
             }
         }

src/app.hpp

@@ -12,6 +12,7 @@ class AppCliArgs {
     char *mode;
     NnUint nThreads;
     NnUint nBatches;
+    bool info;
     bool help;
 
     // inference

src/dllama-api.cpp

@@ -9,6 +9,8 @@
 #include <vector>
 #include <string>
 #include <csignal>
+#include <thread>
+#include <chrono>
 
 #ifdef _WIN32
 #include <winsock2.h>
@@ -532,7 +534,7 @@ void handleModelsRequest(HttpRequest& request, const char* modelPath) {
 }
 
 static void server(AppInferenceContext *context) {
-    int serverSocket = createServerSocket(context->args->port);
+    NnSocket serverSocket(createServerSocket(context->args->port));
 
     TokenizerChatStops stops(context->tokenizer);
     ChatTemplateGenerator templateGenerator(context->args->chatTemplateType, context->tokenizer->chatTemplate, stops.stops[0]);
@@ -556,19 +558,16 @@ static void server(AppInferenceContext *context) {
 
     while (true) {
         try {
-            int clientSocket = acceptSocket(serverSocket);
-            HttpRequest request = HttpRequest::read(clientSocket);
+            NnSocket clientSocket(acceptSocket(serverSocket.fd));
+            HttpRequest request = HttpRequest::read(clientSocket.fd);
             printf("🔷 %s %s\n", request.getMethod().c_str(), request.path.c_str());
             Router::resolve(request, routes);
-            destroySocket(clientSocket);
-        } catch (NnReadNetworkException& ex) {
-            printf("Read socket error: %d %s\n", ex.code, ex.message);
-        } catch (NnWriteNetworkException& ex) {
-            printf("Write socket error: %d %s\n", ex.code, ex.message);
+        } catch (const NnTransferSocketException& e) {
+            printf("Socket error: %d %s\n", e.code, e.what());
+        } catch (const NnExecutorException &e) {
+            throw;
         }
     }
-
-    destroySocket(serverSocket);
 }
 
 #ifdef _WIN32
@@ -601,22 +600,29 @@ int main(int argc, char *argv[]) {
     std::signal(SIGPIPE, SIG_IGN);
 #endif
 
+    AppCliArgs args = AppCliArgs::parse(argc, argv, false);
+    if (args.help) {
+        usage();
+        return EXIT_SUCCESS;
+    }
+
     initQuants();
     initSockets();
 
-    int returnCode = EXIT_SUCCESS;
-    try {
-        AppCliArgs args = AppCliArgs::parse(argc, argv, false);
-        if (args.help) {
-            usage();
-        } else {
+    while (true) {
+        try {
             runInferenceApp(&args, server);
+        } catch (const NnConnectionSocketException &e) {
+            printf("🚨 Connection error: %s\n", e.what());
+        } catch (const NnExecutorException &e) {
+            printf("🚨 Inference error: %s\n", e.what());
         }
-    } catch (std::exception &e) {
-        printf("🚨 Critical error: %s\n", e.what());
-        returnCode = EXIT_FAILURE;
+
+        printf("🔄 Retrying in 3 seconds...\n");
+        std::this_thread::sleep_for(std::chrono::seconds(3));
+        args.info = false;
     }
 
     cleanupSockets();
-    return returnCode;
+    return EXIT_SUCCESS;
 }

src/dllama.cpp

@@ -275,7 +275,7 @@ int main(int argc, char **argv) {
             runWorkerApp(&args);
         else
             throw std::runtime_error("Unsupported mode");
-    } catch (std::exception &e) {
+    } catch (const std::exception &e) {
         printf("🚨 Critical error: %s\n", e.what());
         returnCode = EXIT_FAILURE;
     }

src/nn/nn-core.cpp

@@ -145,6 +145,8 @@ void releaseNetConfig(NnNetConfig *netConfig) {
     for (NnUint pipeIndex = 0; pipeIndex < netConfig->nPipes; pipeIndex++) {
         delete[] netConfig->pipes[pipeIndex].name;
     }
+    if (netConfig->nPreSyncs > 0)
+        delete[] netConfig->preSyncs;
     delete[] netConfig->pipes;
 }
 

src/nn/nn-executor.cpp

@@ -1,6 +1,5 @@
 #include <cassert>
 #include <cstring>
-#include <stdexcept>
 #include "nn-executor.hpp"
 
 void NnFakeNodeSynchronizer::sync(NnUint segmentIndex, NnUint nThreads, NnUint threadIndex) {
@@ -39,6 +38,10 @@ NnExecutorDevice::NnExecutorDevice(NnDevice *device, int segmentFrom, int segmen
     this->segmentTo = segmentTo;
 }
 
+NnExecutorException::NnExecutorException(const std::string message)
+    : std::runtime_error(message) 
+{}
+
 NnExecutor::NnExecutor(NnNetConfig *netConfig, NnNodeConfig *nodeConfig, std::vector<NnExecutorDevice> *devices, NnNetExecution *netExecution, NnNodeSynchronizer *synchronizer, bool benchmark)
     : segments(nodeConfig->nSegments), steps()
 {
@@ -137,13 +140,19 @@ static inline void *executorThreadHandler(void *arg) {
     NnUint nThreads = context->nThreads;
     NnUint doneCount = nThreads - 1;
 
-    while (true) {
+    while (context->isAlive.load()) {
         const unsigned int currentStepIndex = context->currentStepIndex.load();
         if (currentStepIndex == context->nSteps)
             break;
 
         NnExecutorStep *step = &context->steps[currentStepIndex];
-        executeStep(step, nThreads, thread, context);
+        try {
+            executeStep(step, nThreads, thread, context);
+        } catch (const std::runtime_error &e) {
+            context->isAlive.store(false);
+            printf("🚨 Execution error: %s\n", e.what());
+            break;
+        }
 
         NnUint currentCount = context->doneThreadCount.fetch_add(1);
         if (currentCount == doneCount) {
@@ -156,7 +165,10 @@ static inline void *executorThreadHandler(void *arg) {
             context->doneThreadCount.store(0);
             context->currentStepIndex.fetch_add(1);
         } else {
-            while (context->currentStepIndex.load() == currentStepIndex);
+            while (
+                context->currentStepIndex.load() == currentStepIndex &&
+                context->isAlive.load()
+            );
         }
     }
     return nullptr;
@@ -166,6 +178,7 @@ void NnExecutor::forward() {
     assert(netExecution->batchSize > 0);
 
     NnUint nThreads = netExecution->nThreads;
+    context.isAlive.exchange(true);
     context.currentStepIndex.exchange(0);
     context.doneThreadCount.exchange(0);
     context.batchSize = netExecution->batchSize;
@@ -178,12 +191,14 @@ void NnExecutor::forward() {
     NnUint threadIndex;
     for (threadIndex = 1; threadIndex < nThreads; threadIndex++) {
         int result = pthread_create(&threads[threadIndex].handler, NULL, (PthreadFunc)executorThreadHandler, (void *)&threads[threadIndex]);
-        if (result != 0)
-            throw std::runtime_error("Failed to create thread");
+        assert(result == 0 && "Failed to create thread");
     }
     executorThreadHandler((void *)&threads[0]);
     for (threadIndex = 1; threadIndex < nThreads; threadIndex++)
         pthread_join(threads[threadIndex].handler, NULL);
+
+    if (!context.isAlive.load())
+        throw NnExecutorException("Execution failed in one of the threads");
 }
 
 NnUint NnExecutor::getTotalTime(NnExecutorStepType type) {

src/nn/nn-executor.hpp

@@ -4,6 +4,7 @@
 #include "nn-core.hpp"
 #include <atomic>
 #include <vector>
+#include <stdexcept>
 #include "pthread.h"
 
 class NnDeviceSegment {
@@ -73,6 +74,7 @@ typedef struct {
     NnNodeSynchronizer *synchronizer;
     std::atomic_uint currentStepIndex;
     std::atomic_uint doneThreadCount;
+    std::atomic_bool isAlive;
     NnUint batchSize;
     Timer *timer;
     NnUint totalTime[N_STEP_TYPES];
@@ -84,6 +86,11 @@ typedef struct {
     PthreadHandler handler;
 } NnExecutorThread;
 
+class NnExecutorException : public std::runtime_error {
+public:
+    NnExecutorException(const std::string message);
+};
+
 class NnExecutor {
 private:
     NnNetExecution *netExecution;