Merge pull request #23 from ChiSym/shutdown

add async cleanup method to trie

Author: benlebrun

genlm_backend/trie/async_impl.py

@@ -133,10 +133,24 @@ async def _background_loop(self):
                             future.set_exception(e)
                 raise
 
+    async def cleanup(self):
+        """Async cleanup - preferred method"""
+        if self._task and not self._task.done():
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+            self._task = None
+
     def shutdown(self):
         """Stop the background processing task and cleanup resources."""
-        if self._task:
-            self._task.cancel()
+        if self._task is not None:
+            try:
+                self._task.cancel()
+            except RuntimeError:
+                # Ignore runtime errors that might occur if event loop is closed
+                pass
             self._task = None
 
     def __del__(self):

tests/test_trie.py

@@ -174,6 +174,16 @@ async def test_async_trie(mock_llm, backend):
         np.testing.assert_allclose(have, want, rtol=1e-5, atol=1e-8)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["sequential", "parallel"])
+async def test_async_trie_cleanup(mock_llm, backend):
+    async_trie = AsyncTokenCharacterTrie.from_vocab(
+        mock_llm.byte_vocab, backend=backend
+    )
+    await async_trie.cleanup()
+    assert async_trie._task is None
+
+
 def test_sequential_preprocessing(decode):
     trie = TokenCharacterTrie(decode=decode)