diff --git a/src/anonymizer/__init__.py b/src/anonymizer/__init__.py index 64932ab3..4a63168a 100644 --- a/src/anonymizer/__init__.py +++ b/src/anonymizer/__init__.py @@ -9,6 +9,7 @@ __version__ = version("nemo-anonymizer") from data_designer.config.models import ModelProvider as ModelProvider +from data_designer.config.run_config import RunConfig as RunConfig from anonymizer.config.anonymizer_config import ( AnonymizerConfig, @@ -64,6 +65,7 @@ def __getattr__(name: str) -> object: "Redact", "Rewrite", "RiskTolerance", + "RunConfig", "Substitute", "configure_logging", ] diff --git a/src/anonymizer/interface/anonymizer.py b/src/anonymizer/interface/anonymizer.py index ba9af8d1..bcfee0c5 100644 --- a/src/anonymizer/interface/anonymizer.py +++ b/src/anonymizer/interface/anonymizer.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING from data_designer.config.models import ModelProvider +from data_designer.config.run_config import RunConfig from data_designer.config.utils.io_helpers import load_config_file from data_designer.interface.data_designer import DataDesigner @@ -109,6 +110,7 @@ def __init__( model_providers: list[ModelProvider] | str | Path | None = None, artifact_path: str | Path | None = None, data_designer: DataDesigner | None = None, + data_designer_run_config: RunConfig | None = None, detection_workflow: EntityDetectionWorkflow | None = None, replace_runner: ReplacementWorkflow | None = None, rewrite_runner: RewriteWorkflow | None = None, @@ -125,6 +127,10 @@ def __init__( artifact_path: Directory for intermediate artifacts. Defaults to ``.anonymizer-artifacts``. data_designer: Pre-configured DataDesigner instance (advanced usage). + data_designer_run_config: Optional DataDesigner run configuration + applied to the Anonymizer-managed or caller-supplied DataDesigner + instance. Use this for DataDesigner execution knobs such as + ``buffer_size`` and ``max_in_flight_tasks``. detection_workflow: Custom detection workflow (advanced/testing). replace_runner: Custom replacement workflow (advanced/testing). rewrite_runner: Custom rewrite workflow (advanced/testing). @@ -164,6 +170,8 @@ def __init__( model_providers=self._resolved_providers, ) reapply_log_levels() + if data_designer_run_config is not None: + self._data_designer.set_run_config(data_designer_run_config) self._adapter = NddAdapter(data_designer=self._data_designer) self._detection_workflow = detection_workflow or EntityDetectionWorkflow(adapter=self._adapter) self._replace_runner = replace_runner or ReplacementWorkflow( diff --git a/tests/interface/test_anonymizer_interface.py b/tests/interface/test_anonymizer_interface.py index 289c4c74..f307a609 100644 --- a/tests/interface/test_anonymizer_interface.py +++ b/tests/interface/test_anonymizer_interface.py @@ -11,6 +11,7 @@ import pytest from data_designer.config.models import ModelConfig +from anonymizer import RunConfig from anonymizer.config.anonymizer_config import AnonymizerConfig, AnonymizerInput, Rewrite from anonymizer.config.models import ModelSelection, ReplaceModelSelection from anonymizer.config.replace_strategies import Redact, Substitute @@ -169,6 +170,37 @@ def test_anonymizer_custom_model_providers_override_bundled_defaults() -> None: assert passed_providers is custom_providers +def test_anonymizer_applies_data_designer_run_config_to_managed_instance() -> None: + run_config = RunConfig(buffer_size=20, max_in_flight_tasks=64) + + with patch("anonymizer.interface.anonymizer.DataDesigner") as mock_data_designer: + Anonymizer( + data_designer_run_config=run_config, + detection_workflow=Mock(), + replace_runner=Mock(), + rewrite_runner=Mock(), + ) + + mock_data_designer.return_value.set_run_config.assert_called_once_with(run_config) + + +def test_anonymizer_applies_data_designer_run_config_to_supplied_instance() -> None: + from data_designer.interface.data_designer import DataDesigner + + data_designer = Mock(spec=DataDesigner) + run_config = RunConfig(buffer_size=20, max_in_flight_tasks=64) + + Anonymizer( + data_designer=data_designer, + data_designer_run_config=run_config, + detection_workflow=Mock(), + replace_runner=Mock(), + rewrite_runner=Mock(), + ) + + data_designer.set_run_config.assert_called_once_with(run_config) + + def test_anonymizer_rejects_missing_provider_as_invalid_config_error() -> None: yaml_str = """ model_configs: