Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 22 additions & 5 deletions ecc/app/graphrag/graph_rag.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,7 +333,7 @@ async def load(conn: AsyncTigerGraphConnection):


async def embed(
embed_chan: Channel, embedding_store: EmbeddingStore, graphname: str
embed_chan: Channel, embedding_store: EmbeddingStore, graphname: str, progress=None
):
"""
Creates and starts one worker for each embed job
Expand Down Expand Up @@ -366,6 +366,11 @@ async def embed(
n_embed += 1
if n_embed % 100 == 0:
logger.info(f"Embedding Processing: {n_embed} embedded so far")
if progress is not None:
try:
progress(f"Embedding documents ({n_embed} embedded so far)")
except Exception:
pass
Comment on lines +369 to +373

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Suggestion: If progress is an async callback, calling it without awaiting will silently skip the UI update and emit an un-awaited coroutine warning. Capture the return value and await it when it is a coroutine; apply the same pattern to the new progress(...) calls in extract and summarize_communities. [possible issue, importance: 6]

Suggested change
if progress is not None:
try:
progress(f"Embedding documents ({n_embed} embedded so far)")
except Exception:
pass
if progress is not None:
try:
result = progress(f"Embedding documents ({n_embed} embedded so far)")
if asyncio.iscoroutine(result):
await result
except Exception:
pass

except ChannelClosed:
break
except Exception:
Expand All @@ -383,6 +388,7 @@ async def extract(
extractor: BaseExtractor,
conn: AsyncTigerGraphConnection,
num_senders: int,
progress=None,
):
"""
Creates and starts one worker for each extract job
Expand Down Expand Up @@ -412,6 +418,11 @@ async def extract(
logger.info(
f"Entity Extraction: {n_chunks} chunks dispatched"
)
if progress is not None:
try:
progress(f"Extracting entities ({n_chunks} chunks processed)")
except Exception:
pass
except ChannelClosed:
break
except Exception:
Expand Down Expand Up @@ -536,6 +547,7 @@ async def summarize_communities(
comm_process_chan: Channel,
upsert_chan: Channel,
embed_chan: Channel,
progress=None,
):
logger.info("Community summarization started")
n_comm = 0
Expand All @@ -552,6 +564,11 @@ async def summarize_communities(
logger.info(
f"Community summarization: {n_comm} dispatched"
)
if progress is not None:
try:
progress(f"Summarizing communities ({n_comm} done so far)")
except Exception:
pass
except ChannelClosed:
break
except Exception:
Expand Down Expand Up @@ -634,9 +651,9 @@ async def new_doc_pipeline():

grp.create_task(upsert(upsert_chan))
grp.create_task(load(conn))
grp.create_task(embed(embed_chan, embedding_store, graphname))
grp.create_task(embed(embed_chan, embedding_store, graphname, progress))
grp.create_task(
extract(extract_chan, upsert_chan, embed_chan, extractor, conn, num_chunk_senders)
extract(extract_chan, upsert_chan, embed_chan, extractor, conn, num_chunk_senders, progress)
)
logger.info("Join docs_chan")
await docs_chan.join()
Expand Down Expand Up @@ -687,11 +704,11 @@ async def new_doc_pipeline():
grp.create_task(communities(conn, comm_process_chan))
# summarize each community
grp.create_task(
summarize_communities(conn, comm_process_chan, upsert_chan, embed_chan)
summarize_communities(conn, comm_process_chan, upsert_chan, embed_chan, progress)
)
grp.create_task(upsert(upsert_chan))
grp.create_task(load(conn))
grp.create_task(embed(embed_chan, embedding_store, graphname))
grp.create_task(embed(embed_chan, embedding_store, graphname, progress))
logger.info("Join comm_process_chan")
await comm_process_chan.join()
logger.info("Join embed_chan")
Expand Down
24 changes: 21 additions & 3 deletions graphrag/app/routers/ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,10 @@
from common.metrics.prometheus_metrics import metrics as pmetrics
from common.metrics.tg_proxy import TigerGraphConnectionProxy
from common.utils.graph_locks import acquire_graph_lock, release_graph_lock, acquire_rebuild_lock, release_rebuild_lock, get_rebuilding_graph, get_current_operation

# Cache the last successful ECC status response per graph so the UI
# still sees started_at and stage when ECC is too busy to respond.
_last_ecc_status_cache: dict = {}
from supportai import supportai
from common.py_schemas.schemas import (
AgentProgess,
Expand Down Expand Up @@ -2258,7 +2262,7 @@ async def rebuild_and_monitor():
elapsed = 0

while elapsed < max_wait_time:
await asyncio.sleep(poll_interval) # Non-blocking sleep
await asyncio.sleep(poll_interval)
elapsed += poll_interval

try:
Expand All @@ -2272,6 +2276,11 @@ async def rebuild_and_monitor():
status_data = status_response.json()
is_running = status_data.get("is_running", False)
status = status_data.get("status", "unknown")

# Cache last known good response so the UI endpoint
# can fall back to it when ECC is too busy to respond.
if is_running and status_data.get("started_at"):
_last_ecc_status_cache[graphname] = status_data

# Log every minute to avoid spam
if elapsed % 60 == 0:
Expand All @@ -2280,6 +2289,7 @@ async def rebuild_and_monitor():
# Check if ALL stages are complete
if not is_running and status in ["completed", "failed", "idle"]:
LogWriter.info(f"ECC rebuild finished for {graphname} with status: {status} after {elapsed}s")
_last_ecc_status_cache.pop(graphname, None)
break
else:
LogWriter.warning(f"ECC status check returned {status_response.status_code} for {graphname}")
Expand All @@ -2290,12 +2300,17 @@ async def rebuild_and_monitor():

if elapsed >= max_wait_time:
LogWriter.error(f"ECC rebuild monitoring timed out for {graphname} after {max_wait_time}s")
_last_ecc_status_cache.pop(graphname, None)

except Exception as e:
LogWriter.error(f"Error during ECC rebuild monitoring for {graphname}: {e}")
Comment on lines 2305 to 2306

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Suggestion: _last_ecc_status_cache is only cleared on normal completion and timeout, so an unexpected monitor failure can leave stale running status for the graph. Clear the cache in finally before releasing the rebuild lock so later timeout fallbacks do not report an old rebuild as active. [possible issue, importance: 7]

New proposed code:
 except Exception as e:
     LogWriter.error(f"Error during ECC rebuild monitoring for {graphname}: {e}")
     import traceback
     LogWriter.error(traceback.format_exc())
 finally:
+    _last_ecc_status_cache.pop(graphname, None)
     # Release lock only after ALL stages complete (or timeout/error)
     release_rebuild_lock(graphname)
     LogWriter.info(f"Released global rebuild lock for {graphname}")

import traceback
LogWriter.error(traceback.format_exc())
finally:
# Always drop cached status when monitoring ends (success,
# timeout, or unexpected failure) so timeout fallbacks do
# not keep reporting a stale rebuild as active.
_last_ecc_status_cache.pop(graphname, None)
# Release lock only after ALL stages complete (or timeout/error)
release_rebuild_lock(graphname)
LogWriter.info(f"Released global rebuild lock for {graphname}")
Expand Down Expand Up @@ -2342,12 +2357,15 @@ def get_rebuild_status(
"error": f"ECC service returned status {response.status_code}"
}
except httpx.TimeoutException as e:
# ECC is busy (heavy processing) - assume rebuild is still running
# ECC is busy (heavy processing) - assume rebuild is still running.
# Return the last cached status so the UI keeps started_at and stage.
LogWriter.warning(f"ECC status check timed out (ECC may be busy): {str(e)}")
cached = _last_ecc_status_cache.get(graphname, {})
return {
**cached,
"graphname": graphname,
"is_running": True,
"status": "unknown",
"status": cached.get("status", "unknown"),
"error": "ECC is busy processing, status check timed out. Rebuild likely still in progress."
}
except Exception as e:
Expand Down