data-designer 0.1.0__py3-none-any.whl

data_designer/engine/resources/managed_dataset_repository.py

@@ -0,0 +1,194 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from functools import cached_property
+import logging
+from pathlib import Path
+import tempfile
+import threading
+import time
+from typing import Any
+
+import duckdb
+import pandas as pd
+
+from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS
+from data_designer.engine.resources.managed_storage import LocalBlobStorageProvider, ManagedBlobStorage
+
+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

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from contextlib import contextmanager
+import logging
+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

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.models import ModelConfig
+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.registry import ModelRegistry, create_model_registry
+from data_designer.engine.resources.managed_storage import ManagedBlobStorage, init_managed_blob_storage
+from data_designer.engine.resources.seed_dataset_data_store import SeedDatasetDataStore
+from data_designer.engine.secret_resolver import SecretResolver
+
+
+class ResourceType(StrEnum):
+    BLOB_STORAGE = "blob_storage"
+    DATASTORE = "datastore"
+    MODEL_REGISTRY = "model_registry"
+
+
+class ResourceProvider(ConfigBase):
+    artifact_storage: ArtifactStorage
+    blob_storage: ManagedBlobStorage | None = None
+    datastore: SeedDatasetDataStore | None = None
+    model_registry: ModelRegistry | None = None
+
+
+def create_resource_provider(
+    *,
+    artifact_storage: ArtifactStorage,
+    model_configs: list[ModelConfig],
+    secret_resolver: SecretResolver,
+    model_provider_registry: ModelProviderRegistry,
+    datastore: SeedDatasetDataStore | None = None,
+    blob_storage: ManagedBlobStorage | None = None,
+) -> ResourceProvider:
+    return ResourceProvider(
+        artifact_storage=artifact_storage,
+        datastore=datastore,
+        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(),
+    )

data_designer/engine/resources/seed_dataset_data_store.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+
+import duckdb
+from huggingface_hub import HfApi, HfFileSystem
+
+from data_designer.logging import quiet_noisy_logger
+
+quiet_noisy_logger("httpx")
+
+_HF_DATASETS_PREFIX = "hf://datasets/"
+
+
+class MalformedFileIdError(Exception):
+    """Raised when file_id format is invalid."""
+
+
+class SeedDatasetDataStore(ABC):
+    """Abstract base class for dataset storage implementations."""
+
+    @abstractmethod
+    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection: ...
+
+    @abstractmethod
+    def get_dataset_uri(self, file_id: str) -> str: ...
+
+
+class LocalSeedDatasetDataStore(SeedDatasetDataStore):
+    """Local filesystem-based dataset storage."""
+
+    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
+        return duckdb.connect()
+
+    def get_dataset_uri(self, file_id: str) -> str:
+        return file_id
+
+
+class HfHubSeedDatasetDataStore(SeedDatasetDataStore):
+    """Hugging Face and Data Store dataset storage."""
+
+    def __init__(self, endpoint: str, token: str | None):
+        self.hfapi = HfApi(endpoint=endpoint, token=token)
+        self.hffs = HfFileSystem(endpoint=endpoint, token=token)
+
+    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
+        conn = duckdb.connect()
+        conn.register_filesystem(self.hffs)
+        return conn
+
+    def get_dataset_uri(self, file_id: str) -> str:
+        identifier = file_id.removeprefix(_HF_DATASETS_PREFIX)
+        repo_id, filename = self._get_repo_id_and_filename(identifier)
+        return f"{_HF_DATASETS_PREFIX}{repo_id}/{filename}"
+
+    def _get_repo_id_and_filename(self, identifier: str) -> tuple[str, str]:
+        """Extract repo_id and filename from identifier."""
+        parts = identifier.split("/", 2)
+        if len(parts) < 3:
+            raise MalformedFileIdError(
+                "Could not extract repo id and filename from file_id, "
+                "expected 'hf://datasets/{repo-namespace}/{repo-name}/{filename}'"
+            )
+        repo_ns, repo_name, filename = parts
+        return f"{repo_ns}/{repo_name}", filename

data_designer/engine/sampling_gen/column.py

@@ -0,0 +1,89 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+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)

data_designer/engine/sampling_gen/constraints.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from typing import Type
+
+import numpy as np
+from numpy.typing import NDArray
+import pandas as pd
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.sampler_constraints import (
+    ColumnInequalityConstraint,
+    Constraint,
+    ConstraintType,
+    InequalityOperator,
+    ScalarInequalityConstraint,
+)
+
+
+class ConstraintChecker(ConfigBase, ABC):
+    constraint: Constraint
+
+    def get_required_column_names(self) -> tuple[str, ...]:
+        return (self.constraint.target_column,)
+
+    @abstractmethod
+    def check(self, dataframe: pd.DataFrame) -> NDArray[np.bool_]: ...
+
+
+class WithCompareMixin:
+    @property
+    def lhs(self) -> str:
+        return self.constraint.target_column
+
+    def compare(self, lhs: float | int | NDArray, rhs: float | int | NDArray) -> bool | NDArray[np.bool_]:
+        operator = {
+            InequalityOperator.LT: np.less,
+            InequalityOperator.LE: np.less_equal,
+            InequalityOperator.GT: np.greater,
+            InequalityOperator.GE: np.greater_equal,
+        }[InequalityOperator(self.constraint.operator)]
+        return operator(lhs, rhs)
+
+
+class ScalarInequalityChecker(ConstraintChecker, WithCompareMixin):
+    """Compare a column to a scalar value.
+
+    Args:
+        column_name: Name of the constrained column. Will be
+            used as the left-hand side (lhs) of the comparison.
+        operator: Comparison operator.
+        rhs: Scalar value to compare against.
+    """
+
+    constraint: ScalarInequalityConstraint
+
+    def check(self, dataframe: pd.DataFrame) -> NDArray[np.bool_]:
+        return self.compare(dataframe[self.lhs].values, self.constraint.rhs)
+
+
+class ColumnInequalityChecker(ConstraintChecker, WithCompareMixin):
+    """Compare the values of two columns.
+
+    Args:
+        column_name: Name of the constrained column. Will be
+            used as the left-hand side (lhs) of the comparison.
+        operator: Comparison operator.
+        rhs: Name of the column to compare against.
+    """
+
+    constraint: ColumnInequalityConstraint
+
+    def get_required_column_names(self) -> tuple[str, ...]:
+        """Return the names of columns required for the constraint.
+
+        Note that order matters. Edges in the DAG are created as column_names[1], column_names[0].
+        """
+        return (self.lhs, self.constraint.rhs)
+
+    def check(self, dataframe: pd.DataFrame) -> NDArray[np.bool_]:
+        return self.compare(
+            dataframe[self.lhs].values,
+            dataframe[self.constraint.rhs].values,
+        )
+
+
+CONSTRAINT_TYPE_TO_CHECKER = {
+    ConstraintType.SCALAR_INEQUALITY: ScalarInequalityChecker,
+    ConstraintType.COLUMN_INEQUALITY: ColumnInequalityChecker,
+}
+
+
+def get_constraint_checker(constraint_type: ConstraintType) -> Type[ConstraintChecker]:
+    return CONSTRAINT_TYPE_TO_CHECKER[ConstraintType(constraint_type)]