feat: add bounded graph extraction concurrency#370
Conversation
imbajin
left a comment
There was a problem hiding this comment.
Blocking: yes. Summary: The concurrency feature is wired through the demo path, but the REST graph extraction path cannot accept or forward the new worker setting, and backend validation does not enforce the advertised upper bound. Evidence: GraphExtractRequest and GraphExtractService do not include graph_extract_max_workers, while GraphExtractFlow.prepare and graph_index_utils only reject values below 1; the UI slider is capped at 8.
| graph_extract_max_workers = int(graph_extract_max_workers) | ||
| except (TypeError, ValueError) as exc: | ||
| raise ValueError("graph_extract_max_workers must be a positive integer") from exc | ||
| if graph_extract_max_workers < 1: |
There was a problem hiding this comment.
This only rejects values below 1, while the UI advertises a maximum of 8. Programmatic callers can pass a much larger graph_extract_max_workers value, and PropertyGraphExtract will create up to min(max_workers, len(chunks)) threads. Please enforce the same backend cap here, or in a shared validator used by both this flow and graph_index_utils, so the concurrency limit stays bounded outside the Gradio slider path.
There was a problem hiding this comment.
Thanks for the review. I addressed this by adding graph_extract_max_workers to the REST request model and forwarding it through the REST graph extraction path as well.
The same bounded concurrency setting is now available from both the demo UI path and the REST API path, and the backend validates it with a shared cap of 1–8 workers instead of relying only on the UI slider.
imbajin
left a comment
There was a problem hiding this comment.
Blocking: yes. Summary: A malformed per-chunk extraction response can still be silently merged as an empty result instead of surfacing chunk-level failure context. Evidence: _extract_chunk_items only wraps exceptions from extract_property_graph_by_llm, while _extract_and_filter_label returns an empty list for missing JSON, invalid graph shape, or JSON decode errors.
| chunk, | ||
| proceeded_chunk, | ||
| ) | ||
| return self._extract_and_filter_label(schema, proceeded_chunk) |
There was a problem hiding this comment.
This still treats malformed model output as a successful empty chunk in several cases. _extract_chunk_items only wraps exceptions raised by extract_property_graph_by_llm, but _extract_and_filter_label returns [] when no JSON is found, when the graph shape is invalid, or when JSON decoding fails. With multiple chunks, one bad chunk can therefore be silently dropped and the overall extraction still returns success, which loses the chunk-level failure visibility described in the PR. Please make parse/format failures distinguishable from a valid empty extraction and raise or report them with the chunk index.
There was a problem hiding this comment.
Thanks for the review. I updated the chunk extraction path so malformed LLM responses are no longer silently treated as empty results.
Each chunk response is now validated for JSON extraction and expected property-graph shape before filtering. If a chunk returns malformed JSON or an invalid graph format, extraction fails with chunk context such as Graph extraction failed for chunk 2/5: ....
I also added regression tests for malformed chunk output, backend concurrency cap validation, and REST request validation.
imbajin
left a comment
There was a problem hiding this comment.
Blocking: yes. Summary: Concurrent extraction does not stop queued LLM work after a chunk fails, and each successful response is parsed twice. Evidence: static control-flow review of ThreadPoolExecutor submission/shutdown at property_graph_extract.py:125-145; latest-head CI is green.
| } | ||
| for future in as_completed(future_to_index): | ||
| index = future_to_index[future] | ||
| chunk_results[index] = future.result() |
There was a problem hiding this comment.
future.result() raises, leaving the ThreadPoolExecutor context waits for running work while queued calls may still start, delaying the error and consuming unnecessary rate limit or paid requests. Please cancel futures that have not started on the first failure, define the policy for already-running calls, and add a regression test proving later queued chunks are not invoked after an early failure.
| try: | ||
| proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk) | ||
| self._extract_property_graph_json(proceeded_chunk) | ||
| self._extract_property_graph_json(proceeded_chunk) |
There was a problem hiding this comment.
🧹 The same response is parsed and validated twice consecutively, and both return values are discarded. This doubles the regex and JSON-decoding work for every chunk without adding coverage. Please remove the duplicate call or parse once and reuse the result.
VGalaxies
left a comment
There was a problem hiding this comment.
Review summary
- Blocking: no
- Summary: Fractional concurrency values are silently truncated outside the REST boundary.
- Evidence:
- Focused pytest: 50 passed
- Ruff checks passed
- direct validator probe mapped
1.9 -> 1and8.9 -> 8
|
|
||
| def validate_graph_extract_max_workers(value) -> int: | ||
| try: | ||
| workers = int(value) |
There was a problem hiding this comment.
Medium: Reject fractional worker counts instead of truncating them
hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py:21
Evidence
int(value)converts1.9to1and8.9to8, while the REST model rejects fractional values.
Impact
- Internal flow and demo callers can execute with a different concurrency than requested, producing inconsistent validation across entry points.
Requested fix
- Reject booleans and non-integral numeric values before conversion, while accepting integer-valued slider output; add regression tests for fractional values.
imbajin
left a comment
There was a problem hiding this comment.
Blocking: yes. Summary: The latest head fixes the previously reported concurrency issues, but malformed container shapes can still lose chunk-level failure context; the existing malformed-response thread already covers this, so no duplicate inline comment is added. Evidence: static review of property_graph_extract.py lines 170-182 and 203-204; all visible latest-head checks are green.
Thanks for the review. I fixed this by validating the property graph container shape during per-chunk response parsing.
I also added regression coverage for malformed container shapes. |
imbajin
left a comment
There was a problem hiding this comment.
Blocking: yes. Summary: REST worker validation still coerces invalid raw inputs, and the existing cancellation-policy thread remains only partially addressed; two focused follow-ups are included. Evidence: direct Pydantic/shared-validator probe, static review, 17 focused tests and 56 related tests passed, and Ruff passed.
| extract_type: Literal["property_graph"] = Query("property_graph", description="Extraction type.") | ||
| language: Literal["zh", "en"] = Query("zh", description="Language for chunk splitting.") | ||
| split_type: Literal["document", "paragraph", "sentence"] = Query("document", description="Chunk split granularity.") | ||
| graph_extract_max_workers: int = Query( |
There was a problem hiding this comment.
🧹 REST validation coerces some raw inputs before the shared validator can enforce its contract. A direct probe shows graph_extract_max_workers=true and "1.0" are both accepted here as integer 1, while validate_graph_extract_max_workers() rejects those same raw values. This makes REST behavior inconsistent with the flow/helper entry points and can silently turn malformed input into a valid request. Please add a mode="before" field validator (or equivalent strict raw-input validation) that applies the shared rules, and cover boolean and decimal-string request values with regression tests.
|
|
||
|
|
||
| def test_property_graph_extract_preserves_chunk_merge_order_with_concurrency(): | ||
| llm = CountingLLM() |
There was a problem hiding this comment.
🧹 This order-preservation test gives every chunk the same delay, so it does not deterministically make completion order differ from input order. An implementation that accidentally merges futures in completion order could still pass when these calls finish first/second/third. Please coordinate the fake calls with distinct events or delays so they finish in a fixed reverse order, then assert that the merged result remains in input order.
imbajin
left a comment
There was a problem hiding this comment.
Blocking: yes. Summary: The latest head fixes the worker-validation issues, but the existing cancellation-policy thread still covers an unresolved failure-order and in-flight-call lifecycle concern; no duplicate inline comment is added. Evidence: static control-flow review, 51 focused tests passed, and all visible latest-head checks are green.
Purpose
Closes #344.
This PR adds bounded chunk-level concurrency for property graph extraction.
Changes
graph_extract_max_workersoption for graph extraction.1, preserving the existing serial behavior.PropertyGraphExtractusingThreadPoolExecutor.Graph extraction failed for chunk 2/5: ....Graph Extraction Chunk Concurrency.Notes
The default value is
1, so existing behavior remains serial unless users explicitly increase the chunk concurrency limit. The helper also avoids forwarding the default concurrency value to preserve existing scheduler call behavior.Tests
uv run ruff check hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py hugegraph-llm/src/hugegraph_llm/state/ai_state.py hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py hugegraph-llm/src/tests/document/test_graph_extract_concurrency.pyuv run ruff format --check .uv run pytest hugegraph-llm/src/tests/document/test_graph_extract_concurrency.pyuv run pytest hugegraph-llm/src/tests/document