Add default batch size for SQLReader

[chenyushuo]

Description

As the title says.

Checklist

Please check the following items before code is ready to be reviewed.

• Code has passed all tests
• Docstrings have been added/updated in Google Style
• Documentation has been updated
• Code is ready for review

[gemini-code-assist]

Summary of Changes

Hello @chenyushuo, 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 introduces a default batch size mechanism for the SQLReader component, enhancing its configurability. Previously, read operations might have required an explicit batch size. With this change, the SQLReader can be initialized with a default batch size from its configuration, which will then be automatically used for subsequent read operations unless a specific batch size is provided. This streamlines data retrieval from SQL storage by offering a convenient default.

Highlights

• SQLReader Default Batch Size: The SQLReader class now initializes with a read_batch_size attribute, which is populated from the batch_size specified in its StorageConfig.
• Flexible Read Operations: The read and read_async methods within SQLReader have been updated to utilize this self.read_batch_size as a default if no batch_size argument is explicitly provided during the method call.
• Test Configuration Update: The sql_test.py file was modified to explicitly set the batch_size for the SQLWriter during testing by deepcopying the configuration, ensuring the writer's batch size is distinct and correctly configured.

🧠 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.

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. ↩

[chenyushuo]

/unittest-module-buffer

[gemini-code-assist]

Code Review

This pull request introduces a default batch_size for SQLReader, configured during its initialization. This allows read() and read_async() to be called without a batch_size argument, falling back to the configured default. The implementation for setting this default value has a subtle bug: passing batch_size=0 would incorrectly use the default batch size instead of 0. I've suggested fixes for this in both read and read_async methods to ensure 0 is handled correctly. The accompanying test changes are well-structured and effectively validate the new functionality.

[github-actions]

Summary

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Other ❓	Flaky 🍂	Duration ⏱️
48	48	0	0	0	0	1m 58s

Tests

Test Name	Status	Flaky	Duration
tests/buffer/experience_pipeline_test.py::TestExperiencePipeline::test_experience_pipeline	✅		10.6s
tests/buffer/experience_pipeline_test.py::TestExperiencePipeline::test_pass_rate_calculation	✅		4.7s
tests/buffer/experience_storage_test.py::ExperienceStorageTest::test_sql_experience_buffer	✅		2.8s
tests/buffer/experience_storage_test.py::ExperienceStorageTest::test_sql_storage_0_sft	✅		5.6s
tests/buffer/experience_storage_test.py::ExperienceStorageTest::test_sql_storage_1_dpo	✅		5.3s
tests/buffer/file_test.py::TestFileBuffer::test_file_reader	✅		441ms
tests/buffer/file_test.py::TestFileBuffer::test_file_writer	✅		1.9s
tests/buffer/formatter_test.py::TestFormatter::test_dpo_messages_formatter	✅		579ms
tests/buffer/formatter_test.py::TestFormatter::test_dpo_plaintext_formatter	✅		470ms
tests/buffer/formatter_test.py::TestFormatter::test_multi_modal_sft_formatter	✅		1.4s
tests/buffer/formatter_test.py::TestFormatter::test_sft_messages_formatter	✅		1.0s
tests/buffer/formatter_test.py::TestFormatter::test_sft_plaintext_formatter	✅		743ms
tests/buffer/formatter_test.py::TestFormatter::test_task_formatter	✅		230ms
tests/buffer/queue_test.py::TestQueueBuffer::test_priority_queue_buffer_reuse	✅		6.5s
tests/buffer/queue_test.py::TestQueueBuffer::test_priority_queue_capacity	✅		2.7s
tests/buffer/queue_test.py::TestQueueBuffer::test_priority_queue_reuse_count_control	✅		4.5s
tests/buffer/queue_test.py::TestQueueBuffer::test_queue_buffer_0_queue	✅		3.7s
tests/buffer/queue_test.py::TestQueueBuffer::test_queue_buffer_1_priority_queue	✅		3.7s
tests/buffer/queue_test.py::TestQueueBuffer::test_queue_buffer_capacity	✅		4.1s
tests/buffer/reader_test.py::TestBufferReader::test_buffer_reader_registration	✅		547ms
tests/buffer/reward_shaping_mapper_test.py::TestRewardShapingMapper::test_basic_usage	✅		7ms
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_0::test_default_queue_default_sample_strategy	✅		1.9s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_0::test_default_queue_staleness_control_sample_strategy	✅		2.1s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_0::test_priority_queue_default_sample_strategy	✅		2.0s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_0::test_priority_queue_staleness_control_sample_strategy	✅		2.2s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_0::test_sql_staleness_control_sample_strategy	✅		5.0s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_1::test_default_queue_default_sample_strategy	✅		2.0s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_1::test_default_queue_staleness_control_sample_strategy	✅		2.0s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_1::test_priority_queue_default_sample_strategy	✅		2.2s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_1::test_priority_queue_staleness_control_sample_strategy	✅		2.2s
tests/buffer/sample_strategy_test.py::ExperienceStorageTest_1::test_sql_staleness_control_sample_strategy	✅		3.8s
tests/buffer/sql_test.py::TestSQLBuffer::test_sql_exp_buffer_read_write_0	✅		5.8s
tests/buffer/sql_test.py::TestSQLBuffer::test_sql_exp_buffer_read_write_1	✅		3.0s
tests/buffer/sql_test.py::TestSQLBuffer::test_sql_task_buffer_read_write	✅		3.5s
tests/buffer/task_scheduler_test.py::TestTaskScheduler::test_task_scheduler_0	✅		72ms
tests/buffer/task_scheduler_test.py::TestTaskScheduler::test_task_scheduler_1	✅		56ms
tests/buffer/task_scheduler_test.py::TestTaskScheduler::test_task_scheduler_2	✅		88ms
tests/buffer/task_scheduler_test.py::TestTaskScheduler::test_task_scheduler_3	✅		88ms
tests/buffer/task_scheduler_test.py::TestTaskScheduler::test_task_scheduler_4	✅		88ms
tests/buffer/task_scheduler_test.py::TestTaskScheduler::test_task_scheduler_5	✅		92ms
tests/buffer/task_scheduler_test.py::TestTaskScheduler::test_task_scheduler_6	✅		108ms
tests/buffer/task_scheduler_test.py::TestTaskScheduler::test_task_scheduler_simple	✅		44ms
tests/buffer/task_storage_test.py::TaskStorageTest::test_read_task_0_file	✅		60ms
tests/buffer/task_storage_test.py::TaskStorageTest::test_read_task_1_sql	✅		3.1s
tests/buffer/task_storage_test.py::TaskStorageTest::test_read_task_2_file	✅		41ms
tests/buffer/task_storage_test.py::TaskStorageTest::test_read_task_3_sql	✅		3.4s
tests/buffer/task_storage_test.py::TaskStorageTest::test_read_task_4_file	✅		39ms
tests/buffer/task_storage_test.py::TaskStorageTest::test_read_task_5_sql	✅		3.9s

Github Test Reporter by CTRF 💚