Android: Closeable lifecycle, error reporting, thread safety, and test coverage (#19229)

[psiddh]

Summary:
Combines all previously reverted Android improvement PRs (#18669, #19012, #19028, #19092, #19099, #19124) plus new Module lifecycle tests into a single atomic change.

Error reporting: all sync errors throw exceptions instead of returning status codes or logging silently. Module.loadMethod() throws ExecutorchRuntimeException on failure. LlmModule.generate() and load() throw on error. Native methods renamed with "Native" suffix; public wrappers check status and throw.

Lifecycle: Module, LlmModule, and TrainingModule implement Closeable for try-with-resources. LlmModule adds ReentrantLock to serialize access to non-thread-safe native state. stop() remains lock-free (C++ atomic flag). TrainingModule replaces Log.e silent failures with IllegalStateException.

Error consistency: ExecuTorchRuntime.validateFilePath throws IllegalArgumentException. SGD throws IllegalStateException. AsrModule throws ExecutorchRuntimeException. Cause-chaining constructor added to ExecutorchRuntimeException.

JNI safety: Module and LlmModule constructors wrapped in try-catch to surface native initialization failures. LlmModule.load() uses throwExecutorchException with diagnostic detail. All DoNotStrip annotations preserved on JNI-called methods; new DoNotStrip added to renamed private native methods (generateNative, resetContextNative, loadNative).

Test Plan:
fix 4 Ignored Module tests by providing required input tensor. Add 13 new lifecycle/API coverage tests. Add LlmModule use-after-close and idempotent close tests.

This diff updates callers of Module.loadMethod() to use try/catch instead of checking the return code, calling destroy() on failure.

Reviewed By: kirklandsign

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19229

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

❌ 2 New Failures, 4 Unrelated Failures

As of commit 85cc5e4 with merge base fe98297 ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / linux / linux-job (gh)
  exir/tests/test_joint_graph.py::TestJointGraph::test_joint_graph
• pull / unittest-editable / linux / linux-job (gh)
  exir/tests/test_joint_graph.py::TestJointGraph::test_joint_graph

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / unittest / macos / macos-job (gh) (similar failure)
  exir/tests/test_memory_planning.py::TestMisc::test_multiple_pools_1

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / windows / windows-job (gh) (trunk failure)
  backends/xnnpack/test/recipes/test_xnnpack_recipes.py::TestXnnpackRecipes::test_8a4w_recipe
• pull / unittest-editable / macos / macos-job (gh) (trunk failure)
  exir/tests/test_memory_planning.py::TestMisc::test_multiple_pools_1
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  backends/xnnpack/test/recipes/test_xnnpack_recipes.py::TestXnnpackRecipes::test_8a4w_recipe

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[Copilot]

Pull request overview

Consolidates several Android API improvements across the ExecuTorch Java/Kotlin wrappers and JNI layer, focusing on exception-based error reporting, explicit lifecycle management (Closeable), and basic thread-safety guarantees, plus expanded instrumentation test coverage.

Changes:

• Standardize synchronous error handling to throw exceptions (vs returning status codes / silent logging) across Module, LlmModule, training/ASR wrappers, and JNI glue.
• Add/solidify lifecycle management via Closeable (close()/idempotent destroy patterns) and serialize access to non-thread-safe native state in LlmModule via ReentrantLock.
• Update JNI registrations and Java native method naming (*Native), and expand Android instrumentation tests to cover lifecycle and API behavior.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
extension/benchmark/android/benchmark/app/src/main/java/org/pytorch/minibench/ModelRunner.java	Handle Module.loadMethod() throwing by capturing an error code and early-returning safely; ensure destroy() in finally.
extension/android/jni/jni_layer_llama.cpp	Add constructor exception-to-Java conversion; improve load() error reporting; rename registered natives (generateNative, loadNative, resetContextNative).
extension/android/jni/jni_layer.cpp	Wrap native Module construction in exception mapping; throw on unsupported input EValue type codes; rename registered natives (etdumpNative, getMethodsNative).
extension/android/executorch_android/src/main/java/org/pytorch/executorch/training/TrainingModule.java	Implement Closeable; replace silent failures with IllegalStateException on use-after-destroy.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/training/SGD.java	Throw IllegalStateException (vs generic RuntimeException) when optimizer is destroyed.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java	Implement Closeable; add locking + destroyed checks; convert many status-returning APIs to void + exception; rename native methods to *Native.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt	Throw ExecutorchRuntimeException with error code on creation/transcription failures.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/Module.java	Implement Closeable; make loadMethod() throw; add locked wrappers around additional APIs; rename native methods to *Native.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/ExecutorchRuntimeException.java	Add error-code documentation, ALREADY_LOADED, improved message prefix, and cause-chaining constructor.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/ExecuTorchRuntime.java	Tighten file-path validation and switch to IllegalArgumentException.
extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/ModuleInstrumentationTest.kt	Un-ignore and fix tests by providing required dummy input; add new API/lifecycle coverage.
extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleInstrumentationTest.kt	Update tests for exception-based load; add close() lifecycle tests.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

This PR consolidates several previously reverted Android improvements into one atomic update, standardizing ExecuTorch Android APIs around exception-based error reporting, explicit Closeable lifecycles, and improved thread-safety—along with significant new/updated instrumentation test coverage.

Changes:

• Shift Android APIs away from status-code returns/silent failures to consistent exception throwing (Java/Kotlin + JNI).
• Add/align Closeable lifecycle semantics (try-with-resources) and introduce locking around non-thread-safe native state (notably for LlmModule).
• Update JNI registrations to match renamed *Native methods and expand Android instrumentation tests (including new lifecycle tests).

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
extension/benchmark/android/benchmark/app/src/main/java/org/pytorch/minibench/ModelRunner.java	Adapts benchmark load-path to exception-based Module.loadMethod() and ensures destroy on success path.
extension/benchmark/android/benchmark/app/src/main/java/org/pytorch/minibench/LlmModelRunner.java	Adapts load-path to exception-based LlmModule.load() (but generate path still needs updates).
extension/android/jni/jni_layer_llama.cpp	Adds constructor error surfacing, improves load error reporting, and updates native method registrations (*Native).
extension/android/jni/jni_layer.cpp	Adds constructor error surfacing, throws on unsupported EValue input type, and updates native registrations (etdumpNative, getMethodsNative).
extension/android/executorch_android/src/main/java/org/pytorch/executorch/training/TrainingModule.java	Implements Closeable and replaces silent failures with exceptions.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/training/SGD.java	Uses IllegalStateException for destroyed optimizer usage.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java	Implements Closeable, adds locking, converts sync APIs to throw-on-error wrappers, and renames JNI methods to *Native.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt	Switches to ExecutorchRuntimeException for native creation/transcription failures.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/Module.java	Implements Closeable, adds throw-on-error wrappers and lock-guarded public methods, and renames JNI methods to *Native.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/ExecutorchRuntimeException.java	Adds richer docs, new error code constant, improved message format, and a cause-chaining constructor.
extension/android/executorch_android/src/main/java/org/pytorch/executorch/ExecuTorchRuntime.java	Makes file-path validation stricter and consistently IllegalArgumentException-based.
extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/ModuleInstrumentationTest.kt	Un-ignores/fixes Module tests and adds broader API/lifecycle coverage (load modes, log buffer, etdump, destroyed-state).
extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleInstrumentationTest.kt	Updates to new throwing APIs and adds lifecycle tests (use-after-close, idempotent close).

Comments suppressed due to low confidence (1)

extension/benchmark/android/benchmark/app/src/main/java/org/pytorch/minibench/LlmModelRunner.java:102

• LlmModule.generate(...) now throws on failure, but the MESSAGE_GENERATE branch doesn't catch exceptions. An exception here will crash the HandlerThread and can prevent onGenerationStopped() from running / results from being recorded. Wrap the generate call in try/catch (ideally extracting an error code from ExecutorchRuntimeException) and route the failure to the callback (e.g., call onGenerationStopped() and/or add an explicit error callback).

      mLlmModelRunner.mCallback.onModelLoaded(status);
    } else if (msg.what == MESSAGE_GENERATE) {
      mLlmModelRunner.mModule.generate((String) msg.obj, mLlmModelRunner);
      mLlmModelRunner.mCallback.onGenerationStopped();

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[meta-codesync]

@psiddh has imported this pull request. If you are a Meta employee, you can view this in D103233465.

[meta-codesync]

@psiddh has exported this pull request. If you are a Meta employee, you can view the originating Diff in D103233465.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.

Comments suppressed due to low confidence (2)

extension/benchmark/android/benchmark/app/src/main/java/org/pytorch/minibench/LlmModelRunner.java:103

• LlmModule.generate(...) now throws on failure (it no longer returns a status code). In the MESSAGE_GENERATE path, any exception will escape handleMessage, potentially killing the handler thread and skipping onGenerationStopped(). Wrap generate() in try/catch and ensure the callback is notified (and/or propagate an error status) instead of letting the background thread crash.

      mLlmModelRunner.mCallback.onModelLoaded(status);
    } else if (msg.what == MESSAGE_GENERATE) {
      mLlmModelRunner.mModule.generate((String) msg.obj, mLlmModelRunner);
      mLlmModelRunner.mCallback.onGenerationStopped();
    }

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleInstrumentationTest.kt:62

• These tests create a new LlmModule in setUp() for every test case but never close it, so native resources can accumulate across the suite (and potentially cause flakiness/OOM). Add an @After that calls llmModule.close() (or wrap each test body in try/finally) so the module is always released.

    llmModule =
        LlmModule(getTestFilePath(TEST_FILE_NAME), getTestFilePath(TOKENIZER_FILE_NAME), 0.0f)
  }

  @Test
  @Throws(IOException::class, URISyntaxException::class)
  fun testGenerate() {
    llmModule.load()

    llmModule.generate(TEST_PROMPT, SEQ_LEN, this@LlmModuleInstrumentationTest)
    assertEquals(results.size.toLong(), SEQ_LEN.toLong())
    assertTrue(tokensPerSecond[tokensPerSecond.size - 1] > 0)
  }

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[psiddh]

