feat: Elevate non-LLM concurrency limits to RunConfig (#242)

Author: eric-tramel

docs/code_reference/run_config.md

@@ -0,0 +1,6 @@
+# Run Config
+
+The `run_config` module defines runtime settings that control dataset generation behavior,
+including early shutdown thresholds, batch sizing, and non-inference worker concurrency.
+
+:::: data_designer.config.run_config

mkdocs.yml

@@ -43,6 +43,7 @@ nav:
       - column_configs: code_reference/column_configs.md
       - config_builder: code_reference/config_builder.md
       - data_designer_config: code_reference/data_designer_config.md
+      - run_config: code_reference/run_config.md
       - sampler_params: code_reference/sampler_params.md
       - validator_params: code_reference/validator_params.md
       - processors: code_reference/processors.md

src/data_designer/config/run_config.py

@@ -26,6 +26,8 @@ class RunConfig(ConfigBase):
         buffer_size: Number of records to process in each batch during dataset generation.
             A batch is processed end-to-end (column generation, post-batch processors, and writing the batch
             to artifact storage) before moving on to the next batch. Must be > 0. Default is 1000.
+        non_inference_max_parallel_workers: Maximum number of worker threads used for non-inference
+            cell-by-cell generators. Must be >= 1. Default is 4.
         max_conversation_restarts: Maximum number of full conversation restarts permitted when
             generation tasks call `ModelFacade.generate(...)`. Must be >= 0. Default is 5.
         max_conversation_correction_steps: Maximum number of correction rounds permitted within a
@@ -37,6 +39,7 @@ class RunConfig(ConfigBase):
     shutdown_error_rate: float = Field(default=0.5, ge=0.0, le=1.0)
     shutdown_error_window: int = Field(default=10, ge=0)
     buffer_size: int = Field(default=1000, gt=0)
+    non_inference_max_parallel_workers: int = Field(default=4, ge=1)
     max_conversation_restarts: int = Field(default=5, ge=0)
     max_conversation_correction_steps: int = Field(default=0, ge=0)
 

src/data_designer/engine/dataset_builders/column_wise_builder.py

@@ -31,10 +31,7 @@
 from data_designer.engine.dataset_builders.artifact_storage import SDG_CONFIG_FILENAME, ArtifactStorage
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError, DatasetProcessingError
 from data_designer.engine.dataset_builders.multi_column_configs import MultiColumnConfig
-from data_designer.engine.dataset_builders.utils.concurrency import (
-    MAX_CONCURRENCY_PER_NON_LLM_GENERATOR,
-    ConcurrentThreadExecutor,
-)
+from data_designer.engine.dataset_builders.utils.concurrency import ConcurrentThreadExecutor
 from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
 from data_designer.engine.dataset_builders.utils.dataset_batch_manager import DatasetBatchManager
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
@@ -202,7 +199,7 @@ def _run_from_scratch_column_generator(self, generator: ColumnGenerator) -> None
         self.batch_manager.add_records(df.to_dict(orient="records"))
 
     def _run_cell_by_cell_generator(self, generator: ColumnGenerator) -> None:
-        max_workers = MAX_CONCURRENCY_PER_NON_LLM_GENERATOR
+        max_workers = self._resource_provider.run_config.non_inference_max_parallel_workers
         if isinstance(generator, ColumnGeneratorWithModel):
             max_workers = generator.inference_parameters.max_parallel_requests
         self._fan_out_with_threads(generator, max_workers=max_workers)

src/data_designer/engine/dataset_builders/utils/concurrency.py

@@ -16,9 +16,6 @@
 
 logger = logging.getLogger(__name__)
 
-# Constants
-MAX_CONCURRENCY_PER_NON_LLM_GENERATOR = 4
-
 
 class ExecutorResults(BaseModel):
     failure_threshold: float = 0.0  # Error rate threshold

src/data_designer/interface/data_designer.py

@@ -316,7 +316,8 @@ def set_run_config(self, run_config: RunConfig) -> None:
 
         Args:
             run_config: A RunConfig instance containing runtime settings such as
-                early shutdown behavior and batch sizing via `buffer_size`. Import RunConfig from
+                early shutdown behavior, batch sizing via `buffer_size`, and non-inference worker
+                concurrency via `non_inference_max_parallel_workers`. Import RunConfig from
                 data_designer.essentials.
 
         Example:

tests/engine/dataset_builders/test_column_wise_builder.py

@@ -15,10 +15,7 @@
 from data_designer.config.run_config import RunConfig
 from data_designer.config.sampler_params import SamplerType, UUIDSamplerParams
 from data_designer.engine.column_generators.generators.base import GenerationStrategy
-from data_designer.engine.dataset_builders.column_wise_builder import (
-    MAX_CONCURRENCY_PER_NON_LLM_GENERATOR,
-    ColumnWiseDatasetBuilder,
-)
+from data_designer.engine.dataset_builders.column_wise_builder import ColumnWiseDatasetBuilder
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum
 from data_designer.engine.models.usage import ModelUsageStats, TokenUsageStats
@@ -243,8 +240,9 @@ def test_column_wise_dataset_builder_initialize_processors(stub_column_wise_buil
     assert processors[BuildStage.POST_BATCH][0].config.column_names == ["column_to_drop"]
 
 
-def test_constants_max_concurrency_constant():
-    assert MAX_CONCURRENCY_PER_NON_LLM_GENERATOR == 4
+def test_run_config_default_non_inference_max_parallel_workers() -> None:
+    run_config = RunConfig()
+    assert run_config.non_inference_max_parallel_workers == 4
 
 
 @patch("data_designer.engine.dataset_builders.column_wise_builder.TelemetryHandler")