[BREAKING] FEAT: Make scenario dataset sources and limits explicit (#2956)

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: richlundeen

.github/instructions/scenarios.instructions.md

@@ -28,7 +28,7 @@ class MyScenario(Scenario):
         super().__init__(
             version=self.VERSION,
             technique_class=MyTechnique,
-            default_dataset_config=DatasetConfiguration(dataset_names=["my_dataset"]),
+            default_dataset_config=DatasetAttackConfiguration(sources=[DatasetSource(name="my_dataset")]),
             objective_scorer=objective_scorer or self._get_default_objective_scorer(),
             scenario_result_id=scenario_result_id,
         )
@@ -68,7 +68,7 @@ def __init__(
     super().__init__(
         version=self.VERSION,
         technique_class=MyTechnique,
-        default_dataset_config=DatasetConfiguration(dataset_names=["my_dataset"]),
+        default_dataset_config=DatasetAttackConfiguration(sources=[DatasetSource(name="my_dataset")]),
         objective_scorer=objective_scorer,
     )
 ```
@@ -113,30 +113,64 @@ Dropping a common input is not silent: `set_params_from_args` rejects any value
 
 ## Dataset Loading
 
-Datasets are read from `CentralMemory`.
+Datasets are read from `CentralMemory`. New runs call `prepare_async()` before
+reading. Reads, discovery, estimates, and resume must never fetch or store datasets.
+Resolve parameter-dependent source names before preparation, not during seed reads.
 
 ### Basic — named datasets:
 ```python
-DatasetConfiguration(
-    dataset_names=["airt_hate", "airt_violence"],
-    max_dataset_size=10,  # optional: sample up to N per dataset
+DatasetAttackConfiguration(
+    sources=[
+        DatasetSource(name="harmbench"),
+        DatasetSource(name="airt_hate", max_size=10),
+        DatasetSource(name="team_objectives", max_size="all"),
+    ],
+    max_per_dataset=5,
+    max_total=12,
 )
 ```
 
+With filtered populations of 100, 40, and 7 groups, selection has two steps:
+
+| Dataset | Full population | After source cap |
+| --- | --- | --- |
+| `harmbench` | 100 | 5 (inherits `max_per_dataset`) |
+| `airt_hate` | 40 | 10 (explicit `max_size`) |
+| `team_objectives` | 7 | 7 (`"all"` removes the source cap) |
+
+The total cap then samples 12 groups from these 22 groups. The split can be 3, 6,
+and 3, for example; it is not balanced or guaranteed to cover every dataset.
+A dataset with no selected groups keeps its key with an empty list.
+
 ### Advanced — custom subclass for filtering:
 ```python
-class MyDatasetConfiguration(DatasetConfiguration):
-    def get_seed_groups(self) -> dict[str, list[SeedGroup]]:
-        result = super().get_seed_groups()
-        # Filter by selected techniques via self._scenario_techniques
-        return filtered_result
+class MyDatasetConfiguration(DatasetAttackConfiguration):
+    def _build_attack_groups(self, seeds: list[Seed]) -> list[AttackSeedGroup]:
+        return build_custom_groups(seeds)
 ```
 
 Options:
-- `dataset_names` — load by name from memory
-- `seed_groups` — pass explicit groups (mutually exclusive with `dataset_names`)
-- `max_dataset_size` — cap per dataset
-- Override `_load_seed_groups_for_dataset()` for custom loading
+- `sources` selects datasets by name. Source `max_size` overrides `max_per_dataset`;
+  omitted, `None`, empty, and `"default"` mean inherit; `"all"` means no cap.
+- All dataset limits use the same rule: omitted, `None`, empty, and `"default"` use the
+  default; `"all"` removes that limit; a positive integer sets a cap.
+- `max_per_dataset=5` is the named-objective default. `max_total="all"` leaves the
+  combined selection uncapped. Apply source limits before the total limit.
+- `fetch=DatasetFetchPolicy.IF_MISSING` prepares absent registered datasets.
+  `NEVER` requires stored data. A filter miss must never fetch.
+- `seed_groups` and `seeds` are inline alternatives. They never use memory or providers.
+- Validators run on full filtered populations before any sampling.
+- Ingredients must remain complete. Use `max_per_dataset="all"` and no finite source
+  `max_size`; cap the assembled attack groups with `max_total`, not ingredient rows.
+- Use `with_overrides()` instead of reconstructing a subclass or mutating its defaults.
+- Keep custom shaping in `_build_attack_groups()` or `_build_groups_by_dataset_async()`.
+  Reads use `_collect_seeds_for_dataset_async()` and must not fetch.
+- `dataset_names`, `max_dataset_size`, `auto_fetch`, and `per_dataset()` are deprecated.
+
+`max_dataset_size` is an exact alias for `max_total`; it does not disable source caps.
+To keep the old total-only selection of up to 10 groups, use
+`max_per_dataset="all", max_total=10`. To remove both caps, set both to `"all"`;
+`None` now uses the default, not unlimited selection.
 
 ## Technique Enum
 

doc/code/datasets/6_generated_datasets.ipynb

@@ -240,11 +240,11 @@
    ],
    "source": [
     "from pyrit.output import output_scenario_async, output_scenario_attacks_async\n",
-    "from pyrit.scenario import DatasetAttackConfiguration\n",
+    "from pyrit.scenario import DatasetAttackConfiguration, DatasetFetchPolicy, DatasetSource\n",
     "from pyrit.scenario.airt import RapidResponse, RapidResponseTechnique\n",
     "from pyrit.score import SelfAskTrueFalseScorer\n",
     "\n",
-    "dataset_config = DatasetAttackConfiguration(dataset_names=[dataset_name], auto_fetch=False)\n",
+    "dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name=dataset_name)], fetch=DatasetFetchPolicy.NEVER)\n",
     "scenario = RapidResponse(objective_scorer=SelfAskTrueFalseScorer(chat_target=OpenAIChatTarget()))\n",
     "scenario.set_params_from_args(\n",
     "    args={\n",

doc/code/datasets/6_generated_datasets.py

@@ -76,11 +76,11 @@
 
 # %%
 from pyrit.output import output_scenario_async, output_scenario_attacks_async
-from pyrit.scenario import DatasetAttackConfiguration
+from pyrit.scenario import DatasetAttackConfiguration, DatasetFetchPolicy, DatasetSource
 from pyrit.scenario.airt import RapidResponse, RapidResponseTechnique
 from pyrit.score import SelfAskTrueFalseScorer
 
-dataset_config = DatasetAttackConfiguration(dataset_names=[dataset_name], auto_fetch=False)
+dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name=dataset_name)], fetch=DatasetFetchPolicy.NEVER)
 scenario = RapidResponse(objective_scorer=SelfAskTrueFalseScorer(chat_target=OpenAIChatTarget()))
 scenario.set_params_from_args(
     args={

doc/code/scenarios/0_scenarios.ipynb

@@ -65,9 +65,11 @@
     "     Matrix-shaped scenarios delegate to `build_matrix_atomic_attacks(context=...)` in one line.\n",
     "\n",
     "3. **Default Dataset**: Pass `default_dataset_config=` to `super().__init__()` to specify the datasets your scenario uses out of the box.\n",
-    "   - Returns a `DatasetConfiguration` with one or more named datasets (e.g., `DatasetConfiguration(dataset_names=[\"my_dataset\"])`)\n",
+    "   - Returns a `DatasetAttackConfiguration` with named sources (e.g., `sources=[DatasetSource(name=\"my_dataset\")]`)\n",
     "   - Users can override this at runtime via `--dataset-names` in the CLI or by passing a custom `dataset_config` programmatically\n",
-    "   - `DatasetAttackConfiguration` selects at most 5 attack groups unless you set `max_dataset_size`; `max_dataset_size=None` uses all groups\n",
+    "   - Named sources select at most 5 attack groups per dataset by default; `max_total` caps the union. For any limit, omitted/`None`/empty/`\"default\"` uses the default; `\"all\"` removes that limit.\n",
+    "   - The deprecated `max_dataset_size` is an exact alias for `max_total`, not a way to disable source caps. To keep the old total-only selection, use `max_per_dataset=\"all\", max_total=10`; to remove both caps, set both to `\"all\"`. `None` now uses the default, not unlimited selection.\n",
+    "   - New runs prepare missing registered datasets once; reads, estimates, and resume never fetch\n",
     "\n",
     "4. **Constructor**: Use `@apply_defaults` decorator and call `super().__init__()` with scenario metadata:\n",
     "   - `name`: Descriptive name for your scenario\n",
@@ -113,7 +115,8 @@
    "source": [
     "from pyrit.common import apply_defaults\n",
     "from pyrit.scenario import (\n",
-    "    DatasetConfiguration,\n",
+    "    DatasetAttackConfiguration,\n",
+    "    DatasetSource,\n",
     "    Scenario,\n",
     "    ScenarioTechnique,\n",
     ")\n",
@@ -159,7 +162,9 @@
     "            version=self.VERSION,\n",
     "            objective_scorer=self._objective_scorer,\n",
     "            technique_class=MyTechnique,\n",
-    "            default_dataset_config=DatasetConfiguration(dataset_names=[\"dataset_name\"], max_dataset_size=4),\n",
+    "            default_dataset_config=DatasetAttackConfiguration(\n",
+    "                sources=[DatasetSource(name=\"dataset_name\")], max_total=4\n",
+    "            ),\n",
     "            scenario_result_id=scenario_result_id,\n",
     "        )\n",
     "\n",

doc/code/scenarios/0_scenarios.py

@@ -67,9 +67,11 @@
 #      Matrix-shaped scenarios delegate to `build_matrix_atomic_attacks(context=...)` in one line.
 #
 # 3. **Default Dataset**: Pass `default_dataset_config=` to `super().__init__()` to specify the datasets your scenario uses out of the box.
-#    - Returns a `DatasetConfiguration` with one or more named datasets (e.g., `DatasetConfiguration(dataset_names=["my_dataset"])`)
+#    - Returns a `DatasetAttackConfiguration` with named sources (e.g., `sources=[DatasetSource(name="my_dataset")]`)
 #    - Users can override this at runtime via `--dataset-names` in the CLI or by passing a custom `dataset_config` programmatically
-#    - `DatasetAttackConfiguration` selects at most 5 attack groups unless you set `max_dataset_size`; `max_dataset_size=None` uses all groups
+#    - Named sources select at most 5 attack groups per dataset by default; `max_total` caps the union. For any limit, omitted/`None`/empty/`"default"` uses the default; `"all"` removes that limit.
+#    - The deprecated `max_dataset_size` is an exact alias for `max_total`, not a way to disable source caps. To keep the old total-only selection, use `max_per_dataset="all", max_total=10`; to remove both caps, set both to `"all"`. `None` now uses the default, not unlimited selection.
+#    - New runs prepare missing registered datasets once; reads, estimates, and resume never fetch
 #
 # 4. **Constructor**: Use `@apply_defaults` decorator and call `super().__init__()` with scenario metadata:
 #    - `name`: Descriptive name for your scenario
@@ -97,7 +99,8 @@
 # %%
 from pyrit.common import apply_defaults
 from pyrit.scenario import (
-    DatasetConfiguration,
+    DatasetAttackConfiguration,
+    DatasetSource,
     Scenario,
     ScenarioTechnique,
 )
@@ -143,7 +146,9 @@ def __init__(
             version=self.VERSION,
             objective_scorer=self._objective_scorer,
             technique_class=MyTechnique,
-            default_dataset_config=DatasetConfiguration(dataset_names=["dataset_name"], max_dataset_size=4),
+            default_dataset_config=DatasetAttackConfiguration(
+                sources=[DatasetSource(name="dataset_name")], max_total=4
+            ),
             scenario_result_id=scenario_result_id,
         )
 

doc/code/scenarios/1_common_scenario_parameters.ipynb

@@ -75,7 +75,7 @@
     "## Dataset Configuration\n",
     "\n",
     "`DatasetAttackConfiguration` controls which prompts (objectives) are sent to the target.\n",
-    "The simplest approach uses `dataset_names` to load datasets by name from memory.\n",
+    "Use `sources` to select named datasets from memory.\n",
     "By default, `RedTeamAgent` loads four random objectives from HarmBench [@mazeika2024harmbench]."
    ]
   },
@@ -86,9 +86,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from pyrit.scenario import DatasetAttackConfiguration\n",
+    "from pyrit.scenario import DatasetAttackConfiguration, DatasetSource\n",
     "\n",
-    "dataset_config = DatasetAttackConfiguration(dataset_names=[\"harmbench\"], max_dataset_size=2)"
+    "dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name=\"harmbench\")], max_total=2)"
    ]
   },
   {
@@ -129,8 +129,8 @@
     "    if all(seed.value.isascii() and seed.value.isprintable() for seed in group.seeds)\n",
     "]\n",
     "\n",
-    "# Pass explicit seed_groups instead of dataset_names\n",
-    "dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=2)"
+    "# Pass explicit seed_groups instead of named sources\n",
+    "dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_total=2)"
    ]
   },
   {

doc/code/scenarios/1_common_scenario_parameters.py

@@ -42,13 +42,13 @@
 # ## Dataset Configuration
 #
 # `DatasetAttackConfiguration` controls which prompts (objectives) are sent to the target.
-# The simplest approach uses `dataset_names` to load datasets by name from memory.
+# Use `sources` to select named datasets from memory.
 # By default, `RedTeamAgent` loads four random objectives from HarmBench [@mazeika2024harmbench].
 
 # %%
-from pyrit.scenario import DatasetAttackConfiguration
+from pyrit.scenario import DatasetAttackConfiguration, DatasetSource
 
-dataset_config = DatasetAttackConfiguration(dataset_names=["harmbench"], max_dataset_size=2)
+dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name="harmbench")], max_total=2)
 
 # %% [markdown]
 # For more control, use `SeedDatasetProvider` to fetch datasets and pass explicit `seed_groups`.
