Skip to content

feat: add bounded graph extraction concurrency#370

Open
nw9663644-eng wants to merge 7 commits into
apache:mainfrom
nw9663644-eng:feat-bounded-graph-extract-concurrency
Open

feat: add bounded graph extraction concurrency#370
nw9663644-eng wants to merge 7 commits into
apache:mainfrom
nw9663644-eng:feat-bounded-graph-extract-concurrency

Conversation

@nw9663644-eng

@nw9663644-eng nw9663644-eng commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

Purpose

Closes #344.

This PR adds bounded chunk-level concurrency for property graph extraction.

Changes

  • Added a configurable graph_extract_max_workers option for graph extraction.
  • Kept the conservative default as 1, preserving the existing serial behavior.
  • Added bounded chunk-level concurrency in PropertyGraphExtract using ThreadPoolExecutor.
  • Preserved deterministic merge order by collecting concurrent chunk results by chunk index before merging.
  • Kept the same vertex/edge parsing and filtering path as serial extraction.
  • Added clear chunk-level failure context, for example Graph extraction failed for chunk 2/5: ....
  • Passed the concurrency setting through the graph extraction flow path.
  • Exposed the setting in the demo UI as Graph Extraction Chunk Concurrency.
  • Persisted the demo concurrency setting through prompt config.

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.py
  • uv run ruff format --check .
  • uv run pytest hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py
  • uv run pytest hugegraph-llm/src/tests/document

@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. enhancement New feature or request labels Jul 10, 2026
@github-actions github-actions Bot added the llm label Jul 10, 2026

@imbajin imbajin left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 imbajin left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 imbajin left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ A failed chunk does not stop queued LLM calls. Every chunk is submitted above, and when 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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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 VGalaxies left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 -> 1 and 8.9 -> 8


def validate_graph_extract_max_workers(value) -> int:
try:
workers = int(value)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Medium: Reject fractional worker counts instead of truncating them

hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py:21

Evidence

  • int(value) converts 1.9 to 1 and 8.9 to 8, 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 imbajin left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@nw9663644-eng

Copy link
Copy Markdown
Contributor Author

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.阻塞:是。摘要:最新版本修复了之前报告的并发问题,但格式错误的容器结构仍然可能丢失块级故障上下文;现有的 malformed-response 讨论串已涵盖此问题,因此不再添加重复的内联注释。证据:对 property_graph_extract.py 文件第 170-182 行和 203-204 行进行静态审查;所有可见的最新版本检查结果均为绿色。

Thanks for the review. I fixed this by validating the property graph container shape during per-chunk response parsing.

vertices and edges now must both be lists before the response can proceed to filtering. Malformed container shapes such as {"vertices": {}, "edges": []} now fail with chunk context instead of being silently merged as an empty result.

I also added regression coverage for malformed container shapes.

@imbajin imbajin left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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 imbajin left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request llm size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature] Add bounded chunk-level concurrency to graph extraction

4 participants