MAINT: declare scenario required data apart from dataset_config

.github/instructions/scenarios.instructions.md

@@ -182,8 +182,14 @@ Options:
   Use one writer per name; concurrent runs can generate and store different batches.
 - `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.
+- Ingredients must remain complete. Prefer declaring them as required data: override
+  `get_required_data()` to return a `ScenarioRequiredData`, assemble groups from
+  `required_data.load_async()`, and limit them with `dataset_config.select_assembled_groups()`.
+  Required data is prepared and refreshed like any source (fresh runs and explicit loads only,
+  never estimates or resume), is never capped, sampled, filtered, or deselected, and stays in
+  the dataset settings projection and scenario identity. `ApiKey` is the reference. Older
+  scenarios that keep ingredients in `dataset_config` 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.
 - Clients and Python use `ScenarioDatasetRequest` through the common `dataset_overrides`
   parameter. `dataset_config` still replaces the default; overrides apply on top.

doc/scanner/garak.ipynb

@@ -991,8 +991,9 @@
     "\n",
     "**Available techniques:** `GetKey` and `CompleteKey`. `DEFAULT` and `ALL` both select the two\n",
     "techniques. `max_total` samples across all selected technique populations, not per service.\n",
-    "The base scenario persists the sample for resume. Use `ApiKeyDatasetConfiguration` with\n",
-    "`max_total=\"all\"` to run all 348 requests. Standard technique converter stacks are supported."
+    "The base scenario persists the sample for resume. The four corpus datasets are the scenario's\n",
+    "required data and are always read in full; `ApiKeyDatasetConfiguration` only limits the assembled\n",
+    "requests. Use `max_total=\"all\"` to run all 348 requests. Standard technique converter stacks are supported."
    ]
   },
   {
@@ -1051,9 +1052,7 @@
     "    args={\n",
     "        \"objective_target\": objective_target,\n",
     "        \"scenario_techniques\": [ApiKeyTechnique.GetKey],\n",
-    "        \"dataset_config\": ApiKeyDatasetConfiguration(\n",
-    "            sources=[DatasetSource(name=name) for name in ApiKey.required_datasets()], max_total=2\n",
-    "        ),\n",
+    "        \"dataset_config\": ApiKeyDatasetConfiguration(max_total=2),\n",
     "    }\n",
     ")\n",
     "await api_key_scenario.initialize_async()  # type: ignore\n",

doc/scanner/garak.py

@@ -338,18 +338,17 @@
 #
 # **Available techniques:** `GetKey` and `CompleteKey`. `DEFAULT` and `ALL` both select the two
 # techniques. `max_total` samples across all selected technique populations, not per service.
-# The base scenario persists the sample for resume. Use `ApiKeyDatasetConfiguration` with
-# `max_total="all"` to run all 348 requests. Standard technique converter stacks are supported.
+# The base scenario persists the sample for resume. The four corpus datasets are the scenario's
+# required data and are always read in full; `ApiKeyDatasetConfiguration` only limits the assembled
+# requests. Use `max_total="all"` to run all 348 requests. Standard technique converter stacks are supported.
 
 # %%
 api_key_scenario = ApiKey()
 api_key_scenario.set_params_from_args(  # type: ignore
     args={
         "objective_target": objective_target,
         "scenario_techniques": [ApiKeyTechnique.GetKey],
-        "dataset_config": ApiKeyDatasetConfiguration(
-            sources=[DatasetSource(name=name) for name in ApiKey.required_datasets()], max_total=2
-        ),
+        "dataset_config": ApiKeyDatasetConfiguration(max_total=2),
     }
 )
 await api_key_scenario.initialize_async()  # type: ignore

pyrit/registry/components/scenario_registry.py

@@ -232,7 +232,7 @@ def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata:
             )
             for aggregate in technique_class.get_aggregate_techniques()
         )
-        default_datasets = tuple(instance._default_dataset_config.dataset_names)
+        default_datasets = tuple(instance.get_default_dataset_names())
 
         return ScenarioMetadata(
             class_name=cls.__name__,

pyrit/scenario/__init__.py

@@ -31,6 +31,7 @@
         IncompatibleTechniquePolicy,
         ResolvedDataset,
         Scenario,
+        ScenarioRequiredData,
         ScenarioTechnique,
         TechniqueRequirements,
     )
