fix fetch batch

Signed-off-by: Praneeth Bedapudi <[email protected]>

Author: bedapudi6788

fastdeploy/_infer.py

@@ -109,8 +109,8 @@ def create_response(self, unique_id, response, is_compressed, input_type):
             pass
 
         if is_compressed:
-            _utils.logger.debug(f"{unique_id}: compressing response")
             response = self._compressor.compress(response)
+            _utils.logger.debug(f"{unique_id}: response compressed")
 
         return success, response
 
@@ -224,6 +224,7 @@ def get_responses_for_unique_ids(self, unique_ids, is_compresseds, input_types):
         all_responses = {}
 
         updations = {}
+        still_processing = []
 
         for unique_id, is_compressed, input_type in zip(
             unique_ids, is_compresseds, input_types
@@ -232,24 +233,21 @@ def get_responses_for_unique_ids(self, unique_ids, is_compresseds, input_types):
 
             if current_results["timedout_in_queue"]:
                 _utils.logger.warning(f"{unique_id}: timedout in queue")
+                updations[unique_id] = {
+                    "-1.predicted_at": time.time(),
+                }
                 all_responses[unique_id] = self.get_timeout_response(
                     unique_id, is_compressed, input_type
                 )
+                _utils.logger.debug(f"{unique_id}: timedout in queue response created")
 
             elif (
                 current_results["last_predictor_success"] is True
                 and current_results["last_predictor_sequence"]
                 == _utils.LAST_PREDICTOR_SEQUENCE
             ):
                 updations[unique_id] = {
-                    **{
-                        "-1.predicted_at": time.time(),
-                        "-1.outputs": None,
-                    },
-                    **{
-                        f"{__}.outputs": None
-                        for __ in _utils.PREDICTOR_SEQUENCE_TO_FILES
-                    },
+                    "-1.predicted_at": time.time(),
                 }
 
                 all_responses[unique_id] = self.create_response(
@@ -265,10 +263,14 @@ def get_responses_for_unique_ids(self, unique_ids, is_compresseds, input_types):
                     is_compressed,
                     input_type,
                 )
+                _utils.logger.debug(f"{unique_id}: response created")
             elif current_results["last_predictor_success"] is False:
                 _utils.logger.warning(
                     f"{unique_id}: predictor failed at {current_results['last_predictor_sequence']}"
                 )
+                updations[unique_id] = {
+                    "-1.predicted_at": time.time(),
+                }
                 all_responses[unique_id] = self.create_response(
                     unique_id,
                     {
@@ -280,8 +282,15 @@ def get_responses_for_unique_ids(self, unique_ids, is_compresseds, input_types):
                     is_compressed,
                     input_type,
                 )
+                _utils.logger.debug(f"{unique_id}: failed response created")
+
+            else:
+                still_processing.append(unique_id)
 
         if updations:
             _utils.MAIN_INDEX.update(updations)
 
+        if still_processing:
+            _utils.logger.debug(f"Still processing: {still_processing}")
+
         return all_responses

fastdeploy/_loop.py

@@ -75,10 +75,16 @@ def process_batch(predictor, input_batch, optimal_batch_size):
     return results, last_predictor_success, received_at, predicted_at
 
 
-def fetch_batch(main_index, predictor_sequence, optimal_batch_size):
+def fetch_batch(
+    main_index,
+    predictor_sequence,
+    optimal_batch_size,
+    max_wait_time_for_batch_collection,
+):
     unique_id_wise_input_count = {}
     input_batch = []
     current_batch_length = 0
+    batch_collection_started_at = time.time()
 
     while current_batch_length < optimal_batch_size:
         to_process = main_index.search(
@@ -98,9 +104,6 @@ def fetch_batch(main_index, predictor_sequence, optimal_batch_size):
             },
         )
 
-        if not to_process:  # No more items to process
-            break
-
         for unique_id, data in to_process.items():
             outputs = data[f"{predictor_sequence - 1}.outputs"]
 
@@ -110,6 +113,22 @@ def fetch_batch(main_index, predictor_sequence, optimal_batch_size):
             input_batch.extend(outputs)
             current_batch_length += input_count
 
+        if current_batch_length == 0:
+            time.sleep(max_wait_time_for_batch_collection / 2)
+            continue
+
+        elif (
+            time.time() - batch_collection_started_at
+            < max_wait_time_for_batch_collection
+            and current_batch_length / optimal_batch_size < 0.9
+        ):
+            time.sleep(max_wait_time_for_batch_collection / 2)
+            continue
+
+        else:
+            # finished collecting batch
+            break
+
     return unique_id_wise_input_count, input_batch
 
 
@@ -156,7 +175,7 @@ def start_loop(
 
     optimal_batch_size = predictor_info["optimal_batch_size"]
     time_per_example = predictor_info["time_per_example"]
-    max_wait_time_for_batch_collection = max(0.003, time_per_example * 0.25)
+    max_wait_time_for_batch_collection = max(0.003, time_per_example * 0.51)
 
     _utils.logger.info(
         f"""{predictor_name}
@@ -167,23 +186,27 @@ def start_loop(
     """
     )
 
-    last_batch_collection_started_at = 0
-
     while True:
         """
         Set timedout_in_queue to True for all the predictions that have been in the queue for more than timeout_time seconds
         and delete older than 7 seconds predictions that have finished prediction
         """
 
-        _utils.MAIN_INDEX.search(
+        timedout_in_queue_unique_ids = _utils.MAIN_INDEX.search(
             query={
                 "-1.predicted_at": 0,
                 "-1.received_at": {"$lt": time.time() - timeout_time},
                 "timedout_in_queue": {"$ne": True},
             },
             update={"timedout_in_queue": True},
+            select_keys=[],
         )
 
+        if timedout_in_queue_unique_ids:
+            _utils.logger.warning(
+                f"{_utils.MAIN_INDEX.count()} in queue, set timedout_in_queue to True for {list(timedout_in_queue_unique_ids)} unique_ids"
+            )
+
         _utils.MAIN_INDEX.delete(
             query={
                 "$and": [
@@ -194,23 +217,12 @@ def start_loop(
         )
 
         unique_id_wise_input_count, input_batch = fetch_batch(
-            _utils.MAIN_INDEX, predictor_sequence, optimal_batch_size
+            _utils.MAIN_INDEX,
+            predictor_sequence,
+            optimal_batch_size,
+            max_wait_time_for_batch_collection,
         )
 
-        current_batch_length = len(input_batch)
-
-        if current_batch_length == 0:
-            time.sleep(max_wait_time_for_batch_collection)
-            continue
-
-        if (
-            time.time() - last_batch_collection_started_at
-            < max_wait_time_for_batch_collection
-            and current_batch_length / optimal_batch_size < 0.9
-        ):
-            time.sleep(max_wait_time_for_batch_collection / 2)
-            continue
-
         _utils.logger.debug(f"Processing batch {unique_id_wise_input_count}")
 
         results, last_predictor_success, received_at, predicted_at = process_batch(
@@ -223,11 +235,12 @@ def start_loop(
             last_predictor_success,
             received_at,
             predicted_at,
-            current_batch_length,
+            len(input_batch),
         )
         _utils.MAIN_INDEX.update(unique_id_wise_results)
+
         _utils.logger.debug(
-            f"Updated results predictor {predictor_sequence}: list({unique_id_wise_results})"
+            f"Updated results predictor {predictor_sequence}: {list(unique_id_wise_results)}"
         )
 
         last_batch_collection_started_at = time.time()

fastdeploy/_rest.py

@@ -24,7 +24,7 @@
 
 
 class AsyncResponseHandler:
-    def __init__(self, check_interval=0.01):
+    def __init__(self, check_interval=0.003):
         self.pending_requests = {}
         self.check_interval = check_interval
         self.lock = threading.Lock()
@@ -75,6 +75,9 @@ def _response_checker(self):
                         input_types.append(data["input_type"])
 
                 if unique_ids:
+                    _utils.logger.debug(
+                        f"Checking responses for unique_ids: {unique_ids}"
+                    )
                     try:
                         responses = self.infer.get_responses_for_unique_ids(
                             unique_ids=unique_ids,

fastdeploy/_utils.py

@@ -3,7 +3,7 @@
 logging.basicConfig(
     format="%(asctime)s,%(msecs)d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s",
     datefmt="%Y-%m-%d:%H:%M:%S",
-    level=logging.INFO,
+    level=logging.DEBUG,
 )
 
 logger = logging.getLogger(__name__)