@@ -69,8 +69,8 @@
     if all(seed.value.isascii() and seed.value.isprintable() for seed in group.seeds)
 ]
 
-# Pass explicit seed_groups instead of dataset_names
-dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=2)
+# Pass explicit seed_groups instead of named sources
+dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_total=2)
 
 # %% [markdown]
 # ## Technique Selection and Composition

doc/code/scenarios/3_adaptive_scenarios.ipynb

@@ -113,7 +113,7 @@
     "\n",
     "from pyrit.output.scenario_result.pretty import PrettyScenarioResultMemoryPrinter\n",
     "from pyrit.registry import TargetRegistry\n",
-    "from pyrit.scenario import DatasetAttackConfiguration\n",
+    "from pyrit.scenario import DatasetAttackConfiguration, DatasetSource\n",
     "from pyrit.scenario.scenarios.adaptive import TextAdaptive\n",
     "from pyrit.setup import initialize_from_config_async\n",
     "\n",
@@ -452,8 +452,8 @@
     "        \"objective_target\": objective_target,\n",
     "        \"scenario_techniques\": [technique_class(\"single_turn\")],\n",
     "        \"dataset_config\": DatasetAttackConfiguration(\n",
-    "            dataset_names=[\"airt_hate\", \"airt_violence\"],\n",
-    "            max_dataset_size=4,\n",
+    "            sources=[DatasetSource(name=\"airt_hate\"), DatasetSource(name=\"airt_violence\")],\n",
+    "            max_total=4,\n",
     "        ),\n",
     "    }\n",
     ")\n",
