feat: expose Data Designer run config (#203)

Author: lipikaramaswamy

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",
 ]

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(

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: