Skip to content

♻️ Refactor(data-process): replace Ray with Celery prefork - #3888

Merged
jeffwu-1999 merged 24 commits into
developfrom
codex/data-process-celery-prefork
Oct 8, 2026
Merged

jeffwu-1999 merged 24 commits into
developfrom
codex/data-process-celery-prefork

Conversation

@cj2026-bit

@cj2026-bit cj2026-bit commented Sep 8, 2026 •

Copy link
Copy Markdown
Collaborator

Problem

The data-process parsing stage ran on a Ray cluster. Ray brings a large resident footprint (actor pool + cluster system processes), keeps its actor pool alive permanently, and pulls the full data-process import stack into every Celery worker process. This PR removes Ray entirely and moves the parsing stage to a Celery prefork pool.

Solution

  • Parsing now runs on a Celery prefork pool (--autoscale=DP_PARSE_MAX_PROCESSES,DP_PARSE_MIN_PROCESSES, --max-tasks-per-child=DP_PARSE_MAX_TASKS_PER_CHILD); forwarding, aggregation and ES ingestion keep using the existing Celery thread pools.
  • Lazy, process-local model registry (sdk/nexent/data_process/model_registry.py): parser children preload DP_PRELOAD_MODELS (default unstructured_default) at child init and reuse the loaded models across tasks; other models (e.g. table_transformer) load on demand.
  • Configurable parser sizing via backend/consts/const.py: DP_PARSE_MAX_PROCESSES=3, DP_PARSE_MIN_PROCESSES=1, DP_PARSE_THREADS_PER_PROCESS=2, DP_PARSE_MAX_TASKS_PER_CHILD=1000, DP_PRELOAD_MODELS=unstructured_default, DP_PARSER_STARTUP_TIMEOUT_S=300.
  • Ray removed end to end (runtime, dependencies, dashboard port, deployment/compose/env references, ray_init_failed error code) with no fallback switch.
  • Delete fence, part cancellation and aggregation-failure cleanup were simplified over several review rounds (parsing stops at the delete fence; the aggregation stage owns part cleanup).

Verification

CI

All 17 checks green on the pushed head (including the test job with the full backend pytest suite + coverage.xml + Codecov upload, codecov/patch, build-data-process amd64/arm64 image builds, frontend type-check, CodeQL and SonarCloud).

Live benchmark: Celery prefork (this PR) vs Ray (develop)

Same machine, same ES knowledge-base stack (Docker PG/Redis/ES/MinIO/Supabase + WSL backends), same driver script and assets: three 4.95 MiB GBK TXT files (2237 chunks each, the same file used in the design benchmark), full API flow (sign-in -> create KB -> upload -> /file/process -> poll to terminal state). Prefork runs with default config (min=1 / max=3, unstructured_default preloaded). Host: 16 CPU / 20 GiB WSL2.

1. Idle baseline (no tasks)

Component prefork Ray (develop)
Main service 331 MiB 1185 MiB
5 Celery workers 5 x 183 MiB 5 x 1032 MiB
Parser processes / actors 1 x 714 MiB (model preloaded) 3 x 992 MiB + pool manager 993 MiB (models not loaded yet)
Ray system processes (gcs/raylet/agents) - ~740 MiB
Total 1961 MiB 11094 MiB (5.7x)

Ray workers are ~1 GiB each because ray_actors.py imports nexent.data_process.DataProcessCore at module level, which pulls the torch/unstructured stack into every Celery worker; this PR moves the heavy imports into the parser children (launchers stay at ~183 MiB). Idle numbers reproduced across 3 independent starts on each side.

2. Full ingestion flow (3 x 4.95 MiB TXT)

Metric prefork (4 runs) Ray (clean run)
Total time to all COMPLETED 501 / 504.6 / 507.3 / 514.2 s 504.6 s
Parse per file 83.7 - 88.8 s 87.4 - 87.7 s
Parse CPU, steady state (instantaneous) avg 347-362%, peak 434% avg 351%, peak 441%
Scale-up CPU transient 1 s at ~1701% (two children cold-starting concurrently) n/a (pool pre-warmed)
Peak flow memory 5420 - 5434 MiB 12586 MiB
Idle after the flow 2730 MiB (scaled down to min=1) 12574 MiB (no recycling)

