Kill the generate thread when getting a "None" new request (#218)

If we don't return immediately when getting a None, we may fail the assertion in line 970

Author: Lumosis

jetstream/core/orchestrator.py

@@ -730,7 +730,7 @@ def _insert_if_possible(
       try:
         new_request = my_generate_backlog.get(block=block, timeout=1.0)
         if new_request is None:
-          break
+          return None
         ThreadDebugLog(
             thread_name,
             f"Got a new ActiveRequest from generate backlog {idx}.",
@@ -828,7 +828,7 @@ def _bulk_insert_if_possible(
           break
 
       if my_generate_backlog.queue[0] is None:
-        break
+        return None
 
       expected_slots = my_generate_backlog.queue[0].num_samples
 
@@ -857,7 +857,7 @@ def _bulk_insert_if_possible(
       try:
         new_request = my_generate_backlog.get(block=False)
         if new_request is None:
-          break
+          return None
         ThreadDebugLog(
             thread_name,
             f"Got a new ActiveRequest from generate backlog {idx}.",
@@ -965,6 +965,8 @@ def _generate_thread(self, idx: int):
             generate_engine,
             my_detokenize_backlog,
         )
+      if decode_state is None:
+        break
 
       # At this point, we know that we have at least some slots filled.
       assert (