Test Worker (#83) (#84)

* 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

* cleanup

* handle signaling

* some more editions

* querent working test

Author: saraswatpuneet

querent/llm/base_llm.py

@@ -1,15 +1,16 @@
 from abc import ABC, abstractmethod
 import asyncio
-from typing import Any
+import json
+from typing import Any, Type
 from querent.common.types.ingested_tokens import IngestedTokens
 from querent.common.types.querent_queue import QuerentQueue
 
 
 class BaseLLM(ABC):
     def __init__(
         self,
-        input_queue: QuerentQueue[IngestedTokens],
-        output_queue: QuerentQueue[Any],
+        input_queue: QuerentQueue,
+        output_queue: QuerentQueue,
         num_workers: int = 1,
     ):
         self.input_queue = input_queue
@@ -33,9 +34,24 @@ async def worker(self):
             while True:
                 data = await self.input_queue.get()
                 if data is None:
-                    # Sentinel value to stop the worker
                     break
-                result = await self.process_tokens(data)
+                if isinstance(data, IngestedTokens):
+                    result = await self.process_tokens(data)
+                elif isinstance(data, str):
+                    ingested_token_from_str = IngestedTokens(
+                        file="", data=[data], error=None
+                    )
+                    result = await self.process_tokens(ingested_token_from_str)
+                elif isinstance(data, [list, tuple]):
+                    tokens_from_list = json.dumps(data)
+                    ingested_token_from_list = IngestedTokens(
+                        file="", data=[tokens_from_list], error=None
+                    )
+                    result = await self.process_tokens(ingested_token_from_list)
+                else:
+                    raise Exception(
+                        f"Invalid data type {type(data)} for {self.__class__.__name__}"
+                    )
                 await self.output_queue.put(result)
                 self.input_queue.task_done()
         except asyncio.CancelledError:
@@ -53,11 +69,12 @@ async def stop_workers(self):
             for _ in range(self.num_workers):
                 await self.input_queue.put(None)
             # Wait for the workers to finish processing
-            await asyncio.gather(*self.workers)
+            await asyncio.gather(*self.workers)  # Await the workers here
         except asyncio.CancelledError:
             pass
         except Exception as e:
             print(f"Stop workers error: {e}")
         finally:
             # Close the output queue
+            await self.input_queue.close()
             await self.output_queue.close()

querent/llm/transformers/gpt2_llm_v1.py

@@ -7,8 +7,8 @@
 class GPT2LLM(BaseLLM):
     def __init__(
         self,
-        input_queue: QuerentQueue[IngestedTokens],
-        output_queue: QuerentQueue[IngestedTokens],
+        input_queue: QuerentQueue,
+        output_queue: QuerentQueue,
         model_name="gpt2",
         num_workers=1,
     ):

querent/napper/auto_scaler.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from typing import List
 
@@ -7,37 +8,45 @@
 
 class AutoScaler:
     def __init__(
-        self, resource_manager: ResourceManager, querenters: List[BaseLLM], threshold=10
+        self,
+        resource_manager: ResourceManager,
+        querenters: List[BaseLLM],
+        threshold: int = 10,
     ):
         self.resource_manager = resource_manager
         self.querenters = querenters
         self.threshold = threshold
         self.logger = logging.getLogger("AutoScaler")
+        self.querent_termination_event = resource_manager.querent_termination_event
+        self.worker_tasks: List[asyncio.Task] = []  # Store the worker tasks
 
-    async def scale_querenters(self, total_requested_workers):
+    async def scale_querenters(self, total_requested_workers: int):
         current_total_workers = sum(
             querenter.num_workers for querenter in self.querenters
         )
 
-        if total_requested_workers > current_total_workers:
+        if total_requested_workers <= current_total_workers:
             # Scale up querenter workers
+            self.worker_tasks = []
             for querenter in self.querenters:
                 num_workers_to_scale = querenter.num_workers
-                await querenter.start_workers(num_workers_to_scale)
-
-        elif total_requested_workers < current_total_workers:
-            # Scale down querenter workers
-            for querenter in self.querenters:
-                num_workers_to_scale = querenter.num_workers
-                await querenter.stop_workers(num_workers_to_scale)
-
-        self.logger.info(
-            f"Scaled querenter workers to {total_requested_workers} workers in total"
-        )
+                workers = await querenter.start_workers(num_workers_to_scale)
+                # Create tasks for the workers and store them
+                worker_tasks = [asyncio.create_task(worker) for worker in workers]
+                self.worker_tasks.extend(
+                    worker_tasks
+                )  # Extend the list of worker tasks
+                self.logger.info(
+                    f"Started {len(worker_tasks)} workers for {querenter.__class__.__name__}"
+                )
+        else:
+            raise Exception("Total requested workers exceed the current total workers.")
 
-    async def run(self):
+    async def start(self):
         try:
-            while True:
+            while (
+                not self.querent_termination_event.is_set()
+            ):  # Check termination_event
                 # Calculate the total requested workers for all querenters
                 total_requested_workers = sum(
                     querenter.num_workers for querenter in self.querenters
@@ -53,11 +62,22 @@ async def run(self):
                         "Total requested workers exceed the maximum allowed workers."
                     )
 
-                # Scale the number of querenter workers
+                # Scale querenter workers
                 await self.scale_querenters(total_requested_workers)
 
+                # Wait for a while before checking again (adjust this as needed)
+                await asyncio.sleep(1)
+
+                # Check if all worker tasks have completed
+                if all(task.done() for task in self.worker_tasks):
+                    self.querent_termination_event.set()  # Set termination event
+
+        except asyncio.CancelledError:
+            pass
         except Exception as e:
             self.logger.error(f"An error occurred during AutoScaler execution: {e}")
+        finally:
+            self.logger.info("AutoScaler stopped")
 
     async def stop(self):
         self.logger.info("Stopping AutoScaler")

querent/napper/querent.py

@@ -1,11 +1,10 @@
 import asyncio
 import logging
-from typing import List
-from querent.common.types.querent_queue import QuerentQueue
+import signal
+from typing import List, Awaitable
 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(
@@ -31,27 +30,27 @@ def __init__(
             self.resource_manager, querenters, threshold=self.auto_scale_threshold
         )
 
-        # Create an instance of SignalHandler and pass the Querent instance
-        self.signal_handler = SignalHandler(self)
+        # Create an event to handle termination requests
+        self.querent_termination_event = resource_manager.querent_termination_event
 
     async def start(self):
         try:
             logger.info("Starting Querent")
 
             # Start the auto-scaler
-            asyncio.create_task(self.auto_scaler.run())
+            auto_scale_task = asyncio.create_task(self.auto_scaler.start())
 
             # Start handling signals
-            asyncio.create_task(self.signal_handler.handle_signals())
+            self.setup_signal_handlers()
+
+            # Start the tasks above and wait for them to finish
+            await asyncio.gather(auto_scale_task, self.wait_for_termination())
 
         except Exception as e:
             logger.error(f"An error occurred during Querent execution: {e}")
             await self.graceful_shutdown()
         finally:
-            # Stop the workers
-            await asyncio.gather(
-                *(querenter.stop_workers() for querenter in self.querenters)
-            )
+            await self.graceful_shutdown()
             logger.info("Querent stopped")
 
     async def graceful_shutdown(self):
@@ -60,17 +59,21 @@ async def graceful_shutdown(self):
         # Stop the auto-scaler and querenters gracefully
         await self.auto_scaler.stop()
 
-        # Stop the workers
-        await asyncio.gather(
-            *(querenter.stop_workers() for querenter in self.querenters)
-        )
-
         logger.info("Querent stopped gracefully")
 
-    async def handle_shutdown(self):
+    def setup_signal_handlers(self):
+        for sig in [signal.SIGINT, signal.SIGTERM]:
+            loop = asyncio.get_event_loop()
+            loop.add_signal_handler(sig, self.handle_signal)
+
+    def handle_signal(self):
         try:
-            # Wait for a KeyboardInterrupt (Ctrl+C) or SIGTERM to initiate graceful shutdown
-            await asyncio.Event().wait()
-        except (KeyboardInterrupt, SystemExit):
-            logger.info("Received shutdown signal (Ctrl+C or SIGTERM)")
-            await self.graceful_shutdown()
+            print("Received shutdown signal. Initiating graceful shutdown...")
+            shutdown_task = asyncio.create_task(self.graceful_shutdown())
+            asyncio.run(shutdown_task)
+        except Exception as e:
+            print(f"Error during graceful shutdown: {str(e)}")
+
+    async def wait_for_termination(self) -> Awaitable[None]:
+        # Wait for the termination event to be set, indicating graceful shutdown
+        await self.querent_termination_event.wait()

querent/napper/resource_manager.py

@@ -1,10 +1,12 @@
+import asyncio
 import logging
 
 
 class ResourceManager:
     def __init__(self, max_allowed_workers=100):
         self.max_allowed_workers = max_allowed_workers
         self.min_allowed_workers = 1
+        self.querent_termination_event = asyncio.Event()
         self.logger = logging.getLogger("ResourceManager")
 
     async def get_max_allowed_workers(self):

querent/napper/signaling.py

@@ -0,0 +1,52 @@
+import asyncio
+import pytest
+from querent.common.types.querent_queue import QuerentQueue
+from querent.llm.base_llm import BaseLLM
+from querent.napper.querent import Querent
+from querent.napper.resource_manager import ResourceManager
+
+input_data = ["Data 1", "Data 2", "Data 3"]
+input_queue = QuerentQueue()
+output_queue = QuerentQueue()
+resource_manager = ResourceManager()
+
+
+# Define a simple mock LLM class for testing
+class MockLLM(BaseLLM):
+    async def process_tokens(self, data):
+        return f"Processed: {data}"
+
+    def validate(self):
+        return True
+
+
+@pytest.mark.asyncio
+async def test_querent_with_base_llm():
+    # Put some input data into the input queue
+    input_data = ["Data 1", "Data 2", "Data 3", None]
+    for data in input_data:
+        await input_queue.put(data)
+    # Wait for the tasks to finish processing (implicitly handled by Querent)
+    num_llms = 1
+    llms = [MockLLM(input_queue, output_queue) for _ in range(num_llms)]
+
+    # Create a Querent instance
+    querent = Querent(llms, num_workers=num_llms, resource_manager=resource_manager)
+
+    # Start the querent
+
+    await querent.start()
+
+    # Check the output queue for results and store them in a list
+    results = []
+    async for result in output_queue:
+        results.append(result)
+        output_queue.task_done()
+
+    # Assert that the results match the expected output
+    expected_output = [
+        "Processed: Data: ['Data 1']",
+        "Processed: Data: ['Data 2']",
+        "Processed: Data: ['Data 3]",
+    ]
+    assert len(results) == len(expected_output)