add datasetio to distribution

This commit is contained in:
raspawar 2025-03-26 15:21:38 +05:30 committed by Rashmi Pawar
parent ae973c9595
commit 1e77873a02
7 changed files with 37 additions and 35 deletions

View file

@ -24,11 +24,6 @@ class NvidiaDatasetIOConfig(BaseModel):
description="The NVIDIA dataset namespace.",
)
access_policies: Optional[dict] = Field(
default_factory=lambda: os.getenv("NVIDIA_ACCESS_POLICIES", {}),
description="The NVIDIA access policies.",
)
project_id: Optional[str] = Field(
default_factory=lambda: os.getenv("NVIDIA_PROJECT_ID", "test-project"),
description="The NVIDIA project ID.",
@ -46,8 +41,6 @@ class NvidiaDatasetIOConfig(BaseModel):
default_values.append("project_id='test-project'")
if os.getenv("NVIDIA_DATASET_NAMESPACE") is None:
default_values.append("dataset_namespace='default'")
if os.getenv("NVIDIA_ACCESS_POLICIES") is None:
default_values.append("access_policies='{}'")
if os.getenv("NVIDIA_DATASETS_URL") is None:
default_values.append("datasets_url='http://nemo.test'")
@ -64,8 +57,5 @@ class NvidiaDatasetIOConfig(BaseModel):
"api_key": "${env.NVIDIA_API_KEY:}",
"user_id": "${env.NVIDIA_USER_ID:llama-stack-user}",
"dataset_namespace": "${env.NVIDIA_DATASET_NAMESPACE:default}",
"access_policies": "${env.NVIDIA_ACCESS_POLICIES:}",
"project_id": "${env.NVIDIA_PROJECT_ID:test-project}",
"customizer_url": "${env.NVIDIA_CUSTOMIZER_URL:}",
"output_model_dir": "${env.NVIDIA_OUTPUT_MODEL_DIR:test-example-model@v1}",
}

View file

@ -4,14 +4,13 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Any, Dict, Literal, Optional
from typing import Any, Dict, List, Optional
import aiohttp
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.common.type_system import ParamType
from llama_stack.apis.datasets.datasets import Dataset, ListDatasetsResponse
from llama_stack.apis.resource import ResourceType
from llama_stack.apis.datasetio import IterrowsResponse
from llama_stack.schema_utils import webmethod
from .config import NvidiaDatasetIOConfig
@ -20,8 +19,6 @@ from .config import NvidiaDatasetIOConfig
class NvidiaDatasetIOAdapter:
"""Nvidia NeMo DatasetIO API."""
type: Literal[ResourceType.dataset.value] = ResourceType.dataset.value
def __init__(self, config: NvidiaDatasetIOConfig):
self.config = config
self.headers = {}
@ -80,19 +77,6 @@ class NvidiaDatasetIOAdapter:
"""
...
@webmethod(route="/datasets/{dataset_id:namespace}", method="GET")
async def get_dataset(
self,
dataset_id: str,
) -> Optional[Dataset]:
raise NotImplementedError("Not implemented")
@webmethod(route="/datasets", method="GET")
async def list_datasets(
self,
) -> ListDatasetsResponse:
raise NotImplementedError("Not implemented")
@webmethod(route="/datasets/{dataset_id:path}", method="POST")
async def update_dataset(
self,
@ -112,3 +96,14 @@ class NvidiaDatasetIOAdapter:
namespace: Optional[str] = "default",
) -> None:
raise NotImplementedError("Not implemented")
async def iterrows(
self,
dataset_id: str,
start_index: Optional[int] = None,
limit: Optional[int] = None,
) -> IterrowsResponse:
raise NotImplementedError("Not implemented")
async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:
raise NotImplementedError("Not implemented")