refactor: refactor pipeline engine using ray data #110
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ChenZiHong-Gavin
changed the title
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
…to refactor/refactor-with-ray-data
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
…open-sciencelab/GraphGen into refactor/refactor-with-ray-data
Collaborator
Author
|
/gemini summary |
Contributor
|
Hi @ChenZiHong-Gavin, it looks like you've tried to invoke the |
Collaborator
Author
|
/gemini summary |
Collaborator
Author
|
/gemini review |
Contributor
Summary of ChangesThis pull request introduces a major architectural overhaul to the GraphGen project, transitioning its pipeline execution from a custom threading model to a robust, distributed framework built on Ray Data. This refactoring is designed to enhance the system's scalability, resource efficiency, and maintainability by standardizing operations around Ray's data processing primitives. Key changes include the introduction of a Highlights
Changelog
Activity
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This is an impressive and substantial refactoring of the pipeline execution engine. Migrating from a custom threading-based system to a Ray Data-native framework is a significant improvement for scalability, maintainability, and robustness. The new architecture with a declarative YAML configuration, a Ray Data-based engine, and modular operators is well-designed. The use of Ray actors for managing stateful services like LLMs and storage is a solid pattern for distributed execution.
My review includes a few suggestions to further improve robustness and performance. Specifically, I've commented on improving error handling in data processing, making logging more consistent in storage classes, and optimizing database write patterns.
Overall, this is an excellent piece of work that modernizes the project's core infrastructure.
Comment on lines
+28
to
+46
| def _should_keep_item(self, item: Dict[str, Any]) -> bool: | ||
| """ | ||
| Determine whether to keep the given item based on the text column. | ||
| :param item: Dictionary representing a data entry. | ||
| :return: True if the item should be kept, False otherwise. | ||
| """ | ||
| Filter out entries with empty or missing text in the specified column. | ||
| item_type = item.get("type") | ||
| assert item_type in [ | ||
| "text", | ||
| "image", | ||
| "table", | ||
| "equation", | ||
| "protein", | ||
| ], f"Unsupported item type: {item_type}" | ||
| if item_type == "text": | ||
| content = item.get(self.text_column, "").strip() | ||
| return bool(content) | ||
| return True |
Contributor
There was a problem hiding this comment.
In _should_keep_item, using assert to validate item_type will crash the Ray worker on unsupported types. It's more robust to log a warning and return False to filter out invalid items without halting the pipeline.
def _should_keep_item(self, item: Dict[str, Any]) -> bool:
"""
Determine whether to keep the given item based on the text column.
:param item: Dictionary representing a data entry.
:return: True if the item should be kept, False otherwise.
"""
item_type = item.get("type")
if item_type not in {
"text",
"image",
"table",
"equation",
"protein",
}:
# Consider logging a warning for visibility on data quality issues.
return False
if item_type == "text":
content = item.get(self.text_column, "").strip()
return bool(content)
return True| self._graph.nodes[node_id].update(node_data) | ||
| else: | ||
| logger.warning("Node %s not found in the graph for update.", node_id) | ||
| print(f"Node {node_id} not found in the graph for update.") |
Contributor
There was a problem hiding this comment.
Replacing logger.warning with print is a regression in logging practices. print bypasses logging configuration (levels, formatters, handlers), making it difficult to control log output in different environments. It's better to use the standard logging module. Since the context-aware logger might not be available in this class, importing logging and using logging.warning(...) would be a more robust solution. This feedback applies to other print statements in this file and other storage classes (e.g., kuzu_storage.py, json_storage.py).
Comment on lines
+40
to
+44
| except Exception as e: # pylint: disable=broad-except | ||
| logger.error("Error in judging description: %s", e) | ||
| logger.info("Use default loss 0.1") | ||
| item["loss"] = -math.log(0.1) | ||
| return item |
Contributor
There was a problem hiding this comment.
Catching a broad Exception can hide bugs and make debugging difficult. If the LLM client library defines more specific exceptions (e.g., for API errors, timeouts, or content filtering), it's better to catch those explicitly. This allows for more granular error handling and reporting. If specific exceptions are not available, consider using logger.exception() to automatically include traceback information in the log, which is very helpful for debugging.
Comment on lines
+89
to
+122
| def _pre_tokenize(self) -> None: | ||
| """Pre-tokenize all nodes and edges to add token length information.""" | ||
| logger.info("Starting pre-tokenization of nodes and edges...") | ||
|
|
||
| nodes = self.kg_instance.get_all_nodes() | ||
| edges = self.kg_instance.get_all_edges() | ||
|
|
||
| # Process nodes | ||
| for node_id, node_data in nodes: | ||
| if "length" not in node_data: | ||
| try: | ||
| description = node_data.get("description", "") | ||
| tokens = self.tokenizer_instance.encode(description) | ||
| node_data["length"] = len(tokens) | ||
| self.kg_instance.update_node(node_id, node_data) | ||
| except Exception as e: | ||
| logger.warning("Failed to tokenize node %s: %s", node_id, e) | ||
| node_data["length"] = 0 | ||
|
|
||
| # Process edges | ||
| for u, v, edge_data in edges: | ||
| if "length" not in edge_data: | ||
| try: | ||
| description = edge_data.get("description", "") | ||
| tokens = self.tokenizer_instance.encode(description) | ||
| edge_data["length"] = len(tokens) | ||
| self.kg_instance.update_edge(u, v, edge_data) | ||
| except Exception as e: | ||
| logger.warning("Failed to tokenize edge %s-%s: %s", u, v, e) | ||
| edge_data["length"] = 0 | ||
|
|
||
| # Persist changes | ||
| self.kg_instance.index_done_callback() | ||
| logger.info("Pre-tokenization completed.") |
Contributor
There was a problem hiding this comment.
The _pre_tokenize method updates nodes and edges one by one inside loops. This can lead to a large number of individual write operations to the storage backend, which is inefficient, especially for database-backed storages like KuzuDB. It would be more performant to batch these updates. Consider collecting all node and edge updates into lists and then performing a bulk update operation if the storage backend supports it. This might require adding bulk update methods to the BaseGraphStorage interface for better performance.
CHERRY-ui8
added a commit
to CHERRY-ui8/GraphGen
that referenced
this pull request
Dec 16, 2025
…ing read functionality
CHERRY-ui8
added a commit
to CHERRY-ui8/GraphGen
that referenced
this pull request
Dec 16, 2025
CHERRY-ui8
added a commit
to CHERRY-ui8/GraphGen
that referenced
this pull request
Dec 17, 2025
…ing read functionality
CHERRY-ui8
pushed a commit
to CHERRY-ui8/GraphGen
that referenced
this pull request
Dec 17, 2025
* feat: add config and operator node types * refactor: refactor readers with ray data * fix: delete param parallelism for readers * fix: fix import error * refactor read and chunk operators with no side effects * fix: fix import error * fix: fix return logic * refactor: rename operator split to chunk * refactor: refactor build_kg to accomodate ray data * feat: add StorageFactory & global params * refactor: refactor quiz to accomodata ray data engine * fix: reload graph before quizzing * Potential fix for pull request finding 'Unreachable code' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> * fix: fix quiz params * refactor: refactor quiz&judge to ray actors * fix: fix transferring quizzed data to JudgeService * refactor: refactor partition to accomodate ray data * fix: fix lint problem * refactor: refactor op generate * feat: write results in output folder * fix: raise error when no dataset is created * fix: return generator in ece_partitioner * fix: return generator in ece_partitioner * refactor: refactor data format to support multi-modal input * fix: delete fetching schema to avoid ray's duplicate execution * fix: fix operators' registry * feat: refactor schema_guided_extraction & add examples * feat: seperate ray logs and service logs * feat: use storage actor * feat: add kuzu graph database * feat: add llm as actors * refactor: delete old runner * fix: fix vllm wrapper * docs: update .env.example * fix: use kuzudb in quiz_service * fix: update webui * feat: make storage backend configuragble * docs: update README” --------- Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
CHERRY-ui8
added a commit
to CHERRY-ui8/GraphGen
that referenced
this pull request
Dec 17, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR fundamentally refactors GraphGen's pipeline execution engine from a custom threading-based orchestration system to a Ray Data-native distributed processing framework. The changes improve scalability, resource management, and maintainability while preserving the core pipeline semantics.