data-designer 0.1.0__py3-none-any.whl

data_designer/engine/dataset_builders/utils/dataset_batch_manager.py

@@ -0,0 +1,190 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from pathlib import Path
+import shutil
+from typing import Callable, Container, Iterator
+
+import pandas as pd
+import pyarrow.parquet as pq
+
+from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage, BatchStage
+from data_designer.engine.dataset_builders.utils.errors import DatasetBatchManagementError
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetBatchManager:
+    def __init__(self, artifact_storage: ArtifactStorage):
+        self._buffer: list[dict] = []
+        self._current_batch_number = 0
+        self._num_records_list: list[int] | None = None
+        self._buffer_size: int | None = None
+        self.artifact_storage = artifact_storage
+
+    @property
+    def num_batches(self) -> int:
+        if self._num_records_list is None:
+            return 0
+        return len(self._num_records_list)
+
+    @property
+    def num_records_batch(self) -> int:
+        if self._num_records_list is None or self._current_batch_number >= len(self._num_records_list):
+            raise DatasetBatchManagementError("🛑 Invalid batch number or num_records_list not set.")
+        return self._num_records_list[self._current_batch_number]
+
+    @property
+    def num_records_list(self) -> list[int]:
+        if self._num_records_list is None:
+            raise DatasetBatchManagementError("🛑 `num_records_list` is not set. Call start() first.")
+        return self._num_records_list
+
+    @property
+    def num_records_in_buffer(self) -> int:
+        return len(self._buffer)
+
+    @property
+    def buffer_is_empty(self) -> bool:
+        return len(self._buffer) == 0
+
+    @property
+    def buffer_size(self) -> int:
+        if self._buffer_size is None:
+            raise DatasetBatchManagementError("🛑 `buffer_size` is not set. Call start() first.")
+        return self._buffer_size
+
+    def add_record(self, record: dict) -> None:
+        self.add_records([record])
+
+    def add_records(self, records: list[dict]) -> None:
+        self._buffer.extend(records)
+        if len(self._buffer) > self.buffer_size:
+            raise DatasetBatchManagementError(
+                f"🛑 Buffer size exceeded. Current: {len(self._buffer)}, Max: {self.buffer_size}. "
+                "Flush the batch before adding more records."
+            )
+
+    def drop_records(self, index: Container[int]) -> None:
+        self._buffer = [record for i, record in enumerate(self._buffer) if i not in index]
+
+    def finish_batch(self, on_complete: Callable[[Path], None] | None = None) -> Path:
+        """Finish the batch by moving the results from the partial results path to the final parquet folder.
+
+        Returns:
+            The path to the written parquet file.
+        """
+        if self._current_batch_number >= self.num_batches:
+            raise DatasetBatchManagementError("🛑 All batches have been processed.")
+
+        if not self.write():
+            raise DatasetBatchManagementError("🛑 Batch finished without any results to write.")
+
+        final_file_path = self.artifact_storage.move_partial_result_to_final_file_path(self._current_batch_number)
+
+        self.artifact_storage.write_metadata(
+            {
+                "target_num_records": sum(self.num_records_list),
+                "total_num_batches": self.num_batches,
+                "buffer_size": self._buffer_size,
+                "schema": {field.name: str(field.type) for field in pq.read_schema(final_file_path)},
+                "file_paths": [str(f) for f in sorted(self.artifact_storage.final_dataset_path.glob("*.parquet"))],
+                "num_records": self.num_records_list[: self._current_batch_number + 1],
+                "num_completed_batches": self._current_batch_number + 1,
+                "dataset_name": self.artifact_storage.dataset_name,
+            }
+        )
+        self._current_batch_number += 1
+        self._buffer: list[dict] = []
+
+        if on_complete:
+            on_complete(final_file_path)
+
+        return final_file_path
+
+    def finish(self) -> None:
+        """Finish the dataset writing process by deleting the partial results path if it exists and is empty."""
+
+        # If the partial results path is empty, delete it.
+        if not any(self.artifact_storage.partial_results_path.iterdir()):
+            self.artifact_storage.partial_results_path.rmdir()
+
+        # Otherwise, log a warning, since existing partial results means the dataset is not complete.
+        else:
+            logger.warning("⚠️ Dataset writing finished with partial results.")
+
+        self.reset()
+
+    def get_current_batch_number(self) -> int:
+        return self._current_batch_number
+
+    def get_current_batch(self, *, as_dataframe: bool = False) -> pd.DataFrame | list[dict]:
+        if as_dataframe:
+            return pd.DataFrame(self._buffer)
+        return self._buffer
+
+    def iter_current_batch(self) -> Iterator[tuple[int, dict]]:
+        for i, record in enumerate(self._buffer):
+            yield i, record
+
+    def reset(self, delete_files: bool = False) -> None:
+        self._current_batch_number = 0
+        self._buffer: list[dict] = []
+        if delete_files:
+            for dir_path in [
+                self.artifact_storage.final_dataset_path,
+                self.artifact_storage.partial_results_path,
+                self.artifact_storage.dropped_columns_dataset_path,
+                self.artifact_storage.base_dataset_path,
+            ]:
+                if dir_path.exists():
+                    try:
+                        shutil.rmtree(dir_path)
+                    except OSError as e:
+                        raise DatasetBatchManagementError(f"🛑 Failed to delete directory {dir_path}: {e}")
+
+    def start(self, *, num_records: int, buffer_size: int) -> None:
+        if num_records <= 0:
+            raise DatasetBatchManagementError("🛑 num_records must be positive.")
+        if buffer_size <= 0:
+            raise DatasetBatchManagementError("🛑 buffer_size must be positive.")
+
+        self._buffer_size = buffer_size
+        self._num_records_list = [buffer_size] * (num_records // buffer_size)
+        if remaining_records := num_records % buffer_size:
+            self._num_records_list.append(remaining_records)
+        self.reset()
+
+    def write(self) -> Path | None:
+        """Write the current batch to a parquet file.
+
+        This method always writes results to the partial results path.
+
+        Returns:
+            The path to the written parquet file. If the buffer is empty, returns None.
+        """
+        if len(self._buffer) == 0:
+            return None
+        try:
+            file_path = self.artifact_storage.write_batch_to_parquet_file(
+                batch_number=self._current_batch_number,
+                dataframe=pd.DataFrame(self._buffer),
+                batch_stage=BatchStage.PARTIAL_RESULT,
+            )
+            return file_path
+        except Exception as e:
+            raise DatasetBatchManagementError(f"🛑 Failed to write batch {self._current_batch_number}: {e}")
+
+    def update_record(self, index: int, record: dict) -> None:
+        if index < 0 or index >= len(self._buffer):
+            raise IndexError(f"🛑 Index {index} is out of bounds for buffer of size {len(self._buffer)}.")
+        self._buffer[index] = record
+
+    def update_records(self, records: list[dict]) -> None:
+        if len(records) != len(self._buffer):
+            raise DatasetBatchManagementError(
+                f"🛑 Number of records to update ({len(records)}) must match "
+                f"the number of records in the buffer ({len(self._buffer)})."
+            )
+        self._buffer = records

data_designer/engine/dataset_builders/utils/errors.py

@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.engine.errors import DataDesignerError
+
+
+class DatasetBatchManagementError(DataDesignerError): ...
+
+
+class ConfigCompilationError(DataDesignerError): ...
+
+
+class DAGCircularDependencyError(DataDesignerError): ...

data_designer/engine/errors.py

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from pydantic import BaseModel, Field
+
+from ..errors import DataDesignerError
+
+
+class DataDesignerRuntimeError(DataDesignerError): ...
+
+
+class UnknownModelAliasError(DataDesignerError): ...
+
+
+class UnknownProviderError(DataDesignerError): ...
+
+
+class NoModelProvidersError(DataDesignerError): ...
+
+
+class SecretResolutionError(DataDesignerError): ...
+
+
+class RemoteValidationSchemaError(DataDesignerError): ...
+
+
+class LocalCallableValidationError(DataDesignerError): ...
+
+
+class ErrorTrap(BaseModel):
+    error_count: int = 0
+    task_errors: dict[str, int] = Field(default_factory=dict)
+
+    def _track_error(self, error: DataDesignerError) -> None:
+        """
+        Track a specific error type.
+        """
+        error_type = type(error).__name__
+        if error_type not in self.task_errors:
+            self.task_errors[error_type] = 0
+        self.task_errors[error_type] += 1
+
+    def handle_error(self, error: Exception) -> None:
+        self.error_count += 1
+
+        if not isinstance(error, DataDesignerError):
+            error = DataDesignerError(str(error))
+
+        self._track_error(error)

data_designer/engine/model_provider.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from functools import cached_property
+
+from pydantic import BaseModel, field_validator, model_validator
+from typing_extensions import Self
+
+from data_designer.config.models import ModelProvider
+from data_designer.engine.errors import NoModelProvidersError, UnknownProviderError
+
+
+class ModelProviderRegistry(BaseModel):
+    providers: list[ModelProvider]
+    default: str | None = None
+
+    @field_validator("providers", mode="after")
+    @classmethod
+    def validate_providers_not_empty(cls, v: list[ModelProvider]) -> list[ModelProvider]:
+        if len(v) == 0:
+            raise ValueError("At least one model provider must be defined")
+        return v
+
+    @field_validator("providers", mode="after")
+    @classmethod
+    def validate_providers_have_unique_names(cls, v: list[ModelProvider]) -> list[ModelProvider]:
+        names = set()
+        dupes = set()
+        for provider in v:
+            if provider.name in names:
+                dupes.add(provider.name)
+            names.add(provider.name)
+
+        if len(dupes) > 0:
+            raise ValueError(f"Model providers must have unique names, found duplicates: {dupes}")
+        return v
+
+    @model_validator(mode="after")
+    def check_implicit_default(self) -> Self:
+        if self.default is None and len(self.providers) != 1:
+            raise ValueError("A default provider must be specified if multiple model providers are defined")
+        return self
+
+    @model_validator(mode="after")
+    def check_default_exists(self) -> Self:
+        if self.default and self.default not in self._providers_dict:
+            raise ValueError(f"Specified default {self.default!r} not found in providers list")
+        return self
+
+    def get_default_provider_name(self) -> str:
+        return self.default or self.providers[0].name
+
+    @cached_property
+    def _providers_dict(self) -> dict[str, ModelProvider]:
+        return {p.name: p for p in self.providers}
+
+    def get_provider(self, name: str | None) -> ModelProvider:
+        if name is None:
+            name = self.get_default_provider_name()
+
+        try:
+            return self._providers_dict[name]
+        except KeyError:
+            raise UnknownProviderError(f"No provider named {name!r} registered")
+
+
+def resolve_model_provider_registry(
+    model_providers: list[ModelProvider], default_provider_name: str | None = None
+) -> ModelProviderRegistry:
+    if len(model_providers) == 0:
+        raise NoModelProvidersError("At least one model provider must be defined")
+    return ModelProviderRegistry(
+        providers=model_providers,
+        default=default_provider_name or model_providers[0].name,
+    )

data_designer/engine/models/__init__.py

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

data_designer/engine/models/errors.py

@@ -0,0 +1,308 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from functools import wraps
+import logging
+from typing import Any
+
+from litellm.exceptions import (
+    APIConnectionError,
+    APIError,
+    AuthenticationError,
+    BadRequestError,
+    ContextWindowExceededError,
+    InternalServerError,
+    NotFoundError,
+    PermissionDeniedError,
+    RateLimitError,
+    Timeout,
+    UnprocessableEntityError,
+    UnsupportedParamsError,
+)
+from pydantic import BaseModel
+
+from data_designer.engine.errors import DataDesignerError
+
+logger = logging.getLogger(__name__)
+
+
+def get_exception_primary_cause(exception: BaseException) -> BaseException:
+    """Returns the primary cause of an exception by walking backwards.
+
+    This recursive walkback halts when it arrives at an exception which
+    has no provided __cause__ (e.g. __cause__ is None).
+
+    Args:
+        exception (Exception): An exception to start from.
+
+    Raises:
+        RecursionError: if for some reason exceptions have circular
+            dependencies (seems impossible in practice).
+    """
+    if exception.__cause__ is None:
+        return exception
+    else:
+        return get_exception_primary_cause(exception.__cause__)
+
+
+class GenerationValidationFailureError(Exception): ...
+
+
+class ModelRateLimitError(DataDesignerError): ...
+
+
+class ModelTimeoutError(DataDesignerError): ...
+
+
+class ModelContextWindowExceededError(DataDesignerError): ...
+
+
+class ModelAuthenticationError(DataDesignerError): ...
+
+
+class ModelPermissionDeniedError(DataDesignerError): ...
+
+
+class ModelNotFoundError(DataDesignerError): ...
+
+
+class ModelUnsupportedParamsError(DataDesignerError): ...
+
+
+class ModelBadRequestError(DataDesignerError): ...
+
+
+class ModelInternalServerError(DataDesignerError): ...
+
+
+class ModelAPIError(DataDesignerError): ...
+
+
+class ModelUnprocessableEntityError(DataDesignerError): ...
+
+
+class ModelAPIConnectionError(DataDesignerError): ...
+
+
+class ModelStructuredOutputError(DataDesignerError): ...
+
+
+class ModelGenerationValidationFailureError(DataDesignerError): ...
+
+
+class FormattedLLMErrorMessage(BaseModel):
+    cause: str
+    solution: str
+
+    def __str__(self) -> str:
+        return "\n".join(
+            [
+                "  |----------",
+                f"  | Cause: {self.cause}",
+                f"  | Solution: {self.solution}",
+                "  |----------",
+            ]
+        )
+
+
+def handle_llm_exceptions(
+    exception: Exception, model_name: str, model_provider_name: str, purpose: str | None = None
+) -> None:
+    """Handle LLM-related exceptions and convert them to appropriate DataDesignerError errors.
+
+    This method centralizes the exception handling logic for LLM operations,
+    making it reusable across different contexts.
+
+    Args:
+        exception: The exception that was raised
+        model_name: Name of the model that was being used
+        model_provider_name: Name of the model provider that was being used
+        purpose: The purpose of the model usage to show as context in the error message
+    Raises:
+        DataDesignerError: A more user-friendly error with appropriate error type and message
+    """
+    purpose = purpose or "running generation"
+    authentication_error = FormattedLLMErrorMessage(
+        cause=f"The API key provided for model {model_name!r} was found to be invalid or expired while {purpose}.",
+        solution=f"Verify your API key for model provider and update it in your settings for model provider {model_provider_name!r}.",
+    )
+    err_msg_parser = DownstreamLLMExceptionMessageParser(model_name, model_provider_name, purpose)
+    match exception:
+        # Common errors that can come from LiteLLM
+        case APIError():
+            raise err_msg_parser.parse_api_error(exception, authentication_error) from None
+
+        case APIConnectionError():
+            raise ModelAPIConnectionError(
+                FormattedLLMErrorMessage(
+                    cause=f"Connection to model {model_name!r} hosted on model provider {model_provider_name!r} failed while {purpose}.",
+                    solution="Check your network/proxy/firewall settings.",
+                )
+            ) from None
+
+        case AuthenticationError():
+            raise ModelAuthenticationError(authentication_error) from None
+
+        case ContextWindowExceededError():
+            raise err_msg_parser.parse_context_window_exceeded_error(exception) from None
+
+        case UnsupportedParamsError():
+            raise ModelUnsupportedParamsError(
+                FormattedLLMErrorMessage(
+                    cause=f"One or more of the parameters you provided were found to be unsupported by model {model_name!r} while {purpose}.",
+                    solution=f"Review the documentation for model provider {model_provider_name!r} and adjust your request.",
+                )
+            ) from None
+
+        case BadRequestError():
+            raise err_msg_parser.parse_bad_request_error(exception) from None
+
+        case InternalServerError():
+            raise ModelInternalServerError(
+                FormattedLLMErrorMessage(
+                    cause=f"Model {model_name!r} is currently experiencing internal server issues while {purpose}.",
+                    solution=f"Try again in a few moments. Check with your model provider {model_provider_name!r} if the issue persists.",
+                )
+            ) from None
+
+        case NotFoundError():
+            raise ModelNotFoundError(
+                FormattedLLMErrorMessage(
+                    cause=f"The specified model {model_name!r} could not be found while {purpose}.",
+                    solution=f"Check that the model name is correct and supported by your model provider {model_provider_name!r} and try again.",
+                )
+            ) from None
+
+        case PermissionDeniedError():
+            raise ModelPermissionDeniedError(
+                FormattedLLMErrorMessage(
+                    cause=f"Your API key was found to lack the necessary permissions to use model {model_name!r} while {purpose}.",
+                    solution=f"Use an API key that has the right permissions for the model or use a model the API key in use has access to in model provider {model_provider_name!r}.",
+                )
+            ) from None
+
+        case RateLimitError():
+            raise ModelRateLimitError(
+                FormattedLLMErrorMessage(
+                    cause=f"You have exceeded the rate limit for model {model_name!r} while {purpose}.",
+                    solution="Wait and try again in a few moments.",
+                )
+            ) from None
+
+        case Timeout():
+            raise ModelTimeoutError(
+                FormattedLLMErrorMessage(
+                    cause=f"The request to model {model_name!r} timed out while {purpose}.",
+                    solution="Check your connection and try again. You may need to increase the timeout setting for the model.",
+                )
+            ) from None
+
+        case UnprocessableEntityError():
+            raise ModelUnprocessableEntityError(
+                FormattedLLMErrorMessage(
+                    cause=f"The request to model {model_name!r} failed despite correct request format while {purpose}.",
+                    solution="This is most likely temporary. Try again in a few moments.",
+                )
+            ) from None
+
+        # Parsing and validation errors
+        case GenerationValidationFailureError():
+            raise ModelGenerationValidationFailureError(
+                FormattedLLMErrorMessage(
+                    cause=f"The provided output schema was unable to be parsed from model {model_name!r} responses while {purpose}.",
+                    solution="This is most likely temporary as we make additional attempts. If you continue to see more of this, simplify or modify the output schema for structured output and try again. If you are attempting token-intensive tasks like generations with high-reasoning effort, ensure that max_tokens in the model config is high enough to reach completion.",
+                )
+            ) from None
+
+        case DataDesignerError():
+            raise exception from None
+
+        case _:
+            raise DataDesignerError(
+                FormattedLLMErrorMessage(
+                    cause=f"An unexpected error occurred while {purpose}.",
+                    solution=f"Review the stack trace for more details: {exception}",
+                )
+            ) from exception
+
+
+def catch_llm_exceptions(func: Callable) -> Callable:
+    """This decorator should be used on any `ModelFacade` method that could potentially raise
+    exceptions that should turn into upstream user-facing errors.
+    """
+
+    @wraps(func)
+    def wrapper(model_facade: Any, *args, **kwargs):
+        try:
+            return func(model_facade, *args, **kwargs)
+        except Exception as e:
+            logger.debug(
+                "\n".join(
+                    [
+                        "",
+                        "|----------",
+                        f"| Caught an exception downstream of type {type(e)!r}. Re-raising it below as a custom error with more context.",
+                        "|----------",
+                    ]
+                ),
+                exc_info=True,
+                stack_info=True,
+            )
+            handle_llm_exceptions(
+                e, model_facade.model_name, model_facade.model_provider_name, purpose=kwargs.get("purpose")
+            )
+
+    return wrapper
+
+
+class DownstreamLLMExceptionMessageParser:
+    def __init__(self, model_name: str, model_provider_name: str, purpose: str):
+        self.model_name = model_name
+        self.model_provider_name = model_provider_name
+        self.purpose = purpose
+
+    def parse_bad_request_error(self, exception: BadRequestError) -> DataDesignerError:
+        err_msg = FormattedLLMErrorMessage(
+            cause=f"The request for model {self.model_name!r} was found to be malformed or missing required parameters while {self.purpose}.",
+            solution="Check your request parameters and try again.",
+        )
+        if "is not a multimodal model" in str(exception):
+            err_msg = FormattedLLMErrorMessage(
+                cause=f"Model {self.model_name!r} is not a multimodal model, but it looks like you are trying to provide multimodal context while {self.purpose}.",
+                solution="Check your request parameters and try again.",
+            )
+        return ModelBadRequestError(err_msg)
+
+    def parse_context_window_exceeded_error(self, exception: ContextWindowExceededError) -> DataDesignerError:
+        cause = f"The input data for model '{self.model_name}' was found to exceed its supported context width while {self.purpose}."
+        try:
+            if "OpenAIException - This model's maximum context length is " in str(exception):
+                openai_exception_cause = (
+                    str(exception).split("OpenAIException - ")[1].split("\n")[0].split(" Please reduce ")[0]
+                )
+                cause = f"{cause} {openai_exception_cause}"
+        except Exception:
+            pass
+        finally:
+            return ModelContextWindowExceededError(
+                FormattedLLMErrorMessage(
+                    cause=cause,
+                    solution="Check the model's supported max context width. Adjust the length of your input along with completions and try again.",
+                )
+            )
+
+    def parse_api_error(
+        self, exception: InternalServerError, auth_error_msg: FormattedLLMErrorMessage
+    ) -> DataDesignerError:
+        if "Error code: 403" in str(exception):
+            return ModelAuthenticationError(auth_error_msg)
+
+        return ModelAPIError(
+            FormattedLLMErrorMessage(
+                cause=f"An unexpected API error occurred with model {self.model_name!r} while {self.purpose}.",
+                solution=f"Try again in a few moments. Check with your model provider {self.model_provider_name!r} if the issue persists.",
+            )
+        )