diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 5f45135bfd..a20c8eabd9 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -25,7 +25,7 @@ memory_db_type: sqlite # - target: Registers available prompt targets into the TargetRegistry # - scorer: Registers pre-configured scorers into the ScorerRegistry # - technique: Registers attack techniques into the AttackTechniqueRegistry -# - load_default_datasets: Optionally preloads all registered datasets into memory +# - load_default_datasets: Optionally preloads bounded scenario datasets into memory # - preload_scenario_metadata: Preloads scenario metadata into the registry # # Each initializer can be specified as: @@ -51,13 +51,18 @@ initializers: - scorer - name: scorer - name: technique - # Optional full preload/cache warming for offline or shared environments. + # Optional scenario dataset preload/cache warming for offline or shared environments. + # Large opt-in datasets are skipped unless selected explicitly with + # dataset_names or tags; either parameter replaces the scenario-default selection. # This can take several minutes and may require network access, # provider credentials, or accepted dataset licenses. Scenarios fetch only their # requested datasets on demand without this initializer. # If intentional preload exceeds backend startup time, increase # server.startup_timeout below. # - name: load_default_datasets + # args: + # dataset_names: + # - garak_npm_packages # Operator and Operation Labels # ------------------------------ diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 64235420ea..9efc00e393 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -127,11 +127,11 @@ initializers: - name: technique ``` -#### Optional Full Dataset Preload +#### Optional Dataset Preload `load_default_datasets` is not required for `pyrit_scan`. Scenario `DatasetAttackConfiguration` objects fetch only their requested datasets from registered providers on demand, then add them to memory. -Use `load_default_datasets` only when you intentionally want to preload every registered dataset—for example, to warm a shared cache or prepare an offline environment: +Use `load_default_datasets` only when you intentionally want to preload datasets—for example, to warm a shared cache or prepare an offline environment. By default, it loads the bounded datasets required by registered scenarios and skips large opt-in datasets: ```yaml initializers: @@ -141,7 +141,17 @@ initializers: - name: load_default_datasets ``` -Full preload can take several minutes and may require network access, provider credentials, or acceptance of gated dataset licenses. When preloading during local backend startup, increase `server.startup_timeout` if the configured timeout is not long enough. +Select opt-in datasets explicitly with `dataset_names` or `tags`. Either parameter replaces the scenario-default selection rather than adding to it: + +```yaml +initializers: + - name: load_default_datasets + args: + dataset_names: + - garak_npm_packages +``` + +Preloading can take several minutes and may require network access, provider credentials, or acceptance of gated dataset licenses. When preloading during local backend startup, increase `server.startup_timeout` if the configured timeout is not long enough. ### `initialization_scripts` @@ -366,10 +376,14 @@ initializers: - scorer - name: scorer - name: technique - # Optional full preload/cache warming; scenarios fetch requested datasets on demand. - # Full preload can take several minutes and may require network access, + # Optional scenario dataset preload/cache warming; scenarios fetch requested + # datasets on demand. Large opt-in datasets must be selected explicitly. + # Preloading can take several minutes and may require network access, # provider credentials, accepted dataset licenses, and a larger startup_timeout. # - name: load_default_datasets + # args: + # dataset_names: + # - garak_npm_packages # Custom initialization scripts (optional) # Omit or set to null for no scripts; [] to explicitly load nothing diff --git a/doc/scanner/pyrit_conf.yaml b/doc/scanner/pyrit_conf.yaml index 992122f8d0..490d5bb644 100644 --- a/doc/scanner/pyrit_conf.yaml +++ b/doc/scanner/pyrit_conf.yaml @@ -8,5 +8,5 @@ initializers: - scorer - name: scorer - name: technique - # Optional full preload/cache warming; scenarios fetch requested datasets on demand and add them to memory. + # Optional scenario dataset preloading for offline user or to avoid first-run loading delays # - name: load_default_datasets diff --git a/pyrit/setup/initializers/load_default_datasets.py b/pyrit/setup/initializers/load_default_datasets.py index bdaf053490..b141c034bf 100644 --- a/pyrit/setup/initializers/load_default_datasets.py +++ b/pyrit/setup/initializers/load_default_datasets.py @@ -20,14 +20,23 @@ logger = logging.getLogger(__name__) +OPT_IN_DATASET_NAMES: frozenset[str] = frozenset( + { + "garak_crates_packages", + "garak_npm_packages", + "garak_pypi_packages", + "garak_rubygems_packages", + } +) + class LoadDefaultDatasets(PyRITInitializer): """ Load datasets into memory so scenarios can run. - By default this loads the datasets required by all registered scenarios. - Pass ``dataset_names`` to load specific datasets by name, or ``tags`` to - select datasets by metadata. + By default this loads the bounded datasets required by all registered + scenarios. Large opt-in datasets are only loaded when selected through + ``dataset_names`` or ``tags``. """ @property @@ -36,8 +45,8 @@ def description(self) -> str: return textwrap.dedent( """ Loads datasets into memory so scenarios can run. By default loads the datasets - required by all registered scenarios; use the dataset_names or tags parameters to - select datasets explicitly. + required by all registered scenarios except large opt-in datasets; use the + dataset_names or tags parameters to select datasets explicitly. Note: if you are using persistent memory, avoid calling this every time as datasets can take time to load. @@ -81,7 +90,7 @@ async def initialize_async(self) -> None: logger.info(f"Loading {len(unique_datasets)} dataset(s) matching tags: {sorted(tags)}") else: unique_datasets = self._scenario_default_dataset_names() - logger.info(f"Loading {len(unique_datasets)} unique datasets required by all scenarios") + logger.info(f"Loading {len(unique_datasets)} bounded datasets required by all scenarios") if not unique_datasets: logger.warning("No datasets matched the requested selection") @@ -112,4 +121,11 @@ def _scenario_default_dataset_names() -> list[str]: all_default_datasets.extend(datasets) logger.info(f"Scenario '{metadata.registry_name}' uses datasets: {datasets}") - return list(dict.fromkeys(all_default_datasets)) + selected = list(dict.fromkeys(name for name in all_default_datasets if name not in OPT_IN_DATASET_NAMES)) + skipped = sorted(set(all_default_datasets) & OPT_IN_DATASET_NAMES) + if skipped: + logger.info( + f"Skipping opt-in datasets from automatic loading: {skipped}. " + "Select them explicitly with dataset_names or tags." + ) + return selected diff --git a/tests/integration/datasets/test_load_default_datasets_integration.py b/tests/integration/datasets/test_load_default_datasets_integration.py index 37877e8148..a3362e0032 100644 --- a/tests/integration/datasets/test_load_default_datasets_integration.py +++ b/tests/integration/datasets/test_load_default_datasets_integration.py @@ -4,18 +4,23 @@ """ Integration test for the LoadDefaultDatasets initializer. -Runs the full pipeline: discovers scenario default datasets, fetches them -from real remote sources, and stores them in in-memory CentralMemory. +Runs the full pipeline with a bounded representative selection and stores the +datasets in in-memory CentralMemory. """ import logging from pyrit.memory import CentralMemory from pyrit.setup.initializers.load_default_datasets import LoadDefaultDatasets -from pyrit.setup.initializers.techniques import TechniqueInitializer logger = logging.getLogger(__name__) +BOUNDED_DATASET_NAMES = [ + "garak_package_hallucination_real_tasks", + "garak_package_hallucination_stubs", + "garak_package_hallucination_unreal_tasks", +] + class TestLoadDefaultDatasetsIntegration: """Integration test that LoadDefaultDatasets loads real datasets into memory.""" @@ -26,11 +31,11 @@ async def test_initialize_loads_datasets_into_memory(self, sqlite_instance): real datasets and stores them in CentralMemory. """ initializer = LoadDefaultDatasets() - await TechniqueInitializer().initialize_async() + initializer.params = {"dataset_names": BOUNDED_DATASET_NAMES} await initializer.initialize_async() memory = CentralMemory.get_memory_instance() - dataset_names = memory.get_seed_dataset_names() + dataset_names = set(memory.get_seed_dataset_names()) - assert len(dataset_names) > 0, "No datasets were loaded into memory" + assert dataset_names == set(BOUNDED_DATASET_NAMES) logger.info(f"LoadDefaultDatasets loaded {len(dataset_names)} datasets into memory") diff --git a/tests/unit/setup/test_load_default_datasets.py b/tests/unit/setup/test_load_default_datasets.py index 120a301df8..95726e0c8d 100644 --- a/tests/unit/setup/test_load_default_datasets.py +++ b/tests/unit/setup/test_load_default_datasets.py @@ -16,7 +16,7 @@ from pyrit.prompt_target import PromptTarget from pyrit.registry import ScenarioRegistry, TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry -from pyrit.setup.initializers.load_default_datasets import LoadDefaultDatasets +from pyrit.setup.initializers.load_default_datasets import OPT_IN_DATASET_NAMES, LoadDefaultDatasets from pyrit.setup.initializers.techniques import build_technique_factories @@ -128,6 +128,36 @@ async def test_initialize_async_deduplicates_datasets(self) -> None: assert set(call_kwargs["dataset_names"]) == {"dataset1", "dataset2", "dataset3"} assert len(call_kwargs["dataset_names"]) == 3 + async def test_initialize_async_excludes_opt_in_datasets_from_scenario_defaults(self) -> None: + """Test that implicit loading excludes datasets that must be selected explicitly.""" + initializer = LoadDefaultDatasets() + metadata = [ + _FakeMetadata( + registry_name="package_hallucination", + default_datasets=("bounded_dataset", *sorted(OPT_IN_DATASET_NAMES)), + ) + ] + + with ( + patch.object(ScenarioRegistry, "get_all_registered_class_metadata", return_value=metadata), + patch.object(SeedDatasetProvider, "fetch_datasets_async", new_callable=AsyncMock) as mock_fetch, + patch.object(CentralMemory, "get_memory_instance") as mock_memory, + ): + mock_fetch.return_value = [] + mock_memory_instance = MagicMock() + mock_memory_instance.add_seed_datasets_to_memory_async = AsyncMock() + mock_memory.return_value = mock_memory_instance + + await initializer.initialize_async() + + assert mock_fetch.call_args.kwargs["dataset_names"] == ["bounded_dataset"] + + async def test_opt_in_datasets_are_registered_without_fetching(self) -> None: + """Test that opt-in package registries remain discoverable for explicit loading.""" + available_datasets = set(await SeedDatasetProvider.get_all_dataset_names_async()) + + assert available_datasets >= OPT_IN_DATASET_NAMES + async def test_all_required_datasets_available_in_seed_provider(self, populated_technique_registry) -> None: """ Test that all datasets required by scenarios are available in SeedDatasetProvider. @@ -201,7 +231,8 @@ def test_supported_parameters_defaults(self) -> None: async def test_dataset_names_loads_exact_names(self) -> None: initializer = LoadDefaultDatasets() - initializer.params = {"dataset_names": ["alpha", "beta"]} + requested_names = ["alpha", "garak_npm_packages"] + initializer.params = {"dataset_names": requested_names} with ( patch.object(ScenarioRegistry, "get_all_registered_class_metadata") as mock_list_metadata, @@ -218,7 +249,7 @@ async def test_dataset_names_loads_exact_names(self) -> None: mock_list_metadata.assert_not_called() mock_names.assert_not_called() - assert mock_fetch.call_args.kwargs["dataset_names"] == ["alpha", "beta"] + assert mock_fetch.call_args.kwargs["dataset_names"] == requested_names async def test_tags_selects_via_metadata_filter(self) -> None: initializer = LoadDefaultDatasets()