Improve ray actors resilience

.env.example

@@ -41,6 +41,8 @@ RAY_DEDUP_LOGS=0 # turns off ray log deduplication that appear across multiple p
 RAY_ENABLE_RECORD_ACTOR_TASK_LOGGING=1 # # to enable logs at task level in ray dashboard
 RAY_task_retry_delay_ms=3000
 RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 # critical with the newest version of UV
+# # To disable worker killing
+# RAY_memory_monitor_refresh_ms=0
 
 # Indexer UI 
 ## 1. replace X.X.X.X with localhost if launching local or with your server IP

.hydra_config/config.yaml

@@ -159,8 +159,9 @@ ray:
       update: ${oc.decode:${oc.env:INDEXER_UPDATE_CONCURRENCY, 100}}
       search: ${oc.decode:${oc.env:INDEXER_SEARCH_CONCURRENCY, 100}}
       delete: ${oc.decode:${oc.env:INDEXER_DELETE_CONCURRENCY, 100}}
+      serialize: ${oc.decode:${oc.env:INDEXER_SERIALIZE_CONCURRENCY, 50}}
       chunk: ${oc.decode:${oc.env:INDEXER_CHUNK_CONCURRENCY, 1000}}
-      insert: ${oc.decode:${oc.env:INDEXER_INSERT_CONCURRENCY, 10}}
+      insert: ${oc.decode:${oc.env:INDEXER_INSERT_CONCURRENCY, 100}}
   semaphore:
     concurrency: ${oc.decode:${oc.env:RAY_SEMAPHORE_CONCURRENCY, 100000}}
   serve:

docs/content/docs/documentation/env_vars.md

@@ -314,6 +314,7 @@ The following environment variables control Ray's logging behavior, task retry s
 | `RAY_ENABLE_RECORD_ACTOR_TASK_LOGGING` | `number` | `1` | Enables logs at task level in the Ray dashboard for better debugging and monitoring. |
 | `RAY_task_retry_delay_ms` | `number` | `3000` | Delay (in milliseconds) before retrying a failed task. Controls the wait time between retry attempts. |
 | `RAY_ENABLE_UV_RUN_RUNTIME_ENV` | `number` | `0` | Controls UV runtime environment integration. **Critical**: Must be set to `0` when using the newest version of UV to avoid compatibility issues. |
+|`RAY_memory_monitor_refresh_ms`| `number` | 250 ms | To control the frequency of memory usage checks and task or actor termination if needed. If you set this value to 0, task killing is disabled. |
 
 #### Indexer Configuration
 
@@ -330,6 +331,7 @@ Controls the maximum number of concurrent operations for different indexer tasks
 |----------|------|---------|-------------|
 | `INDEXER_DEFAULT_CONCURRENCY` | int | 1000 | Default concurrency limit for general operations |
 | `INDEXER_UPDATE_CONCURRENCY` | int | 100 | Maximum concurrent document update operations |
+| `INDEXER_SERIALIZE_CONCURRENCY` | int | 50 | Maximum concurrent serialization operations |
 | `INDEXER_SEARCH_CONCURRENCY` | int | 100 | Maximum concurrent search/retrieval operations |
 | `INDEXER_DELETE_CONCURRENCY` | int | 100 | Maximum concurrent document deletion operations |
 | `INDEXER_CHUNK_CONCURRENCY` | int | 1000 | Maximum concurrent document chunking operations |

openrag/components/indexer/indexer.py

@@ -30,6 +30,7 @@
         "delete": config.ray.indexer.concurrency_groups["delete"],
         "insert": config.ray.indexer.concurrency_groups["insert"],
         "chunk": config.ray.indexer.concurrency_groups["chunk"],
+        "serialize": config.ray.indexer.concurrency_groups["serialize"],
     },
 )
 class Indexer:
@@ -55,6 +56,17 @@ async def chunk(
         chunks = await self.chunker.split_document(doc, task_id)
         return chunks
 
+    @ray.method(concurrency_group="serialize")
+    async def serialize_file(
+        self,
+        path: Union[str, List[str]],
+        metadata: Optional[Dict] = {},
+        task_id: str = None,
+    ):
+        # Serialize
+        doc = await serialize_file(task_id, path, metadata=metadata)
+        return doc
+
     async def add_file(
         self,
         path: Union[str, List[str]],
@@ -86,11 +98,19 @@ async def add_file(
             metadata = {**metadata, "partition": partition}
 
             # Serialize
-            doc = await serialize_file(task_id, path, metadata=metadata)
+            doc = await self.serialize_file(
+                path=path, metadata=metadata, task_id=task_id
+            )
 
             # Chunk
-            await task_state_manager.set_state.remote(task_id, "CHUNKING")
-            chunks = await self.handle.chunk.remote(doc, str(path), task_id)
+            if doc:
+                await task_state_manager.set_state.remote(task_id, "CHUNKING")
+                chunks = await self.handle.chunk.remote(doc, str(path), task_id)
+            else:
+                log.warning(
+                    "No document returned from serialization; skipping indexing."
+                )
+                chunks = []
 
             if self.enable_insertion:
                 if chunks:

openrag/components/indexer/loaders/pdf_loaders/marker.py

@@ -24,7 +24,7 @@
     MARKER_NUM_GPUS = 0
 
 
-@ray.remote(num_gpus=MARKER_NUM_GPUS)
+@ray.remote(num_gpus=MARKER_NUM_GPUS, max_restarts=5)
 class MarkerWorker:
     def __init__(self):
         import os
@@ -156,7 +156,7 @@ def __del__(self):
                 pass  # Best effort cleanup
 
 
-@ray.remote
+@ray.remote(max_restarts=5)
 class MarkerPool:
     def __init__(self):
         from config import load_config

openrag/components/indexer/loaders/serializer.py

@@ -1,4 +1,3 @@
-import asyncio
 import gc
 from pathlib import Path
 from typing import Dict, Optional, Union
@@ -21,7 +20,7 @@
 DICT_MIMETYPES = dict(config.loader["mimetypes"])
 
 
-@ray.remote
+@ray.remote(max_restarts=5)
 class DocSerializer:
     def __init__(self, data_dir=None, **kwargs) -> None:
         from config import load_config
@@ -85,6 +84,6 @@ async def serialize_document(
                 torch.cuda.ipc_collect()
             log.info("Document serialized successfully")
             return doc
-        except Exception:
-            log.exception("Failed to serialize document")
-            raise
+        except Exception as e:
+            log.exception("Failed to serialize document", error=str(e))
+            raise