Unified error reporting for Android ExecuTorch JNI layers (pytorch#18128)

All four Android JNI modules (Generic, LLM, ASR, Training) previously
used inconsistent exception types — a mix of RuntimeException,
IllegalStateException, IllegalArgumentException, and the structured
ExecutorchRuntimeException. This unifies them so every JNI error path
throws ExecutorchRuntimeException with a proper error code, giving
callers a single type to catch and programmatic access to the underlying
runtime error.

Key changes per module:

LLM (jni_layer_llama.cpp) — generate() now captures the Error return
from runner_->generate(), wraps the call in try/catch, reports failures
via a new onError(errorCode, message) callback, and returns the actual
error code instead of always returning 0.

ASR (jni_layer_asr.cpp) — replaced six env->ThrowNew(...) calls with
setExecutorchPendingException (for pure JNI path)


Training (jni_layer_training.cpp) — added jni_helper.h include; replaced
five throwNewJavaException("java/lang/Exception", ...) calls with
throwExecutorchException, preserving the actual error codes from the
failed operations.

Additionally: added default onError callbacks to LlmCallback (Java) and
AsrCallback (Kotlin);

This PR was authored with the assistance of Claude.



cc @kirklandsign @cbilgin

Author: psiddh

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrCallback.kt

@@ -25,4 +25,12 @@ interface AsrCallback {
    * @param token The decoded text token
    */
   fun onToken(token: String)
+
+  /**
+   * Called when an error occurs during transcription.
+   *
+   * @param errorCode Error code from the ExecuTorch runtime
+   * @param message Human-readable error description
+   */
+  fun onError(errorCode: Int, message: String) {}
 }

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmCallback.java

@@ -37,4 +37,14 @@ public interface LlmCallback {
    */
   @DoNotStrip
   default void onStats(String stats) {}
+
+  /**
+   * Called when an error occurs during generate().
+   *
+   * @param errorCode Error code from the ExecuTorch runtime (see {@link
+   *     org.pytorch.executorch.ExecutorchRuntimeException})
+   * @param message Human-readable error description
+   */
+  @DoNotStrip
+  default void onError(int errorCode, String message) {}
 }

extension/android/jni/jni_helper.cpp

@@ -35,6 +35,115 @@ void throwExecutorchException(uint32_t errorCode, const std::string& details) {
   facebook::jni::throwNewJavaException(exception.get());
 }
 
+void setExecutorchPendingException(
+    JNIEnv* env,
+    uint32_t errorCode,
+    const std::string& details) {
+  if (!env) {
+    return;
+  }
+  if (env->ExceptionCheck()) {
+    // Preserve any preexisting pending exception; do not overwrite it here.
+    return;
+  }
+
+  jclass exceptionClass =
+      env->FindClass("org/pytorch/executorch/ExecutorchRuntimeException");
+  if (env->ExceptionCheck()) {
+    if (exceptionClass) {
+      env->DeleteLocalRef(exceptionClass);
+    }
+    // Preserve the original exception; do not clear or overwrite it.
+    return;
+  }
+
+  if (!exceptionClass) {
+    // FindClass failed. It should have set a pending exception; leave it as is.
+    return;
+  }
+
+  jmethodID factoryMethod = env->GetStaticMethodID(
+      exceptionClass,
+      "makeExecutorchException",
+      "(ILjava/lang/String;)Ljava/lang/RuntimeException;");
+  if (env->ExceptionCheck()) {
+    // Preserve the original exception; do not overwrite it.
+    env->DeleteLocalRef(exceptionClass);
+    return;
+  }
+
+  if (!factoryMethod) {
+    // Factory unavailable; try the (int, String) constructor directly so
+    // callers still receive a structured ExecutorchRuntimeException.
+    jmethodID ctor =
+        env->GetMethodID(exceptionClass, "<init>", "(ILjava/lang/String;)V");
+    if (!env->ExceptionCheck() && ctor) {
+      jstring jDetails = env->NewStringUTF(details.c_str());
+      if (!env->ExceptionCheck() && jDetails) {
+        jobject exObj = env->NewObject(
+            exceptionClass, ctor, static_cast<jint>(errorCode), jDetails);
+        if (!env->ExceptionCheck() && exObj) {
+          env->Throw(reinterpret_cast<jthrowable>(exObj));
+          env->DeleteLocalRef(exObj);
+        }
+        env->DeleteLocalRef(jDetails);
+      }
+    } else {
+      // Clear any NoSuchMethodError from GetMethodID before falling back.
+      if (env->ExceptionCheck()) {
+        env->ExceptionClear();
+      }
+      jclass runtimeExClass = env->FindClass("java/lang/RuntimeException");
+      if (!env->ExceptionCheck() && runtimeExClass) {
+        env->ThrowNew(runtimeExClass, details.c_str());
+        env->DeleteLocalRef(runtimeExClass);
+      }
+    }
+    env->DeleteLocalRef(exceptionClass);
+    return;
+  }
+
+  jstring jDetails = env->NewStringUTF(details.c_str());
+  if (env->ExceptionCheck()) {
+    // Preserve the original exception; do not overwrite it.
+    if (jDetails) {
+      env->DeleteLocalRef(jDetails);
+    }
+    env->DeleteLocalRef(exceptionClass);
+    return;
+  }
+
+  if (!jDetails) {
+    // NewStringUTF returned null without setting an exception; fall back to
+    // a standard RuntimeException since ExecutorchRuntimeException lacks a
+    // (String) ctor.
+    jclass runtimeExClass = env->FindClass("java/lang/RuntimeException");
+    if (!env->ExceptionCheck() && runtimeExClass) {
+      env->ThrowNew(runtimeExClass, details.c_str());
+      env->DeleteLocalRef(runtimeExClass);
+    }
+    env->DeleteLocalRef(exceptionClass);
+    return;
+  }
+
+  auto exception = static_cast<jthrowable>(env->CallStaticObjectMethod(
+      exceptionClass, factoryMethod, static_cast<jint>(errorCode), jDetails));
+  if (env->ExceptionCheck() || !exception) {
+    // If a Java exception was thrown, it is already pending; just clean up.
+    if (exception) {
+      env->DeleteLocalRef(exception);
+    }
+    env->DeleteLocalRef(jDetails);
+    env->DeleteLocalRef(exceptionClass);
+    return;
+  }
+
+  env->Throw(exception);
+  env->DeleteLocalRef(exception);
+  env->DeleteLocalRef(jDetails);
+  env->DeleteLocalRef(exceptionClass);
+}
+
 bool utf8_check_validity(const char* str, size_t length) {
   for (size_t i = 0; i < length; ++i) {
     uint8_t byte = static_cast<uint8_t>(str[i]);

extension/android/jni/jni_helper.h

@@ -20,11 +20,35 @@ namespace executorch::jni_helper {
  * code and details. Uses the Java factory method
  * ExecutorchRuntimeException.makeExecutorchException(int, String).
  *
+ * IMPORTANT: This attempts to throw a C++ exception (via fbjni). Only use in
+ * fbjni HybridClass methods where fbjni catches it at the JNI boundary.
+ * For plain extern "C" JNIEXPORT functions, use setExecutorchPendingException.
+ *
+ * Note: If there is no current JNI environment (for example, if
+ * facebook::jni::Environment::current() returns null), this function is a
+ * no-op and does not throw. Callers must not rely on this always aborting
+ * control flow.
+ *
  * @param errorCode The error code from the C++ Executorch runtime.
  * @param details Additional details to include in the exception message.
  */
 void throwExecutorchException(uint32_t errorCode, const std::string& details);
 
+/**
+ * Sets a pending Java ExecutorchRuntimeException without throwing a C++
+ * exception. Safe to call from plain extern "C" JNIEXPORT functions.
+ * After calling this, the caller must return from the JNI function promptly;
+ * the Java exception will be delivered when control returns to the JVM.
+ *
+ * @param env The JNI environment pointer.
+ * @param errorCode The error code from the C++ Executorch runtime.
+ * @param details Additional details to include in the exception message.
+ */
+void setExecutorchPendingException(
+    JNIEnv* env,
+    uint32_t errorCode,
+    const std::string& details);
+
 // Define the JavaClass wrapper
 struct JExecutorchRuntimeException
     : public facebook::jni::JavaClass<JExecutorchRuntimeException> {

extension/android/jni/jni_layer_asr.cpp

@@ -128,8 +128,9 @@ Java_org_pytorch_executorch_extension_asr_AsrModule_nativeCreate(
       auto load_error = handle->preprocessor->load();
       if (load_error != Error::Ok) {
         ET_LOG(Error, "Failed to load preprocessor module");
-        env->ThrowNew(
-            env->FindClass("java/lang/RuntimeException"),
+        executorch::jni_helper::setExecutorchPendingException(
+            env,
+            static_cast<uint32_t>(load_error),
             "Failed to load preprocessor module");
         return 0;
       }
@@ -138,9 +139,10 @@ Java_org_pytorch_executorch_extension_asr_AsrModule_nativeCreate(
     return reinterpret_cast<jlong>(handle.release());
   } catch (const std::exception& e) {
     ET_LOG(Error, "Failed to create AsrModule: %s", e.what());
-    env->ThrowNew(
-        env->FindClass("java/lang/RuntimeException"),
-        ("Failed to create AsrModule: " + std::string(e.what())).c_str());
+    executorch::jni_helper::setExecutorchPendingException(
+        env,
+        static_cast<uint32_t>(Error::Internal),
+        "Failed to create AsrModule: " + std::string(e.what()));
     return 0;
   }
 }
@@ -172,8 +174,9 @@ Java_org_pytorch_executorch_extension_asr_AsrModule_nativeLoad(
     jobject /* this */,
     jlong nativeHandle) {
   if (nativeHandle == 0) {
-    env->ThrowNew(
-        env->FindClass("java/lang/IllegalStateException"),
+    executorch::jni_helper::setExecutorchPendingException(
+        env,
+        static_cast<uint32_t>(Error::InvalidState),
         "Module has been destroyed");
     return -1;
   }
@@ -218,15 +221,17 @@ Java_org_pytorch_executorch_extension_asr_AsrModule_nativeTranscribe(
     jlong decoderStartTokenId,
     jobject callback) {
   if (nativeHandle == 0) {
-    env->ThrowNew(
-        env->FindClass("java/lang/IllegalStateException"),
+    executorch::jni_helper::setExecutorchPendingException(
+        env,
+        static_cast<uint32_t>(Error::InvalidState),
         "Module has been destroyed");
     return -1;
   }
 
   if (wavPath == nullptr) {
-    env->ThrowNew(
-        env->FindClass("java/lang/IllegalArgumentException"),
+    executorch::jni_helper::setExecutorchPendingException(
+        env,
+        static_cast<uint32_t>(Error::InvalidArgument),
         "WAV path cannot be null");
     return -1;
   }
@@ -239,15 +244,17 @@ Java_org_pytorch_executorch_extension_asr_AsrModule_nativeTranscribe(
   try {
     audioData = ::executorch::extension::llm::load_wav_audio_data(wavPathStr);
   } catch (const std::exception& e) {
-    env->ThrowNew(
-        env->FindClass("java/lang/RuntimeException"),
-        ("Failed to load WAV file: " + std::string(e.what())).c_str());
+    executorch::jni_helper::setExecutorchPendingException(
+        env,
+        static_cast<uint32_t>(Error::AccessFailed),
+        "Failed to load WAV file: " + std::string(e.what()));
     return -1;
   }
 
   if (audioData.empty()) {
-    env->ThrowNew(
-        env->FindClass("java/lang/IllegalArgumentException"),
+    executorch::jni_helper::setExecutorchPendingException(
+        env,
+        static_cast<uint32_t>(Error::InvalidArgument),
         "WAV file contains no audio data");
     return -1;
   }
@@ -267,16 +274,18 @@ Java_org_pytorch_executorch_extension_asr_AsrModule_nativeTranscribe(
     auto processedResult =
         handle->preprocessor->execute("forward", audioTensor);
     if (processedResult.error() != Error::Ok) {
-      env->ThrowNew(
-          env->FindClass("java/lang/RuntimeException"),
+      executorch::jni_helper::setExecutorchPendingException(
+          env,
+          static_cast<uint32_t>(processedResult.error()),
           "Audio preprocessing failed");
       return -1;
     }
 
     auto outputs = std::move(processedResult.get());
     if (outputs.empty() || !outputs[0].isTensor()) {
-      env->ThrowNew(
-          env->FindClass("java/lang/RuntimeException"),
+      executorch::jni_helper::setExecutorchPendingException(
+          env,
+          static_cast<uint32_t>(Error::Internal),
           "Preprocessor returned unexpected output");
       return -1;
     }

extension/android/jni/jni_layer_llama.cpp

@@ -71,6 +71,14 @@ class ExecuTorchLlmCallbackJni
         facebook::jni::make_jstring(
             executorch::extension::llm::stats_to_json_string(result)));
   }
+
+  void onError(int errorCode, const std::string& message) const {
+    static auto cls = ExecuTorchLlmCallbackJni::javaClassStatic();
+    static const auto on_error_method =
+        cls->getMethod<void(jint, facebook::jni::local_ref<jstring>)>(
+            "onError");
+    on_error_method(self(), errorCode, facebook::jni::make_jstring(message));
+  }
 };
 
 class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
@@ -198,6 +206,17 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       jfloat temperature,
       jint num_bos,
       jint num_eos) {
+    Error err = Error::Ok;
+    if (!prompt) {
+      err = Error::InvalidArgument;
+      if (callback) {
+        callback->onError(
+            static_cast<int>(err),
+            "generate() failed: prompt must not be null");
+      }
+      return static_cast<jint>(err);
+    }
+
     float effective_temperature = temperature >= 0 ? temperature : temperature_;
     std::string token_buffer;
     auto token_callback = [callback, &token_buffer](const std::string& token) {
@@ -209,26 +228,53 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       }
       std::string result = token_buffer;
       token_buffer.clear();
-      callback->onResult(result);
+      if (callback) {
+        callback->onResult(result);
+      }
     };
 
     if (!runner_) {
-      return static_cast<jint>(Error::InvalidState);
-    }
-    executorch::extension::llm::GenerationConfig config{
-        .echo = static_cast<bool>(echo),
-        .seq_len = seq_len,
-        .temperature = effective_temperature,
-        .num_bos = needs_bos_ ? num_bos_ : 0,
-        .num_eos = num_eos,
-    };
-    auto err = runner_->generate(
-        prompt->toStdString(),
-        config,
-        token_callback,
-        [callback](const llm::Stats& result) { callback->onStats(result); });
-    if (err == Error::Ok) {
-      needs_bos_ = false;
+      err = Error::InvalidState;
+      if (callback) {
+        callback->onError(
+            static_cast<int>(err), "generate() failed: runner not initialized");
+      }
+      return static_cast<jint>(err);
+    }
+
+    try {
+      executorch::extension::llm::GenerationConfig config{
+          .echo = static_cast<bool>(echo),
+          .seq_len = seq_len,
+          .temperature = effective_temperature,
+          .num_bos = needs_bos_ ? num_bos_ : 0,
+          .num_eos = num_eos,
+      };
+      err = runner_->generate(
+          prompt->toStdString(),
+          config,
+          token_callback,
+          [callback](const llm::Stats& result) {
+            if (callback) {
+              callback->onStats(result);
+            }
+          });
+      if (err == Error::Ok) {
+        needs_bos_ = false;
+      }
+      if (err != Error::Ok && callback) {
+        callback->onError(
+            static_cast<int>(err),
+            "generate() failed with error code " +
+                std::to_string(static_cast<int>(err)));
+      }
+    } catch (const std::exception& e) {
+      if (callback) {
+        callback->onError(
+            static_cast<int>(Error::Internal),
+            std::string("generate() threw: ") + e.what());
+      }
+      return static_cast<jint>(Error::Internal);
     }
     return static_cast<jint>(err);
   }