@claude review this PR and the comments

[claude]

Claude finished @psiddh's task in 2m 14s —— View job

---

PR Review: Android Closeable lifecycle, error reporting, thread safety, and test coverage

• Gather context and understand the request
• Review core Java/Kotlin changes (Module, LlmModule, error handling)
• Review JNI layer changes
• Review benchmark/caller updates
• Review test coverage
• Assess existing review comments and provide consolidated feedback

---

Overall Assessment

This is a well-structured consolidation of several improvements. The move from status-code returns to exception-based error reporting is the right direction for Android APIs, and the Closeable lifecycle + ReentrantLock serialization are meaningful safety improvements. The test coverage additions are solid.

That said, there are a few issues worth addressing — some flagged by Copilot that I agree with, and some additional ones.

---

Issues I'd Recommend Fixing

1. stop() has no use-after-close guard (High Priority)

File: LlmModule.java:782

stop() is a bare public native void stop() with no mDestroyed check. Every other public method checks mDestroyed and throws IllegalStateException, but stop() would call into released native memory after close(). Since stop() intentionally skips the lock (it uses a C++ atomic flag), a simple volatile check is sufficient:

public void stop() {
    if (mDestroyed) throw new IllegalStateException("LlmModule has been destroyed");
    stopNative();
}

@DoNotStrip
private native void stopNative();

