diff --git a/backend/consts/const.py b/backend/consts/const.py index 7500a83150..715f824be3 100644 --- a/backend/consts/const.py +++ b/backend/consts/const.py @@ -317,10 +317,17 @@ class VectorDatabaseType(str, Enum): # Worker Configuration RAY_ADDRESS = os.getenv("RAY_ADDRESS", "auto") -QUEUES = os.getenv("QUEUES", "process_q,process_part_q,forward_q") +QUEUES = os.getenv( + "QUEUES", + "process_q,process_part_q,forward_q,forward_part_q,forward_aggregate_q", +) # Will be dynamically set based on PID if not provided WORKER_NAME = os.getenv("WORKER_NAME") -WORKER_CONCURRENCY = DP_PART_PROCESSOR_COUNT + 1 +# The data-process service sets a queue-specific value for each child worker. +# Keep the historical default when the variable is not provided. +WORKER_CONCURRENCY = int( + os.getenv("WORKER_CONCURRENCY", str(DP_PART_PROCESSOR_COUNT + 1)) +) RAY_WARM_ACTOR_POOL_SIZE_PART = int( os.getenv("RAY_WARM_ACTOR_POOL_SIZE_PART", "2")) RAY_WARM_ACTOR_POOL_SIZE_PROCESS = int( diff --git a/backend/data_process/app.py b/backend/data_process/app.py index c3346c0882..3ec2ad8d16 100644 --- a/backend/data_process/app.py +++ b/backend/data_process/app.py @@ -44,11 +44,14 @@ # Explicitly set result backend broker_url=REDIS_URL, result_backend=REDIS_BACKEND_URL, - # Two task queues for processing and forward steps + # Explicitly route the newly isolated forward child and aggregate tasks. + # Other tasks keep their queue from the @app.task declaration. task_routes={ f'{import_path}.process': {'queue': 'process_q'}, f'{import_path}.forward': {'queue': 'forward_q'}, f'{import_path}.process_and_forward': {'queue': 'process_q'}, + f'{import_path}.forward_part': {'queue': 'forward_part_q'}, + f'{import_path}.aggregate_forward_parts': {'queue': 'forward_aggregate_q'}, }, task_serializer='json', accept_content=['json'], diff --git a/backend/data_process/tasks.py b/backend/data_process/tasks.py index 9135b1e263..ed6ba8a539 100644 --- a/backend/data_process/tasks.py +++ b/backend/data_process/tasks.py @@ -1088,7 +1088,7 @@ def aggregate_store_chunks( } -@app.task(bind=True, base=LoggingTask, name='data_process.tasks.forward_part', queue='forward_q') +@app.task(bind=True, base=LoggingTask, name='data_process.tasks.forward_part', queue='forward_part_q') @trace_knowledge_operation("knowledge.forward.batch", "forward.batch") def forward_part( self, @@ -1210,14 +1210,19 @@ def forward_part( ) -@app.task(bind=True, base=LoggingTask, name='data_process.tasks.aggregate_forward_parts', queue='forward_q') +@app.task( + bind=True, + base=LoggingTask, + name='data_process.tasks.aggregate_forward_parts', + queue='forward_aggregate_q', +) @trace_knowledge_operation("knowledge.forward.aggregate", "forward.aggregate") def aggregate_forward_parts( self, parts_results: List[Dict[str, Any]], source: Optional[str] = None, index_name: Optional[str] = None, - original_filename: Optional[str] = None + original_filename: Optional[str] = None, ) -> Dict[str, Any]: """ Aggregate forward_part results. @@ -2058,13 +2063,13 @@ def forward( total_batches=total_batches, # If request was split into multiple groups, force all groups to use large path. large_mode=True, - ).set(queue='forward_q') for idx, batch in enumerate(batches) + ).set(queue='forward_part_q') for idx, batch in enumerate(batches) ) callback = aggregate_forward_parts.s( source=original_source, index_name=original_index_name, - original_filename=original_filename - ).set(queue='forward_q') + original_filename=original_filename, + ).set(queue='forward_aggregate_q') result = chord(group_tasks)(callback) with allow_join_result(): es_result = result.get() diff --git a/backend/data_process_service.py b/backend/data_process_service.py index 1d955ebf80..961b10d8e2 100644 --- a/backend/data_process_service.py +++ b/backend/data_process_service.py @@ -179,6 +179,45 @@ def start_ray_cluster(self): logger.error(traceback.format_exc()) return False + @staticmethod + def _build_worker_configs(total_cpus: int) -> list[dict[str, Any]]: + """Build isolated Celery worker pools for each processing stage.""" + total_cpus = max(1, int(total_cpus)) + ray_actor_num_cpus = max(1, int(RAY_ACTOR_NUM_CPUS)) + process_worker_concurrency = min( + DP_PART_PROCESSOR_COUNT, + max(1, total_cpus // ray_actor_num_cpus), + ) + forward_worker_concurrency = min(8, total_cpus * 2) + forward_aggregate_worker_concurrency = min(2, total_cpus) + return [ + { + 'name': 'process-worker', + 'queue': 'process_q', + 'concurrency': process_worker_concurrency, + }, + { + 'name': 'process-part-worker', + 'queue': 'process_part_q', + 'concurrency': process_worker_concurrency, + }, + { + 'name': 'forward-worker', + 'queue': 'forward_q', + 'concurrency': forward_worker_concurrency, + }, + { + 'name': 'forward-part-worker', + 'queue': 'forward_part_q', + 'concurrency': forward_worker_concurrency, + }, + { + 'name': 'forward-aggregate-worker', + 'queue': 'forward_aggregate_q', + 'concurrency': forward_aggregate_worker_concurrency, + }, + ] + def start_workers(self): """Start Celery workers for process and forward queues""" if not self.config.get('start_workers', True): @@ -194,47 +233,23 @@ def start_workers(self): # Fallback to 1 if os.cpu_count() is None. total_cpus = int(RAY_NUM_CPUS) if RAY_NUM_CPUS else (os.cpu_count() or 1) - # Get the number of CPUs requested by each actor. - ray_actor_num_cpus = RAY_ACTOR_NUM_CPUS - - # Calculate concurrency for the process-worker. Each worker will spawn an actor, - # so we limit concurrency to avoid oversubscribing Ray's CPU resources. - process_worker_concurrency = min( - DP_PART_PROCESSOR_COUNT, - max(1, total_cpus // ray_actor_num_cpus), - ) - - # For forward-worker, it's I/O bound. A higher concurrency is fine, but we can cap it - # relative to CPU count to avoid creating excessive threads on small machines. - forward_worker_concurrency = min(8, total_cpus * 2) + workers_config = self._build_worker_configs(total_cpus) + concurrency_by_name = { + config['name']: config['concurrency'] for config in workers_config + } + process_worker_concurrency = concurrency_by_name['process-worker'] + forward_worker_concurrency = concurrency_by_name['forward-worker'] + forward_aggregate_worker_concurrency = concurrency_by_name['forward-aggregate-worker'] + ray_actor_num_cpus = max(1, int(RAY_ACTOR_NUM_CPUS)) logger.debug(f"Total available CPUs: {total_cpus}") logger.debug(f"CPUs per processing actor (RAY_ACTOR_NUM_CPUS): {ray_actor_num_cpus}") logger.debug(f"Process-worker concurrency set to: {process_worker_concurrency}") logger.debug(f"Forward-worker concurrency set to: {forward_worker_concurrency}") + logger.debug( + f"Forward-aggregate-worker concurrency set to: {forward_aggregate_worker_concurrency}" + ) - # Define worker configurations based on split architecture: - # - process-worker handles orchestration (process_q) - # - process-part-worker handles split sub-tasks (process_part_q) - # - forward-worker handles vectorization/storage (forward_q) - workers_config = [ - { - 'name': 'process-worker', - 'queue': 'process_q', - 'concurrency': process_worker_concurrency - }, - { - 'name': 'process-part-worker', - 'queue': 'process_part_q', - 'concurrency': process_worker_concurrency - }, - { - 'name': 'forward-worker', - 'queue': 'forward_q', - 'concurrency': forward_worker_concurrency - } - ] - # Start each worker in a separate process for config in workers_config: # Use full Python path and correct module path diff --git a/test/backend/data_process/test_tasks.py b/test/backend/data_process/test_tasks.py index 89aeb8c818..fe82418a4a 100644 --- a/test/backend/data_process/test_tasks.py +++ b/test/backend/data_process/test_tasks.py @@ -2167,8 +2167,10 @@ def is_task_cancelled(self, *args, **kwargs): class _Sig: def __init__(self, kwargs): self.kwargs = kwargs + self.queue = None - def set(self, **_kw): + def set(self, **kw): + self.queue = kw.get("queue") return self captured = {"group_sigs": None} @@ -2217,6 +2219,77 @@ def _fake_allow_join_result(): assert len(captured["group_sigs"]) == 2 assert all(sig.kwargs.get("large_mode") is True for sig in captured["group_sigs"]) + assert all(sig.queue == "forward_part_q" for sig in captured["group_sigs"]) + + +def test_forward_large_chunks_routes_aggregate_to_dedicated_queue(monkeypatch): + tasks, _ = import_tasks_with_fake_ray(monkeypatch) + monkeypatch.setattr(tasks, "get_file_size", lambda *args, **kwargs: 0) + + class _RedisSvc: + def save_progress_info(self, *args, **kwargs): + return True + + def is_task_cancelled(self, *args, **kwargs): + return False + + monkeypatch.setattr(tasks, "get_redis_service", lambda: _RedisSvc()) + + captured = {} + + class _Sig: + def __init__(self, kwargs): + self.kwargs = kwargs + self.queue = None + + def set(self, **kwargs): + self.queue = kwargs.get("queue") + return self + + monkeypatch.setattr(tasks, "forward_part", types.SimpleNamespace( + s=lambda **kwargs: _Sig(kwargs))) + monkeypatch.setattr(tasks, "aggregate_forward_parts", types.SimpleNamespace( + s=lambda **kwargs: _Sig(kwargs))) + + def _fake_group(signatures): + captured["parts"] = list(signatures) + return captured["parts"] + + def _fake_chord(group_tasks): + def _runner(callback): + captured["callback"] = callback + total = sum(len(sig.kwargs["chunks"]) for sig in group_tasks) + return types.SimpleNamespace(get=lambda: { + "success": True, + "total_indexed": total, + "total_submitted": total, + }) + return _runner + + @contextmanager + def _fake_allow_join_result(): + yield + + monkeypatch.setattr(tasks, "group", _fake_group) + monkeypatch.setattr(tasks, "chord", _fake_chord) + monkeypatch.setattr(tasks, "allow_join_result", _fake_allow_join_result) + monkeypatch.setattr(tasks, "_send_chunks_to_es", lambda **kwargs: { + "success": True, + "total_indexed": len(kwargs["chunks"]), + "total_submitted": len(kwargs["chunks"]), + }) + + out = tasks.forward( + FakeSelf("forward-aggregate-queue"), + processed_data={"chunks": [{"content": f"c-{i}", "metadata": {}} for i in range(70)]}, + index_name="idx", + source="/big.txt", + source_type="local", + file_id="file-1", + ) + + assert out["chunks_stored"] == 70 + assert captured["callback"].queue == "forward_aggregate_q" def test_process_sync_unsupported_raises_and_updates_state(monkeypatch): diff --git a/test/backend/test_data_process_service_entrypoint.py b/test/backend/test_data_process_service_entrypoint.py index f2eb74fbdc..3401902895 100644 --- a/test/backend/test_data_process_service_entrypoint.py +++ b/test/backend/test_data_process_service_entrypoint.py @@ -93,6 +93,50 @@ def test_start_ray_cluster_returns_when_disabled(service_module): service_module.RayConfig.init_ray_for_service.assert_not_called() +def test_worker_configs_isolate_forward_parent_parts_and_aggregate(service_module): + configs = service_module.ServiceManager._build_worker_configs(4) + + assert [config["queue"] for config in configs] == [ + "process_q", + "process_part_q", + "forward_q", + "forward_part_q", + "forward_aggregate_q", + ] + assert configs[0]["concurrency"] == configs[1]["concurrency"] == 2 + assert configs[2]["concurrency"] == configs[3]["concurrency"] == 8 + assert configs[4]["concurrency"] == 2 + + +def test_start_workers_launches_each_isolated_queue(service_module, monkeypatch): + launched = [] + + class _Process: + def __init__(self, command, **kwargs): + self.pid = len(launched) + 100 + self.stdout = types.SimpleNamespace(readline=lambda: "") + launched.append((command, kwargs)) + + monkeypatch.setattr(service_module, "RAY_NUM_CPUS", "4") + monkeypatch.setattr(service_module, "RAY_ACTOR_NUM_CPUS", 2) + monkeypatch.setattr(service_module.subprocess, "Popen", _Process) + monkeypatch.setattr(service_module.threading, "Thread", lambda **kwargs: types.SimpleNamespace(start=lambda: None)) + + service_module.service_processes["workers"] = [] + manager = service_module.ServiceManager({"start_workers": True}) + + assert manager.start_workers() is True + assert [row["queue"] for row in service_module.service_processes["workers"]] == [ + "process_q", + "process_part_q", + "forward_q", + "forward_part_q", + "forward_aggregate_q", + ] + assert len(launched) == 5 + service_module.service_processes["workers"] = [] + + def test_start_all_services_starts_enabled_services_in_order(service_module, monkeypatch): scheduler = types.SimpleNamespace(start=MagicMock()) scheduler_module = types.ModuleType("services.auto_summary_scheduler")