Fix: set chunk_num before set_progress, avoid race condition (#17724)

This commit is contained in:
Lynn
2026-08-03 15:52:23 +08:00
committed by GitHub
parent 59b682e6a4
commit 51af9e2eae
2 changed files with 16 additions and 4 deletions

View File

@@ -934,8 +934,12 @@ async def run_dataflow(task: dict):
time_cost = timer() - start_ts
task_time_cost = timer() - task_start_ts
try:
ret = DocumentService.increment_chunk_num(doc_id, task_dataset_id, embedding_token_consumption, len(chunks), task_time_cost)
except Exception:
logging.exception("increment_chunk_num failed for doc %s", doc_id)
ret = None
set_progress(task_id, prog=1.0, msg="Indexing done ({:.2f}s). Task done ({:.2f}s)".format(time_cost, task_time_cost))
ret = DocumentService.increment_chunk_num(doc_id, task_dataset_id, embedding_token_consumption, len(chunks), task_time_cost)
get_recording_context().save_func_return_value("DocumentService.increment_chunk_num", ret)
logging.info("[Done], chunks({}), token({}), elapsed:{:.2f}".format(len(chunks), embedding_token_consumption, task_time_cost))
get_recording_context().record("dataflow_chunks", chunks)

View File

@@ -166,13 +166,21 @@ class DataflowService:
time_cost = timer() - start_ts
task_time_cost = timer() - task_start_ts
self._progress(prog=1.0, msg="Indexing done ({:.2f}s). Task done ({:.2f}s)".format(time_cost, task_time_cost))
# Update document stats
# Update document stats (chunk counters) BEFORE marking the task
# as done, so the async _sync_progress loop never observes DONE
# with stale chunk_num=0. If the stats update fails, the task
# is still marked DONE — chunk counters are not critical to the
# parse result.
if ctx.write_interceptor:
ctx.write_interceptor.intercept("DocumentService.increment_chunk_num")
else:
DocumentService.increment_chunk_num(doc_id, task_dataset_id, embedding_token_consumption, len(chunks), task_time_cost)
try:
DocumentService.increment_chunk_num(doc_id, task_dataset_id, embedding_token_consumption, len(chunks), task_time_cost)
except Exception:
logging.exception("increment_chunk_num failed for doc %s", doc_id)
self._progress(prog=1.0, msg="Indexing done ({:.2f}s). Task done ({:.2f}s)".format(time_cost, task_time_cost))
logging.info("[Done], chunks({}), token({}), elapsed:{:.2f}".format(len(chunks), embedding_token_consumption, task_time_cost))
ctx.recording_context.record("dataflow_chunks", chunks)