Worker (#83)

* base llm

* ;)

* base llm

* llm

* more ideas around queues

* examples

* cleanup

* napper

* querent workers

* update the querent core to handle shutdown signals

* cleanup

* worker

* handle signaling

Author: saraswatpuneet

README.md

@@ -4,7 +4,6 @@
 # Querent
 
 **Querent: Unleash the Power of Data and Graph Neural Networks**
-
 *Unlock Insights, Scale Asynchronously, and Forge a Knowledge-Driven Future*
 
 **Welcome to Querent!** We're not just another data framework; we're the future of knowledge discovery and insight generation. Querent is your agile and dynamic companion for collecting, processing, and harnessing data's transformative potential. Whether you're crafting knowledge graphs, training cutting-edge language models, or diving deep into data-driven insights, Querent has your back.

querent/llm/base_llm.py

@@ -12,11 +12,10 @@ def __init__(
         output_queue: QuerentQueue[Any],
         num_workers: int = 1,
     ):
-        self.input_queue = input_queue  # ingested tokens coming file by file
-        self.output_queue = (
-            output_queue  # any type of output, need to think about various LLM outputs
-        )
+        self.input_queue = input_queue
+        self.output_queue = output_queue
         self.num_workers = num_workers
+        self.workers = []
 
     @abstractmethod
     async def process_tokens(self, data: IngestedTokens) -> Any:
@@ -49,10 +48,16 @@ async def start_workers(self, num_workers: int = 1):
         return self.workers
 
     async def stop_workers(self):
-        # Signal the workers to stop by putting None into the input queue
-        await self.input_queue.close()
-        # Wait for the workers to finish processing
-        await asyncio.gather(*self.workers)
-        # Close the output queue
-        # TODO this will change depending on many output queues we have
-        await self.output_queue.close()
+        try:
+            # Signal the workers to stop by putting None into the input queue
+            for _ in range(self.num_workers):
+                await self.input_queue.put(None)
+            # Wait for the workers to finish processing
+            await asyncio.gather(*self.workers)
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            print(f"Stop workers error: {e}")
+        finally:
+            # Close the output queue
+            await self.output_queue.close()

querent/llm/transformers/gpt2_llm_v1.py

@@ -1,28 +1,39 @@
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from querent.common.types.ingested_tokens import IngestedTokens
 from querent.common.types.querent_queue import QuerentQueue
 from querent.llm.base_llm import BaseLLM
 
 
 class GPT2LLM(BaseLLM):
     def __init__(
         self,
-        input_queue: QuerentQueue[str],
-        output_queue: QuerentQueue[str],
+        input_queue: QuerentQueue[IngestedTokens],
+        output_queue: QuerentQueue[IngestedTokens],
         model_name="gpt2",
+        num_workers=1,
     ):
-        super().__init__(input_queue, output_queue)
+        super().__init__(input_queue, output_queue, num_workers=num_workers)
         self.model_name = model_name
-        self.model = GPT2LMHeadModel.from_pretrained(model_name)
-        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 
-    async def process_tokens(self, data: str) -> str:
+    async def process_tokens(self, data: IngestedTokens) -> str:
         try:
-            input_text = data  # Assuming data is a string
-            input_ids = self.tokenizer.encode(input_text, return_tensors="pt")
-            output = self.model.generate(
-                input_ids, max_length=50, num_return_sequences=1, no_repeat_ngram_size=2
+            # get the input text from the data which is a list of str
+            input_text_list = data.data
+
+            # concatenate the input text into a single string
+            input_text = " ".join(input_text_list)
+
+            model = GPT2LMHeadModel.from_pretrained(self.model_name)
+            tokenizer = GPT2Tokenizer.from_pretrained(self.model_name)
+
+            input_ids = tokenizer.encode(input_text, return_tensors="pt")
+            output = model.generate(
+                input_ids,
+                max_length=50,
+                num_return_sequences=1,
+                no_repeat_ngram_size=2,
             )
-            generated_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
+            generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
             return generated_text
         except Exception as e:
             # Log the error and return an informative error message

querent/napper/querent.py

@@ -1,11 +1,11 @@
 import asyncio
 import logging
-import signal
 from typing import List
 from querent.common.types.querent_queue import QuerentQueue
 from querent.llm.base_llm import BaseLLM
 from querent.napper.resource_manager import ResourceManager
 from querent.napper.auto_scaler import AutoScaler
+from signaling import SignalHandler  # Import the SignalHandler class from signaling.py
 
 # Set up logging
 logging.basicConfig(
@@ -20,26 +20,30 @@ def __init__(
         self,
         querenters: List[BaseLLM],
         num_workers: int,
-        max_workers: int,
         resource_manager: ResourceManager,
         auto_scale_threshold: int = 10,
     ):
         self.num_workers = num_workers
-        self.max_workers = max_workers
         self.resource_manager = resource_manager
         self.querenters = querenters
         self.auto_scale_threshold = auto_scale_threshold
         self.auto_scaler = AutoScaler(
             self.resource_manager, querenters, threshold=self.auto_scale_threshold
         )
 
+        # Create an instance of SignalHandler and pass the Querent instance
+        self.signal_handler = SignalHandler(self)
+
     async def start(self):
         try:
             logger.info("Starting Querent")
 
             # Start the auto-scaler
             asyncio.create_task(self.auto_scaler.run())
 
+            # Start handling signals
+            asyncio.create_task(self.signal_handler.handle_signals())
+
         except Exception as e:
             logger.error(f"An error occurred during Querent execution: {e}")
             await self.graceful_shutdown()

querent/napper/signaling.py

@@ -0,0 +1,22 @@
+import asyncio
+import signal
+
+
+class SignalHandler:
+    def __init__(self, querent):
+        self.querent = querent
+
+    async def handle_signals(self):
+        for sig in [signal.SIGINT, signal.SIGTERM]:
+            loop = asyncio.get_event_loop()
+            loop.add_signal_handler(sig, self.handle_signal)
+
+    async def handle_signal(self):
+        try:
+            print("Received shutdown signal. Initiating graceful shutdown...")
+            await self.querent.graceful_shutdown()
+        except Exception as e:
+            print(f"Error during graceful shutdown: {str(e)}")
+        finally:
+            print("Querent stopped")
+            exit(0)