@@ -101,6 +102,7 @@ def find_spec(
     "Parameter": "pyrit.models.parameter",
     "ResolvedDataset": "pyrit.scenario.core.dataset_configuration",
     "Scenario": "pyrit.scenario.core.scenario",
+    "ScenarioRequiredData": "pyrit.scenario.core.required_data",
     "ScenarioTechnique": "pyrit.scenario.core.scenario_technique",
     "TechniqueRequirements": "pyrit.scenario.core.technique_requirements",
     "IncompatibleTechniqueError": "pyrit.scenario.core.technique_requirements",

pyrit/scenario/core/__init__.py

@@ -30,6 +30,7 @@
         ResolvedDataset,
         require_nonempty,
     )
+    from pyrit.scenario.core.required_data import ScenarioRequiredData
     from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario
     from pyrit.scenario.core.scenario_target_defaults import (
         get_default_adversarial_target,
@@ -70,6 +71,7 @@
     "ResolvedDataset": "pyrit.scenario.core.dataset_configuration",
     "require_nonempty": "pyrit.scenario.core.dataset_configuration",
     "Scenario": "pyrit.scenario.core.scenario",
+    "ScenarioRequiredData": "pyrit.scenario.core.required_data",
     "ScenarioTechnique": "pyrit.scenario.core.scenario_technique",
     "TechniqueRequirements": "pyrit.scenario.core.technique_requirements",
     "IncompatibleTechniqueError": "pyrit.scenario.core.technique_requirements",

pyrit/scenario/core/dataset_configuration.py

@@ -28,6 +28,7 @@
 
 import copy
 import random
+from collections import Counter
 from contextvars import ContextVar
 from dataclasses import dataclass
 from enum import Enum
@@ -52,7 +53,7 @@
 from pyrit.models.scenario_dataset import DatasetFetchPolicy
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Sequence
+    from collections.abc import Awaitable, Callable, Sequence
 
     from pyrit.datasets import SeedDatasetProvider
     from pyrit.memory import MemoryInterface
@@ -71,6 +72,70 @@
 )
 
 
+async def prepare_configurations_async(
+    *,
+    configurations: Sequence[DatasetConfiguration],
+    fetch: DatasetFetchPolicy | None = None,
+    validate_assembly: Callable[[], Awaitable[object]] | None = None,
+) -> None:
+    """
+    Load or refresh several configurations' datasets as one batch.
+
+    Every source is checked and fetched before anything is written. When any dataset is refreshed,
+    each configuration that fetched data and then ``validate_assembly`` run while all fetched
+    candidates are visible to reads, so a rejected refresh leaves every stored dataset unchanged.
+
+    Args:
+        configurations (Sequence[DatasetConfiguration]): The configurations to prepare together.
+        fetch (DatasetFetchPolicy | None): Optional policy for this operation only.
+        validate_assembly (Callable[[], Awaitable[object]] | None): A check over everything the
+            configurations build together, such as a scenario's full assembly.
+
+    Raises:
+        DatasetConstraintError: If a dataset is named by more than one configuration, a source cannot
+            be fetched, or the candidates fail validation.
+    """
+    # Candidates are keyed by name, so a shared name would validate one batch and commit both.
+    counts = Counter(name for config in configurations for name in dict.fromkeys(config.dataset_names))
+    shared = sorted(name for name, count in counts.items() if count > 1)
+    if shared:
+        raise DatasetConstraintError(
+            f"Datasets {shared} are named by more than one configuration prepared together, such as "
+            "required data and the population. List each dataset in one place."
+        )
+    fetches = [(config, *await config._fetch_for_preparation_async(fetch=fetch)) for config in configurations]
+    if not any(policies[name] is DatasetFetchPolicy.ALWAYS for _, fetched, policies in fetches for name in fetched):
+        for config, fetched, _ in fetches:
+            for dataset in fetched.values():
+                await config._memory.add_seed_datasets_to_memory_async(
+                    datasets=[dataset], added_by="DatasetConfiguration"
+                )
+        return
+    prepared: list[tuple[DatasetConfiguration, dict[str, PreparedSeedBatch], dict[str, DatasetFetchPolicy]]] = []
+    for config, fetched, policies in fetches:
+        batches = {
+            name: await config._memory.prepare_seed_batch_async(seeds=dataset.seeds, added_by="DatasetConfiguration")
+            for name, dataset in fetched.items()
+        }
+        prepared.append((config, batches, policies))
+    candidates = {name: batch.seeds for _, batches, _ in prepared for name, batch in batches.items()}
+    token = _candidate_populations.set(candidates)
+    try:
+        for config, batches, _ in prepared:
+            if batches:
+                await config.with_overrides()._validate_candidate_population_async()
+        if validate_assembly is not None:
+            await validate_assembly()
+    finally:
+        _candidate_populations.reset(token)
+    for config, batches, policies in prepared:
+        for name, batch in batches.items():
+            if policies[name] is DatasetFetchPolicy.ALWAYS:
+                await config._memory.replace_prepared_seed_batch_async(dataset_name=name, batch=batch)
+            else:
+                await config._memory.add_prepared_seed_batch_async(batch)
+
+
 class _Unset(Enum):
     VALUE = "unset"
 