This also requires updating the JNI registration in jni_layer_llama.cpp to map "stopNative" → ExecuTorchLlmJni::stop.

Fix this →

2. Missing std::move for etdump_gen in profiling build (Medium Priority)

File: jni_layer.cpp:288-289

module_ = std::make_unique<Module>(
    modelPath->toStdString(), load_mode, etdump_gen);

The Module constructor takes std::unique_ptr<runtime::EventTracer> by value. In the EXECUTORCH_ANDROID_PROFILING path, etdump_gen is a std::unique_ptr<ETDumpGen> — passing it without std::move won't compile. This should be:

module_ = std::make_unique<Module>(
    modelPath->toStdString(), load_mode, std::move(etdump_gen));

The non-profiling path (auto etdump_gen = nullptr) works fine since nullptr implicitly converts to a null unique_ptr.

Fix this →

3. LlmModelRunner generate path lacks exception handling (Medium Priority)

File: LlmModelRunner.java:101-103

} else if (msg.what == MESSAGE_GENERATE) {
  mLlmModelRunner.mModule.generate((String) msg.obj, mLlmModelRunner);
  mLlmModelRunner.mCallback.onGenerationStopped();
}

generate() now throws ExecutorchRuntimeException on failure, but this path doesn't catch it. An exception here kills the HandlerThread and skips onGenerationStopped(). Should be wrapped like the load path:

} else if (msg.what == MESSAGE_GENERATE) {
  try {
    mLlmModelRunner.mModule.generate((String) msg.obj, mLlmModelRunner);
  } catch (Exception e) {
    // Optionally surface error to callback
  }
  mLlmModelRunner.mCallback.onGenerationStopped();
}

