data-designer 0.3.8rc1__py3-none-any.whl → 0.4.0__py3-none-any.whl

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -1,338 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import functools
-import importlib.metadata
-import json
-import logging
-import time
-import uuid
-from pathlib import Path
-from typing import TYPE_CHECKING, Callable
-
-from data_designer.config.column_types import ColumnConfigT
-from data_designer.config.config_builder import BuilderConfig
-from data_designer.config.data_designer_config import DataDesignerConfig
-from data_designer.config.dataset_builders import BuildStage
-from data_designer.config.processors import (
-    DropColumnsProcessorConfig,
-    ProcessorConfig,
-    ProcessorType,
-)
-from data_designer.engine.column_generators.generators.base import (
-    ColumnGenerator,
-    ColumnGeneratorWithModel,
-    GenerationStrategy,
-)
-from data_designer.engine.column_generators.utils.generator_classification import column_type_is_model_generated
-from data_designer.engine.compiler import compile_data_designer_config
-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.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
-from data_designer.engine.processing.processors.base import Processor
-from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
-from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry
-from data_designer.engine.resources.resource_provider import ResourceProvider
-from data_designer.lazy_heavy_imports import pd
-
-if TYPE_CHECKING:
-    import pandas as pd
-
-    from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModelRegistry
-    from data_designer.engine.models.usage import ModelUsageStats
-
-logger = logging.getLogger(__name__)
-
-_CLIENT_VERSION: str = importlib.metadata.version("data_designer")
-
-
-class ColumnWiseDatasetBuilder:
-    def __init__(
-        self,
-        data_designer_config: DataDesignerConfig,
-        resource_provider: ResourceProvider,
-        registry: DataDesignerRegistry | None = None,
-    ):
-        self.batch_manager = DatasetBatchManager(resource_provider.artifact_storage)
-        self._resource_provider = resource_provider
-        self._records_to_drop: set[int] = set()
-        self._registry = registry or DataDesignerRegistry()
-
-        self._data_designer_config = compile_data_designer_config(data_designer_config, resource_provider)
-        self._column_configs = compile_dataset_builder_column_configs(self._data_designer_config)
-        self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(
-            self._data_designer_config.processors or []
-        )
-        self._validate_column_configs()
-
-    @property
-    def artifact_storage(self) -> ArtifactStorage:
-        return self._resource_provider.artifact_storage
-
-    @functools.cached_property
-    def single_column_configs(self) -> list[ColumnConfigT]:
-        configs = []
-        for config in self._column_configs:
-            if isinstance(config, MultiColumnConfig):
-                configs.extend(config.columns)
-            else:
-                configs.append(config)
-        return configs
-
-    @functools.cached_property
-    def llm_generated_column_configs(self) -> list[ColumnConfigT]:
-        return [config for config in self.single_column_configs if column_type_is_model_generated(config.column_type)]
-
-    def build(
-        self,
-        *,
-        num_records: int,
-        on_batch_complete: Callable[[Path], None] | None = None,
-    ) -> Path:
-        self._run_model_health_check_if_needed()
-        self._write_builder_config()
-        generators = self._initialize_generators()
-        start_time = time.perf_counter()
-        group_id = uuid.uuid4().hex
-
-        buffer_size = self._resource_provider.run_config.buffer_size
-        self.batch_manager.start(num_records=num_records, buffer_size=buffer_size)
-        for batch_idx in range(self.batch_manager.num_batches):
-            logger.info(f"⏳ Processing batch {batch_idx + 1} of {self.batch_manager.num_batches}")
-            self._run_batch(generators, batch_mode="batch", group_id=group_id)
-            df_batch = self._run_processors(
-                stage=BuildStage.POST_BATCH,
-                dataframe=self.batch_manager.get_current_batch(as_dataframe=True),
-                current_batch_number=batch_idx,
-            )
-            self._write_processed_batch(df_batch)
-            self.batch_manager.finish_batch(on_batch_complete)
-        self.batch_manager.finish()
-
-        model_usage_stats = self._resource_provider.model_registry.get_model_usage_stats(
-            time.perf_counter() - start_time
-        )
-        logger.info(f"📊 Model usage summary:\n{json.dumps(model_usage_stats, indent=4)}")
-
-        return self.artifact_storage.final_dataset_path
-
-    def build_preview(self, *, num_records: int) -> pd.DataFrame:
-        self._run_model_health_check_if_needed()
-
-        generators = self._initialize_generators()
-        group_id = uuid.uuid4().hex
-        start_time = time.perf_counter()
-        self.batch_manager.start(num_records=num_records, buffer_size=num_records)
-        self._run_batch(generators, batch_mode="preview", save_partial_results=False, group_id=group_id)
-        dataset = self.batch_manager.get_current_batch(as_dataframe=True)
-        self.batch_manager.reset()
-
-        model_usage_stats = self._resource_provider.model_registry.get_model_usage_stats(
-            time.perf_counter() - start_time
-        )
-        logger.info(f"📊 Model usage summary:\n{json.dumps(model_usage_stats, indent=4)}")
-
-        return dataset
-
-    def process_preview(self, dataset: pd.DataFrame) -> pd.DataFrame:
-        return self._run_processors(
-            stage=BuildStage.POST_BATCH,
-            dataframe=dataset.copy(),
-            current_batch_number=None,  # preview mode does not have a batch number
-        )
-
-    def _initialize_generators(self) -> list[ColumnGenerator]:
-        return [
-            self._registry.column_generators.get_for_config_type(type(config))(
-                config=config, resource_provider=self._resource_provider
-            )
-            for config in self._column_configs
-        ]
-
-    def _write_builder_config(self) -> None:
-        self.artifact_storage.mkdir_if_needed(self.artifact_storage.base_dataset_path)
-        BuilderConfig(data_designer=self._data_designer_config).to_json(
-            self.artifact_storage.base_dataset_path / SDG_CONFIG_FILENAME
-        )
-
-    def _run_batch(
-        self, generators: list[ColumnGenerator], *, batch_mode: str, save_partial_results: bool = True, group_id: str
-    ) -> None:
-        pre_batch_snapshot = self._resource_provider.model_registry.get_model_usage_snapshot()
-        for generator in generators:
-            generator.log_pre_generation()
-            try:
-                generation_strategy = generator.get_generation_strategy()
-                if generator.can_generate_from_scratch and self.batch_manager.buffer_is_empty:
-                    self._run_from_scratch_column_generator(generator)
-                elif generation_strategy == GenerationStrategy.CELL_BY_CELL:
-                    self._run_cell_by_cell_generator(generator)
-                elif generation_strategy == GenerationStrategy.FULL_COLUMN:
-                    self._run_full_column_generator(generator)
-                else:
-                    logger.error(f"❌ Unknown generation strategy: {generation_strategy}")
-                    raise DatasetGenerationError(f"🛑 Unknown generation strategy: {generation_strategy}")
-                if save_partial_results:
-                    self.batch_manager.write()
-            except Exception as e:
-                column_error_str = (
-                    f"columns {generator.config.column_names}"
-                    if hasattr(generator.config, "column_names")
-                    else f"column {generator.config.name!r}"
-                )
-                raise DatasetGenerationError(f"🛑 Failed to process {column_error_str}:\n{e}")
-
-        try:
-            usage_deltas = self._resource_provider.model_registry.get_usage_deltas(pre_batch_snapshot)
-            self._emit_batch_inference_events(batch_mode, usage_deltas, group_id)
-        except Exception:
-            pass
-
-    def _run_from_scratch_column_generator(self, generator: ColumnGenerator) -> None:
-        df = generator.generate_from_scratch(self.batch_manager.num_records_batch)
-        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
-        if isinstance(generator, ColumnGeneratorWithModel):
-            max_workers = generator.inference_parameters.max_parallel_requests
-        self._fan_out_with_threads(generator, max_workers=max_workers)
-
-    def _run_full_column_generator(self, generator: ColumnGenerator) -> None:
-        df = generator.generate(self.batch_manager.get_current_batch(as_dataframe=True))
-        self.batch_manager.update_records(df.to_dict(orient="records"))
-
-    def _run_model_health_check_if_needed(self) -> bool:
-        if any(column_type_is_model_generated(config.column_type) for config in self.single_column_configs):
-            self._resource_provider.model_registry.run_health_check(
-                list(set(config.model_alias for config in self.llm_generated_column_configs))
-            )
-
-    def _fan_out_with_threads(self, generator: ColumnGeneratorWithModelRegistry, max_workers: int) -> None:
-        if generator.get_generation_strategy() != GenerationStrategy.CELL_BY_CELL:
-            raise DatasetGenerationError(
-                f"Generator {generator.name} is not a {GenerationStrategy.CELL_BY_CELL} "
-                "generator so concurrency through threads is not supported."
-            )
-
-        logger.info(
-            f"🐙 Processing {generator.config.column_type} column '{generator.config.name}' "
-            f"with {max_workers} concurrent workers"
-        )
-        settings = self._resource_provider.run_config
-        with ConcurrentThreadExecutor(
-            max_workers=max_workers,
-            column_name=generator.config.name,
-            result_callback=self._worker_result_callback,
-            error_callback=self._worker_error_callback,
-            shutdown_error_rate=settings.shutdown_error_rate,
-            shutdown_error_window=settings.shutdown_error_window,
-            disable_early_shutdown=settings.disable_early_shutdown,
-        ) as executor:
-            for i, record in self.batch_manager.iter_current_batch():
-                executor.submit(lambda record: generator.generate(record), record, context={"index": i})
-
-        if len(self._records_to_drop) > 0:
-            self.batch_manager.drop_records(self._records_to_drop)
-            self._records_to_drop.clear()
-
-    def _write_processed_batch(self, dataframe: pd.DataFrame) -> None:
-        self.batch_manager.update_records(dataframe.to_dict(orient="records"))
-        self.batch_manager.write()
-
-    def _validate_column_configs(self) -> None:
-        if len(self._column_configs) == 0:
-            raise DatasetGenerationError("🛑 No column configs provided.")
-
-        if not self._registry.column_generators.get_for_config_type(
-            type(self._column_configs[0])
-        ).can_generate_from_scratch:
-            raise DatasetGenerationError("🛑 The first column config must be a from-scratch column generator.")
-
-    def _initialize_processors(self, processor_configs: list[ProcessorConfig]) -> dict[BuildStage, list[Processor]]:
-        # Check columns marked for drop
-        columns_to_drop = [config.name for config in self.single_column_configs if config.drop]
-
-        processors: dict[BuildStage, list[Processor]] = {stage: [] for stage in BuildStage}
-        for config in processor_configs:
-            processors[config.build_stage].append(
-                self._registry.processors.get_for_config_type(type(config))(
-                    config=config,
-                    resource_provider=self._resource_provider,
-                )
-            )
-
-            # Manually included "drop columns" processor takes precedence (can e.g., pick stages other than post-batch)
-            if config.processor_type == ProcessorType.DROP_COLUMNS:
-                for column in config.column_names:
-                    if column in columns_to_drop:
-                        columns_to_drop.remove(column)
-
-        # If there are still columns marked for drop, add the "drop columns" processor to drop them
-        if len(columns_to_drop) > 0:
-            processors[BuildStage.POST_BATCH].append(  # as post-batch by default
-                DropColumnsProcessor(
-                    config=DropColumnsProcessorConfig(
-                        name="default_drop_columns_processor",
-                        column_names=columns_to_drop,
-                        build_stage=BuildStage.POST_BATCH,
-                    ),
-                    resource_provider=self._resource_provider,
-                )
-            )
-
-        return processors
-
-    def _run_processors(
-        self, stage: BuildStage, dataframe: pd.DataFrame, current_batch_number: int | None = None
-    ) -> pd.DataFrame:
-        for processor in self._processors[stage]:
-            try:
-                dataframe = processor.process(dataframe, current_batch_number=current_batch_number)
-            except Exception as e:
-                raise DatasetProcessingError(
-                    f"🛑 Failed to process dataset with processor {processor.name} in stage {stage}: {e}"
-                ) from e
-        return dataframe
-
-    def _worker_error_callback(self, exc: Exception, *, context: dict | None = None) -> None:
-        """If a worker fails, we can handle the exception here."""
-        logger.warning(
-            f"⚠️ Generation for record at index {context['index']} failed. "
-            f"Will omit this record from the dataset.\n{exc}"
-        )
-        self._records_to_drop.add(context["index"])
-
-    def _worker_result_callback(self, result: dict, *, context: dict | None = None) -> None:
-        self.batch_manager.update_record(context["index"], result)
-
-    def _emit_batch_inference_events(
-        self, batch_mode: str, usage_deltas: dict[str, ModelUsageStats], group_id: str
-    ) -> None:
-        if not usage_deltas:
-            return
-
-        events = [
-            InferenceEvent(
-                nemo_source=NemoSourceEnum.DATADESIGNER,
-                task=batch_mode,
-                task_status=TaskStatusEnum.SUCCESS,
-                model=model_name,
-                input_tokens=delta.token_usage.input_tokens,
-                output_tokens=delta.token_usage.output_tokens,
-            )
-            for model_name, delta in usage_deltas.items()
-        ]
-
-        with TelemetryHandler(source_client_version=_CLIENT_VERSION, session_id=group_id) as telemetry_handler:
-            for event in events:
-                telemetry_handler.enqueue(event)

