Fix resharding bug when dataset variable is reassigned

[m-laprise]

Problem: flame/data.py reuses dataset for both the dataset path string and the loaded IterableDataset. After load_dataset(), dataset is no longer a string, so when resharding, load_dataset(path=dataset) receives the dataset object instead of the path, causing it to error out.

Fix: Preserve the original dataset path string in dataset_path and pass dataset_path to resharding.

Edited to add: This happens when num_shards is smaller than dp_degree * num_workers, e.g. for datasets like wikitext-103-v1 that expose a small number of shards.

Summary by CodeRabbit

• Refactor
  • Internal code optimization to improve maintainability without affecting end-user functionality.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[gemini-code-assist]

Summary of Changes

Hello @m-laprise, 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 resolves a critical bug in the data loading and resharding logic by ensuring that the original dataset path is consistently used when reloading or re-sharding datasets. The change prevents errors that occurred due to a variable being overwritten with a dataset object, thereby improving the robustness of the data pipeline.

Highlights

• Bug Fix: Dataset Variable Reassignment: Addresses a bug where the dataset variable in flame/data.py was being reassigned from a string path to an IterableDataset object. This caused issues during resharding when load_dataset() was called again with the object instead of the original path string.
• Introduced dataset_path: A new variable, dataset_path, has been introduced to explicitly store the original dataset path string. This ensures that the correct path is always passed to load_dataset() during initial loading and subsequent resharding operations.

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

[coderabbitai]

Note

Other AI code review bot(s) detected

CodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review.

📝 Walkthrough

Walkthrough

A minimal refactor in flame/data.py introduces a dataset_path variable alias to replace direct references to dataset in load_dataset calls within the single-dataset code path. No functional behavior changes.

Changes

Cohort / File(s)	Summary
Variable alias refactor flame/data.py	Introduced dataset_path variable as an alias for dataset in the single-dataset branch; updated two load_dataset calls to use dataset_path instead of dataset directly. No control flow modifications.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

Poem

🐰 A hop, a skip, a name so clean,
Dataset becomes the path we've seen,
No logic bent, just variables renamed,
The flow remains, none of it changed! 🌿

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: fixing a bug where the dataset variable is reassigned, causing issues with resharding. The title is specific and directly related to the core problem being addressed.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug in flame/data.py where the dataset variable was being reused, causing an error during resharding. The fix preserves the original dataset path in a new dataset_path variable. The change is correct and addresses the issue. I've added one minor suggestion related to code maintainability to consider.

[gemini-code-assist]

Introducing dataset_path correctly fixes the bug. On a related note for maintainability, the load_dataset call on the next line is duplicated in the resharding logic (lines 588-597). Consider refactoring these calls to reduce duplication, for example by using a shared dictionary of arguments.