Fix this →

4. LlmModuleInstrumentationTest leaks native resources (Low Priority)

File: LlmModuleInstrumentationTest.kt

setUp() creates a new LlmModule for each test but there's no @After method calling close(). Tests that don't explicitly close (e.g., testGenerate, testGenerateAndStop, all the prefill validation tests) leak native memory. Add:

@After
fun tearDown() {
    if (::llmModule.isInitialized) {
        llmModule.close()
    }
}

---

Minor Observations (Non-blocking)

5. testMethodMetadata() is a no-op test

File: ModuleInstrumentationTest.kt:56-59

fun testMethodMetadata() {
    val module = Module.load(getTestFilePath(TEST_FILE_NAME))
    module.destroy()
}

This test loads and immediately destroys without asserting anything about metadata. The newer testGetMethodMetadata() at line 216 does the actual work. The old test is dead code — consider removing it.

6. testNonPteFile doesn't clean up on unexpected success

File: ModuleInstrumentationTest.kt:112-117

fun testNonPteFile() {
    val module = Module.load(getTestFilePath(NON_PTE_FILE_NAME))
    module.loadMethod(FORWARD_METHOD)  // expected to throw
}

If loadMethod doesn't throw (test expectation fails), the module is never destroyed. Consider using assertThrows with a finally block for cleanup.

7. ReentrantLock callback re-entrancy in LlmModule

File: LlmModule.java:322-331

The generate() method holds the ReentrantLock while native code invokes LlmCallback methods. Since the lock is reentrant, callbacks running on the same thread could re-enter other LlmModule methods (e.g., resetContext(), close()). The close() has a getHoldCount() > 1 guard, but other methods don't. This is an inherent limitation of using ReentrantLock — documenting that callbacks must not call back into LlmModule (except stop()) would help.

---

What Looks Good

• ExecutorchRuntimeException: The error code mapping with human-readable messages and hex codes is excellent for debugging. The makeExecutorchException factory pattern is clean.
• ExecuTorchRuntime.validateFilePath: Clear, specific error messages for each failure case (null, doesn't exist, not a file, not readable). This addresses the earlier Copilot feedback about !! — the current implementation is already clean.
• Module.close()/destroy() pattern: Using tryLock() to detect concurrent execution and throwing rather than blocking is the right call for lifecycle methods.
• TrainingModule: Simple and correct — checkNotDestroyed() + Closeable with idempotent close().
• AsrModule: Good use of AtomicLong + getAndSet(0L) for thread-safe idempotent close without needing a lock.
• JNI error surfacing: Both jni_layer.cpp and jni_layer_llama.cpp now catch native exceptions during construction and convert them to Java exceptions — this was a significant gap before.
• Test coverage: The multi-threaded forward test, destroyed-state tests for all public methods, and the LLM lifecycle tests provide good confidence.

---

Summary

The PR is in good shape overall. The 3 higher-priority items (#1 stop() guard, #2 std::move, #3 generate exception handling) are worth fixing before merge as they represent potential crashes or UB. The rest are polish items.

---

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.