add a bunch of stuff and now were passing pytest

shortfin/python/shortfin_apps/llm/components/kvcache/trie_attention_cache.py

@@ -1,7 +1,5 @@
-from typing import Dict, Set, List, Tuple, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from .page_pool import PageInfo
+from typing import Dict, Set, List, Tuple
+from .page_pool import PageInfo
 import heapq
 from dataclasses import dataclass
 import time
@@ -36,8 +34,12 @@ def parent(self, val: "TrieNode"):
             val.children[hash(tuple(self.tokens))] = self
         self._parent = val
 
-    def __lt__(self, other):
-        return True
+    # nodes are uniquely identified by their memory address
+    def __hash__(self):
+        return id(self)
+
+    def __eq__(self, other):
+        return self is other
 
 
 class TriePagedAttentionCache(BasePagedAttentionCache):

shortfin/python/shortfin_apps/llm/components/service.py

@@ -212,7 +212,7 @@ def board_flights(self):
         self.pending_prefills.clear()
         logger.debug("Post boarding cache state: %r", cache)
 
-    def board_prefills(self, cache: BasePagedAttentionCache):
+    def board_prefills(self, cache: TriePagedAttentionCache):
         # Fill prefill flights.
         pending_prefills = self.pending_prefills
         if len(pending_prefills) == 0:
@@ -252,7 +252,7 @@ def board_prefills(self, cache: BasePagedAttentionCache):
             # And takeoff.
             exec_process.launch()
 
-    def board_decodes(self, cache: BasePagedAttentionCache):
+    def board_decodes(self, cache: TriePagedAttentionCache):
         # Fill decode flights.
         pending_decodes = self.pending_decodes
         if len(pending_decodes) == 0: