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

data_designer/engine/resources/managed_dataset_generator.py

@@ -1,39 +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 typing import TYPE_CHECKING, Any
-
-from data_designer.engine.resources.managed_dataset_repository import ManagedDatasetRepository
-from data_designer.lazy_heavy_imports import pd
-
-if TYPE_CHECKING:
-    import pandas as pd
-
-
-class ManagedDatasetGenerator:
-    def __init__(self, managed_datasets: ManagedDatasetRepository, dataset_name: str):
-        self.managed_datasets = managed_datasets
-        self.dataset_name = dataset_name
-
-    def generate_samples(
-        self,
-        size: int = 1,
-        evidence: dict[str, Any | list[Any]] = {},
-    ) -> pd.DataFrame:
-        parameters = []
-        query = f"select * from {self.dataset_name}"
-        if evidence:
-            where_conditions = []
-            for column, values in evidence.items():
-                if values:
-                    values = values if isinstance(values, list) else [values]
-                    formatted_values = ["?"] * len(values)
-                    condition = f"{column} IN ({', '.join(formatted_values)})"
-                    where_conditions.append(condition)
-                    parameters.extend(values)
-            if where_conditions:
-                query += " where " + " and ".join(where_conditions)
-        query += f" order by random() limit {size}"
-        return self.managed_datasets.query(query, parameters)

data_designer/engine/resources/managed_dataset_repository.py