@@ -417,9 +482,10 @@ def __init__(
         self.max_total = None if isinstance(max_total, _Unset) else max_total
         self.fetch = DatasetFetchPolicy.IF_MISSING if isinstance(fetch, _Unset) else fetch
         self._filters: dict[str, list[str]] = dict(filters or {})
+        self._user_validators: list[Callable[[ResolvedDataset], None]] = list(validators) if validators else []
         self._validators: list[Callable[[ResolvedDataset], None]] = [
             *self._default_validators(),
-            *(list(validators) if validators else []),
+            *self._user_validators,
         ]
         self._validate_selection_options()
 
@@ -494,6 +560,7 @@ def with_overrides(
         result = copy.copy(self)
         result._filters = {key: list(values) for key, values in self._filters.items()}
         result._validators = list(self._validators)
+        result._user_validators = list(self._user_validators)
         if not isinstance(sources, _Unset):
             if self.source_kind is DatasetSourceKind.INLINE:
                 raise DatasetConstraintError("Cannot replace inline seeds with named sources through overrides.")
@@ -533,14 +600,29 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non
         Args:
             fetch: Optional policy for this operation only. Source settings remain unchanged.
 
+        Raises:
+            DatasetConstraintError: If a source cannot be fetched or provider data is invalid.
+        """
+        await prepare_configurations_async(configurations=[self], fetch=fetch)
+
+    async def _fetch_for_preparation_async(
+        self, *, fetch: DatasetFetchPolicy | None
+    ) -> tuple[dict[str, SeedDataset], dict[str, DatasetFetchPolicy]]:
+        """
+        Check all sources and fetch the datasets that preparation would write, without writing.
+
+        Returns:
+            tuple[dict[str, SeedDataset], dict[str, DatasetFetchPolicy]]: The fetched datasets and
+            every source's effective policy, keyed by dataset name.
+
         Raises:
             DatasetConstraintError: If a source cannot be fetched or provider data is invalid.
         """
         self.validate_configuration()
         if fetch is not None and not isinstance(fetch, DatasetFetchPolicy):
             raise DatasetConstraintError("'fetch' must be a DatasetFetchPolicy.")
         if self.source_kind is DatasetSourceKind.INLINE:
-            return
+            return {}, {}
         policies: dict[str, DatasetFetchPolicy] = {}
         configured_providers: dict[str, SeedDatasetProvider | None] = {}
         for source, policy in self._preparation_sources():
@@ -553,7 +635,7 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non
             policies[name] = policy
             configured_providers[name] = source.provider
         if not policies:
-            return
+            return {}, {}
         stored_names = set(await self._memory.get_seed_dataset_names_async())
         to_fetch: list[str] = []
         for name, policy in policies.items():
@@ -563,7 +645,7 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non
                 raise DatasetConstraintError(f"Dataset '{name}' is missing and fetch is 'never'. Import it first.")
             to_fetch.append(name)
         if not to_fetch:
-            return
+            return {}, policies
         from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider
 
         providers = {name: provider for name in to_fetch if (provider := configured_providers[name]) is not None}
@@ -597,34 +679,7 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non
                 )
             SeedDataset.group_seed_prompts_by_prompt_group_id(dataset.seeds)
             fetched[name] = dataset
-        if any(policies[name] is DatasetFetchPolicy.ALWAYS for name in fetched):
-            await self._commit_refresh_candidates_async(fetched=fetched, policies=policies)
-            return
-        for dataset in fetched.values():
-            await self._memory.add_seed_datasets_to_memory_async(datasets=[dataset], added_by="DatasetConfiguration")
-
-    async def _commit_refresh_candidates_async(
-        self, *, fetched: dict[str, SeedDataset], policies: dict[str, DatasetFetchPolicy]
-    ) -> None:
-        prepared = {
-            name: await self._memory.prepare_seed_batch_async(seeds=dataset.seeds, added_by="DatasetConfiguration")
-            for name, dataset in fetched.items()
-        }
-        await self._validate_refresh_candidates_async(prepared=prepared)
-        for name, batch in prepared.items():
-            if policies[name] is DatasetFetchPolicy.ALWAYS:
-                await self._memory.replace_prepared_seed_batch_async(dataset_name=name, batch=batch)
-            else:
-                await self._memory.add_prepared_seed_batch_async(batch)
-
-    async def _validate_refresh_candidates_async(self, *, prepared: dict[str, PreparedSeedBatch]) -> None:
-        """Resolve task-local copies of storage-ready snapshots without modifying them."""
-        candidates = {name: batch.seeds for name, batch in prepared.items()}
-        token = _candidate_populations.set(candidates)
-        try:
-            await self.with_overrides()._validate_candidate_population_async()
-        finally:
-            _candidate_populations.reset(token)
+        return fetched, policies
 
     async def _validate_candidate_population_async(self) -> None:
         """Validate raw configuration populations before replacement."""
@@ -1021,11 +1076,48 @@ async def get_attack_groups_by_dataset_async(
                 dataset fails validation, or no attack groups could be built.
         """
         selection = await self._resolve_attack_groups_async()
-        sampled = selection.select(apply_sampling=apply_sampling)
-        if not any(sampled.values()):
+        return self._require_selected_groups(selection.select(apply_sampling=apply_sampling))
+
+    def select_assembled_groups(
+        self, *, groups_by_population: dict[str, list[AttackSeedGroup]], apply_sampling: bool = True
+    ) -> dict[str, list[AttackSeedGroup]]:
+        """
+        Validate and limit attack groups a scenario assembled from its required data.
+
+        The configuration has no sources of its own here: it validates the assembled
+        population and applies ``max_total`` across it, keeping each population's key.
+
+        Args:
+            groups_by_population (dict[str, list[AttackSeedGroup]]): Assembled groups keyed by population.
+            apply_sampling (bool): When True (default), apply ``max_total``. Resume passes False.
+
+        Returns:
+            dict[str, list[AttackSeedGroup]]: The validated, limited groups.
+
+        Raises:
+            DatasetConstraintError: If the configuration has its own sources or inline seeds, the
+                population fails validation, or no attack groups remain.
+        """
+        self.validate_configuration()
+        if self.sources or self.source_kind is DatasetSourceKind.INLINE:
+            raise DatasetConstraintError(
+                "Assembled populations can't be combined with dataset sources or inline seeds."
+            )
+        self.validate(
+            ResolvedDataset(
+                seeds=[seed for groups in groups_by_population.values() for group in groups for seed in group.seeds],
+                source_kind=self.source_kind,
+                dataset_names=tuple(groups_by_population),
+            )
+        )
+        selection = _ResolvedAttackGroups(configuration=self, groups=groups_by_population)
+        return self._require_selected_groups(selection.select(apply_sampling=apply_sampling))
+
+    def _require_selected_groups(self, selected: dict[str, list[AttackSeedGroup]]) -> dict[str, list[AttackSeedGroup]]:
+        if not any(selected.values()):
             names = ", ".join(self.dataset_names) if self.dataset_names else "<inline>"
             raise DatasetConstraintError(f"Resolved attack-group dataset is empty (datasets: {names}).")
-        return sampled
+        return selected
 
     async def _resolve_attack_groups_async(self) -> _ResolvedAttackGroups:
         self.validate_configuration()

pyrit/scenario/core/dataset_overrides.py

@@ -20,6 +20,7 @@
     DatasetSource,
     DatasetSourceKind,
 )
+from pyrit.scenario.core.required_data import ScenarioRequiredData
 
 
 def apply_dataset_overrides(
@@ -96,16 +97,23 @@ def validate_dataset_selection(
         raise DatasetConstraintError(f"Select at most {options.max_selected} dataset(s).")
 
 
-def project_dataset_settings(config: DatasetAttackConfiguration) -> ScenarioDatasetSettings:
+def project_dataset_settings(
+    config: DatasetAttackConfiguration, *, required_data: ScenarioRequiredData | None = None
+) -> ScenarioDatasetSettings:
     """
     Project configuration settings without memory reads or live wire objects.
 
+    Required datasets come first as uncapped sources, so clients keep seeing and
+    refreshing every dataset a run loads.
+
     Returns:
         ScenarioDatasetSettings: Effective settings for response data.
     """
     children = config.configurations if isinstance(config, CompoundDatasetAttackConfiguration) else (config,)
+    required_sources = required_data.project_sources() if required_data is not None else []
     return ScenarioDatasetSettings(
-        sources=[
+        sources=required_sources
+        + [
             ScenarioDatasetSourceSettings(
                 name=source.name,
                 max_size=child.source_limit(source.name),