ignore if not needed

Author: johnnygreco

src/data_designer/engine/column_generators/generators/samplers.py

@@ -3,6 +3,7 @@
 
 from functools import partial
 import logging
+from pathlib import Path
 import random
 from typing import Callable
 
@@ -16,8 +17,8 @@
 )
 from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig
 from data_designer.engine.processing.utils import concat_datasets
+from data_designer.engine.resources.errors import ManagedAssetMissingError, ManagedAssetsPathNotSetError
 from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator
-from data_designer.engine.resources.resource_provider import ResourceType
 from data_designer.engine.sampling_gen.data_sources.sources import SamplerType
 from data_designer.engine.sampling_gen.entities.person import load_person_data_sampler
 from data_designer.engine.sampling_gen.generator import DatasetGenerator as SamplingDatasetGenerator
@@ -32,7 +33,7 @@ def metadata() -> GeneratorMetadata:
             name="sampler_column_generator",
             description="Generate columns using sampling-based method.",
             generation_strategy=GenerationStrategy.FULL_COLUMN,
-            required_resources=[ResourceType.BLOB_STORAGE],
+            required_resources=None,
         )
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:
@@ -52,7 +53,32 @@ def _needs_person_generator(self) -> bool:
     def _person_generator_loader(self) -> Callable[[bool], ManagedDatasetGenerator]:
         return partial(load_person_data_sampler, blob_storage=self.resource_provider.blob_storage)
 
+    def _check_managed_assets_exist_if_needed(self) -> None:
+        if self._needs_person_generator:
+            if (
+                self.resource_provider.blob_storage is None
+                or not self.resource_provider.blob_storage.root_path.exists()
+            ):
+                raise ManagedAssetsPathNotSetError(
+                    "🛑 The managed assets path does not exist. If you are using the Person Sampler, "
+                    "You must have a managed assets directory that contains the Nemotron-Personas dataset "
+                    "for each locale you want to sample from."
+                )
+            is_missing = []
+            for c in [c for c in self.config.columns if c.sampler_type == SamplerType.PERSON]:
+                locale_file_path = self.resource_provider.blob_storage.root_path / f"datasets/{c.params.locale}.parquet"
+                if not Path(locale_file_path).exists() or not Path(locale_file_path).is_file():
+                    is_missing.append([c.params.locale, locale_file_path])
+            if len(is_missing) > 0:
+                raise ManagedAssetMissingError(
+                    "🛑 The Nemotron-Personas dataset is missing for the following locales: "
+                    f"{', '.join([f'{locale}' for locale, _ in is_missing])}. "
+                    "Please ensure the files exist at the following paths: "
+                    f"{', '.join([f'{str(file_path)!r}' for _, file_path in is_missing])}"
+                )
+
     def _create_sampling_dataset_generator(self) -> SamplingDatasetGenerator:
+        self._check_managed_assets_exist_if_needed()
         return SamplingDatasetGenerator(
             sampler_columns=self.config,
             person_generator_loader=(self._person_generator_loader if self._needs_person_generator else None),

src/data_designer/engine/resources/errors.py

@@ -0,0 +1,13 @@
+from data_designer.engine.errors import DataDesignerError
+
+
+class ResourceError(DataDesignerError):
+    """Base exception for all errors related to resources."""
+
+
+class ManagedAssetsPathNotSetError(ResourceError):
+    """Exception for all errors related to the managed assets path not being set."""
+
+
+class ManagedAssetMissingError(ResourceError):
+    """Exception for all errors related to the managed assets missing."""

src/data_designer/engine/resources/managed_storage.py

@@ -43,15 +43,15 @@ class LocalBlobStorageProvider(ManagedBlobStorage):
     """
 
     def __init__(self, root_path: Path) -> None:
-        self._root_path = root_path
+        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}"
+        return f"{self.root_path}/{key}"
 
 
 def init_managed_blob_storage(assets_storage: str) -> ManagedBlobStorage:

src/data_designer/interface/data_designer.py

@@ -339,7 +339,9 @@ def _create_resource_provider(
                 model_provider_registry=self._model_provider_registry,
                 secret_resolver=self._secret_resolver,
             ),
-            blob_storage=init_managed_blob_storage(str(self._resolve_managed_assets_path())),
+            blob_storage=None
+            if not self._managed_assets_path.exists()
+            else init_managed_blob_storage(str(self._managed_assets_path)),
             datastore=(
                 LocalSeedDatasetDataStore()
                 if (settings := config_builder.get_seed_datastore_settings()) is None
@@ -349,9 +351,3 @@ def _create_resource_provider(
                 )
             ),
         )
-
-    def _resolve_managed_assets_path(self) -> Path:
-        if not self._managed_assets_path.exists():
-            logger.info(f"🏗️ Creating your managed assets path at: {str(self._managed_assets_path)!r}.")
-            self._managed_assets_path.mkdir(parents=True, exist_ok=True)
-        return self._managed_assets_path

tests/engine/resources/test_managed_storage.py

@@ -45,7 +45,7 @@ def test_uri_for_key_normalization(stub_concrete_storage, test_key, expected_uri
 )
 def test_local_blob_storage_provider_init(test_case, root_path):
     provider = LocalBlobStorageProvider(root_path)
-    assert provider._root_path == root_path
+    assert provider.root_path == root_path
 
 
 @pytest.mark.parametrize(