@@ -1,197 +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 logging
-import tempfile
-import threading
-import time
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from functools import cached_property
-from pathlib import Path
-from typing import TYPE_CHECKING, Any
-
-from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS
-from data_designer.engine.resources.managed_storage import LocalBlobStorageProvider, ManagedBlobStorage
-from data_designer.lazy_heavy_imports import duckdb, pd
-
-if TYPE_CHECKING:
-    import duckdb
-    import pandas as pd
-
-logger = logging.getLogger(__name__)
-
-DATASETS_ROOT = "datasets"
-"""
-Path in object storage to managed datasets
-"""
-
-
-@dataclass
-class Table:
-    """
-    Managed datasets are organized by dataset by table under a root
-    table path in object storage.
-    """
-
-    source: str
-    """
-    Table source path
-    """
-
-    schema: str = "main"
-    """
-    Specifies the schema to use when registering the table.
-
-    Note: this is not the schema of the table, but rather the _database_
-    schema to associated with the table.
-    """
-
-    @cached_property
-    def name(self) -> str:
-        return Path(self.source).stem
-
-
-DataCatalog = list[Table]
-
-# For now we hardcode the remote data catalog in code. This make it easier
-# initialize the data catalog. Eventually we can make this work more
-# dynamically once this data catalog pattern becomes more widely adopted.
-DEFAULT_DATA_CATALOG: DataCatalog = [Table(f"{locale}.parquet") for locale in LOCALES_WITH_MANAGED_DATASETS]
-
-
-class ManagedDatasetRepository(ABC):
-    @abstractmethod
-    def query(self, sql: str, parameters: list[Any]) -> pd.DataFrame: ...
-
-    @property
-    @abstractmethod
-    def data_catalog(self) -> DataCatalog: ...
-
-
-class DuckDBDatasetRepository(ManagedDatasetRepository):
-    """
-    Provides a duckdb based sql interface over Gretel managed datasets.
-    """
-
-    _default_config = {"threads": 2, "memory_limit": "4 gb"}
-
-    def __init__(
-        self,
-        blob_storage: ManagedBlobStorage,
-        config: dict | None = None,
-        data_catalog: DataCatalog = DEFAULT_DATA_CATALOG,
-        datasets_root: str = DATASETS_ROOT,
-        use_cache: bool = True,
-    ):
-        """
-        Create a new DuckDB backed dataset repository
-
-        Args:
-            blob_storage: A managed blob storage provider
-            config: DuckDB configuration options,
-            https://duckdb.org/docs/configuration/overview.html#configuration-reference
-            data_catalog: A list of tables to register with the DuckDB instance
-            datasets_root: The root path in blob storage to managed datasets
-            use_cache: Whether to cache datasets locally. Trades off disk memory
-            and startup time for faster queries.
-        """
-        self._data_catalog = data_catalog
-        self._data_sets_root = datasets_root
-        self._blob_storage = blob_storage
-        self._config = self._default_config if config is None else config
-        self._use_cache = use_cache
-
-        # Configure database and register tables
-        self.db = duckdb.connect(config=self._config)
-
-        # Dataset registration completion is tracked with an event. Consumers can
-        # wait on this event to ensure the catalog is ready.
-        self._registration_event = threading.Event()
-        self._register_lock = threading.Lock()
-
-        # Kick off dataset registration in a background thread so that IO-heavy
-        # caching and view creation can run asynchronously without blocking the
-        # caller that constructs this repository instance.
-        self._register_thread = threading.Thread(target=self._register_datasets, daemon=True)
-        self._register_thread.start()
-
-    def _register_datasets(self):
-        # Just in case this method gets called from inside a thread.
-        # This operation isn't thread-safe by default, so we
-        # synchronize the registration process.
-        if self._registration_event.is_set():
-            return
-        with self._register_lock:
-            # check once more to see if the catalog is ready it's possible a
-            # previous thread already registered the dataset.
-            if self._registration_event.is_set():
-                return
-            try:
-                for table in self.data_catalog:
-                    key = table.source if table.schema == "main" else f"{table.schema}/{table.source}"
-                    if self._use_cache:
-                        tmp_root = Path(tempfile.gettempdir()) / "dd_cache"
-                        local_path = tmp_root / key
-                        local_path.parent.mkdir(parents=True, exist_ok=True)
-                        if not local_path.exists():
-                            start = time.time()
-                            logger.debug("Caching database %s to %s", table.name, local_path)
-                            with self._blob_storage.get_blob(f"{self._data_sets_root}/{key}") as src_fd:
-                                with open(local_path, "wb") as dst_fd:
-                                    dst_fd.write(src_fd.read())
-                            logger.debug(
-                                "Cached database %s in %.2f s",
-                                table.name,
-                                time.time() - start,
-                            )
-                        data_path = local_path.as_posix()
-                    else:
-                        data_path = self._blob_storage.uri_for_key(f"{self._data_sets_root}/{key}")
-                    if table.schema != "main":
-                        self.db.sql(f"CREATE SCHEMA IF NOT EXISTS {table.schema}")
-                    logger.debug(f"Registering dataset {table.name} from {data_path}")
-                    self.db.sql(f"CREATE VIEW {table.schema}.{table.name} AS FROM '{data_path}'")
-
-                logger.debug("DuckDBDatasetRepository registration complete")
-
-            except Exception as e:
-                logger.exception(f"Failed to register datasets: {str(e)}")
-
-            finally:
-                # Signal that registration is complete so any waiting queries can proceed.
-                self._registration_event.set()
-
-    def query(self, sql: str, parameters: list[Any]) -> pd.DataFrame:
-        # Ensure dataset registration has completed. Possible future optimization:
-        # pull datasets in parallel and only wait here if the query requires a
-        # table that isn't cached.
-        if not self._registration_event.is_set():
-            logger.debug("Waiting for dataset caching and registration to finish...")
-            self._registration_event.wait()
-
-        # the duckdb connection isn't thread-safe, so we create a new
-        # connection per query using cursor().
-        # more details here: https://duckdb.org/docs/stable/guides/python/multiple_threads.html
-        cursor = self.db.cursor()
-        try:
-            df = cursor.execute(sql, parameters).df()
-        finally:
-            cursor.close()
-        return df
-
-    @property
-    def data_catalog(self) -> DataCatalog:
-        return self._data_catalog
-
-
-def load_managed_dataset_repository(blob_storage: ManagedBlobStorage, locales: list[str]) -> ManagedDatasetRepository:
-    return DuckDBDatasetRepository(
-        blob_storage,
-        config={"threads": 1, "memory_limit": "2 gb"},
-        data_catalog=[Table(f"{locale}.parquet") for locale in locales],
-        # Only cache if not using local storage.
-        use_cache=not isinstance(blob_storage, LocalBlobStorageProvider),
-    )

data_designer/engine/resources/managed_storage.py