data_designer/engine/dataset_builders/errors.py

@@ -1,15 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from data_designer.engine.errors import DataDesignerError
-
-
-class ArtifactStorageError(DataDesignerError): ...
-
-
-class DatasetGenerationError(DataDesignerError): ...
-
-
-class DatasetProcessingError(DataDesignerError): ...

data_designer/engine/dataset_builders/multi_column_configs.py

@@ -1,46 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from abc import ABC
-from typing import TypeAlias
-
-from pydantic import Field, field_validator
-
-from data_designer.config.base import ConfigBase
-from data_designer.config.column_configs import SamplerColumnConfig, SeedDatasetColumnConfig, SingleColumnConfig
-from data_designer.config.column_types import ColumnConfigT, DataDesignerColumnType
-from data_designer.config.sampler_constraints import ColumnConstraintT
-from data_designer.config.seed import SeedConfig
-
-
-class MultiColumnConfig(ConfigBase, ABC):
-    columns: list[SingleColumnConfig] = Field(..., min_length=1)
-
-    @property
-    def column_names(self) -> list[str]:
-        return [c.name for c in self.columns]
-
-    @property
-    def column_type(self) -> DataDesignerColumnType:
-        return self.columns[0].column_type
-
-    @field_validator("columns", mode="after")
-    def validate_column_types_are_the_same(cls, v: list[SingleColumnConfig]) -> list[SingleColumnConfig]:
-        if len(set([c.column_type for c in v])) != 1:
-            raise ValueError("All column types must be of the same type")
-        return v
-
-
-class SamplerMultiColumnConfig(MultiColumnConfig):
-    columns: list[SamplerColumnConfig]
-    constraints: list[ColumnConstraintT] = []
-    max_rejections_factor: int = 5
-
-
-class SeedDatasetMultiColumnConfig(SeedConfig, MultiColumnConfig):
-    columns: list[SeedDatasetColumnConfig]
-
-
-DatasetBuilderColumnConfigT: TypeAlias = ColumnConfigT | SeedDatasetMultiColumnConfig | SamplerMultiColumnConfig