Throughput and CPU are equivalent within noise; the embedding API is the flow bottleneck on both sides (~390-400 s of forwarding).

3. Model cold start

prefork Ray
Mechanism preload at child creation (startup bootstrap / scale-up) lazy load on the actor's first task
Single model, in service 1.53 s (startup) / 1.87-1.93 s (scale-up child) ~1.18 s (A/B: cold 1.238 s vs warm 0.056 s on a 1.2 KB file)
Both models per process 4.72-5.80 s; child RSS 714 -> 1256 MiB (+543) not preloaded; a PDF parse lazily added +133 MiB to one actor
Standalone process (same machine) unstructured_default 1.95 s / table_transformer 3.83 s same (library-level cost)

4. Recycling and scale-up

prefork Ray
Recycling rule Celery autoscaler: checked every 1 s, scale-down allowed once >= 30 s (AUTOSCALE_KEEPALIVE) have passed since the last scale-up and the pool is idle; shrinks stepwise to min=1 none (static actor pool)
Observed "Scaling down" ~8 s after the last parse task; consistent across runs actors still resident 60 s after the flow
Scale-up cost submit -> child ready ~2 s (1 model) / ~6 s (2 models) n/a (3 warm actors kept alive)
Memory return partial (surviving child 714 -> 1354 MiB after work) none (actors 992 -> 1349 MiB and stay)
c29a54ab-2ec6-470b-83df-6629dca91eb3 0563f37a-dc6a-4fd8-8085-ae1df19b13da aa958710-832a-41cc-8b42-2ea75138ac90 image

Run parsing in process-local Celery prefork children, keep orchestration and forwarding on thread workers, and remove the Ray runtime, dependency, and deployment surface.

Co-authored-by: Codex <noreply@openai.com>

Generated-by: gpt-5-codex
# Conflicts:
#	backend/data_process/app.py
#	backend/data_process/tasks.py
#	backend/data_process_service.py
#	doc/docs/en/backend/overview.md
#	doc/docs/en/developer-guide/overview.md
#	doc/docs/zh/backend/overview.md
#	doc/docs/zh/developer-guide/overview.md
#	test/backend/data_process/test_tasks.py
#	test/backend/test_data_process_service_entrypoint.py
@cj2026-bit cj2026-bit changed the title Codex/data process celery prefork ♻️ Refactor(data-process): replace Ray with Celery prefork Sep 22, 2026
Problem: the stub_modules fixture in test_data_process_app.py assigned
stub modules directly to sys.modules without restoring them, so the
stubs leaked into every later test file sharing the same pytest
process. Running test_data_process_app.py before
test_data_process_service.py failed all 52 service tests in setUp with
AttributeError: '_CeleryAppStub' object has no attribute 'reset_mock'.
The leak was invisible to CI because run_all_test.py runs each test
file in its own subprocess.

Fix: replace the six direct sys.modules assignments with
monkeypatch.setitem so the real modules are restored when the fixture
tears down. Behavior during each test is unchanged.

Verification:
- pytest test/backend/app/test_data_process_app.py
  test/backend/services/test_data_process_service.py
  -> 76 passed (was 24 passed / 52 failed)
- each file alone: 24 passed / 52 passed (no regression)
- wider combo (app + services + tasks + worker + coverage):
  app+services now pass together; the remaining failures are a
  pre-existing module-level mock leak in test_data_process_service.py
  that exists on develop and is out of scope here.
@cj2026-bit
cj2026-bit force-pushed the codex/data-process-celery-prefork branch from 083dfca to 5a9718a Compare September 24, 2026 08:20
@cj2026-bit
cj2026-bit force-pushed the codex/data-process-celery-prefork branch from 5a9718a to 53024c4 Compare September 24, 2026 08:31
Comment thread sdk/nexent/data_process/model_registry.py
@jeffwu-1999
jeffwu-1999 merged commit 8fe2381 into develop Oct 8, 2026
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants