diff --git a/.env.example b/.env.example index bbe8360d0..79c2c2be0 100644 --- a/.env.example +++ b/.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 diff --git a/.hydra_config/config.yaml b/.hydra_config/config.yaml index eded6d2c2..3e743db9c 100644 --- a/.hydra_config/config.yaml +++ b/.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: diff --git a/docs/content/docs/documentation/env_vars.md b/docs/content/docs/documentation/env_vars.md index 1c5613404..55a665f7f 100644 --- a/docs/content/docs/documentation/env_vars.md +++ b/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 | diff --git a/openrag/components/indexer/indexer.py b/openrag/components/indexer/indexer.py index 471439ec9..df4e337cb 100644 --- a/openrag/components/indexer/indexer.py +++ b/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: diff --git a/openrag/components/indexer/loaders/pdf_loaders/marker.py b/openrag/components/indexer/loaders/pdf_loaders/marker.py index 3c9233ff5..e7446d925 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/marker.py +++ b/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 diff --git a/openrag/components/indexer/loaders/serializer.py b/openrag/components/indexer/loaders/serializer.py index df755e011..9270c1b22 100644 --- a/openrag/components/indexer/loaders/serializer.py +++ b/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 \ No newline at end of file + except Exception as e: + log.exception("Failed to serialize document", error=str(e)) + raise