@@ -1,65 +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 logging
-from abc import ABC, abstractmethod
-from collections.abc import Iterator
-from contextlib import contextmanager
-from pathlib import Path
-from typing import IO
-
-logger = logging.getLogger(__name__)
-
-
-class ManagedBlobStorage(ABC):
-    """
-    Provides a low-level interface for access object in blob storage. This interface
-    can be used to access model weights, raw datasets, or any artifact in blob
-    storage.
-
-    If you want a high-level interface for accessing datasets, use the `ManagedDatasetRepository`
-    which provides a high-level SQL interface over each dataset.
-    """
-
-    @abstractmethod
-    @contextmanager
-    def get_blob(self, blob_key: str) -> Iterator[IO]: ...
-
-    @abstractmethod
-    def _key_uri_builder(self, key: str) -> str: ...
-
-    def uri_for_key(self, key: str) -> str:
-        """
-        Returns a qualified storage URI for a given a key. `key` is
-        normalized to ensure that and leading path components ("/")  are removed.
-        """
-        return self._key_uri_builder(key.lstrip("/"))
-
-
-class LocalBlobStorageProvider(ManagedBlobStorage):
-    """
-    Provide a local blob storage service. Useful for running
-    tests that don't require access to external infrastructure
-    """
-
-    def __init__(self, root_path: Path) -> None:
-        self._root_path = root_path
-
-    @contextmanager
-    def get_blob(self, blob_key: str) -> Iterator[IO]:
-        with open(self._key_uri_builder(blob_key), "rb") as fd:
-            yield fd
-
-    def _key_uri_builder(self, key: str) -> str:
-        return f"{self._root_path}/{key}"
-
-
-def init_managed_blob_storage(assets_storage: str) -> ManagedBlobStorage:
-    path = Path(assets_storage)
-    if not path.exists():
-        raise RuntimeError(f"Local storage path {assets_storage!r} does not exist.")
-
-    logger.debug(f"Using local storage for managed datasets: {assets_storage!r}")
-    return LocalBlobStorageProvider(Path(assets_storage))

data_designer/engine/resources/resource_provider.py

@@ -1,77 +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.base import ConfigBase
-from data_designer.config.dataset_metadata import DatasetMetadata
-from data_designer.config.models import ModelConfig
-from data_designer.config.run_config import RunConfig
-from data_designer.config.seed_source import SeedSource
-from data_designer.config.utils.type_helpers import StrEnum
-from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
-from data_designer.engine.model_provider import ModelProviderRegistry
-from data_designer.engine.models.factory import create_model_registry
-from data_designer.engine.models.registry import ModelRegistry
-from data_designer.engine.resources.managed_storage import ManagedBlobStorage, init_managed_blob_storage
-from data_designer.engine.resources.seed_reader import SeedReader, SeedReaderRegistry
-from data_designer.engine.secret_resolver import SecretResolver
-
-
-class ResourceType(StrEnum):
-    BLOB_STORAGE = "blob_storage"
-    MODEL_REGISTRY = "model_registry"
-    SEED_READER = "seed_reader"
-
-
-class ResourceProvider(ConfigBase):
-    artifact_storage: ArtifactStorage
-    blob_storage: ManagedBlobStorage | None = None
-    model_registry: ModelRegistry | None = None
-    run_config: RunConfig = RunConfig()
-    seed_reader: SeedReader | None = None
-
-    def get_dataset_metadata(self) -> DatasetMetadata:
-        """Get metadata about the dataset being generated.
-
-        Returns:
-            DatasetMetadata with seed column names and other metadata.
-        """
-        seed_column_names = []
-        if self.seed_reader is not None:
-            seed_column_names = self.seed_reader.get_column_names()
-        return DatasetMetadata(seed_column_names=seed_column_names)
-
-
-def create_resource_provider(
-    *,
-    artifact_storage: ArtifactStorage,
-    model_configs: list[ModelConfig],
-    secret_resolver: SecretResolver,
-    model_provider_registry: ModelProviderRegistry,
-    seed_reader_registry: SeedReaderRegistry,
-    blob_storage: ManagedBlobStorage | None = None,
-    seed_dataset_source: SeedSource | None = None,
-    run_config: RunConfig | None = None,
-) -> ResourceProvider:
-    """Factory function for creating a ResourceProvider instance.
-    This function triggers lazy loading of heavy dependencies like litellm.
-    """
-    seed_reader = None
-    if seed_dataset_source:
-        seed_reader = seed_reader_registry.get_reader(
-            seed_dataset_source,
-            secret_resolver,
-        )
-
-    return ResourceProvider(
-        artifact_storage=artifact_storage,
-        model_registry=create_model_registry(
-            model_configs=model_configs,
-            secret_resolver=secret_resolver,
-            model_provider_registry=model_provider_registry,
-        ),
-        blob_storage=blob_storage or init_managed_blob_storage(),
-        seed_reader=seed_reader,
-        run_config=run_config or RunConfig(),
-    )

data_designer/engine/resources/seed_reader.py

