move nvidia to remote datasetio

llama_stack/providers/registry/datasetio.py

@@ -36,4 +36,15 @@ def available_providers() -> List[ProviderSpec]:
                 config_class="llama_stack.providers.remote.datasetio.huggingface.HuggingfaceDatasetIOConfig",
             ),
         ),
+        remote_provider_spec(
+            api=Api.datasetio,
+            adapter=AdapterSpec(
+                adapter_type="nvidia",
+                pip_packages=[
+                    "datasets",
+                ],
+                module="llama_stack.providers.remote.datasetio.nvidia",
+                config_class="llama_stack.providers.remote.datasetio.nvidia.NvidiaDatasetIOConfig",
+            ),
+        ),
     ]

llama_stack/providers/registry/datasets.py

@@ -1,30 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# 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
-
-from llama_stack.providers.datatypes import (
-    AdapterSpec,
-    Api,
-    ProviderSpec,
-    remote_provider_spec,
-)
-
-
-def available_providers() -> List[ProviderSpec]:
-    return [
-        remote_provider_spec(
-            api=Api.datasets,
-            adapter=AdapterSpec(
-                adapter_type="nvidia",
-                pip_packages=[
-                    "datasets",
-                ],
-                module="llama_stack.providers.remote.datasets.nvidia",
-                config_class="llama_stack.providers.remote.datasets.nvidia.NvidiaDatasetConfig",
-            ),
-        ),
-    ]

llama_stack/providers/remote/datasetio/nvidia/__init__.py

@@ -4,20 +4,20 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from .config import NvidiaDatasetConfig
+from .config import NvidiaDatasetIOConfig
 
 
 async def get_adapter_impl(
-    config: NvidiaDatasetConfig,
+    config: NvidiaDatasetIOConfig,
     _deps,
 ):
-    from .datasets import NvidiaDatasetAdapter
+    from .datasetio import NvidiaDatasetIOAdapter
 
-    if not isinstance(config, NvidiaDatasetConfig):
+    if not isinstance(config, NvidiaDatasetIOConfig):
         raise RuntimeError(f"Unexpected config type: {type(config)}")
 
-    impl = NvidiaDatasetAdapter(config)
+    impl = NvidiaDatasetIOAdapter(config)
     return impl
 
 
-__all__ = ["get_adapter_impl", "NvidiaDatasetAdapter"]
+__all__ = ["get_adapter_impl", "NvidiaDatasetIOAdapter"]

llama_stack/providers/remote/datasetio/nvidia/config.py

@@ -11,8 +11,8 @@ from typing import Any, Dict, Optional
 from pydantic import BaseModel, Field
 
 
-class NvidiaDatasetConfig(BaseModel):
-    """Configuration for NVIDIA Dataset implementation."""
+class NvidiaDatasetIOConfig(BaseModel):
+    """Configuration for NVIDIA DatasetIO implementation."""
 
     api_key: Optional[str] = Field(
         default_factory=lambda: os.getenv("NVIDIA_API_KEY"),

llama_stack/providers/remote/datasetio/nvidia/datasetio.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from datetime import datetime
 from typing import Any, Dict, Literal, Optional
 
 import aiohttp
@@ -15,15 +14,15 @@ from llama_stack.apis.datasets.datasets import Dataset, ListDatasetsResponse
 from llama_stack.apis.resource import ResourceType
 from llama_stack.schema_utils import webmethod
 
-from .config import NvidiaDatasetConfig
+from .config import NvidiaDatasetIOConfig
 
 
-class NvidiaDatasetAdapter:
-    """Nvidia NeMo Dataset API."""
+class NvidiaDatasetIOAdapter:
+    """Nvidia NeMo DatasetIO API."""
 
     type: Literal[ResourceType.dataset.value] = ResourceType.dataset.value
 
-    def __init__(self, config: NvidiaDatasetConfig):
+    def __init__(self, config: NvidiaDatasetIOConfig):
         self.config = config
         self.headers = {}
         if config.api_key:
@@ -86,48 +85,13 @@ class NvidiaDatasetAdapter:
         self,
         dataset_id: str,
     ) -> Optional[Dataset]:
-        dataset_id, namespace = dataset_id.split(":")
-        dataset = await self._make_request(
-            method="GET",
-            path=f"/v1/datasets/{namespace}/{dataset_id}",
-        )
-        created_at = datetime.fromisoformat(dataset.pop("created_at")) if "created_at" in dataset else datetime.now()
-        identifier = dataset.pop("name")
-        url = URL(uri=dataset.pop("files_url"))
-        return Dataset(
-            identifier=identifier,
-            provider_id="nvidia",  # confirm this
-            url=url,
-            dataset_schema={},  # ToDo: get schema from the dataset
-            created_at=created_at,
-            metadata=dataset,
-        )
+        raise NotImplementedError("Not implemented")
 
     @webmethod(route="/datasets", method="GET")
     async def list_datasets(
         self,
     ) -> ListDatasetsResponse:
-        ## ToDo: add pagination
-        response = await self._make_request(method="GET", path="/v1/datasets")
-        datasets = []
-        for dataset in response["data"]:
-            created_at = (
-                datetime.fromisoformat(dataset.pop("created_at")) if "created_at" in dataset else datetime.now()
-            )
-            identifier = dataset.pop("name")
-            url = URL(uri=dataset.pop("files_url"))
-            datasets.append(
-                Dataset(
-                    identifier=identifier,
-                    provider_id="nvidia",  # confirm this
-                    url=url,
-                    dataset_schema={},
-                    created_at=created_at,
-                    metadata=dataset,
-                )
-            )  # add remaining fields as metadata
-
-        return ListDatasetsResponse(data=datasets)
+        raise NotImplementedError("Not implemented")
 
     @webmethod(route="/datasets/{dataset_id:path}", method="POST")
     async def update_dataset(
@@ -138,11 +102,13 @@ class NvidiaDatasetAdapter:
         provider_dataset_id: Optional[str] = None,
         provider_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
-    ) -> None: ...
+    ) -> None:
+        raise NotImplementedError("Not implemented")
 
     @webmethod(route="/datasets/{dataset_id:path}", method="DELETE")
     async def unregister_dataset(
         self,
         dataset_id: str,
         namespace: Optional[str] = "default",
-    ) -> None: ...
+    ) -> None:
+        raise NotImplementedError("Not implemented")