data-designer-engine 0.4.0__py3-none-any.whl

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -0,0 +1,354 @@
+# 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 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.dataset_builders.utils.progress_tracker import ProgressTracker
+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-engine")
+
+
+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 = 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)
+
+    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."
+            )
+
+        progress_tracker = ProgressTracker(
+            total_records=self.batch_manager.num_records_batch,
+            label=f"{generator.config.column_type} column '{generator.config.name}'",
+        )
+        progress_tracker.log_start(max_workers)
+
+        settings = self._resource_provider.run_config
+        with ConcurrentThreadExecutor(
+            max_workers=max_workers,
+            column_name=generator.config.name,
+            result_callback=self._make_result_callback(progress_tracker),
+            error_callback=self._make_error_callback(progress_tracker),
+            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})
+
+        progress_tracker.log_final()
+
+        if len(self._records_to_drop) > 0:
+            self.batch_manager.drop_records(self._records_to_drop)
+            self._records_to_drop.clear()
+
+    def _make_result_callback(self, progress_tracker: ProgressTracker) -> Callable[[dict], None]:
+        def callback(result: dict, *, context: dict | None = None) -> None:
+            self._worker_result_callback(result, context=context)
+            progress_tracker.record_success()
+
+        return callback
+
+    def _make_error_callback(self, progress_tracker: ProgressTracker) -> Callable[[Exception], None]:
+        def callback(exc: Exception, *, context: dict | None = None) -> None:
+            self._worker_error_callback(exc, context=context)
+            progress_tracker.record_failure()
+
+        return callback
+
+    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

@@ -0,0 +1,15 @@
+# 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

@@ -0,0 +1,46 @@
+# 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

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

data_designer/engine/dataset_builders/utils/concurrency.py

@@ -0,0 +1,212 @@
+# 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__)
+
+
+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)