data_designer/engine/dataset_builders/utils/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0

data_designer/engine/dataset_builders/utils/concurrency.py

@@ -1,215 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import contextvars
-import json
-import logging
-from concurrent.futures import Future, ThreadPoolExecutor
-from threading import Lock, Semaphore
-from typing import Any, Protocol
-
-from pydantic import BaseModel, Field
-
-from data_designer.engine.errors import DataDesignerRuntimeError, ErrorTrap
-
-logger = logging.getLogger(__name__)
-
-# Constants
-MAX_CONCURRENCY_PER_NON_LLM_GENERATOR = 4
-
-
-class ExecutorResults(BaseModel):
-    failure_threshold: float = 0.0  # Error rate threshold
-    completed_count: int = 0  # How many tasks/jobs completed
-    success_count: int = 0  # How many tasks/jobs were successful
-    early_shutdown: bool = False  # Did we shutdown early due to errors?
-    error_trap: ErrorTrap = Field(default_factory=ErrorTrap)
-
-    @property
-    def summary(self) -> dict:
-        summary = self.model_dump(exclude={"error_trap"})
-        summary |= self.error_trap.model_dump()
-        return summary
-
-    def get_error_rate(self, window: int) -> float:
-        # We don't start actually tracking until our minimum window size is met
-        if self.completed_count < window:
-            return 0.0
-        return self.error_trap.error_count / max(1, self.completed_count)
-
-    def is_error_rate_exceeded(self, window: int) -> bool:
-        return self.get_error_rate(window) >= self.failure_threshold
-
-
-class CallbackWithContext(Protocol):
-    """Executor callback functions must accept a context kw argument."""
-
-    def __call__(self, result: Any, *, context: dict | None = None) -> Any: ...
-
-
-class ErrorCallbackWithContext(Protocol):
-    """Error callbacks take the Exception instance and context."""
-
-    def __call__(self, exc: Exception, *, context: dict | None = None) -> Any: ...
-
-
-class ConcurrentThreadExecutor:
-    """
-    Interface for executing multiple concurrent tasks with error rate monitoring.
-
-    This interface should be used exclusively as
-    a context manager. New tasks can be submitted to the executor using the `submit`
-    method. This submit method functions similarly to the
-    submit method of a ThreadPoolExecutor.
-
-    The underlying queue of tasks is bounded by the `max_workers`
-    parameter. This means that only `max_workers` number of
-    tasks can be queued up for execution. As tasks complete,
-    if there are errors, those are tracked and counted. If
-    a certain error rate is exceeded, the executor will shutdown
-    early. All queued and running tasks will complete.
-
-    The reason we bound the underlying task queue is to ensure that when
-    a certain error threshold is met there aren't an unbounded
-    number of tasks that need to complete. Generally speaking,
-    tasks should not be sitting in the queue for long at all since
-    the queue size == `max_workers`. The side effect of this is that
-    the `submit()` method will block, however this should not matter
-    because upstream Tasks need to wait for all jobs to complete
-    before the Task can be considered complete.
-
-    ContextVars from the main parent thread are automatically propagated
-    to all child threads.
-
-    When a task is completed, the user provided `result_callback`
-    function will be called with the task result as the only argument.
-    """
-
-    def __init__(
-        self,
-        *,
-        max_workers: int,
-        column_name: str,
-        result_callback: CallbackWithContext | None = None,
-        error_callback: ErrorCallbackWithContext | None = None,
-        shutdown_error_rate: float = 0.50,
-        shutdown_error_window: int = 10,
-        disable_early_shutdown: bool = False,
-    ):
-        self._executor = None
-        self._column_name = column_name
-        self._max_workers = max_workers
-        self._lock = Lock()
-        self._semaphore = Semaphore(self._max_workers)
-        self._result_callback = result_callback
-        self._error_callback = error_callback
-        self._shutdown_error_rate = shutdown_error_rate
-        self._shutdown_window_size = shutdown_error_window
-        self._disable_early_shutdown = disable_early_shutdown
-        self._results = ExecutorResults(failure_threshold=shutdown_error_rate)
-
-    @property
-    def results(self) -> ExecutorResults:
-        return self._results
-
-    @property
-    def max_workers(self) -> int:
-        return self._max_workers
-
-    @property
-    def shutdown_error_rate(self) -> float:
-        return self._shutdown_error_rate
-
-    @property
-    def shutdown_window_size(self) -> int:
-        return self._shutdown_window_size
-
-    @property
-    def semaphore(self) -> Semaphore:
-        return self._semaphore
-
-    def __enter__(self) -> ConcurrentThreadExecutor:
-        self._executor = ThreadPoolExecutor(
-            max_workers=self._max_workers,
-            thread_name_prefix="ConcurrentThreadExecutor",
-            initializer=_set_worker_contextvars,
-            initargs=(contextvars.copy_context(),),
-        )
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self._shutdown_executor()
-        if not self._disable_early_shutdown and self._results.early_shutdown is True:
-            self._raise_task_error()
-
-    def _shutdown_executor(self) -> None:
-        if self._executor is not None:
-            self._executor.shutdown()
-
-    def _raise_task_error(self):
-        raise DataDesignerRuntimeError(
-            "\n".join(
-                [
-                    "  |-- Data generation was terminated early due to error rate exceeding threshold.",
-                    f"  |-- The summary of encountered errors is: \n{json.dumps(self._results.summary, indent=4)}",
-                ]
-            )
-        )
-
-    def submit(self, fn, *args, context: dict | None = None, **kwargs) -> None:
-        if self._executor is None:
-            raise RuntimeError("Executor is not initialized, this class should be used as a context manager.")
-
-        if not self._disable_early_shutdown and self._results.early_shutdown:
-            self._shutdown_executor()
-            self._raise_task_error()
-
-        def _handle_future(future: Future) -> None:
-            try:
-                result = future.result()
-                if self._result_callback is not None:
-                    self._result_callback(result, context=context)
-                with self._lock:
-                    self._results.completed_count += 1
-                    self._results.success_count += 1
-            except Exception as err:
-                with self._lock:
-                    self._results.completed_count += 1
-                    self._results.error_trap.handle_error(err)
-                    if not self._disable_early_shutdown and self._results.is_error_rate_exceeded(
-                        self._shutdown_window_size
-                    ):
-                        # Signal to shutdown early on the next submission (if received).
-                        # We cannot trigger shutdown from within this thread as it can
-                        # cause a deadlock.
-                        if not self._results.early_shutdown:
-                            self._results.early_shutdown = True
-                if self._error_callback is not None:
-                    self._error_callback(err, context=context)
-            finally:
-                self._semaphore.release()
-
-        try:
-            self._semaphore.acquire()
-            future = self._executor.submit(fn, *args, **kwargs)
-            future.add_done_callback(_handle_future)
-        except Exception as err:
-            # If we get here, the pool is shutting down (likely due to early termination from errors)
-            # We'll re-raise a custom error that can be handled at the call-site and the summary
-            # can also be inspected.
-            self._semaphore.release()
-            is_shutdown_error = isinstance(err, RuntimeError) and (
-                "after shutdown" in str(err) or "Pool shutdown" in str(err)
-            )
-            if not is_shutdown_error:
-                raise err
-            if self._disable_early_shutdown:
-                raise err
-            self._raise_task_error()
-
-
-def _set_worker_contextvars(context: contextvars.Context):
-    for var, value in context.items():
-        var.set(value)

