[Android] Output error message and stack trace in chat (#866)

* [Android] Output error message and stack trace in chat This PR enables the output of error message and stack trace in chat, when the chat backend fails. Co-authored-by: spectrometerHBH <[email protected]> * apply code review suggestions --------- Co-authored-by: spectrometerHBH <[email protected]>
mlc-ai · Sep 5, 2023 · 8a4b4d4 · 8a4b4d4
1 parent 42d683c
commit 8a4b4d4
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 15 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -82,7 +82,7 @@ set(
   ${TVM_HOME}/3rdparty/picojson
 )
 
-set(MLC_LLM_COMPILE_DEFS DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>)
+set(MLC_LLM_COMPILE_DEFS ${MLC_LLM_COMPILE_DEFS} DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>)
 
 target_include_directories(mlc_llm_objs PRIVATE ${MLC_LLM_INCLUDES})
 target_compile_definitions(mlc_llm_objs PRIVATE ${MLC_LLM_COMPILE_DEFS})

diff --git a/android/CMakeLists.txt b/android/CMakeLists.txt
@@ -7,6 +7,7 @@ set(ANDROID_BIN_DIR ${CMAKE_CURRENT_BINARY_DIR})
 
 set(MLC_LLM_DIR ${ANDROID_DIR}/..)
 set(MLC_LLM_BINARY_DIR mlc_llm)
+set(MLC_LLM_COMPILE_DEFS TVM_LOG_CUSTOMIZE=1)
 add_subdirectory(${MLC_LLM_DIR} ${MLC_LLM_BINARY_DIR} EXCLUDE_FROM_ALL)
 
 if (NOT DEFINED TVM_HOME)

diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/AppViewModel.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/AppViewModel.kt
@@ -498,7 +498,7 @@ class AppViewModel(application: Application) : AndroidViewModel(application) {
 
         private fun mainResetChat() {
             executorService.submit {
-                backend.resetChat()
+                callBackend { backend.resetChat() }
                 viewModelScope.launch {
                     clearHistory()
                     switchToReady()
@@ -528,6 +528,28 @@ class AppViewModel(application: Application) : AndroidViewModel(application) {
             modelChatState.value = ModelChatState.Ready
         }
 
+        private fun switchToFailed() {
+            modelChatState.value = ModelChatState.Falied
+        }
+
+        private fun callBackend(callback: () -> Unit): Boolean {
+            try {
+                callback()
+            } catch (e: Exception) {
+                viewModelScope.launch {
+                    val stackTrace = e.stackTraceToString()
+                    val errorMessage = e.localizedMessage
+                    appendMessage(
+                        MessageRole.Bot,
+                        "MLCChat failed\n\nStack trace:\n$stackTrace\n\nError message:\n$errorMessage"
+                    )
+                    switchToFailed()
+                }
+                return false
+            }
+            return true
+        }
+
         fun requestResetChat() {
             require(interruptable())
             interruptChat(
@@ -571,7 +593,7 @@ class AppViewModel(application: Application) : AndroidViewModel(application) {
 
         private fun mainTerminateChat(callback: () -> Unit) {
             executorService.submit {
-                backend.unload()
+                callBackend { backend.unload() }
                 viewModelScope.launch {
                     clearHistory()
                     switchToReady()
@@ -609,8 +631,10 @@ class AppViewModel(application: Application) : AndroidViewModel(application) {
                 viewModelScope.launch {
                     Toast.makeText(application, "Initialize...", Toast.LENGTH_SHORT).show()
                 }
-                backend.unload()
-                backend.reload(modelLib, modelPath)
+                if (!callBackend {
+                        backend.unload()
+                        backend.reload(modelLib, modelPath)
+                    }) return@submit
                 viewModelScope.launch {
                     Toast.makeText(application, "Ready to chat", Toast.LENGTH_SHORT).show()
                     switchToReady()
@@ -624,11 +648,13 @@ class AppViewModel(application: Application) : AndroidViewModel(application) {
             executorService.submit {
                 appendMessage(MessageRole.User, prompt)
                 appendMessage(MessageRole.Bot, "")
-                backend.prefill(prompt)
+                if (!callBackend { backend.prefill(prompt) }) return@submit
                 while (!backend.stopped()) {
-                    backend.decode()
-                    val newText = backend.getMessage()
-                    viewModelScope.launch { updateMessage(MessageRole.Bot, newText) }
+                    if (!callBackend {
+                            backend.decode()
+                            val newText = backend.message
+                            viewModelScope.launch { updateMessage(MessageRole.Bot, newText) }
+                        }) return@submit
                     if (modelChatState.value != ModelChatState.Generating) return@submit
                 }
                 val runtimeStats = backend.runtimeStatsText()
@@ -653,7 +679,9 @@ class AppViewModel(application: Application) : AndroidViewModel(application) {
         }
 
         fun interruptable(): Boolean {
-            return modelChatState.value == ModelChatState.Ready || modelChatState.value == ModelChatState.Generating
+            return modelChatState.value == ModelChatState.Ready
+                    || modelChatState.value == ModelChatState.Generating
+                    || modelChatState.value == ModelChatState.Falied
         }
     }
 }
@@ -674,7 +702,8 @@ enum class ModelChatState {
     Resetting,
     Reloading,
     Terminating,
-    Ready
+    Ready,
+    Falied
 }
 
 enum class MessageRole {

diff --git a/android/prepare_libs.sh b/android/prepare_libs.sh
@@ -23,8 +23,8 @@ cmake .. \
       -DUSE_HEXAGON_SDK=OFF \
       -DMLC_LLM_INSTALL_STATIC_LIB=ON \
       -DCMAKE_SKIP_INSTALL_ALL_DEPENDENCY=ON \
-      -DUSE_OPENCL=ON
+      -DUSE_OPENCL=ON \
+      -DUSE_CUSTOM_LOGGING=ON \
 
-make tvm4j_runtime_packed -j8
+make tvm4j_runtime_packed -j${nproc}
 cmake --build . --target install --config release -j
-
diff --git a/android/src/cpp/tvm_runtime.h b/android/src/cpp/tvm_runtime.h
@@ -4,4 +4,26 @@
 #include <dlfcn.h>
 #include <dmlc/logging.h>
 #include <dmlc/thread_local.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/c_runtime_api.h>
+
+#include <android/log.h>
+
+static_assert(TVM_LOG_CUSTOMIZE == 1, "TVM_LOG_CUSTOMIZE must be 1");
+
+namespace tvm {
+namespace runtime {
+namespace detail {
+// Override logging mechanism
+[[noreturn]] void LogFatalImpl(const std::string& file, int lineno, const std::string& message) {
+  std::string m = file + ":" + std::to_string(lineno) + ": " + message;
+  __android_log_write(ANDROID_LOG_FATAL, "TVM_RUNTIME", m.c_str());
+  throw InternalError(file, lineno, message);
+}
+void LogMessageImpl(const std::string& file, int lineno, int level, const std::string& message) {
+  std::string m = file + ":" + std::to_string(lineno) + ": " + message;
+  __android_log_write(ANDROID_LOG_DEBUG + level, "TVM_RUNTIME", m.c_str());
+}
+
+}  // namespace detail
+}  // namespace runtime
+}  // namespace tvm