@@ -577,8 +577,8 @@
     "        \"objective_target\": objective_target,\n",
     "        \"scenario_techniques\": [technique_class(\"single_turn\")],\n",
     "        \"dataset_config\": DatasetAttackConfiguration(\n",
-    "            dataset_names=[\"airt_hate\", \"airt_violence\"],\n",
-    "            max_dataset_size=4,\n",
+    "            sources=[DatasetSource(name=\"airt_hate\"), DatasetSource(name=\"airt_violence\")],\n",
+    "            max_total=4,\n",
     "        ),\n",
     "    }\n",
     ")\n",

doc/code/scenarios/3_adaptive_scenarios.py

@@ -48,7 +48,7 @@
 
 from pyrit.output.scenario_result.pretty import PrettyScenarioResultMemoryPrinter
 from pyrit.registry import TargetRegistry
-from pyrit.scenario import DatasetAttackConfiguration
+from pyrit.scenario import DatasetAttackConfiguration, DatasetSource
 from pyrit.scenario.scenarios.adaptive import TextAdaptive
 from pyrit.setup import initialize_from_config_async
 
@@ -102,8 +102,8 @@
         "objective_target": objective_target,
         "scenario_techniques": [technique_class("single_turn")],
         "dataset_config": DatasetAttackConfiguration(
-            dataset_names=["airt_hate", "airt_violence"],
-            max_dataset_size=4,
+            sources=[DatasetSource(name="airt_hate"), DatasetSource(name="airt_violence")],
+            max_total=4,
         ),
     }
 )