data_designer/engine/dataset_builders/utils/config_compiler.py

@@ -1,62 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from data_designer.config.column_types import DataDesignerColumnType
-from data_designer.config.data_designer_config import DataDesignerConfig
-from data_designer.config.processors import ProcessorConfig
-from data_designer.engine.dataset_builders.multi_column_configs import (
-    DatasetBuilderColumnConfigT,
-    SamplerMultiColumnConfig,
-    SeedDatasetMultiColumnConfig,
-)
-from data_designer.engine.dataset_builders.utils.dag import topologically_sort_column_configs
-from data_designer.engine.dataset_builders.utils.errors import ConfigCompilationError
-
-
-def compile_dataset_builder_column_configs(config: DataDesignerConfig) -> list[DatasetBuilderColumnConfigT]:
-    seed_column_configs = []
-    sampler_column_configs = []
-    generated_column_configs = []
-
-    for column_config in topologically_sort_column_configs(config.columns):
-        if column_config.column_type == DataDesignerColumnType.SEED_DATASET:
-            seed_column_configs.append(column_config)
-        elif column_config.column_type == DataDesignerColumnType.SAMPLER:
-            sampler_column_configs.append(column_config)
-        else:
-            generated_column_configs.append(column_config)
-
-    compiled_column_configs = []
-
-    if len(seed_column_configs) > 0:
-        if config.seed_config is None:
-            raise ConfigCompilationError("🛑 Seed column configs require a seed configuration.")
-        compiled_column_configs.append(
-            SeedDatasetMultiColumnConfig(
-                columns=seed_column_configs,
-                source=config.seed_config.source,
-                sampling_strategy=config.seed_config.sampling_strategy,
-                selection_strategy=config.seed_config.selection_strategy,
-            )
-        )
-
-    if len(sampler_column_configs) > 0:
-        compiled_column_configs.append(
-            SamplerMultiColumnConfig(
-                columns=sampler_column_configs,
-                constraints=config.constraints or [],
-            )
-        )
-
-    if len(generated_column_configs) > 0:
-        compiled_column_configs.extend(generated_column_configs)
-
-    return compiled_column_configs
-
-
-def compile_dataset_builder_processor_configs(
-    config: DataDesignerConfig,
-) -> list[ProcessorConfig]:
-    return config.processors or []