@@ -1,154 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Generic, TypeVar, get_args, get_origin
-
-from huggingface_hub import HfFileSystem
-from typing_extensions import Self
-
-from data_designer.config.seed_source import (
-    DataFrameSeedSource,
-    HuggingFaceSeedSource,
-    LocalFileSeedSource,
-    SeedSource,
-)
-from data_designer.engine.secret_resolver import SecretResolver
-from data_designer.errors import DataDesignerError
-from data_designer.lazy_heavy_imports import duckdb
-
-if TYPE_CHECKING:
-    import duckdb
-
-
-class SeedReaderError(DataDesignerError): ...
-
-
-SourceT = TypeVar("ConfigT", bound=SeedSource)
-
-
-class SeedReader(ABC, Generic[SourceT]):
-    """Base class for reading a seed dataset.
-
-    Seeds are read using duckdb. Reader implementations define duckdb connection setup details
-    and how to get a URI that can be queried with duckdb (i.e. "... FROM <uri> ...").
-
-    The Data Designer engine automatically supplies the appropriate SeedSource
-    and a SecretResolver to use for any secret fields in the config.
-    """
-
-    source: SourceT
-    secret_resolver: SecretResolver
-
-    @abstractmethod
-    def get_dataset_uri(self) -> str: ...
-
-    @abstractmethod
-    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection: ...
-
-    def attach(self, source: SourceT, secret_resolver: SecretResolver):
-        """Attach a source and secret resolver to the instance.
-
-        This is called internally by the engine so that these objects do not
-        need to be provided in the reader's constructor.
-        """
-        self.source = source
-        self.secret_resolver = secret_resolver
-
-    def get_column_names(self) -> list[str]:
-        """Returns the seed dataset's column names"""
-        conn = self.create_duckdb_connection()
-        describe_query = f"DESCRIBE SELECT * FROM '{self.get_dataset_uri()}'"
-        column_descriptions = conn.execute(describe_query).fetchall()
-        return [col[0] for col in column_descriptions]
-
-    def get_seed_type(self) -> str:
-        """Return the seed_type of the source class this reader is generic over."""
-        # Get the generic type arguments from the reader class
-        # Check __orig_bases__ for the generic base class
-        for base in getattr(type(self), "__orig_bases__", []):
-            origin = get_origin(base)
-            if origin is SeedReader:
-                args = get_args(base)
-                if args:
-                    source_cls = args[0]
-                    # Extract seed_type from the source class
-                    if hasattr(source_cls, "model_fields") and "seed_type" in source_cls.model_fields:
-                        field = source_cls.model_fields["seed_type"]
-                        default_value = field.default
-                        if isinstance(default_value, str):
-                            return default_value
-
-        raise SeedReaderError("Reader does not have a valid generic source type with seed_type")
-
-
-class LocalFileSeedReader(SeedReader[LocalFileSeedSource]):
-    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
-        return duckdb.connect()
-
-    def get_dataset_uri(self) -> str:
-        return self.source.path
-
-
-class HuggingFaceSeedReader(SeedReader[HuggingFaceSeedSource]):
-    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
-        token = self.secret_resolver.resolve(self.source.token) if self.source.token else None
-
-        # Use skip_instance_cache to avoid fsspec-level caching
-        hffs = HfFileSystem(endpoint=self.source.endpoint, token=token, skip_instance_cache=True)
-
-        # Clear all internal caches to avoid stale metadata issues
-        # HfFileSystem caches file metadata (size, etc.) which can become stale when files are re-uploaded
-        if hasattr(hffs, "dircache"):
-            hffs.dircache.clear()
-
-        conn = duckdb.connect()
-        conn.register_filesystem(hffs)
-        return conn
-
-    def get_dataset_uri(self) -> str:
-        return f"hf://{self.source.path}"
-
-
-class DataFrameSeedReader(SeedReader[DataFrameSeedSource]):
-    # This is a "magic string" that gets registered in the duckdb connection to make the dataframe directly queryable.
-    _table_name = "df"
-
-    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
-        conn = duckdb.connect()
-        conn.register(self._table_name, self.source.df)
-        return conn
-
-    def get_dataset_uri(self) -> str:
-        return self._table_name
-
-
-class SeedReaderRegistry:
-    def __init__(self, readers: Sequence[SeedReader]):
-        self._readers: dict[str, SeedReader] = {}
-        for reader in readers:
-            self.add_reader(reader)
-
-    def add_reader(self, reader: SeedReader) -> Self:
-        seed_type = reader.get_seed_type()
-
-        if seed_type in self._readers:
-            raise SeedReaderError(f"A reader for seed_type {seed_type!r} already exists")
-
-        self._readers[seed_type] = reader
-        return self
-
-    def get_reader(self, seed_dataset_source: SeedSource, secret_resolver: SecretResolver) -> SeedReader:
-        reader = self._get_reader_for_source(seed_dataset_source)
-        reader.attach(seed_dataset_source, secret_resolver)
-        return reader
-
-    def _get_reader_for_source(self, seed_dataset_source: SeedSource) -> SeedReader:
-        seed_type = seed_dataset_source.seed_type
-        try:
-            return self._readers[seed_type]
-        except KeyError:
-            raise SeedReaderError(f"No reader found for seed_type {seed_type!r}")

