migrate dataset to resource (#420)

* migrate dataset to resource

* remove auto discovery

* remove listing of providers's datasets

* fix after rebase

---------

Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>

llama_stack/providers/adapters/datasetio/huggingface/huggingface.py

@@ -3,7 +3,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from typing import List, Optional
+from typing import Optional
 
 from llama_stack.apis.datasetio import *  # noqa: F403
 
@@ -15,7 +15,7 @@ from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_u
 from .config import HuggingfaceDatasetIOConfig
 
 
-def load_hf_dataset(dataset_def: DatasetDef):
+def load_hf_dataset(dataset_def: Dataset):
     if dataset_def.metadata.get("path", None):
         return hf_datasets.load_dataset(**dataset_def.metadata)
 
@@ -41,13 +41,10 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
 
     async def register_dataset(
         self,
-        dataset_def: DatasetDef,
+        dataset_def: Dataset,
     ) -> None:
         self.dataset_infos[dataset_def.identifier] = dataset_def
 
-    async def list_datasets(self) -> List[DatasetDef]:
-        return list(self.dataset_infos.values())
-
     async def get_rows_paginated(
         self,
         dataset_id: str,

llama_stack/providers/datatypes.py

@@ -11,7 +11,7 @@ from urllib.parse import urlparse
 from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
-from llama_stack.apis.datasets import DatasetDef
+from llama_stack.apis.datasets import Dataset
 from llama_stack.apis.eval_tasks import EvalTaskDef
 from llama_stack.apis.memory_banks.memory_banks import MemoryBank
 from llama_stack.apis.models import Model
@@ -57,9 +57,7 @@ class MemoryBanksProtocolPrivate(Protocol):
 
 
 class DatasetsProtocolPrivate(Protocol):
-    async def list_datasets(self) -> List[DatasetDef]: ...
-
-    async def register_dataset(self, dataset_def: DatasetDef) -> None: ...
+    async def register_dataset(self, dataset: Dataset) -> None: ...
 
 
 class ScoringFunctionsProtocolPrivate(Protocol):

llama_stack/providers/inline/datasetio/localfs/datasetio.py

@@ -3,7 +3,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from typing import List, Optional
+from typing import Optional
 
 import pandas
 from llama_models.llama3.api.datatypes import *  # noqa: F403
@@ -37,12 +37,12 @@ class BaseDataset(ABC):
 
 @dataclass
 class DatasetInfo:
-    dataset_def: DatasetDef
+    dataset_def: Dataset
     dataset_impl: BaseDataset
 
 
 class PandasDataframeDataset(BaseDataset):
-    def __init__(self, dataset_def: DatasetDef, *args, **kwargs) -> None:
+    def __init__(self, dataset_def: Dataset, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.dataset_def = dataset_def
         self.df = None
@@ -60,9 +60,9 @@ class PandasDataframeDataset(BaseDataset):
 
     def _validate_dataset_schema(self, df) -> pandas.DataFrame:
         # note that we will drop any columns in dataset that are not in the schema
-        df = df[self.dataset_def.dataset_schema.keys()]
+        df = df[self.dataset_def.schema.keys()]
         # check all columns in dataset schema are present
-        assert len(df.columns) == len(self.dataset_def.dataset_schema)
+        assert len(df.columns) == len(self.dataset_def.schema)
         # TODO: type checking against column types in dataset schema
         return df
 
@@ -89,17 +89,14 @@ class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
 
     async def register_dataset(
         self,
-        dataset_def: DatasetDef,
+        dataset: Dataset,
     ) -> None:
-        dataset_impl = PandasDataframeDataset(dataset_def)
-        self.dataset_infos[dataset_def.identifier] = DatasetInfo(
-            dataset_def=dataset_def,
+        dataset_impl = PandasDataframeDataset(dataset)
+        self.dataset_infos[dataset.identifier] = DatasetInfo(
+            dataset_def=dataset,
             dataset_impl=dataset_impl,
         )
 
-    async def list_datasets(self) -> List[DatasetDef]:
-        return [i.dataset_def for i in self.dataset_infos.values()]
-
     async def get_rows_paginated(
         self,
         dataset_id: str,

llama_stack/providers/tests/datasetio/test_datasetio.py

@@ -55,15 +55,11 @@ async def register_dataset(
             "generated_answer": StringType(),
         }
 
-    dataset = DatasetDefWithProvider(
-        identifier=dataset_id,
-        provider_id="",
-        url=URL(
-            uri=test_url,
-        ),
-        dataset_schema=dataset_schema,
+    await datasets_impl.register_dataset(
+        dataset_id=dataset_id,
+        schema=dataset_schema,
+        url=URL(uri=test_url),
     )
-    await datasets_impl.register_dataset(dataset)
 
 
 class TestDatasetIO: