feat: add bounded graph extraction concurrency by nw9663644-eng · Pull Request #370 · apache/hugegraph-ai

nw9663644-eng · 2026-07-10T10:39:39Z

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

imbajin

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.

imbajin · 2026-07-10T21:24:04Z

+            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:


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.

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

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.

imbajin · 2026-07-10T21:24:58Z

+            chunk,
+            proceeded_chunk,
+        )
+        return self._extract_and_filter_label(schema, proceeded_chunk)


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.

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

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.

imbajin · 2026-07-11T13:25:32Z

+            }
+            for future in as_completed(future_to_index):
+                index = future_to_index[future]
+                chunk_results[index] = future.result()


⚠️ 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.

imbajin · 2026-07-11T13:25:32Z

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


🧹 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

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

VGalaxies · 2026-07-11T17:41:22Z

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


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

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 · 2026-07-14T06:20:07Z

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

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.

imbajin · 2026-07-14T09:47:01Z

    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(


🧹 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.

imbajin · 2026-07-14T09:47:01Z

+
+
+def test_property_graph_extract_preserves_chunk_merge_order_with_concurrency():
+    llm = CountingLLM()


🧹 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

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.

imbajin

Blocking: yes. Summary: The bounded failure handling is substantially improved, but item-level normalization can still lose chunk context, and the new tests leave lifecycle, REST forwarding, and concurrency behavior incompletely or nondeterministically covered. Evidence: static control-flow and test review; PR diff check passed; all visible latest-head checks passed; focused local pytest could not start because the configured dependency mirror timed out.

imbajin · 2026-07-18T11:45:49Z

+            chunk,
+            proceeded_chunk,
+        )
+        return self._extract_and_filter_label(schema, property_graph)


⚠️ Item-level normalization failures still escape without the promised chunk context. The try above ends before this call, so a valid top-level graph whose vertex has a non-mapping properties value can fail later at properties.get(...) with a bare AttributeError, making it impossible to identify the offending chunk. Please include filtering and normalization in the chunk-level exception wrapper and add a regression case for malformed item fields.

imbajin · 2026-07-18T11:45:49Z

+
+            if "FAIL_CHUNK" in prompt:
+                assert self.slow_started.wait(timeout=2)
+                self.release_slow.set()


⚠️ Releasing the slow call from the failing worker makes the lifecycle assertion racy: an implementation that returns immediately after observing the failure can still pass if the slow thread reaches slow_finished before the assertion runs. Please execute extractor.run() in a separate test thread, assert it has not returned while the slow call remains blocked, then release the slow call and verify the extraction returns only afterward.

imbajin · 2026-07-18T11:45:49Z



+@pytest.mark.parametrize("raw_value", [True, False, "1.0"])
+def test_graph_extract_rejects_raw_invalid_worker_values(raw_value):


⚠️ The new REST tests only cover rejected raw values; none sends a valid non-default worker count through a successful POST and inspects the scheduler call. Removing or misrouting the forwarding in GraphExtractService would therefore remain green. Please add a mocked-scheduler request with graph_extract_max_workers: 4 and assert that exact value reaches schedule_flow (plus the default if useful).

imbajin · 2026-07-18T11:45:49Z

+    result = extractor.run({"schema": SCHEMA, "chunks": ["a", "b", "c", "d"]})
+
+    assert llm.max_active <= 2
+    assert llm.max_active > 1


⚠️ This concurrency proof depends on the 20 ms sleep in CountingLLM. On a loaded runner, the second executor worker can start after the first call finishes, causing a correct max_workers=2 implementation to report max_active == 1. Please coordinate the first two calls with a barrier or started/release events so their overlap is deterministic while retaining the upper-bound assertion.

imbajin

Blocking: yes. Summary: The latest head resolves the cancellation lifecycle, but persisted worker configuration still lacks round-trip and upgrade-default regression coverage; the existing malformed-output thread also remains only partially addressed for nested item shapes. Evidence: six independent review lanes, static lifecycle review, a direct validator probe, Ruff on all changed files, and all visible latest-head checks are green; full local pytest was blocked before collection by unavailable dependencies.

imbajin · 2026-07-18T11:46:36Z

            "gremlin_generate_prompt": to_literal(self.gremlin_generate_prompt),
            "doc_input_text": to_literal(self.doc_input_text),
            "graph_extract_split_type": to_literal(self.graph_extract_split_type),
+            "graph_extract_max_workers": self.graph_extract_max_workers,


⚠️ This adds graph_extract_max_workers to the persisted prompt configuration, but the persistence tests still cover only graph_extract_split_type: there is no round-trip assertion for a non-default worker count and no old-YAML fixture without this key. A serialization or upgrade-default regression could therefore silently reset the UI concurrency setting while the new concurrency tests still pass. Please extend the prompt-config tests to save/reload a non-default value and to load an older config without this field, asserting the fallback remains 1.

feat: add bounded graph extraction concurrency

3719fb3

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 Bot added the llm label Jul 10, 2026

imbajin reviewed Jul 10, 2026

View reviewed changes

fix: address graph extraction concurrency review comments

fd572e3

imbajin reviewed Jul 11, 2026

View reviewed changes

VGalaxies reviewed Jul 11, 2026

View reviewed changes

wn12222 added 3 commits July 12, 2026 15:31

fix: address graph extraction concurrency review comments

182e1a3

fix: restore graph extraction schema filtering

5e3ac82

fix: preserve graph extraction filter compatibility

fea7c35

imbajin reviewed Jul 12, 2026

View reviewed changes

fix: reject malformed graph extraction containers

e80f434

imbajin reviewed Jul 14, 2026

View reviewed changes

fix: tighten graph extraction worker validation

43d9bd6

imbajin reviewed Jul 15, 2026

View reviewed changes

fix: handle graph extraction concurrency failures deterministically

8f43b44

imbajin reviewed Jul 18, 2026

View reviewed changes



		def test_property_graph_extract_preserves_chunk_merge_order_with_concurrency():
		llm = CountingLLM()



		@pytest.mark.parametrize("raw_value", [True, False, "1.0"])
		def test_graph_extract_rejects_raw_invalid_worker_values(raw_value):

Uh oh!

Conversation

nw9663644-eng commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Changes

Notes

Tests

Uh oh!

imbajin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

imbajin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

imbajin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VGalaxies left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

imbajin left a comment

Choose a reason for hiding this comment

Uh oh!

nw9663644-eng commented Jul 14, 2026

Uh oh!

imbajin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

imbajin left a comment

Choose a reason for hiding this comment

Uh oh!

imbajin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

imbajin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

nw9663644-eng commented Jul 10, 2026 •

edited

Loading