Skip to content

Commit 54b0148

Browse files
authored
Android: LlmModule Closeable lifecycle and ModelRunner crash fix (#19012)
## Summary - **LlmModule**: Implement `Closeable` so modules can be used with try-with-resources. Add `close()` with explicit caller contract. Deprecate `resetNative()` in favor of `close()`. Add `ReentrantLock` to serialize access to non-thread-safe native state, matching `Module.java`'s pattern. - **ModelRunner**: Add early return after `loadMethod()` failure to prevent cascading crash in warmup/inference loops. Fix native memory leak by wrapping benchmark body in try/finally for `module.destroy()`. ### LlmModule lifecycle design - `ReentrantLock` serializes all operations — prevents accidental concurrent access (e.g. double-tap, rotation during generation) from causing native crashes - `generate()`/`prefill*()`/`load()`/`resetContext()` acquire the lock - `close()` uses `tryLock()` — fails fast with `IllegalStateException` instead of blocking (no ANR risk) - `stop()` remains a bare native method — lock-free, uses C++ atomic flag for cross-thread interrupt - Caller contract: call `stop()`, wait for `generate()` to return, then `close()` ### Tests added - `testUseAfterCloseThrows`: `generate()` after `close()` throws `IllegalStateException` - `testCloseIsIdempotent`: double `close()` is safe ## Test plan - [ ] Existing `LlmModuleInstrumentationTest` tests pass - [ ] New lifecycle tests pass - [ ] ModelRunner benchmark handles load failures gracefully This PR was authored with the help of Claude.
1 parent 239f7cc commit 54b0148

3 files changed

Lines changed: 277 additions & 159 deletions

File tree

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,22 @@ class LlmModuleInstrumentationTest : LlmCallback {
271271
}
272272
}
273273

274+
// --- Lifecycle tests ---
275+
276+
@Test
277+
fun testUseAfterCloseThrows() {
278+
llmModule.close()
279+
assertThrows(IllegalStateException::class.java) {
280+
llmModule.generate(TEST_PROMPT, SEQ_LEN, this@LlmModuleInstrumentationTest)
281+
}
282+
}
283+
284+
@Test
285+
fun testCloseIsIdempotent() {
286+
llmModule.close()
287+
llmModule.close()
288+
}
289+
274290
companion object {
275291
private const val TEST_FILE_NAME = "/stories.pte"
276292
private const val TOKENIZER_FILE_NAME = "/tokenizer.bin"

0 commit comments

Comments
 (0)