perf(rag): 사전 계산 query embedding 기반 KB fan-out 병렬화 - #632
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d6a0ac7465
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| executor_factory=GeventNativeThreadRAGRetrievalExecutor, | ||
| cancellation_factory=NativeThreadRAGRetrievalCancellation, | ||
| max_workers=MAX_RAG_FANOUT_CONCURRENCY, |
There was a problem hiding this comment.
동시 RAG 실행마다 여기서 새 scheduler를 만들고 execute()가 호출별로 최대 5-thread ThreadPool을 별도 생성합니다. 운영 entrypoint는 Celery --concurrency=100이므로 RAG burst에서 프로세스 하나가 최대 500개 native worker를 만들 수 있지만 SessionLocal의 DB pool은 overflow를 포함해 50 connection뿐이라, 대량의 thread가 connection을 기다리며 메모리와 스케줄링 자원을 소모하고 retrieval timeout/failure를 유발할 수 있습니다. 호출별 제한 외에 process-wide executor 또는 admission limit을 공유해야 합니다.
AGENTS.md reference: apps/AGENTS.md:L10-L15
Useful? React with 👍 / 👎.
| break | ||
| state = queued.pop(0) | ||
| try: | ||
| job = executor.submit(lambda state=state: invoke(state)) |
There was a problem hiding this comment.
RAG_CROSS_ENCODER_RERANK_ENABLED=true인 상태에서 첫 multi-KB 검색을 실행하면 이 병렬 submit으로 여러 native worker가 RetrievalService._get_cross_encoder_model()의 잠금 없는 class cache를 동시에 통과할 수 있습니다. 그 결과 최대 5개의 대형 CrossEncoder가 중복 로드되어 초기 요청의 메모리 사용량이 급증하거나 OOM 후 rerank fallback이 발생하므로, 모델을 fan-out 전에 한 번 초기화하거나 cache 초기화를 lock으로 보호해야 합니다.
AGENTS.md reference: apps/AGENTS.md:L10-L15
Useful? React with 👍 / 👎.
| for callback in callbacks: | ||
| self._invoke_safely(callback) |
There was a problem hiding this comment.
PostgreSQL 장애나 packet loss 중 timeout이 발생하면 이 callback은 RAGRetrievalSessionRunner가 등록한 DBAPI connection.cancel()인데, scheduler의 gevent greenlet에서 동기 호출됩니다. 이 취소 요청 자체가 block되면 scheduler가 hard deadline을 다시 확인하지 못할 뿐 아니라 monkey-patched Celery 프로세스의 hub까지 멈춰 다른 workflow 실행도 지연되므로, DB 취소를 별도의 bounded native control 경로에서 실행하고 caller deadline과 독립적으로 수렴시켜야 합니다.
AGENTS.md reference: apps/AGENTS.md:L12-L15
Useful? React with 👍 / 👎.
| if not running: | ||
| continue |
There was a problem hiding this comment.
실행 중인 job이 모두 끝난 시점에 queued task가 남아 있고 search deadline까지 남은 시간이 minimum_start_budget_ms 이하이면 submit loop는 아무 task도 꺼내지 않지만, 이 분기가 계속 continue하여 deadline까지 yield 없이 busy-spin합니다. 기본 설정에서도 최대 약 250ms 동안 Celery의 단일 gevent hub를 점유해 같은 프로세스의 다른 workflow greenlet을 전부 지연시킬 수 있으므로, 이 상태에서는 loop를 종료해 queued task를 timeout으로 수렴시켜야 합니다.
AGENTS.md reference: apps/AGENTS.md:L12-L15
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7b3f5bf782
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| finally: | ||
| self._cancel_safely(stop_signal) |
There was a problem hiding this comment.
LLM 노드의 설정된 timeout이 fan-out의 30초보다 짧아 workflow_engine.py의 gevent.Timeout이 발생하면, 이는 Exception이 아닌 BaseException 계열이라 위의 except Exception을 건너뛰고 이 finally로 바로 진입합니다. 여기서는 stop_signal만 취소하고 이미 실행 중인 각 state.cancellation은 취소하지 않으므로 DBAPI cancel callback이 호출되지 않으며, 노드는 timeout으로 실패한 뒤에도 native worker가 최대 statement timeout까지 DB 연결을 점유하거나 rerank를 계속할 수 있습니다. finally에서도 모든 running state의 cancellation을 요청해야 합니다.
AGENTS.md reference: apps/AGENTS.md:L12-L15
Useful? React with 👍 / 👎.
| trace_summary.update( | ||
| self._rag_stage_latency_summary( | ||
| candidate_resolution_latency_ms=candidate_resolution_latency_ms, | ||
| ) |
There was a problem hiding this comment.
Collection의 모든 child가 권한·정책 필터에서 제외되는 경우처럼 후보가 0개이면 이 helper는 hidden/resource-hidden 상태와 정상 empty 상태를 동일한 safe-no-result로 합치지만, 새 코드가 정확한 candidate_resolution_latency_ms를 durable RAG trace에 추가합니다. docs/features/audit-tracing/requirements.md의 hidden 경로 allowlist는 safe reason/scope 등으로 제한되고 stage latency는 authorized retrieval에만 허용되므로, trace 조회자가 반복 실행의 해석 시간으로 hidden 후보의 존재나 규모를 비교하는 timing oracle이 생깁니다. 이 구분 불가능한 경로에서는 stage latency를 생략하거나 hidden allowlist에 맞게 coarse 값으로 축소해야 합니다.
AGENTS.md reference: AGENTS.md:L197-L198
Useful? React with 👍 / 👎.
|
CI 실패 원인을 조치했습니다.
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 51e3378124
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ) | ||
| if cancellation.cancelled: | ||
| raise RAGRetrievalSessionTimeout() | ||
| result = operation(session) |
There was a problem hiding this comment.
search_documents_sync()는 KB 확인과 hierarchy/vector/keyword 조회 등 여러 SQL을 순차 실행하지만, 여기서는 cancellation을 operation 전후에만 확인하고 DBAPI cancel() callback도 한 번만 dispatch합니다. 취소가 두 statement 사이 또는 직전 cancellation 확인 직후에 도착하면 idle connection에 대한 cancel은 이후 쿼리를 막지 못하므로, 취소된 작업이 다음 statement들을 각각 기존 statement_timeout 한도까지 계속 실행하며 KB별 10초 budget과 cleanup reserve가 지난 뒤에도 프로세스 공용 5-thread pool과 DB connection을 점유할 수 있습니다. 각 statement 전에 cancellation과 남은 절대 deadline을 확인하거나 transaction 전체에 적용되는 deadline을 강제해야 합니다.
AGENTS.md reference: apps/AGENTS.md:L12-L15
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 24c5e4ddbb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| failure: Exception | None = None | ||
| cleanup_failed = False | ||
| try: | ||
| connection = session.connection() |
There was a problem hiding this comment.
DB pool이 소진되거나 pool_pre_ping이 지연되는 경우, 이 session.connection()은 cancel callback을 등록하기 전에 실행되므로 child cancellation과 절대 deadline으로 중단할 수 없습니다. 공용 engine의 pool_timeout은 60초(apps/shared/db/session.py)라 per-KB 10초 및 aggregate 30초가 끝난 뒤에도 native worker가 checkout에 남을 수 있고, 이런 작업 5개면 프로세스 공용 RAG pool을 모두 점유해 후속 RAG 요청까지 timeout시킵니다. Connection checkout 자체를 남은 budget으로 제한하거나 cancellation-aware한 bounded acquisition 경계 안으로 옮겨야 합니다.
AGENTS.md reference: apps/AGENTS.md:L12-L15
Useful? React with 👍 / 👎.
| override = getattr(self, "_rag_retrieval_session_runner_override", None) | ||
| if override is not None: | ||
| return override | ||
| return RAGRetrievalSessionRunner(session_factory=SessionLocal) |
There was a problem hiding this comment.
WorkflowEngine은 caller가 제공한 execution_context["db_session_factory"]를 보존하고 candidate resolution과 query-embedding runtime도 이를 사용하지만, 새 child retrieval만 여기서 전역 SessionLocal로 고정됩니다. 별도 DB에 바인딩한 embedded/child 실행이나 테스트 runtime에서는 후보를 주입된 DB에서 승인한 뒤 검색은 환경변수 기반 기본 DB로 보내므로 결과가 전부 누락되거나, 동일 UUID가 존재하면 다른 DB의 데이터를 읽게 됩니다. 각 worker가 독립 session을 갖는 성질은 유지하되 runner의 factory는 execution context에 주입된 factory를 사용해야 합니다.
AGENTS.md reference: apps/AGENTS.md:L7-L7
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7c8f8d60b4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| def close(self) -> None: | ||
| if self._closed: | ||
| return | ||
| self._closed = True |
There was a problem hiding this comment.
CrossEncoder처럼 cancellation에 협조하지 않는 native 작업이 hard deadline 이후에도 process pool을 점유하면, 후속 invocation마다 submit()이 만든 greenlet들이 _pool.apply()의 빈 slot을 기다립니다. Scheduler가 deadline에 반환해도 이 close()는 flag만 변경하므로 해당 greenlet과 callback/task state가 계속 남고, 반복 요청 시 thread 수는 5개로 제한되어도 대기 greenlet은 누적되어 메모리와 RAG 처리 대기열을 고갈시킬 수 있습니다. 기존 리뷰 후 최신 HEAD에서는 pool thread가 공유되지만 제출된 job을 종료 시 추적·취소하지 않는 것이 새 근거이므로, 아직 시작하지 않은 job을 close에서 취소하거나 process-wide bounded admission을 적용해야 합니다.
AGENTS.md reference: apps/AGENTS.md:L12-L15
Useful? React with 👍 / 👎.
| "apps/shared/tests/domain/test_knowledge_runtime_candidates.py", | ||
| "apps/shared/tests/services/test_knowledge_permission_runtime_bulk.py", | ||
| "apps/workflow_engine/adapters/knowledge_runtime_candidates.py", | ||
| "apps/workflow_engine/adapters/rag_retrieval_session.py", |
There was a problem hiding this comment.
새 checkout deadline과 late-session 정리를 실제 PostgreSQL에서 검증하는 테스트는 rag_retrieval_connection_acquirer.py의 동작에 직접 의존하지만, 여기에는 session과 scheduler 경로만 추가되어 acquirer 파일 자체가 _KNOWLEDGE_POSTGRES_PATTERNS에서 빠져 있습니다. 따라서 이후 acquirer만 수정하는 PR은 일반 Workflow 단위 테스트만 실행하고 test_rag_connection_checkout_returns_before_shared_pool_timeout을 포함한 disposable PostgreSQL job을 건너뛰며, .github/workflows/test-knowledge-runtime-postgres.yml의 dev push 경로도 동일하게 누락되어 checkout/pool 회귀가 병합될 수 있으므로 두 selector에 해당 파일을 추가해야 합니다.
AGENTS.md reference: AGENTS.md:L117-L119
Useful? React with 👍 / 👎.
|
@codex review |
bab912e to
7c8f8d6
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bdc2504882
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| "apps/workflow_engine/adapters/rag_retrieval_session.py", | ||
| "apps/workflow_engine/application/rag_retrieval_fanout.py", |
There was a problem hiding this comment.
rag_retrieval_connection_acquirer.py 또는 rag_retrieval_executor.py만 변경하는 후속 PR에서는 _KNOWLEDGE_POSTGRES_PATTERNS와 워크플로의 push.paths가 모두 일치하지 않아, 이번에 추가한 실제 PostgreSQL checkout·취소·connection 복구 테스트가 실행되지 않습니다. 특히 connection acquirer는 test_rag_retrieval_session_postgres.py가 검증하는 pool 고갈 경로의 직접 구현이므로 두 어댑터 경로를 selector와 워크플로 양쪽에 추가하고 해당 분류 테스트에도 포함해야 합니다.
AGENTS.md reference: AGENTS.md:L206-L207
Useful? React with 👍 / 👎.
| def test_node_factory_import_does_not_require_gevent_runtime() -> None: | ||
| """Gateway-side graph validation must not require worker-only gevent.""" |
There was a problem hiding this comment.
apps/workflow_engine/workflow/nodes/llm/__init__.py나 새 RAG 어댑터만 수정하는 후속 PR은 classify_paths()에서 workflow_tests만 선택하므로, apps/gateway/tests에 둔 이 회귀 테스트는 실행되지 않습니다. 그 결과 LLMNode가 다시 eager import되어 Gateway에 worker 전용 gevent 의존성이 유입되는 변경도 선택 CI를 통과할 수 있으므로, 이 테스트를 Workflow Engine 선택 대상에 두거나 해당 import-boundary 소스가 바뀌면 Gateway 테스트도 선택해야 합니다.
AGENTS.md reference: AGENTS.md:L206-L207
Useful? React with 👍 / 👎.
bdc2504 to
c47f179
Compare
c47f179 to
2b485b9
Compare
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
/recheck-ci-control |
2 participants
변경 사항
statement_timeout, DBAPI cancellation, rollback/close 경계를 추가했습니다.organization_id + knowledge_base_id로 제한하고, 조직 컨텍스트가 없으면 DB 접근 전에 종료합니다.관련 이슈
Closes #589
Linear: MBA-354
변경 유형
테스트
로컬 검증:
CI 위임:
pg_sleep취소보호 리소스·외부 실행 경계 (해당 시)
Test Knowledge Runtime PostgreSQLCI에서 확인합니다.스크린샷 (UI 변경 시)
UI 변경 없음