data_designer/engine/sampling_gen/column.py

@@ -1,91 +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 typing import Any
-
-from pydantic import field_serializer, model_validator
-from typing_extensions import Self
-
-from data_designer.config.column_configs import SamplerColumnConfig
-from data_designer.config.sampler_params import SamplerParamsT, SamplerType
-from data_designer.engine.sampling_gen.data_sources.base import DataSource
-from data_designer.engine.sampling_gen.data_sources.sources import SamplerRegistry
-from data_designer.engine.sampling_gen.jinja_utils import extract_column_names_from_expression
-
-
-class ConditionalDataColumn(SamplerColumnConfig):
-    @property
-    def _negative_condition(self) -> str:
-        conditions = list(self.conditional_params.keys())
-        return "not (" + " or ".join([f"({c})" for c in conditions]) + ")"
-
-    @property
-    def conditions(self) -> list[str]:
-        c = list(self.conditional_params.keys())
-        return c + [self._negative_condition] if len(c) > 0 else ["..."]
-
-    @property
-    def conditional_column_names(self) -> set[str]:
-        names = set()
-        for condition in self.conditional_params.keys():
-            names.update(extract_column_names_from_expression(condition))
-        return names
-
-    @field_serializer("sampler_type")
-    def serialize_sampler_type(self, sampler_type: SamplerType) -> str:
-        return SamplerType(sampler_type).value
-
-    @field_serializer("params")
-    def serialize_params(self, params: SamplerParamsT) -> dict:
-        return params.model_dump()
-
-    @field_serializer("conditional_params")
-    def serialize_conditional_params(self, conditional_params: dict[str, SamplerParamsT]) -> dict:
-        for condition, params in conditional_params.items():
-            conditional_params[condition] = params.model_dump()
-        return conditional_params
-
-    @model_validator(mode="before")
-    @classmethod
-    def validate_params_with_type(cls, data: Any) -> Any:
-        if not isinstance(data, dict) or "sampler_type" not in data:
-            return data
-        if isinstance(data["sampler_type"], str):
-            if not SamplerRegistry.is_registered(data["sampler_type"]):
-                raise ValueError(
-                    f"Invalid sampler type: {data['sampler_type']}. Available samplers: {[s.value for s in SamplerType]}"
-                )
-        if "params" in data:
-            data["params"] = SamplerRegistry.get_sampler(data["sampler_type"])(params=data["params"]).params
-        if "conditional_params" in data:
-            for condition, params in data["conditional_params"].items():
-                data["conditional_params"][condition] = SamplerRegistry.get_sampler(data["sampler_type"])(
-                    params=params
-                ).params
-        return data
-
-    @model_validator(mode="after")
-    def validate_params(self) -> Self:
-        self.params = SamplerRegistry.validate_sampler_type(self.sampler_type)(params=self.params).params
-        return self
-
-    @model_validator(mode="after")
-    def validate_data_conversion(self) -> Self:
-        self.get_default_sampler().validate_data_conversion(self.convert_to)
-        return self
-
-    @model_validator(mode="after")
-    def validate_conditional_params(self) -> Self:
-        for condition, params in self.conditional_params.items():
-            self.conditional_params[condition] = SamplerRegistry.get_sampler(self.sampler_type)(params=params).params
-        return self
-
-    def get_default_sampler(self, **kwargs) -> DataSource:
-        return self.get_sampler("...", **kwargs)
-
-    def get_sampler(self, condition: str, **kwargs) -> DataSource:
-        if condition in ["...", self._negative_condition]:
-            return SamplerRegistry.get_sampler(self.sampler_type)(self.params, **kwargs)
-        return SamplerRegistry.get_sampler(self.sampler_type)(self.conditional_params[condition], **kwargs)