perf: reduce system bubble time

[ChenZiHong-Gavin]

Before:

After:

[gemini-code-assist]

Summary of Changes

Hello @ChenZiHong-Gavin, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request focuses on optimizing the performance and robustness of the graph generation pipeline, particularly concerning knowledge graph building and quiz generation. Key changes include refactoring the quiz generation process to leverage batch processing, enhancing the output of knowledge graph building functions to return explicit nodes and edges, and improving error handling in graph partitioning. These modifications aim to reduce "system bubble time" by streamlining data flow and making operations more resilient.

Highlights

• Refactored Quiz Generation Workflow: The quiz operator's type in aggregated_config.yaml changed from aggregate to map_batch, and its internal logic was significantly refactored to process incoming data batches instead of fetching all graph data. This includes removing an explicit concurrency_limit and changing the return mechanism from yield to return a DataFrame.
• Enhanced Knowledge Graph Building Output: The build_kg service and its sub-functions (build_text_kg, build_mm_kg) now explicitly return the generated nodes and edges as a tuple, providing more direct access to the graph components.
• Improved Robustness in Graph Partitioner: The ece_partitioner.py file now uses dict.get() with a default loss value when sorting units, preventing potential KeyError exceptions and improving the stability of the partitioning process.
• Standardized Logging: Logging calls in evaluate_service.py were updated to use a module-level logger instead of self.logger, promoting consistency and potentially simplifying logger management.
• Added Partitioning Metrics: The partition_service.py now logs the total number of communities successfully partitioned, providing better visibility into the partitioning process.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the data processing pipeline to reduce system bubble time, primarily by changing the quiz operator from an aggregate type to map_batch. This allows for better pipelining and parallelism with the build_kg step. The changes in build_kg and related functions to pass nodes and edges forward are well-implemented and support this new architecture.

However, I've identified a few potential issues in the new quiz_service.py implementation that should be addressed to ensure the new implementation is both efficient and correct:

• A caching mechanism that prevented re-generation of existing quizzes has been removed, which could lead to redundant work and increased costs.
• The logic for processing nodes and edges from the input batch appears to handle only the first data entry for each item, potentially ignoring other valid data and leading to incomplete quiz generation.

[gemini-code-assist]

The nodes dictionary maps a node ID to a list of data dictionaries. This code only processes the first element of that list (node_data[0]), potentially ignoring other data associated with the same node. This could result in incomplete processing if a node is associated with multiple descriptions (e.g., from different source chunks). Consider iterating through the entire list of data dictionaries to process all available descriptions.

Suggested change

	for node_id, node_data in nodes.items():
	node_data = node_data[0]
	desc = node_data["description"]
	items.append((node_id, desc))
	for node_id, node_data_list in nodes.items():
	for node_data in node_data_list:
	desc = node_data["description"]
	items.append((node_id, desc))

[gemini-code-assist]

Similar to the loop for nodes, this loop for edges only processes the first data dictionary (edge_data[0]) for each edge. If an edge can have multiple associated data dictionaries, this will lead to data loss. To ensure all descriptions are processed, you should iterate through the list of data dictionaries for each edge.

Suggested change

	for edge_key, edge_data in edges.items():
	edge_data = edge_data[0]
	desc = edge_data["description"]
	items.append((edge_key, desc))
	for edge_key, edge_data_list in edges.items():
	for edge_data in edge_data_list:
	desc = edge_data["description"]
	items.append((edge_key, desc))

[gemini-code-assist]

The check for an existing quiz in quiz_storage has been removed. This will cause quizzes to be regenerated every time _process_single_quiz is called for the same item, even if a quiz already exists. This could lead to redundant processing and increased costs. Consider adding the check back:

if self.quiz_storage.get_by_id(_quiz_id):
    return None