better parallelism in partitioned ray executor

[cyruszhang]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello, 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 significantly enhances the parallelism capabilities of the PartitionedRayExecutor by enabling concurrent processing of data partitions. It introduces a mechanism to dynamically scope the concurrency of individual operators within each parallel partition, optimizing GPU utilization and preventing resource over-subscription. This change aims to reduce overall data processing time, especially for GPU-intensive workflows, by leveraging available resources more effectively.

Highlights

• Concurrent Partition Processing: The PartitionedRayExecutor now supports processing multiple data partitions in parallel as Ray remote tasks, improving overall throughput for data processing pipelines.
• Dynamic Operator Concurrency Scoping: A new utility function, scope_op_concurrency, was introduced to dynamically adjust the num_proc for individual operators within each concurrent partition, ensuring efficient resource allocation and preventing GPU over-subscription.
• Flexible Concurrent Partition Configuration: Users can now configure max_concurrent_partitions to automatically detect available GPUs in the Ray cluster or specify an explicit integer, with automatic adjustment of num_partitions to ensure efficient resource usage.
• Improved Dataset Splitting Robustness: The dataset is now explicitly repartitioned before splitting to prevent the creation of empty partitions, ensuring all partitions contain data for processing.
• New Design Document: A design document, parallel_partition_actor_reuse.md, was added, outlining the problem of sequential GPU operator execution and proposing solutions for improved parallelism and actor reuse.
• Performance Benchmark Script: A new perf-test.py script was added to benchmark the performance of DataJuicer's Ray executors, including the new partitioned mode, and direct Ray Data parallelism.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• data_juicer/core/data/ray_dataset.py
  • Updated Ray map_batches calls to directly specify operator concurrency (concurrency=op.num_proc) instead of using get_compute_strategy, aiming for better GPU utilization.
• data_juicer/core/executor/concurrency_scoping.py
  • Added a new utility module defining scope_op_concurrency to dynamically adjust operator concurrency based on the number of concurrently running partitions.
• data_juicer/core/executor/ray_executor_partitioned.py
  • Introduced max_concurrent_partitions configuration, allowing 'auto' detection of GPUs or explicit integer setting.
  • Implemented a new _process_partitions_concurrent method to execute partitions in parallel as Ray remote tasks.
  • Ensured num_partitions is automatically increased if it's less than max_concurrent_partitions to fully utilize concurrent slots.
  • Added repartitioning logic (dataset.data.repartition(self.num_partitions)) before splitting the dataset to prevent empty partitions.
  • Modified _configure_partitioning to resolve max_concurrent_partitions and adjust num_partitions.
  • Modified _process_with_simple_partitioning to branch to the new concurrent processing logic if max_concurrent_partitions is greater than 1.
• docs/design/parallel_partition_actor_reuse.md
  • Added a new design document outlining the problem of sequential GPU operator execution and proposing solutions for improved parallelism and actor reuse in the partitioned Ray executor.
• perf-test.py
  • Added a new benchmark script to evaluate the performance of DataJuicer's Ray executors, including direct Ray Data parallelism and partitioned executor modes.
• pyproject.toml
  • Added a [tool.uv] section to specify constraint dependencies for the kaleido package.
• tests/core/executor/test_ray_executor_partitioned.py
  • Expanded unit tests to cover the new concurrency scoping logic in ConcurrencyScopingTest.
  • Added new test cases in ConcurrentPartitionConfigTest to validate configuration handling for max_concurrent_partitions and end-to-end concurrent execution scenarios, including checkpointing and backward compatibility.

Activity

• The author, cyruszhang, implemented changes to improve parallelism in the partitioned Ray executor.
• A new design document was added to explain the motivation and proposed architecture for parallel partition processing and actor reuse.
• A new benchmark script was introduced to test the performance of the updated executor.

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 introduces concurrent partition processing in the partitioned Ray executor to improve parallelism. The changes are mostly in ray_executor_partitioned.py, with supporting utilities and tests. My review found a couple of critical bugs in the new concurrent processing logic that could lead to errors when partitions are skipped. I've also pointed out a minor issue with exception handling and a significant discrepancy between the new design document and the actual implementation. Overall, the direction is good, but the identified bugs need to be addressed.