feat: remove use_scenario, add DatasetScenario, make traffic scenario optional

[CatherineSue]

Motivation

• The use_scenario flag leaked loader policy into samplers/CLI and blocked legitimate use-cases (e.g., using scenarios with CSV/JSON/HF datasets).
• Users expect “no scenarios provided + dataset present” to run the dataset directly without token shaping.
• The legacy “F” sentinel was unclear; we needed a user-friendly, first-class representation.

Modifications

• Traffic scenarios are optional end-to-end

  • When no --traffic-scenario is passed and a dataset is provided, we default to dataset mode (“dataset”) without token shaping.
  • We keep the existing scenario-driven benchmark loop unchanged by injecting a single dataset sentinel (“dataset”) when needed.

• use_scenario removed

  • Deleted the use_scenario flag from loaders, CLI, and samplers.
  • DataLoaderFactory now only returns loaded data (no policy-bit alongside).
  • Sampler and TextSampler no longer accept/use use_scenario.

• Introduced DatasetScenario as a first-class scenario

  • Added SpecialScenario.DATASET enum.
  • Removed the opaque “F” alias everywhere.

• Sampler refactor for dataset mode

  • Added Sampler.is_dataset_mode(scenario) to detect dataset/direct sampling without try/except or registry hacks.

• CLI and validation

  • CLI: if no --traffic-scenario and a dataset is present, we inject ["dataset"] and log that we’ll sample raw entries.
  • Validation: if dataset provided and no scenarios passed, return ["dataset"] so CLI can inject dataset mode; otherwise use defaults for that task.

• Embeddings restriction removed

  • Removed “sonnet-only” override for text-to-embeddings. Users can now use any dataset for embeddings.

• Tests updated

• Remove pin_to_core in DistributedRunner as it is not supported on all platforms. And Linux should have a better kernel to handle this.

[gemini-code-assist]

Summary of Changes

Hello @CatherineSue, 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!

I've refactored the data loading and sampling mechanisms to provide a more intuitive and flexible user experience. The primary goal was to decouple data loading policy from samplers and the CLI, allowing for direct dataset usage without explicit scenario definitions. This involved removing the problematic use_scenario flag, introducing a dedicated DatasetScenario to clearly indicate raw dataset sampling, and making traffic scenarios optional. Now, if you provide a dataset but no scenario, the system will intelligently default to sampling directly from your data. Additionally, I've removed an unnecessary restriction on embeddings, allowing any dataset to be used, and made a minor adjustment to CPU pinning defaults for broader platform compatibility.

Highlights

• Removal of use_scenario flag: The use_scenario flag, which previously dictated data loading and sampling policy, has been completely removed from the codebase, simplifying data flow and sampler logic.
• Introduction of DatasetScenario: A new DatasetScenario enum and class have been introduced to provide a clear, first-class representation for direct dataset sampling, replacing the opaque 'F' sentinel.
• Traffic scenarios are now optional: Traffic scenarios are now optional end-to-end. If no --traffic-scenario is provided but a dataset is present, the system automatically defaults to dataset mode, sampling raw entries without token shaping.
• Sampler refactoring for dataset mode: The Sampler and TextSampler classes have been refactored to use an internal _is_dataset_mode check, allowing them to adapt their sampling behavior based on whether a scenario is provided or if direct dataset sampling is intended.
• Embeddings restriction removed: The previous restriction that limited text-to-embeddings tasks to a specific 'sonnet' dataset has been removed, enabling users to utilize any dataset for embeddings.
• Disabled CPU pinning by default: The default for pin_to_cores in DistributedRunner has been changed to False due to platform compatibility issues, with a note that CPU affinity is primarily supported on Linux.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 or fill out our survey to provide feedback.

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 significantly improves the handling of scenarios and datasets by removing the ambiguous use_scenario flag and introducing a first-class DatasetScenario. The changes make the data loading and sampling logic more explicit and maintainable. The refactoring is comprehensive and includes updates to the CLI, validation, data loaders, and samplers, along with corresponding tests. My review identified a potential bug in the token sampling logic, a redundant code block, and a misleading docstring. Addressing these points will further strengthen this already solid contribution.