Addressing concurrency issues (#105)

1- passes the max_workers argument to ThreadPoolExecutor to enable large
number of threads
2- Removes asyncio usage, solely relies on ThreadPoolExecutor for
simplicity.
3- Passes a separate instance of Model to each thread for data safety,
even though through extensive testing I was not able to produce any data
corruption issues.

---------

Co-authored-by: Safoora Yousefi <[email protected]>

Author: safooray

eureka_ml_insights/core/inference.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 import os
 import time
@@ -170,7 +169,7 @@ def retrieve_exisiting_result(self, data, pre_inf_results_df):
 
     def run(self):
         if self.max_concurrent > 1:
-            asyncio.run(self._run_par())
+            self._run_par()
         else:
             self._run()
 
@@ -205,23 +204,7 @@ def _run(self):
                     data.update(response_dict)
                     writer.write(data)
 
-    from functools import partial
-
-    async def run_in_excutor(self, model_inputs, executor):
-        """Run model.generate in a ThreadPoolExecutor.
-        args:
-            model_inputs (tuple): args and kwargs to be passed to the model.generate function.
-            executor (ThreadPoolExecutor): ThreadPoolExecutor instance.
-        """
-        loop = asyncio.get_event_loop()
-
-        # function to run in executor with args and kwargs
-        def sub_func(model_inputs):
-            return self.model.generate(*model_inputs[0], **model_inputs[1])
-
-        return await loop.run_in_executor(executor, sub_func, model_inputs)
-
-    async def _run_par(self):
+    def _run_par(self):
         """parallel inference"""
         concurrent_inputs = []
         concurrent_metadata = []
@@ -240,30 +223,33 @@ async def _run_par(self):
 
                     # if batch is ready for concurrent inference
                     elif len(concurrent_inputs) >= self.max_concurrent:
-                        with ThreadPoolExecutor() as executor:
-                            await self.run_batch(concurrent_inputs, concurrent_metadata, writer, executor)
+                        with ThreadPoolExecutor(max_workers=self.max_concurrent) as executor:
+                            self.run_batch(concurrent_inputs, concurrent_metadata, writer, executor)
                         concurrent_inputs = []
                         concurrent_metadata = []
                     # add data to batch for concurrent inference
                     concurrent_inputs.append((model_args, model_kwargs))
                     concurrent_metadata.append(data)
                 # if data loader is exhausted but there are remaining data points that did not form a full batch
                 if concurrent_inputs:
-                    with ThreadPoolExecutor() as executor:
-                        await self.run_batch(concurrent_inputs, concurrent_metadata, writer, executor)
+                    with ThreadPoolExecutor(max_workers=self.max_concurrent) as executor:
+                        self.run_batch(concurrent_inputs, concurrent_metadata, writer, executor)
 
-    async def run_batch(self, concurrent_inputs, concurrent_metadata, writer, executor):
+    def run_batch(self, concurrent_inputs, concurrent_metadata, writer, executor):
         """Run a batch of inferences concurrently using ThreadPoolExecutor.
         args:
             concurrent_inputs (list): list of inputs to the model.generate function.
             concurrent_metadata (list): list of metadata corresponding to the inputs.
             writer (JsonLinesWriter): JsonLinesWriter instance to write the results.
             executor (ThreadPoolExecutor): ThreadPoolExecutor instance.
         """
-        tasks = [asyncio.create_task(self.run_in_excutor(input_data, executor)) for input_data in concurrent_inputs]
-        results = await asyncio.gather(*tasks)
-        for i in range(len(concurrent_inputs)):
-            data, response_dict = concurrent_metadata[i], results[i]
+
+        def sub_func(model_inputs):
+            return self.model.generate(*model_inputs[0], **model_inputs[1])
+
+        results = executor.map(sub_func, concurrent_inputs)
+        for i, result in enumerate(results):
+            data, response_dict = concurrent_metadata[i], result
             self.validate_response_dict(response_dict)
             # prepare results for writing
             data.update(response_dict)