@@ -132,8 +132,8 @@
         "objective_target": objective_target,
         "scenario_techniques": [technique_class("single_turn")],
         "dataset_config": DatasetAttackConfiguration(
-            dataset_names=["airt_hate", "airt_violence"],
-            max_dataset_size=4,
+            sources=[DatasetSource(name="airt_hate"), DatasetSource(name="airt_violence")],
+            max_total=4,
         ),
     }
 )

doc/scanner/1_pyrit_scan.ipynb

@@ -1287,7 +1287,7 @@
     "\n",
     "from pyrit.common import apply_defaults\n",
     "from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget\n",
-    "from pyrit.scenario import DatasetAttackConfiguration, Scenario, ScenarioTechnique\n",
+    "from pyrit.scenario import DatasetAttackConfiguration, DatasetSource, Scenario, ScenarioTechnique\n",
     "from pyrit.score import SelfAskRefusalScorer, TrueFalseInverterScorer\n",
     "from pyrit.setup import initialize_pyrit_async\n",
     "\n",
@@ -1311,8 +1311,8 @@
     "            version=1,\n",
     "            objective_scorer=TrueFalseInverterScorer(scorer=SelfAskRefusalScorer(chat_target=OpenAIChatTarget())),\n",
     "            technique_class=MyCustomTechnique,\n",
-    "            # DatasetAttackConfiguration selects at most 5 attack groups by default; set max_dataset_size to change it.\n",
-    "            default_dataset_config=DatasetAttackConfiguration(dataset_names=[\"harmbench\"]),\n",
+    "            # Named sources default to 5 groups each; use max_per_dataset and max_total to set caps.\n",
+    "            default_dataset_config=DatasetAttackConfiguration(sources=[DatasetSource(name=\"harmbench\")]),\n",
     "            scenario_result_id=scenario_result_id,\n",
     "        )\n",
     "        # ... your scenario-specific initialization code\n",