diff --git a/e2e/files/test_files.py b/e2e/files/test_files.py index 121097d710..9db3b2bdc3 100644 --- a/e2e/files/test_files.py +++ b/e2e/files/test_files.py @@ -11,7 +11,7 @@ import pytest from nemo_platform import NeMoPlatform -from nemo_platform.types.files import Fileset +from nemo_platform_plugin.files.types import FilesetOutput as Fileset @pytest.fixture diff --git a/packages/data_designer_nemo/src/data_designer_nemo/fileset_file_seed_reader.py b/packages/data_designer_nemo/src/data_designer_nemo/fileset_file_seed_reader.py index 068cc70469..239768a91a 100644 --- a/packages/data_designer_nemo/src/data_designer_nemo/fileset_file_seed_reader.py +++ b/packages/data_designer_nemo/src/data_designer_nemo/fileset_file_seed_reader.py @@ -26,7 +26,7 @@ def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection: if self._sdk is None: raise RuntimeError("FilesetFileSeedReader requires an injected NeMo Platform SDK") - filesystem = FilesetFileSystem(self._sdk) + filesystem = FilesetFileSystem(sdk=self._sdk) conn = duckdb.connect() conn.register_filesystem(filesystem) diff --git a/packages/data_designer_nemo/src/data_designer_nemo/person_reader.py b/packages/data_designer_nemo/src/data_designer_nemo/person_reader.py index 330e7545ce..d7ce8d0aab 100644 --- a/packages/data_designer_nemo/src/data_designer_nemo/person_reader.py +++ b/packages/data_designer_nemo/src/data_designer_nemo/person_reader.py @@ -36,7 +36,7 @@ def __init__(self, sdk: NeMoPlatform | AsyncNeMoPlatform): self._sdk = sdk def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection: - filesystem = FilesetFileSystem(self._sdk) + filesystem = FilesetFileSystem(sdk=self._sdk) conn = duckdb.connect() conn.register_filesystem(filesystem) return conn diff --git a/packages/filesets/src/filesets/filesystem/filesystem.py b/packages/filesets/src/filesets/filesystem/filesystem.py index e7e62eab1d..27f01cc2ce 100644 --- a/packages/filesets/src/filesets/filesystem/filesystem.py +++ b/packages/filesets/src/filesets/filesystem/filesystem.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""FilesetFileSystem - fsspec filesystem for NeMo Platform filesets.""" +"""FilesetFileSystem - fsspec filesystem for NeMo Platform fileset storage.""" from __future__ import annotations @@ -18,13 +18,8 @@ from fsspec.callbacks import DEFAULT_CALLBACK, Callback from fsspec.spec import AbstractBufferedFile from nemo_platform import AsyncNeMoPlatform, NeMoPlatform -from nemo_platform.types.files import FilesetFile as SDKFilesetFile - -# Conditional import for TestClient detection -try: - from starlette.testclient import TestClient -except ImportError: - TestClient = None +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.types import FilesetFileOutput, ListFilesQueryParams T = TypeVar("T") @@ -105,6 +100,18 @@ async def run_one(coro: Coroutine[Any, Any, T], idx: int) -> None: fsspec.asyn._run_coros_in_chunks = run_coros_in_chunks +def _detect_async_transport(sync_client: Any) -> httpx.AsyncBaseTransport | None: + """Detect if a sync httpx client wraps a TestClient and return ASGITransport.""" + try: + from starlette.testclient import TestClient + + if isinstance(sync_client, TestClient): + return httpx.ASGITransport(app=sync_client.app) + except ImportError: + pass + return None + + class FileInfo(TypedDict): """File or directory info returned by fsspec methods.""" @@ -299,30 +306,31 @@ def build_fileset_ref( class FilesetFileSystem(AsyncFileSystem): """ - fsspec filesystem for NeMo Platform filesets. + fsspec filesystem for NeMo Platform fileset storage. URL format: fileset://[workspace/]fileset_name[#path] The optional `#` separator distinguishes the fileset name from the file path. If omitted, assumes root of fileset. Workspace is optional - if omitted, - uses the SDK's default workspace. + uses the client's default workspace. Examples: - >>> from nemo_platform import NeMoPlatform - >>> sdk = NeMoPlatform(base_url="http://localhost:8000", workspace="default") - >>> fs = FilesetFileSystem(sdk=sdk) - >>> fs.ls("my-fileset") # root of fileset, workspace from SDK default + >>> from nemo_platform_plugin.files.client import AsyncFilesClient + >>> client = AsyncFilesClient(base_url="http://localhost:8000", workspace="default") + >>> fs = FilesetFileSystem(client=client) + >>> fs.ls("my-fileset") # root of fileset, workspace from client default >>> fs.ls("my-fileset#data/") # specific path within fileset >>> fs.ls("default/my-fileset#data/") # explicit workspace """ protocol = "fileset" + _client: AsyncFilesClient @classmethod def register_fsspec(cls) -> None: """Register the fileset protocol with fsspec. - After calling this, you can use fsspec.filesystem("fileset", sdk=sdk). + After calling this, you can use fsspec.filesystem("fileset", client=client). """ from fsspec import register_implementation @@ -341,60 +349,81 @@ def register_fsspec(cls) -> None: def __init__( self, - sdk: NeMoPlatform | AsyncNeMoPlatform, + *, + client: FilesClient | AsyncFilesClient | None = None, + sdk: NeMoPlatform | AsyncNeMoPlatform | None = None, batch_size: int | None = None, blocksize: int | None = None, **kwargs, ): + if client is None and sdk is None: + raise TypeError("Either 'client' or 'sdk' must be provided") + + # Normalize: convert sdk to a FilesClient so there's one code path. + # AsyncNeMoPlatform → AsyncFilesClient (already async, _ensure_async is a no-op). + # NeMoPlatform → FilesClient (sync, _ensure_async converts to async). + if sdk is not None: + from nemo_platform_plugin.client.adapter import client_from_platform + + if isinstance(sdk, AsyncNeMoPlatform): + client = client_from_platform(sdk, AsyncFilesClient) + else: + client = client_from_platform(sdk, FilesClient) + + async_client = self._ensure_async(client) + is_async = isinstance(client, AsyncFilesClient) + if batch_size is None: batch_size = self.default_batch_size if blocksize is None: blocksize = self.blocksize - # Set asynchronous mode based on SDK type. When asynchronous=False, - # fsspec creates a global daemon event loop (self.loop) that callers - # can use for sync-to-async bridging via fsspec.asyn.sync(). - is_async_sdk = isinstance(sdk, AsyncNeMoPlatform) - super().__init__(asynchronous=is_async_sdk, batch_size=batch_size, blocksize=blocksize, **kwargs) - self._sdk: AsyncNeMoPlatform = self._get_sdk(sdk) + super().__init__(asynchronous=is_async, batch_size=batch_size, blocksize=blocksize, **kwargs) + self._client = async_client - def _get_sdk( - self, - sdk: NeMoPlatform | AsyncNeMoPlatform, - ) -> AsyncNeMoPlatform: - # If already an async SDK, use it as-is to preserve custom transports (e.g., test clients) - if isinstance(sdk, AsyncNeMoPlatform): - return sdk - - # Convert sync SDK to async SDK - transport: httpx.AsyncBaseTransport | None = None - if TestClient is not None and isinstance(sdk._client, TestClient): - # If using a synchronous test client, we should use the ASGITransport - transport = httpx.ASGITransport(app=sdk._client.app) - - return AsyncNeMoPlatform( - workspace=sdk.workspace, - base_url=sdk.base_url, - timeout=sdk.timeout, - max_retries=sdk.max_retries, - default_headers=sdk._custom_headers, - default_query=sdk.default_query, + @staticmethod + def _ensure_async(client: FilesClient | AsyncFilesClient) -> AsyncFilesClient: + """Ensure we have an AsyncFilesClient, converting from sync if needed. + + Preserves subclass behavior: if the sync client has ``_async_cls`` + (e.g. a remapping subclass), that class is used for the async client. + """ + if isinstance(client, AsyncFilesClient): + return client + + import httpx + + # Use _async_cls if the sync client defines one (e.g. _RemappingFilesClient + # → _RemappingAsyncFilesClient), otherwise plain AsyncFilesClient. + async_cls = getattr(client, "_async_cls", None) or AsyncFilesClient + + transport = _detect_async_transport(client._http) + return async_cls( + base_url=client.base_url, + workspace=client.workspace, + auth=client._auth, + default_headers=client._default_headers or None, + retry=client._retry, http_client=httpx.AsyncClient( transport=transport, - base_url=sdk.base_url, - headers=sdk._custom_headers, + base_url=client.base_url, + headers=dict(client._default_headers) if client._default_headers else None, ), ) - def to_fileset_files(self, results: dict[str, Any]) -> list[SDKFilesetFile]: - """Convert fsspec find results to FilesetFile objects. + @property + def _workspace(self) -> str | None: + return self._client.workspace + + def to_fileset_files(self, results: dict[str, Any]) -> list[FilesetFileOutput]: + """Convert fsspec find results to FilesetFileOutput objects. Args: results: Dict from find(detail=True) mapping paths to file info. Returns: - List of FilesetFile objects with path, size, and file_ref. + List of FilesetFileOutput objects with path, size, and file_ref. """ files = [] for name, info in results.items(): @@ -402,7 +431,7 @@ def to_fileset_files(self, results: dict[str, Any]) -> list[SDKFilesetFile]: continue workspace, fileset, file_path = parse_fileset_ref(name, workspace_fallback=None) files.append( - SDKFilesetFile( + FilesetFileOutput( file_ref=f"{workspace}/{fileset}#{file_path}", file_url=f"/apis/files/v2/workspaces/{workspace}/filesets/{fileset}/-/{file_path}", path=file_path, @@ -485,7 +514,7 @@ async def _info(self, path: str, **kwargs) -> FileInfo: Checks dircache first to avoid redundant API calls. For cache misses, uses _ls which populates the cache for all directory levels. """ - _, _, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + _, _, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) path_key = build_fileset_ref(path) parent_path = self._parent(path_key) @@ -527,17 +556,15 @@ async def _info(self, path: str, **kwargs) -> FileInfo: async def _cat_file(self, path: str, start: int | None = None, end: int | None = None, **kwargs) -> bytes: """Fetch file content with optional byte range.""" - workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) if not file_path: raise IsADirectoryError(path) - extra_headers = {} + client = self._client if start is not None or end is not None: - extra_headers["Range"] = f"bytes={start or 0}-{(end - 1) if end else ''}" + client = client.with_headers({"Range": f"bytes={start or 0}-{(end - 1) if end else ''}"}) - response = await self._sdk.files._download_file( - file_path, workspace=workspace, name=fileset, extra_headers=extra_headers or None - ) + response = await client.download_file(workspace=workspace, name=fileset, path=file_path) return await response.read() @classmethod @@ -569,7 +596,7 @@ async def _ls(self, path: str, detail: bool = True, refresh: bool = False, **kwa detail: If True, return list of dicts. If False, return list of paths. refresh: If True, bypass cache and fetch fresh listing. """ - workspace, fileset, prefix = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, prefix = parse_fileset_ref(path, workspace_fallback=self._workspace) prefix = prefix.rstrip("/") path_key = build_fileset_ref(prefix, workspace=workspace, fileset=fileset) @@ -582,7 +609,13 @@ async def _ls(self, path: str, detail: bool = True, refresh: bool = False, **kwa pass # Fetch from backend and populate cache for all directory levels - response = await self._sdk.files._list_files(fileset, workspace=workspace, path=prefix or None) + query_params: ListFilesQueryParams | None = {"path": prefix} if prefix else None + response = await self._client.list_files( + workspace=workspace, + name=fileset, + query_params=query_params, + ) + response = response.data() dir_contents = self._populate_dircache_from_response(response, workspace, fileset, prefix) # Return the listing for the requested path @@ -591,19 +624,19 @@ async def _ls(self, path: str, detail: bool = True, refresh: bool = False, **kwa async def _rm_file(self, path: str, **kwargs) -> None: """Delete a single file.""" - workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) if not file_path: raise ValueError("Cannot delete fileset root via rm") - await self._sdk.files._delete_file(file_path, workspace=workspace, name=fileset) + await self._client.delete_file(workspace=workspace, name=fileset, path=file_path) # Invalidate parent directory's cache since file info is stored there self.invalidate_cache(self._parent(build_fileset_ref(path))) async def _pipe_file(self, path: str, value: bytes, **kwargs) -> None: """Write bytes to a file.""" - workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) if not file_path: raise ValueError("File path required for upload") - await self._sdk.files._upload_file(file_path, body=value, workspace=workspace, name=fileset) + await self._client.upload_file(workspace=workspace, name=fileset, path=file_path, content=value) # Invalidate parent directory's cache since file info is stored there self.invalidate_cache(self._parent(build_fileset_ref(path))) @@ -627,7 +660,7 @@ async def _pipe_stream( content_length: Optional content length for Content-Length header. If not provided, uses chunked transfer encoding. """ - workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) if not file_path: raise ValueError("File path required for upload") @@ -635,15 +668,11 @@ async def _pipe_stream( if not hasattr(stream, "__anext__"): stream = to_async_iterator(stream) - extra_headers = {"Content-Length": str(content_length)} if content_length is not None else None + client = self._client + if content_length is not None: + client = client.with_headers({"Content-Length": str(content_length)}) - await self._sdk.files._upload_file( - path=file_path, - body=stream, - workspace=workspace, - name=fileset, - extra_headers=extra_headers, - ) + await client.upload_file(workspace=workspace, name=fileset, path=file_path, content=stream) # Invalidate parent directory's cache since file info is stored there self.invalidate_cache(self._parent(build_fileset_ref(path))) @@ -663,7 +692,7 @@ async def _put_file(self, lpath: str, rpath: str, callback: Callback = DEFAULT_C Uses streaming upload to avoid buffering the entire file in memory. Supports per-chunk progress via callback.relative_update(chunk_size). """ - workspace, fileset, file_path = parse_fileset_ref(rpath, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(rpath, workspace_fallback=self._workspace) if not file_path: raise ValueError("File path required for upload") @@ -672,18 +701,17 @@ async def _put_file(self, lpath: str, rpath: str, callback: Callback = DEFAULT_C callback.set_size(file_size) # Create async generator that streams file content with progress - async def stream_file(): + async def stream_file() -> AsyncIterator[bytes]: async with await anyio.open_file(lpath, "rb") as f: while chunk := await f.read(self.blocksize): callback.relative_update(len(chunk)) yield chunk - await self._sdk.files._upload_file( - path=file_path, - body=stream_file(), + await self._client.with_headers({"Content-Length": str(file_size)}).upload_file( workspace=workspace, name=fileset, - extra_headers={"Content-Length": str(file_size)}, + path=file_path, + content=stream_file(), ) # Invalidate parent directory's cache since file info is stored there self.invalidate_cache(self._parent(build_fileset_ref(rpath))) @@ -705,9 +733,15 @@ async def _find( Also populates the dircache so subsequent _ls calls benefit. """ - workspace, fileset, prefix = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, prefix = parse_fileset_ref(path, workspace_fallback=self._workspace) prefix = prefix.rstrip("/") - response = await self._sdk.files._list_files(fileset, workspace=workspace, path=prefix or None) + query_params: ListFilesQueryParams | None = {"path": prefix} if prefix else None + response = await self._client.list_files( + workspace=workspace, + name=fileset, + query_params=query_params, + ) + response = response.data() # Populate dircache for all directory levels (benefits subsequent _ls calls) self._populate_dircache_from_response(response, workspace, fileset, prefix) @@ -718,7 +752,7 @@ async def _find( # Add root path if withdirs requested if withdirs: - root_path = build_fileset_ref(path, workspace=self._sdk.workspace) + root_path = build_fileset_ref(path, workspace=self._workspace) out[root_path] = {"name": root_path, "size": 0, "type": "directory"} for file_info in response.data: @@ -748,28 +782,27 @@ async def _find( async def _get_file(self, rpath: str, lpath: str, callback: Callback = DEFAULT_CALLBACK, **kwargs) -> None: """Download a file to local path. - Uses with_streaming_response to avoid buffering the entire response in memory. - Uses http_response.aiter_raw() for maximum throughput (bypasses httpx chunking overhead). + Uses streaming response to avoid buffering the entire response in memory. Supports per-chunk progress via callback.relative_update(chunk_size). """ - workspace, fileset, file_path = parse_fileset_ref(rpath, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(rpath, workspace_fallback=self._workspace) if not file_path: return - # Use with_streaming_response to not buffer the data in memory. - async with self._sdk.files.with_streaming_response._download_file( - file_path, workspace=workspace, name=fileset - ) as response: - # Set callback size from Content-Length if available - content_length = response.headers.get("content-length") + response = await self._client.download_file( + workspace=workspace, + name=fileset, + path=file_path, + ) + + async with response.stream() as chunks: + content_length = response.http_response.headers.get("content-length") if content_length: callback.set_size(int(content_length)) - await anyio.Path(lpath).parent.mkdir(parents=True, exist_ok=True) async with await anyio.open_file(lpath, "wb") as f: - # Use aiter_raw() instead of iter_bytes() to bypass httpx chunking overhead. - async for chunk in response.http_response.aiter_raw(self.blocksize): + async for chunk in chunks: await f.write(chunk) callback.relative_update(len(chunk)) @@ -810,7 +843,7 @@ async def _get( return # Normalize rpath to new format for comparison (since _find returns new format paths) - rpath_normalized = build_fileset_ref(rpath, workspace=self._sdk.workspace).rstrip("/") + rpath_normalized = build_fileset_ref(rpath, workspace=self._workspace).rstrip("/") lpath_stripped = lpath.rstrip("/") source_is_file = len(source_files) == 1 and self._strip_protocol(source_files[0]) == rpath_normalized @@ -842,7 +875,7 @@ async def _get( # SPECIAL CASE: Fileset root (workspace/fileset with no file path) always # copies contents directly, matching HuggingFace Hub behavior. Users who want # to preserve the fileset name can include it in local_path. - _, _, file_path = parse_fileset_ref(rpath, workspace_fallback=self._sdk.workspace) + _, _, file_path = parse_fileset_ref(rpath, workspace_fallback=self._workspace) copy_contents_directly = rpath.endswith("/") or not file_path # Extract directory name from the file path portion (e.g., "subdir" from "a/b/subdir") diff --git a/packages/filesets/src/filesets/resources.py b/packages/filesets/src/filesets/resources.py index 02bbafe463..154f600b47 100644 --- a/packages/filesets/src/filesets/resources.py +++ b/packages/filesets/src/filesets/resources.py @@ -1,27 +1,35 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Extended FilesResource classes with FilesetFileSystem support. +"""FilesResource classes with FilesetFileSystem support. -These classes extend the SDK's generated FilesResource classes to add -high-level file operations (upload, download, list, delete) and fsspec -filesystem access. +These classes provide high-level file operations (upload, download, list, delete) +backed by the NemoClient typed HTTP client and fsspec filesystem access. """ import uuid from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass +from functools import cached_property from pathlib import PurePath -from typing import Protocol, runtime_checkable +from typing import Any, Protocol, runtime_checkable +import nemo_platform from fsspec.callbacks import Callback from fsspec.core import has_magic -from nemo_platform import ConflictError -from nemo_platform._compat import cached_property -from nemo_platform.resources.files import AsyncFilesResource as BaseAsyncFilesResource -from nemo_platform.resources.files import FilesResource as BaseFilesResource -from nemo_platform.types.files import CacheStatus, FilesetFile -from nemo_platform.types.files.fileset import Fileset +from nemo_platform_plugin.client.errors import NemoHTTPError +from nemo_platform_plugin.client.response import AsyncNemoPaginatedResponse, NemoPaginatedResponse +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.types import ( + CacheStatus, + CreateFilesetRequest, + FilesetFileOutput, + FilesetMetadata, + FilesetOutput, + FilesetPurpose, + StorageConfig, + UpdateFilesetRequest, +) from filesets.filesystem.filesystem import ( FilesetFileSystem, @@ -30,6 +38,85 @@ ) +def _build_error_map() -> dict[type[NemoHTTPError], type[nemo_platform.APIStatusError]]: + """Build a mapping from NemoClient errors to Stainless SDK errors. + + Lazy import to avoid hard-coding the Stainless error classes at module level. + This mapping is temporary — remove when all consumers import errors from + nemo_platform_plugin.client.errors instead of nemo_platform (AIRCORE-840). + """ + from nemo_platform_plugin.client import errors + + return { + errors.BadRequestError: nemo_platform.BadRequestError, + errors.AuthenticationError: nemo_platform.AuthenticationError, + errors.PermissionDeniedError: nemo_platform.PermissionDeniedError, + errors.NotFoundError: nemo_platform.NotFoundError, + errors.ConflictError: nemo_platform.ConflictError, + errors.UnprocessableEntityError: nemo_platform.UnprocessableEntityError, + errors.RateLimitError: nemo_platform.RateLimitError, + errors.InternalServerError: nemo_platform.InternalServerError, + } + + +_ERROR_MAP: dict[type[NemoHTTPError], type[nemo_platform.APIStatusError]] | None = None + + +def _get_error_map() -> dict[type[NemoHTTPError], type[nemo_platform.APIStatusError]]: + global _ERROR_MAP + if _ERROR_MAP is None: + _ERROR_MAP = _build_error_map() + return _ERROR_MAP + + +def _raise_as_stainless(e: NemoHTTPError) -> None: + """Re-raise a NemoClient error as its Stainless SDK equivalent. + + Preserves backward compatibility for consumers that catch + ``nemo_platform.NotFoundError`` etc. Remove with AIRCORE-840. + """ + error_map = _get_error_map() + stainless_cls = error_map.get(type(e)) + if stainless_cls is not None: + raise stainless_cls( + message=str(e), + response=e.http_response, + body=e.body, + ) from e + raise + + +class _RemappingFilesClient(FilesClient): + """FilesClient that re-raises NemoClient errors as Stainless SDK errors. + + Wraps ``send()`` so ALL operations through this client (filesets, files, + fsspec) raise Stainless-compatible exceptions. Remove with AIRCORE-840. + """ + + # Used by FilesetFileSystem._ensure_async to create the matching async + # remapping client when converting sync → async. + _async_cls: type[AsyncFilesClient] | None = None + + def send(self, request, *, headers=None, retry=None): # type: ignore[override] + try: + return super().send(request, headers=headers, retry=retry) + except NemoHTTPError as e: + _raise_as_stainless(e) + + +class _RemappingAsyncFilesClient(AsyncFilesClient): + """AsyncFilesClient that re-raises NemoClient errors as Stainless SDK errors.""" + + async def send(self, request, *, headers=None, retry=None): # type: ignore[override] + try: + return await super().send(request, headers=headers, retry=retry) + except NemoHTTPError as e: + _raise_as_stainless(e) + + +_RemappingFilesClient._async_cls = _RemappingAsyncFilesClient + + @dataclass class ListFilesResponse: """Response from listing files in a fileset. @@ -46,7 +133,7 @@ class ListFilesResponse: - None if no cache information is available """ - data: list[FilesetFile] + data: list[FilesetFileOutput] @property def cache_status(self) -> CacheStatus | None: @@ -129,24 +216,250 @@ def _matches_glob(filepath: str, pattern: str) -> bool: return PurePath(filepath).match(pattern) -class FilesResource(BaseFilesResource): - """Extended FilesResource with high-level file operations. +class FilesetsSubResource: + """Fileset CRUD operations (create, retrieve, update, list, delete). + + Wraps ``FilesClient`` methods with higher-level convenience signatures + (unwrapped params, ``exist_ok`` support). + + .. deprecated:: + Temporary shim for the ``sdk.files`` fileset interface. + New code should use ``FilesClient`` directly. + Once all callers are migrated, this class will be removed. + """ + + def __init__(self, client: FilesClient) -> None: + self._client = client + + def create( + self, + *, + name: str, + workspace: str | None = None, + exist_ok: bool = False, + description: str | None = None, + project: str | None = None, + purpose: FilesetPurpose | None = None, + metadata: FilesetMetadata | None = None, + storage: StorageConfig | None = None, + custom_fields: dict[str, Any] | None = None, + cache: bool = False, + ) -> FilesetOutput: + body = CreateFilesetRequest( + name=name, + description=description, + project=project, + purpose=purpose or FilesetPurpose.GENERIC, + metadata=metadata or FilesetMetadata(), + storage=storage, + custom_fields=custom_fields or {}, + cache=cache, + ) + # The server returns an error body on 409, not the entity, so + # exist_ok is handled here with a follow-up GET rather than at + # the endpoint/client level. + try: + return self._client.create_fileset(workspace=workspace, body=body).data() + except nemo_platform.APIStatusError as e: + if e.status_code == 409 and exist_ok: + return self.retrieve(name=name, workspace=workspace) + raise + + def retrieve(self, name: str, *, workspace: str | None = None) -> FilesetOutput: + return self._client.get_fileset(workspace=workspace, name=name).data() + + def update( + self, + name: str, + *, + workspace: str | None = None, + description: str | None = None, + project: str | None = None, + purpose: FilesetPurpose | None = None, + metadata: FilesetMetadata | None = None, + custom_fields: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> FilesetOutput: + # Only include explicitly provided fields so exclude_unset works correctly + kwargs = { + k: v + for k, v in dict( + description=description, + project=project, + purpose=purpose, + metadata=metadata, + custom_fields=custom_fields, + ).items() + if v is not None + } + body = UpdateFilesetRequest(**kwargs) + client = self._client.with_options(timeout=timeout) if timeout is not None else self._client + return client.update_fileset(workspace=workspace, name=name, body=body).data() + + def list( + self, + *, + workspace: str | None = None, + page: int | None = None, + page_size: int | None = None, + sort: str | None = None, + filter: str | dict | None = None, + ) -> NemoPaginatedResponse[FilesetOutput]: + query_params = { + k: v + for k, v in dict( + page=page, + page_size=page_size, + sort=sort, + filter=filter, + ).items() + if v is not None + } + return self._client.list_filesets(workspace=workspace, query_params=query_params or None) + + def delete(self, name: str, *, workspace: str | None = None) -> FilesetOutput: + return self._client.delete_fileset(workspace=workspace, name=name).data() + + +class AsyncFilesetsSubResource: + """Async fileset CRUD operations (create, retrieve, update, list, delete). + + Wraps ``AsyncFilesClient`` methods with higher-level convenience signatures + (unwrapped params, ``exist_ok`` support). + + .. deprecated:: + Temporary shim for the ``sdk.files`` fileset interface. + New code should use ``AsyncFilesClient`` directly. + Once all callers are migrated, this class will be removed. + """ + + def __init__(self, client: AsyncFilesClient) -> None: + self._client = client + + async def create( + self, + *, + name: str, + workspace: str | None = None, + exist_ok: bool = False, + description: str | None = None, + project: str | None = None, + purpose: FilesetPurpose | None = None, + metadata: FilesetMetadata | None = None, + storage: StorageConfig | None = None, + custom_fields: dict[str, Any] | None = None, + cache: bool = False, + ) -> FilesetOutput: + body = CreateFilesetRequest( + name=name, + description=description, + project=project, + purpose=purpose or FilesetPurpose.GENERIC, + metadata=metadata or FilesetMetadata(), + storage=storage, + custom_fields=custom_fields or {}, + cache=cache, + ) + # The server returns an error body on 409, not the entity, so + # exist_ok is handled here with a follow-up GET rather than at + # the endpoint/client level. + try: + return (await self._client.create_fileset(workspace=workspace, body=body)).data() + except nemo_platform.APIStatusError as e: + if e.status_code == 409 and exist_ok: + return await self.retrieve(name=name, workspace=workspace) + raise + + async def retrieve(self, name: str, *, workspace: str | None = None) -> FilesetOutput: + return (await self._client.get_fileset(workspace=workspace, name=name)).data() + + async def update( + self, + name: str, + *, + workspace: str | None = None, + description: str | None = None, + project: str | None = None, + purpose: FilesetPurpose | None = None, + metadata: FilesetMetadata | None = None, + custom_fields: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> FilesetOutput: + kwargs = { + k: v + for k, v in dict( + description=description, + project=project, + purpose=purpose, + metadata=metadata, + custom_fields=custom_fields, + ).items() + if v is not None + } + body = UpdateFilesetRequest(**kwargs) + client = self._client.with_options(timeout=timeout) if timeout is not None else self._client + return (await client.update_fileset(workspace=workspace, name=name, body=body)).data() + + async def list( + self, + *, + workspace: str | None = None, + page: int | None = None, + page_size: int | None = None, + sort: str | None = None, + filter: str | dict | None = None, + ) -> AsyncNemoPaginatedResponse[FilesetOutput]: + query_params = { + k: v + for k, v in dict( + page=page, + page_size=page_size, + sort=sort, + filter=filter, + ).items() + if v is not None + } + return await self._client.list_filesets(workspace=workspace, query_params=query_params or None) + + async def delete(self, name: str, *, workspace: str | None = None) -> FilesetOutput: + return (await self._client.delete_fileset(workspace=workspace, name=name)).data() + + +class FilesResource: + """FilesResource with high-level file operations. Provides convenient methods for uploading, downloading, and listing files. - For fsspec filesystem access, use `sdk.files.fsspec`. + For fsspec filesystem access, use ``resource.fsspec``. """ + def __init__(self, client) -> None: + # _raw_client kept for otlp delegation (Stainless SDK), removed by AIRCORE-840. + self._raw_client = client + + from nemo_platform_plugin.client.adapter import client_from_platform + + self._client = client_from_platform(client, _RemappingFilesClient) + + @cached_property + def filesets(self) -> FilesetsSubResource: + """Access fileset CRUD operations (create, retrieve, update, list, delete).""" + return FilesetsSubResource(self._client) + + @cached_property + def otlp(self): + """Access OTLP log operations (delegated to Stainless SDK resource).""" + from nemo_platform.resources.files.otlp.otlp import OtlpResource + + return OtlpResource(self._raw_client) + @cached_property def fsspec(self) -> FilesetFileSystem: """Access the underlying fsspec filesystem.""" - return FilesetFileSystem(sdk=self._client) + return FilesetFileSystem(client=self._client) def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None: """Create fileset if it doesn't exist (idempotent).""" - try: - self.filesets.create(name=fileset, workspace=workspace) - except ConflictError: - pass # Already exists + self.filesets.create(name=fileset, workspace=workspace, exist_ok=True) def download( self, @@ -170,7 +483,7 @@ def download( local_path: Local destination path (directory). fileset: Fileset name. If not provided, inferred from remote_path (str only). workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. callback: Optional progress callback (e.g., RichProgressCallback). max_workers: Maximum number of concurrent file transfers. @@ -281,7 +594,7 @@ def upload( callback: Callback | None = None, max_workers: int | None = None, fileset_auto_create: bool = False, - ) -> Fileset: + ) -> FilesetOutput: """Upload files from a local path to a fileset. Args: @@ -292,7 +605,7 @@ def upload( Defaults to "" (root of fileset). fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. callback: Optional progress callback (e.g., RichProgressCallback). max_workers: Maximum number of concurrent file transfers. fileset_auto_create: If True, create the fileset if it doesn't exist. @@ -300,7 +613,7 @@ def upload( a unique name is generated (e.g., "fileset-a1b2c3d4"). Returns: - Fileset: The fileset that was uploaded to. Check `fileset.name` to see + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see the generated name when using fileset_auto_create without specifying a fileset. @@ -373,7 +686,7 @@ def upload_content( fileset: str | None = None, workspace: str | None = None, fileset_auto_create: bool = False, - ) -> Fileset: + ) -> FilesetOutput: """Upload in-memory content to a fileset. Args: @@ -384,13 +697,13 @@ def upload_content( - Iterator[bytes]: Generator or iterator yielding byte chunks remote_path: Destination path within the fileset. fileset: Fileset name. If not provided, inferred from remote_path. - workspace: Workspace name. If not provided, uses SDK default. + workspace: Workspace name. If not provided, uses client default. fileset_auto_create: If True, create the fileset if it doesn't exist. When no fileset is specified (neither as param nor in remote_path), a unique name is generated (e.g., "fileset-a1b2c3d4"). Returns: - Fileset: The fileset that was uploaded to. Check `fileset.name` to see + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see the generated name when using fileset_auto_create without specifying a fileset. @@ -478,7 +791,7 @@ def download_content( Args: remote_path: Path of the file within the fileset. fileset: Fileset name. If not provided, inferred from remote_path. - workspace: Workspace name. If not provided, uses SDK default. + workspace: Workspace name. If not provided, uses client default. Returns: bytes: The file content. @@ -532,12 +845,12 @@ def list( Defaults to "" (root of fileset). fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. include_cache_status: Check and return cache status for each file. When False (default), external storage files return None for cache_status. Returns: - ListFilesResponse with data (list of FilesetFile) and cache_status property. + ListFilesResponse with data (list of FilesetFileOutput) and cache_status property. Examples: # List all files in a fileset @@ -585,12 +898,18 @@ def list( # For path prefixes, the API handles filtering server-side api_path = None if has_magic(path) else (path or None) - response = self._list_files( - fileset, + query_params = {} + if api_path is not None: + query_params["path"] = api_path + if include_cache_status: + query_params["include_cache_status"] = True + + response = self._client.list_files( workspace=ws, - include_cache_status=include_cache_status, - path=api_path, + name=fileset, + query_params=query_params or None, ) + response = response.data() files = list(response.data) # Apply glob filtering if needed @@ -613,7 +932,7 @@ def delete( or a relative path (e.g., "data/file.txt") if fileset is provided. fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. Examples: # Delete a file with explicit fileset @@ -638,28 +957,41 @@ def delete( self.fsspec.rm(fileset_ref) -class AsyncFilesResource(BaseAsyncFilesResource): - """Extended AsyncFilesResource with high-level file operations. +class AsyncFilesResource: + """Async FilesResource with high-level file operations. Provides convenient methods for uploading, downloading, and listing files. - For fsspec filesystem access, use `sdk.files.fsspec`. + For fsspec filesystem access, use ``resource.fsspec``. """ + def __init__(self, client) -> None: + # _raw_client kept for otlp delegation (Stainless SDK), removed by AIRCORE-840. + self._raw_client = client + + from nemo_platform_plugin.client.adapter import client_from_platform + + self._client = client_from_platform(client, _RemappingAsyncFilesClient) + @cached_property - def fsspec(self) -> FilesetFileSystem: - """Get a FilesetFileSystem instance pre-configured with this SDK client. + def filesets(self) -> AsyncFilesetsSubResource: + """Access fileset CRUD operations (create, retrieve, update, list, delete).""" + return AsyncFilesetsSubResource(self._client) - This provides fsspec filesystem access. For high-level file - operations, use `sdk.files` instead. - """ - return FilesetFileSystem(sdk=self._client) + @cached_property + def otlp(self): + """Access OTLP log operations (delegated to Stainless SDK resource).""" + from nemo_platform.resources.files.otlp.otlp import AsyncOtlpResource + + return AsyncOtlpResource(self._raw_client) + + @cached_property + def fsspec(self) -> FilesetFileSystem: + """Access the underlying fsspec filesystem.""" + return FilesetFileSystem(client=self._client) async def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None: """Create fileset if it doesn't exist (idempotent).""" - try: - await self.filesets.create(name=fileset, workspace=workspace) - except ConflictError: - pass # Already exists + await self.filesets.create(name=fileset, workspace=workspace, exist_ok=True) async def download( self, @@ -683,7 +1015,7 @@ async def download( local_path: Local destination path (directory). fileset: Fileset name. If not provided, inferred from remote_path (str only). workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. callback: Optional progress callback (e.g., RichProgressCallback). max_workers: Maximum number of concurrent file transfers. @@ -779,7 +1111,7 @@ async def upload( callback: Callback | None = None, max_workers: int | None = None, fileset_auto_create: bool = False, - ) -> Fileset: + ) -> FilesetOutput: """Upload files from a local path to a fileset (async). Args: @@ -790,7 +1122,7 @@ async def upload( Defaults to "" (root of fileset). fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. callback: Optional progress callback (e.g., RichProgressCallback). max_workers: Maximum number of concurrent file transfers. fileset_auto_create: If True, create the fileset if it doesn't exist. @@ -798,7 +1130,7 @@ async def upload( a unique name is generated (e.g., "fileset-a1b2c3d4"). Returns: - Fileset: The fileset that was uploaded to. Check `fileset.name` to see + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see the generated name when using fileset_auto_create without specifying a fileset. @@ -865,7 +1197,7 @@ async def upload_content( fileset: str | None = None, workspace: str | None = None, fileset_auto_create: bool = False, - ) -> Fileset: + ) -> FilesetOutput: """Upload in-memory data to a fileset (async). Args: @@ -876,13 +1208,13 @@ async def upload_content( - AsyncIterator[bytes]: Async iterator yielding byte chunks (streamed) remote_path: Destination path within the fileset. fileset: Fileset name. If not provided, inferred from remote_path. - workspace: Workspace name. If not provided, uses SDK default. + workspace: Workspace name. If not provided, uses client default. fileset_auto_create: If True, create the fileset if it doesn't exist. When no fileset is specified (neither as param nor in remote_path), a unique name is generated (e.g., "fileset-a1b2c3d4"). Returns: - Fileset: The fileset that was uploaded to. Check `fileset.name` to see + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see the generated name when using fileset_auto_create without specifying a fileset. @@ -974,7 +1306,7 @@ async def download_content( Args: remote_path: Path of the file within the fileset. fileset: Fileset name. If not provided, inferred from remote_path. - workspace: Workspace name. If not provided, uses SDK default. + workspace: Workspace name. If not provided, uses client default. Returns: bytes: The file content. @@ -1020,12 +1352,12 @@ async def list( Defaults to "" (root of fileset). fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. include_cache_status: Check and return cache status for each file. When False (default), external storage files return None for cache_status. Returns: - ListFilesResponse with data (list of FilesetFile) and cache_status property. + ListFilesResponse with data (list of FilesetFileOutput) and cache_status property. Examples: # List all files in a fileset @@ -1070,12 +1402,18 @@ async def list( # For path prefixes, the API handles filtering server-side api_path = None if has_magic(path) else (path or None) - response = await self._list_files( - fileset, + query_params = {} + if api_path is not None: + query_params["path"] = api_path + if include_cache_status: + query_params["include_cache_status"] = True + + response = await self._client.list_files( workspace=ws, - include_cache_status=include_cache_status, - path=api_path, + name=fileset, + query_params=query_params or None, ) + response = response.data() files = list(response.data) # Apply glob filtering if needed @@ -1098,7 +1436,7 @@ async def delete( or a relative path (e.g., "data/file.txt") if fileset is provided. fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. Examples: # Delete a file with explicit fileset diff --git a/packages/nemo_platform/pyproject.toml b/packages/nemo_platform/pyproject.toml index 0762dd69f5..57adab6ccc 100644 --- a/packages/nemo_platform/pyproject.toml +++ b/packages/nemo_platform/pyproject.toml @@ -311,6 +311,7 @@ nemo-guardrails-plugin = [ nemo-platform-plugin = [ "anthropic>=0.88.0", "fastapi>=0.115.4", + "jsonschema>=4.0.0", "lark>=1.1.0", "nemo-platform-sdk", "openai>=1.109.1", diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/formatters.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/formatters.py index bc2d85b0ae..80efdb2564 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/formatters.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/formatters.py @@ -62,16 +62,18 @@ def display_name(self) -> str: "{namespace}/{name}", # Computed OpenAI path field } -# Common field names for list items in API responses -LIST_ITEM_FIELDS = ["data"] +# Common field names for list items in API responses. +# Checked in order via getattr. Callables (e.g. NemoPaginatedResponse.items) +# are invoked; plain attributes/properties are returned directly. +LIST_ITEM_FIELDS = ["data", "items"] def _extract_items_from_response(data: Any) -> list[Any]: """Extract list items from an API response. Handles various response formats: - - Paginated responses with .data attribute - - Non-paginated responses with .files, .items, etc. + - Paginated responses with .items() method (NemoPaginatedResponse) + - Paginated responses with .data attribute (legacy Page) - Plain lists - Dict responses @@ -81,17 +83,19 @@ def _extract_items_from_response(data: Any) -> list[Any]: Returns: List of items from the response """ - # Try common field names for list items - for field in LIST_ITEM_FIELDS: - if hasattr(data, field): - return getattr(data, field) - if isinstance(data, dict) and field in data: - return data[field] - - # Handle plain lists if isinstance(data, list): return data + for field in LIST_ITEM_FIELDS: + if isinstance(data, dict): + if field in data: + return data[field] + else: + value = getattr(data, field, None) + if value is None: + continue + return list(value()) if callable(value) else value + return [] diff --git a/packages/nemo_platform_plugin/pyproject.toml b/packages/nemo_platform_plugin/pyproject.toml index 886e672b5d..fb2b6c66d8 100644 --- a/packages/nemo_platform_plugin/pyproject.toml +++ b/packages/nemo_platform_plugin/pyproject.toml @@ -20,6 +20,7 @@ classifiers = [ dependencies = [ "anthropic>=0.88.0", "fastapi>=0.115.4", + "jsonschema>=4.0.0", "lark>=1.1.0", "nemo-platform-sdk", "openai>=1.109.1", diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py index da6cf6092c..e62fb3443c 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py @@ -40,9 +40,19 @@ def client_from_platform( The overloads ensure callers get the correct concrete return type. """ + # Prefer _custom_headers (set via with_options/set_default_headers), + # fall back to the httpx client's actual headers (set at construction, + # e.g. TestClient(headers={...})), filtering out httpx defaults. + # _custom_headers and _client are private Stainless SDK attrs present on both + # NeMoPlatform and AsyncNeMoPlatform but not visible to the type checker. + headers = platform._custom_headers # type: ignore[union-attr] + if not headers: + _skip = {"accept", "accept-encoding", "connection", "user-agent", "host"} + headers = {k: v for k, v in platform._client.headers.items() if k.lower() not in _skip} # type: ignore[union-attr] + return client_cls( base_url=str(platform.base_url).rstrip("/"), workspace=platform.workspace, - default_headers=platform._custom_headers, # type: ignore[arg-type] + default_headers=headers or None, http_client=platform._client, # type: ignore[arg-type] ) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py index a2d34d403f..f7d3dc63fe 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py @@ -17,12 +17,13 @@ from __future__ import annotations import asyncio +import copy import inspect import json import time from collections.abc import Mapping from pathlib import Path -from typing import Any, TypeVar, get_args, get_origin, overload +from typing import Any, Self, TypeVar, get_args, get_origin, overload import httpx from nemo_platform_plugin.client.auth import ( @@ -124,6 +125,7 @@ def __init__( self._auth: TokenProvider | None = StaticToken(auth) if isinstance(auth, str) else auth self._retry = retry self._default_headers = dict(default_headers) if default_headers else {} + self._timeout: float | None = None @property def base_url(self) -> str: @@ -179,6 +181,39 @@ def _is_stream(self, request: PreparedRequest) -> bool: def _is_paginated(self, request: PreparedRequest) -> bool: return get_origin(request.response_type) is Paginated + def with_options( + self, + *, + headers: Mapping[str, str] | None = None, + retry: RetryPolicy | None = None, + timeout: float | None = None, + ) -> Self: + """Return a copy of this client with the given options merged in. + + The returned client shares the underlying HTTP transport, so it is + cheap to create. Useful for one-off header, retry, or timeout + overrides when calling ``method()``-bound endpoints:: + + client.with_headers({"Range": "bytes=0-99"}).download_file(...) + client.with_options(timeout=300).update_fileset(...) + """ + clone = copy.copy(self) + if headers: + clone._default_headers = {**self._default_headers, **headers} + if retry is not None: + clone._retry = retry + if timeout is not None: + clone._timeout = timeout + return clone + + def with_headers(self, headers: Mapping[str, str]) -> Self: + """Shorthand for ``with_options(headers=...)``.""" + return self.with_options(headers=headers) + + def with_retry(self, retry: RetryPolicy) -> Self: + """Shorthand for ``with_options(retry=...)``.""" + return self.with_options(retry=retry) + def _resolve_query_params(self, request: PreparedRequest) -> dict[str, str | int | bool] | None: """Filter out None values and JSON-serialize dicts/lists in query params.""" if request.query_params is None: @@ -216,6 +251,18 @@ def __init__( timeout=timeout, ) + @classmethod + def from_client(cls, client: NemoClient) -> Self: + """Create an instance of this subclass sharing the transport of *client*.""" + return cls( + base_url=client.base_url, + workspace=client.workspace, + auth=client._auth, + default_headers=client._default_headers or None, + retry=client._retry, + http_client=client._http, + ) + @overload def send( self, @@ -354,7 +401,10 @@ def _request_with_retry( last_response: httpx.Response | None = None for attempt in range(retry.max_retries + 1 if retry else 1): try: - raw = self._http.request(request.method, url, content=request.content, headers=headers, params=params) + kwargs: dict = {"content": request.content, "headers": headers, "params": params} + if self._timeout is not None: + kwargs["timeout"] = self._timeout + raw = self._http.request(request.method, url, **kwargs) except httpx.TransportError as exc: backoff = _should_retry(None, exc, attempt, retry) if retry else None if backoff is not None: @@ -413,6 +463,18 @@ def __init__( timeout=timeout, ) + @classmethod + def from_client(cls, client: AsyncNemoClient) -> Self: + """Create an instance of this subclass sharing the transport of *client*.""" + return cls( + base_url=client.base_url, + workspace=client.workspace, + auth=client._auth, + default_headers=client._default_headers or None, + retry=client._retry, + http_client=client._http, + ) + @overload async def send( self, @@ -543,9 +605,10 @@ async def _request_with_retry( last_response: httpx.Response | None = None for attempt in range(retry.max_retries + 1 if retry else 1): try: - raw = await self._http.request( - request.method, url, content=request.content, headers=headers, params=params - ) + kwargs: dict = {"content": request.content, "headers": headers, "params": params} + if self._timeout is not None: + kwargs["timeout"] = self._timeout + raw = await self._http.request(request.method, url, **kwargs) except httpx.TransportError as exc: backoff = _should_retry(None, exc, attempt, retry) if retry else None if backoff is not None: diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py index 7b0d89e048..0bdc675709 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py @@ -84,8 +84,16 @@ class NemoBinaryResponse: def __init__(self, stream_ctx: AbstractContextManager[httpx.Response], request: PreparedRequest) -> None: self._stream_ctx = stream_ctx + self._http_response: httpx.Response | None = None self.request = request + @property + def http_response(self) -> httpx.Response: + """The underlying httpx response. Available after entering ``stream()``.""" + if self._http_response is None: + raise RuntimeError("http_response is only available inside a stream() context") + return self._http_response + def read(self) -> bytes: """Read and return the entire response body as bytes.""" with self._stream_ctx as raw: @@ -94,11 +102,25 @@ def read(self) -> bytes: return data @contextmanager - def stream(self) -> Iterator[Iterator[bytes]]: - """Yield an iterator of byte chunks.""" + def stream(self, chunk_size: int | None = None) -> Iterator[Iterator[bytes]]: + """Yield an iterator of raw byte chunks. + + Args: + chunk_size: Maximum number of bytes per chunk. If None, uses the + transport's default chunking. + + The underlying httpx response is available as ``http_response`` + after entering the context, e.g. for reading ``Content-Length``:: + + with resp.stream() as chunks: + size = resp.http_response.headers.get("content-length") + for chunk in chunks: + ... + """ with self._stream_ctx as raw: + self._http_response = raw raise_for_status(raw) - yield raw.iter_bytes() + yield raw.iter_raw(chunk_size) if chunk_size else raw.iter_raw() class NemoStreamResponse(Generic[ModelT]): @@ -163,8 +185,16 @@ class AsyncNemoBinaryResponse: def __init__(self, stream_ctx: AbstractAsyncContextManager[httpx.Response], request: PreparedRequest) -> None: self._stream_ctx = stream_ctx + self._http_response: httpx.Response | None = None self.request = request + @property + def http_response(self) -> httpx.Response: + """The underlying httpx response. Available after entering ``stream()``.""" + if self._http_response is None: + raise RuntimeError("http_response is only available inside a stream() context") + return self._http_response + async def read(self) -> bytes: """Read and return the entire response body as bytes.""" async with self._stream_ctx as raw: @@ -173,11 +203,25 @@ async def read(self) -> bytes: return data @asynccontextmanager - async def stream(self) -> AsyncIterator[AsyncIterator[bytes]]: - """Yield an async iterator of byte chunks.""" + async def stream(self, chunk_size: int | None = None) -> AsyncIterator[AsyncIterator[bytes]]: + """Yield an async iterator of raw byte chunks. + + Args: + chunk_size: Maximum number of bytes per chunk. If None, uses the + transport's default chunking. + + The underlying httpx response is available as ``http_response`` + after entering the context, e.g. for reading ``Content-Length``:: + + async with resp.stream() as chunks: + size = resp.http_response.headers.get("content-length") + async for chunk in chunks: + ... + """ async with self._stream_ctx as raw: + self._http_response = raw raise_for_status(raw) - yield raw.aiter_bytes() + yield raw.aiter_raw(chunk_size) if chunk_size else raw.aiter_raw() class AsyncNemoStreamResponse(Generic[ModelT]): @@ -233,11 +277,11 @@ async def _iter() -> AsyncIterator[ModelT]: class PageResult(Generic[ModelT]): """A single page of results with pagination metadata. - Returned by :meth:`NemoPaginatedResponse.data` for callers who want + Returned by :meth:`NemoPaginatedResponse.page` for callers who want one page at a time rather than auto-iterating all pages:: resp = client.send(list_items()) - page = resp.data() + page = resp.page() print(f"Page {page.page} of {page.total_pages} ({page.total_results} total)") for item in page.items: print(item.name) @@ -251,18 +295,23 @@ class PageResult(Generic[ModelT]): class NemoPaginatedResponse(Generic[ModelT]): - """Sync iterable over all items across paginated API responses. + """Sync paginated API response. - Lazily fetches subsequent pages using the pagination strategy configured - on the endpoint's ``Paginated[T, Strategy]`` return type. Iterating - yields individual ``ModelT`` items, not page envelopes:: + Provides two iteration modes:: - for item in client.send(list_items()): + # Iterate all items across all pages + for item in response.items(): print(item.name) - For single-page access with metadata, use :meth:`data`:: + # Iterate page by page with metadata + for page in response.pages(): + print(f"Page {page.page}/{page.total_pages}") + for item in page.items: + process(item) - page = client.send(list_items()).data() + For single-page access, use :meth:`page`:: + + page = response.page() print(f"{page.total_results} total across {page.total_pages} pages") """ @@ -291,13 +340,14 @@ def _parse_page(self, raw: httpx.Response) -> tuple[list[ModelT], dict]: items = [self._model_type.model_validate(item) for item in self._strategy.extract_items(body)] return items, body - def data(self) -> PageResult[ModelT]: + def page(self) -> PageResult[ModelT]: """Return the first page as a :class:`PageResult` with metadata.""" items, body = self._parse_page(self._first_response) metadata = self._strategy.extract_metadata(body) return PageResult(items=items, **metadata) - def __iter__(self) -> Iterator[ModelT]: + def items(self) -> Iterator[ModelT]: + """Iterate all items across all pages, fetching subsequent pages lazily.""" items, body = self._parse_page(self._first_response) yield from items @@ -308,14 +358,31 @@ def __iter__(self) -> Iterator[ModelT]: current = next_page next_page = self._strategy.next_page(body, current) + def pages(self) -> Iterator[PageResult[ModelT]]: + """Iterate page by page, yielding :class:`PageResult` objects with metadata.""" + items, body = self._parse_page(self._first_response) + metadata = self._strategy.extract_metadata(body) + yield PageResult(items=items, **metadata) + + next_page = self._strategy.next_page(body, 1) + while next_page is not None: + items, body = self._parse_page(self._fetch_page(self.request, next_page)) + metadata = self._strategy.extract_metadata(body) + yield PageResult(items=items, **metadata) + current = next_page + next_page = self._strategy.next_page(body, current) + class AsyncNemoPaginatedResponse(Generic[ModelT]): - """Async iterable over all items across paginated API responses. + """Async paginated API response. Async twin of :class:`NemoPaginatedResponse`:: - async for item in await client.send(list_items()): + async for item in response.items(): print(item.name) + + async for page in response.pages(): + print(f"Page {page.page}/{page.total_pages}") """ def __init__( @@ -343,13 +410,14 @@ def _parse_page(self, raw: httpx.Response) -> tuple[list[ModelT], dict]: items = [self._model_type.model_validate(item) for item in self._strategy.extract_items(body)] return items, body - def data(self) -> PageResult[ModelT]: + def page(self) -> PageResult[ModelT]: """Return the first page as a :class:`PageResult` with metadata.""" items, body = self._parse_page(self._first_response) metadata = self._strategy.extract_metadata(body) return PageResult(items=items, **metadata) - async def __aiter__(self) -> AsyncIterator[ModelT]: + async def items(self) -> AsyncIterator[ModelT]: + """Iterate all items across all pages, fetching subsequent pages lazily.""" items, body = self._parse_page(self._first_response) for item in items: yield item @@ -362,3 +430,18 @@ async def __aiter__(self) -> AsyncIterator[ModelT]: yield item current = next_page next_page = self._strategy.next_page(body, current) + + async def pages(self) -> AsyncIterator[PageResult[ModelT]]: + """Iterate page by page, yielding :class:`PageResult` objects with metadata.""" + items, body = self._parse_page(self._first_response) + metadata = self._strategy.extract_metadata(body) + yield PageResult(items=items, **metadata) + + next_page = self._strategy.next_page(body, 1) + while next_page is not None: + raw = await self._fetch_page(self.request, next_page) + items, body = self._parse_page(raw) + metadata = self._strategy.extract_metadata(body) + yield PageResult(items=items, **metadata) + current = next_page + next_page = self._strategy.next_page(body, current) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/client.py new file mode 100644 index 0000000000..656b08a79e --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/client.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed HTTP clients for the Files service. + +Wraps the endpoint functions from ``files.endpoints`` as direct methods +using the ``method()`` descriptor, following the example-plugin pattern. + +Usage:: + + client = FilesClient(base_url="...", workspace="default") + resp = client.create_fileset(body=CreateFilesetRequest(name="my-fs")) + fileset = resp.data() + + files_resp = client.upload_file(name="my-fs", path="data.txt", content=b"hello") +""" + +from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.method import method +from nemo_platform_plugin.files import endpoints + + +class _FilesMethods: + create_fileset = method(endpoints.create_fileset) + list_filesets = method(endpoints.list_filesets) + get_fileset = method(endpoints.get_fileset) + update_fileset = method(endpoints.update_fileset) + delete_fileset = method(endpoints.delete_fileset) + list_files = method(endpoints.list_files) + upload_file = method(endpoints.upload_file) + download_file = method(endpoints.download_file) + delete_file = method(endpoints.delete_file) + + +class FilesClient(_FilesMethods, NemoClient): + """Sync client for the Files service API.""" + + +class AsyncFilesClient(_FilesMethods, AsyncNemoClient): + """Async client for the Files service API.""" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/endpoints.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/endpoints.py new file mode 100644 index 0000000000..d6cd16a041 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/endpoints.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed endpoint definitions for the Files service. + +These are the single source of truth for the HTTP contract. +""" + +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import AsyncIterable, Iterable + +from nemo_platform_plugin.client.endpoint import delete, get, patch, post, put +from nemo_platform_plugin.client.types import BinaryContent, Paginated +from nemo_platform_plugin.files.types import ( + CreateFilesetRequest, + FilesetFileOutput, + FilesetOutput, + ListFilesetFilesResponse, + ListFilesetsQueryParams, + ListFilesQueryParams, + UpdateFilesetRequest, +) + +# --------------------------------------------------------------------------- +# Fileset CRUD +# --------------------------------------------------------------------------- + + +@post("/apis/files/v2/workspaces/{workspace}/filesets") +@abstractmethod +def create_fileset(*, workspace: str | None = None, body: CreateFilesetRequest) -> FilesetOutput: ... + + +@get("/apis/files/v2/workspaces/{workspace}/filesets") +@abstractmethod +def list_filesets( + *, workspace: str | None = None, query_params: ListFilesetsQueryParams | None = None +) -> Paginated[FilesetOutput]: ... + + +@get("/apis/files/v2/workspaces/{workspace}/filesets/{name}") +@abstractmethod +def get_fileset(*, workspace: str | None = None, name: str) -> FilesetOutput: ... + + +@patch("/apis/files/v2/workspaces/{workspace}/filesets/{name}") +@abstractmethod +def update_fileset(*, workspace: str | None = None, name: str, body: UpdateFilesetRequest) -> FilesetOutput: ... + + +@delete("/apis/files/v2/workspaces/{workspace}/filesets/{name}") +@abstractmethod +def delete_fileset(*, workspace: str | None = None, name: str) -> FilesetOutput: ... + + +# --------------------------------------------------------------------------- +# File operations +# --------------------------------------------------------------------------- + + +@get("/apis/files/v2/workspaces/{workspace}/filesets/{name}/files") +@abstractmethod +def list_files( + *, workspace: str | None = None, name: str, query_params: ListFilesQueryParams | None = None +) -> ListFilesetFilesResponse: ... + + +@put("/apis/files/v2/workspaces/{workspace}/filesets/{name}/-/{path}") +@abstractmethod +def upload_file( + *, workspace: str | None = None, name: str, path: str, content: bytes | Iterable[bytes] | AsyncIterable[bytes] +) -> FilesetFileOutput: ... + + +@get("/apis/files/v2/workspaces/{workspace}/filesets/{name}/-/{path}") +@abstractmethod +def download_file(*, workspace: str | None = None, name: str, path: str) -> BinaryContent: ... + + +@delete("/apis/files/v2/workspaces/{workspace}/filesets/{name}/-/{path}") +@abstractmethod +def delete_file(*, workspace: str | None = None, name: str, path: str) -> FilesetFileOutput: ... diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/metadata.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/metadata.py new file mode 100644 index 0000000000..028beac46e --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/metadata.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Metadata types for filesets. + +The metadata uses a tagged/keyed structure where the key indicates the type: + metadata: {dataset: {schema: {...}}} + +The key in metadata should match the fileset's purpose field. +""" + +from jsonschema.exceptions import SchemaError +from jsonschema.validators import validator_for +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class DatasetMetadataContent(BaseModel): + """Content for dataset-type filesets.""" + + # Use `schema_` because `schema` is a BaseModel method. + model_config = ConfigDict(serialize_by_alias=True) + + schema_: dict | str | None = Field( + default=None, + alias="schema", + description="Default row schema for files in this fileset, either inline JSON Schema or a schema_defs key.", + ) + schema_defs: dict[str, dict] = Field( + default_factory=dict, + description="Reusable JSON Schema definitions keyed by name for deduplicating per-file dataset schemas.", + ) + schemas_by_path: dict[str, dict | str] = Field( + default_factory=dict, + description=( + "Optional per-file row schemas keyed by relative path within the fileset. " + "Each value may be inline JSON Schema or a schema_defs key." + ), + ) + + @model_validator(mode="after") + def validate_schema_refs(self) -> "DatasetMetadataContent": + for ref_name, ref_value in [("schema", self.schema_), *self.schemas_by_path.items()]: + if isinstance(ref_value, str) and ref_value not in self.schema_defs: + raise ValueError(f"dataset metadata reference '{ref_name}' points to unknown schema_def '{ref_value}'") + return self + + @model_validator(mode="after") + def validate_json_schemas(self) -> "DatasetMetadataContent": + def _validate_schema_document(schema: dict, ref_name: str) -> None: + validator = validator_for(schema) + try: + validator.check_schema(schema) + except SchemaError as e: + raise ValueError( + f"dataset metadata field '{ref_name}' contains invalid JSON Schema: {e.message}" + ) from e + + if isinstance(self.schema_, dict): + _validate_schema_document(self.schema_, "schema") + + for schema_name, schema in self.schema_defs.items(): + _validate_schema_document(schema, f"schema_defs.{schema_name}") + + for path, schema in self.schemas_by_path.items(): + if isinstance(schema, dict): + _validate_schema_document(schema, f"schemas_by_path.{path}") + + return self + + +class ToolCallingMetadataContent(BaseModel): + """Content for tool-calling configuration on model filesets. + + Stores chat template and tool calling settings that are merged into + the ModelSpec during checkpoint analysis. + """ + + chat_template: str | None = Field( + default=None, + description="Jinja2 chat template for the model.", + ) + tool_call_parser: str | None = Field( + default=None, + description="Name of the tool call parser (e.g., 'openai', 'hermes', 'pythonic', 'llama3_json', 'mistral').", + ) + tool_call_plugin: str | None = Field( + default=None, + description="Reference to a fileset containing a custom tool call plugin Python file. " + "Expected format: '{workspace}/{fileset_name}'.", + ) + auto_tool_choice: bool | None = Field( + default=None, + description="Whether to enable automatic tool choice.", + ) + + +class ModelMetadataContent(BaseModel): + """Content for model-type filesets. + + Contains tool calling configuration that is merged into the ModelSpec + during checkpoint analysis. + """ + + tool_calling: ToolCallingMetadataContent | None = None + + +class FilesetMetadata(BaseModel): + """Tagged metadata container - the key indicates the type. + + Example: + metadata = FilesetMetadata( + dataset=DatasetMetadataContent( + schema={"columns": ["id", "name"]}, + ) + ) + """ + + dataset: DatasetMetadataContent | None = None + model: ModelMetadataContent | None = None diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/storage_config.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/storage_config.py new file mode 100644 index 0000000000..51aa44afe1 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/storage_config.py @@ -0,0 +1,247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Storage configuration classes for various backends. + +These configs can be used by any service that needs to interact with storage backends. +""" + +from __future__ import annotations + +import os +from enum import StrEnum +from pathlib import Path +from typing import ( + Annotated, + Literal, + Self, +) + +from nemo_platform_plugin.schema import SecretRef +from pydantic import BaseModel, Field, field_validator, model_validator + + +class StorageConfigType(StrEnum): + LOCAL = "local" + NGC = "ngc" + HUGGINGFACE = "huggingface" + S3 = "s3" + # AZURE_BLOB = "azure_blob" + # GCS = "gcs" + # HTTP = "http" + + +# Default chunk size for reading/streaming files (1MB) +DEFAULT_READ_CHUNK_SIZE = 1 * 1024 * 1024 + + +class BaseStorageConfig(BaseModel): + read_chunk_size: int = Field( + default=DEFAULT_READ_CHUNK_SIZE, + description="Chunk size in bytes for reading/streaming files. " + "Larger chunks reduce async overhead but increase memory per concurrent download. " + "Default: 1MB.", + ) + + def get_secret_references(self) -> dict[str, SecretRef]: + """Get the secret references for the storage config.""" + return {} + + @property + def owns_storage_data(self) -> bool: + """Whether the platform owns the underlying source data for this backend. + + When True, deleting a fileset must also delete the underlying source + data (e.g. local files, S3 objects under our prefix). When False, the + backend points at source data the platform does not own and must not + delete (e.g. read-only external registries like NGC or HuggingFace). + + Defaults to False so external backends are safe by default. + """ + return False + + def copy_config(self, path: str) -> Self: + """ + This method is necessary for when we're using a storage config + as the default storage config. We will create a new fileset that takes + the config-defined storage config and create a fileset within a subpath of + that storage config. + + Only specific backends will be able to support this functionality, + so by default we should raise an error. + """ + raise NotImplementedError() + + +class LocalStorageConfig(BaseStorageConfig): + type: Literal[StorageConfigType.LOCAL] = StorageConfigType.LOCAL + path: str + + # These flags below will likely never be used by end-users, but they're useful + # during iteration to fine-tune performance. + write_buffer_size: int = Field( + default=16 * 1024 * 1024, + description="How many bytes to buffer before flushing to disk", + ) + + @field_validator("path") + @classmethod + def make_path_relative_to_program(cls, v: str) -> str: + """ + This allows the config to pass in absolute paths, ``~``-prefixed + paths (expanded against the running user's home dir), or relative + paths like ``./files_storage`` (joined against cwd). + """ + return str(Path.cwd() / Path(v).expanduser()) + + @property + def owns_storage_data(self) -> bool: + # Deleting a local-backed fileset removes the underlying directory + # (see LocalStorageImpl.delete_all), so we own that data. + return True + + def copy_config(self, path: str) -> Self: + new_subpath = os.path.join(self.path, path) + return self.model_copy(deep=True, update={"path": new_subpath}) + + +class HuggingfaceStorageConfig(BaseStorageConfig): + type: Literal[StorageConfigType.HUGGINGFACE] = StorageConfigType.HUGGINGFACE + repo_id: str = Field(description="Huggingface repository ID (e.g., 'meta-llama/Llama-2-7b')") + repo_type: Literal["model", "dataset", "space"] = Field( + default="model", + description="Type of Huggingface repository: 'model', 'dataset', or 'space'", + ) + revision: str = Field( + default="main", + description="Branch, tag, or commit SHA. Defaults to 'main'", + ) + original_revision: str | None = Field( + default=None, + description="The original revision requested by the user before resolution (e.g., 'main'). " + "The 'revision' field contains the resolved commit SHA.", + ) + + token_secret: SecretRef | None = Field( + default=None, + description="Huggingface API `token` secret name for private repositories", + ) + + endpoint: str = Field( + default="https://huggingface.co", + description="Huggingface Hub endpoint URL. Use for self-hosted instances.", + ) + + def get_secret_references(self) -> dict[str, SecretRef]: + return {"token": self.token_secret} if self.token_secret else {} + + +class NGCStorageConfig(BaseStorageConfig): + type: Literal[StorageConfigType.NGC] = StorageConfigType.NGC + org: str = Field(description="NGC organization name") + team: str = Field(description="NGC team name") + target: str = Field(description="NGC asset name (model or resource)") + target_type: Literal["resource", "model"] = Field( + default="resource", + description="Type of NGC asset: 'resource' or 'model'", + ) + version: str | None = Field( + default=None, + description="NGC asset version. If not provided, defaults to latest version", + ) + original_version: str | None = Field( + default=None, + description="The original version requested by the user before resolution (e.g., 'latest' or None). " + "The 'version' field contains the resolved version ID.", + ) + + api_key_secret: SecretRef = Field(description="NGC API key secret name") + + host: str = Field( + default="https://api.ngc.nvidia.com", + description="NGC API host URL", + ) + + def get_secret_references(self) -> dict[str, SecretRef]: + return {"api_key": self.api_key_secret} + + +class S3StorageConfig(BaseStorageConfig): + type: Literal[StorageConfigType.S3] = StorageConfigType.S3 + bucket: str = Field(description="S3 bucket name") + prefix: str = Field( + default="", + description="Optional prefix (folder path) within the bucket. All operations will be relative to this prefix.", + ) + region: str | None = Field( + default=None, + description="AWS region. If not specified, uses SDK default (env vars, instance metadata, etc.)", + ) + endpoint_url: str | None = Field( + default=None, + description="Custom endpoint URL for S3-compatible storage (e.g., MinIO, Garage, RustFS). " + "If not specified, uses AWS S3.", + ) + use_sdk_auth: bool = Field( + default=False, + description="Use AWS SDK credential chain for authentication (env vars like AWS_ACCESS_KEY_ID, " + "IAM roles, instance profiles, etc.). This option is only available for the platform's default " + "storage backend. User-provided S3 storage must use explicit credentials via " + "access_key_id_secret and secret_access_key_secret.", + ) + access_key_id_secret: SecretRef | None = Field( + default=None, + description="Secret reference for AWS access key ID. Requires use_sdk_auth=False.", + ) + secret_access_key_secret: SecretRef | None = Field( + default=None, + description="Secret reference for AWS secret access key. Requires use_sdk_auth=False.", + ) + signature_version: Literal["s3v4", "s3"] = Field( + default="s3v4", + description="AWS signature version for request signing. " + "Use 's3' for legacy systems that only support signature v2.", + ) + + @model_validator(mode="after") + def validate_auth_config(self) -> Self: + """Validate auth configuration is consistent.""" + has_secrets = self.access_key_id_secret is not None or self.secret_access_key_secret is not None + + if self.use_sdk_auth and has_secrets: + raise ValueError( + "use_sdk_auth=True is mutually exclusive with access_key_id_secret and " + "secret_access_key_secret. Set use_sdk_auth=False to use explicit credentials." + ) + + if not self.use_sdk_auth: + if self.access_key_id_secret is None or self.secret_access_key_secret is None: + raise ValueError( + "Both access_key_id_secret and secret_access_key_secret must be provided when use_sdk_auth=False." + ) + + return self + + def get_secret_references(self) -> dict[str, SecretRef]: + refs: dict[str, SecretRef] = {} + if self.access_key_id_secret: + refs["access_key_id"] = self.access_key_id_secret + if self.secret_access_key_secret: + refs["secret_access_key"] = self.secret_access_key_secret + return refs + + @property + def owns_storage_data(self) -> bool: + # Deleting an S3-backed fileset removes the objects under our prefix + # (see S3StorageImpl.delete_all), so we own that source data. + return True + + def copy_config(self, path: str) -> Self: + """Create a copy with an extended prefix for subpath filesets.""" + new_prefix = f"{self.prefix.rstrip('/')}/{path}" if self.prefix else path + return self.model_copy(deep=True, update={"prefix": new_prefix}) + + +StorageConfig = LocalStorageConfig | NGCStorageConfig | HuggingfaceStorageConfig | S3StorageConfig + +StorageConfigField = Annotated[StorageConfig, Field(discriminator="type")] diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/types.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/types.py new file mode 100644 index 0000000000..4f7730b646 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/types.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared request/response types for the Files service. + +These types define the HTTP contract for filesets and file operations. +Both the server (FastAPI routes) and the client (NemoClient endpoints) +import from here — one source of truth, no Stainless-generated duplicates. +""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Any, NotRequired, TypedDict + +from nemo_platform_plugin.files.metadata import FilesetMetadata +from nemo_platform_plugin.files.storage_config import StorageConfig +from nemo_platform_plugin.schema import Page +from pydantic import BaseModel, Field + + +class FilesetPurpose(StrEnum): + DATASET = "dataset" + GENERIC = "generic" + MODEL = "model" + + +class CacheStatus(StrEnum): + """Cache status for files in external storage backends.""" + + CACHED = "cached" + CACHING = "caching" + NOT_CACHED = "not_cached" + NOT_CACHEABLE = "not_cacheable" + + +# --------------------------------------------------------------------------- +# Response types +# --------------------------------------------------------------------------- + + +class FilesetOutput(BaseModel): + """Response DTO for fileset operations.""" + + id: str + name: str + workspace: str + description: str + purpose: FilesetPurpose + storage: StorageConfig + metadata: FilesetMetadata + custom_fields: dict[str, Any] + project: str + created_at: str + updated_at: str + + +class FilesetFileOutput(BaseModel): + file_ref: str + file_url: str + path: str + size: int + cache_status: CacheStatus | None = None + + +class ListFilesetFilesResponse(BaseModel): + data: list[FilesetFileOutput] + + +FilesetPage = Page[FilesetOutput] + + +# --------------------------------------------------------------------------- +# Request types +# --------------------------------------------------------------------------- + +NAME_PATTERN = r"^[\w\-.]+$" +MAX_LENGTH = 255 + + +class CreateFilesetRequest(BaseModel): + name: str = Field( + description="The name of the fileset. Allowed characters: letters (a-z, A-Z), digits (0-9), underscores, hyphens, and dots.", + max_length=MAX_LENGTH, + pattern=NAME_PATTERN, + examples=["training-data-v1", "llama-checkpoint"], + ) + description: str | None = Field( + default=None, + description="The description of the fileset.", + max_length=MAX_LENGTH, + ) + project: str | None = Field( + default=None, + description="The name of the project associated with this fileset.", + ) + storage: StorageConfig | None = Field( + default=None, + description="The storage configuration for the fileset. If not provided, uses default storage.", + ) + purpose: FilesetPurpose = Field( + default=FilesetPurpose.GENERIC, + description="The purpose of the fileset.", + ) + metadata: FilesetMetadata = Field( + default_factory=FilesetMetadata, + description="Purpose-specific metadata. Use the purpose as the key (e.g., {dataset: {...}}).", + ) + custom_fields: dict[str, Any] = Field( + default_factory=dict, + description="Custom fields for the fileset.", + ) + cache: bool = Field( + default=False, + description="Cache all files after creation. Only applies to external storage.", + ) + + +class UpdateFilesetRequest(BaseModel): + description: str | None = Field( + default=None, + description="The description of the fileset.", + max_length=MAX_LENGTH, + ) + project: str | None = Field( + default=None, + description="The name of the project associated with this fileset.", + ) + purpose: FilesetPurpose | None = Field( + default=None, + description="The purpose of the fileset.", + ) + metadata: FilesetMetadata | None = Field( + default=None, + description="Purpose-specific metadata. Use the purpose as the key (e.g., {dataset: {...}}).", + ) + custom_fields: dict[str, Any] | None = Field( + default=None, + description="Custom fields for the fileset.", + ) + + +# --------------------------------------------------------------------------- +# Query parameter types +# --------------------------------------------------------------------------- + + +class ListFilesetsQueryParams(TypedDict, total=False): + page: NotRequired[int] + page_size: NotRequired[int] + sort: NotRequired[str] + filter: NotRequired[str] + + +class ListFilesQueryParams(TypedDict, total=False): + path: NotRequired[str] + include_cache_status: NotRequired[bool] diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/schema.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/schema.py index 072c7e4e44..1dd3a82d4b 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/schema.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/schema.py @@ -102,8 +102,6 @@ class Value(BaseModel): class SecretRef(RootModel): - """Reference to a platform secret by name.""" - root: str = Field( description="Reference to a secret. Format: 'secret_name' (uses request workspace) or 'workspace/secret_name' (explicit workspace).", pattern=r"^[a-z0-9_-]+(/[a-z0-9_-]+)?$", diff --git a/packages/nemo_platform_plugin/tests/client/test_client_options.py b/packages/nemo_platform_plugin/tests/client/test_client_options.py index c45e9bbadb..0bf20f4032 100644 --- a/packages/nemo_platform_plugin/tests/client/test_client_options.py +++ b/packages/nemo_platform_plugin/tests/client/test_client_options.py @@ -231,11 +231,6 @@ def test_no_retry_without_policy(self) -> None: assert mock_http.request.call_count == 1 -# --------------------------------------------------------------------------- -# Async: exist_ok -# --------------------------------------------------------------------------- - - # --------------------------------------------------------------------------- # Async: RetryPolicy # --------------------------------------------------------------------------- diff --git a/packages/nemo_platform_plugin/tests/client/test_pagination.py b/packages/nemo_platform_plugin/tests/client/test_pagination.py index f19bc28705..5e5caf0003 100644 --- a/packages/nemo_platform_plugin/tests/client/test_pagination.py +++ b/packages/nemo_platform_plugin/tests/client/test_pagination.py @@ -64,7 +64,7 @@ def test_single_page_iteration(self) -> None: resp = client.send(LIST_ITEMS()) assert isinstance(resp, NemoPaginatedResponse) - items = list(resp) + items = list(resp.items()) assert len(items) == 2 assert items[0].name == "a" assert items[1].name == "b" @@ -85,7 +85,7 @@ def test_multi_page_iteration(self) -> None: client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) resp = client.send(LIST_ITEMS()) - items = list(resp) + items = list(resp.items()) assert len(items) == 5 assert [i.name for i in items] == ["a", "b", "c", "d", "e"] assert mock_http.request.call_count == 3 @@ -98,7 +98,7 @@ def test_data_returns_page_result_with_metadata(self) -> None: client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) resp = client.send(LIST_ITEMS()) - page = resp.data() + page = resp.page() assert len(page.items) == 1 assert page.items[0].name == "a" assert page.page == 1 @@ -116,7 +116,7 @@ def test_empty_page(self) -> None: client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) resp = client.send(LIST_ITEMS()) - items = list(resp) + items = list(resp.items()) assert items == [] def test_no_pagination_metadata(self) -> None: @@ -131,7 +131,7 @@ def test_no_pagination_metadata(self) -> None: client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) resp = client.send(LIST_ITEMS()) - items = list(resp) + items = list(resp.items()) assert len(items) == 1 assert mock_http.request.call_count == 1 @@ -145,7 +145,7 @@ def test_page_query_param_passed_on_subsequent_pages(self) -> None: client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) resp = client.send(LIST_ITEMS()) - list(resp) # consume all pages + list(resp.items()) # consume all pages # Second call should have page=2 in params second_call_params = mock_http.request.call_args_list[1][1]["params"] @@ -173,7 +173,7 @@ class TestClient(_Methods, NemoClient): resp = client.list_items() # Client options are applied but shouldn't break pagination - items = list(resp) + items = list(resp.items()) assert len(items) == 1 assert items[0].name == "a" @@ -197,7 +197,7 @@ async def test_async_multi_page_iteration(self) -> None: resp = await client.send(LIST_ITEMS()) assert isinstance(resp, AsyncNemoPaginatedResponse) - items = [item async for item in resp] + items = [item async for item in resp.items()] assert len(items) == 2 assert items[0].name == "a" assert items[1].name == "b" @@ -211,7 +211,7 @@ async def test_async_data_returns_page_result(self) -> None: client = AsyncNemoClient(base_url=BASE, workspace="default", http_client=mock_http) resp = await client.send(LIST_ITEMS()) - page = resp.data() + page = resp.page() assert len(page.items) == 1 assert page.page == 1 assert page.total_pages == 3 @@ -245,7 +245,7 @@ def test_retry_on_subsequent_page_503(self) -> None: http_client=mock_http, retry=RetryPolicy(max_retries=2, backoff_base=0.0), ) - items = list(client.send(LIST_ITEMS())) + items = list(client.send(LIST_ITEMS()).items()) assert len(items) == 2 assert [i.name for i in items] == ["a", "b"] @@ -291,7 +291,7 @@ def test_custom_items_field(self) -> None: client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) resp = client.send(LIST_THINGS()) - items = list(resp) + items = list(resp.items()) assert len(items) == 2 assert items[0].name == "a" @@ -330,7 +330,7 @@ def test_custom_page_param(self) -> None: ] client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) - items = list(client.send(LIST_THINGS())) + items = list(client.send(LIST_THINGS()).items()) assert len(items) == 2 # Verify the second call used "offset" not "page" diff --git a/packages/nemo_platform_plugin/tests/files/test_endpoints.py b/packages/nemo_platform_plugin/tests/files/test_endpoints.py new file mode 100644 index 0000000000..c3f72fc060 --- /dev/null +++ b/packages/nemo_platform_plugin/tests/files/test_endpoints.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for Files service endpoint definitions.""" + +from __future__ import annotations + +from typing import get_origin + +from nemo_platform_plugin.client.types import BinaryContent, Paginated, PreparedRequest +from nemo_platform_plugin.files import endpoints +from nemo_platform_plugin.files.types import ( + CreateFilesetRequest, + FilesetFileOutput, + FilesetOutput, + ListFilesetFilesResponse, + UpdateFilesetRequest, +) + + +def test_create_fileset() -> None: + body = CreateFilesetRequest(name="my-fileset") + prepared = endpoints.create_fileset(workspace="default", body=body) + + assert isinstance(prepared, PreparedRequest) + assert prepared.method == "POST" + assert prepared.path_template == "/apis/files/v2/workspaces/{workspace}/filesets" + assert prepared.path_params == {"workspace": "default"} + assert prepared.content == body.model_dump_json(exclude_unset=True).encode() + assert prepared.content_type == "application/json" + assert prepared.response_type is FilesetOutput + + +def test_create_fileset_workspace_optional() -> None: + body = CreateFilesetRequest(name="my-fileset") + prepared = endpoints.create_fileset(body=body) + + assert prepared.path_params == {} + + +def test_list_filesets() -> None: + prepared = endpoints.list_filesets(workspace="default") + + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default"} + assert prepared.content is None + assert get_origin(prepared.response_type) is Paginated + + +def test_list_filesets_with_query_params() -> None: + prepared = endpoints.list_filesets(workspace="default", query_params={"page": 2, "page_size": 10}) + + assert prepared.query_params == {"page": 2, "page_size": 10} + + +def test_get_fileset() -> None: + prepared = endpoints.get_fileset(workspace="default", name="my-fileset") + + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default", "name": "my-fileset"} + assert prepared.response_type is FilesetOutput + + +def test_update_fileset() -> None: + body = UpdateFilesetRequest(description="updated desc") + prepared = endpoints.update_fileset(workspace="default", name="my-fileset", body=body) + + assert prepared.method == "PATCH" + assert prepared.path_params == {"workspace": "default", "name": "my-fileset"} + assert prepared.content == body.model_dump_json(exclude_unset=True).encode() + assert prepared.response_type is FilesetOutput + + +def test_delete_fileset() -> None: + prepared = endpoints.delete_fileset(workspace="default", name="my-fileset") + + assert prepared.method == "DELETE" + assert prepared.path_params == {"workspace": "default", "name": "my-fileset"} + assert prepared.content is None + assert prepared.response_type is FilesetOutput + + +def test_list_files() -> None: + prepared = endpoints.list_files(workspace="default", name="my-fileset") + + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default", "name": "my-fileset"} + assert prepared.response_type is ListFilesetFilesResponse + + +def test_list_files_with_query_params() -> None: + prepared = endpoints.list_files( + workspace="default", name="my-fileset", query_params={"path": "data/", "include_cache_status": True} + ) + + assert prepared.query_params == {"path": "data/", "include_cache_status": True} + + +def test_upload_file() -> None: + prepared = endpoints.upload_file(workspace="default", name="my-fileset", path="data/file.txt", content=b"hello") + + assert prepared.method == "PUT" + assert prepared.path_params == {"workspace": "default", "name": "my-fileset", "path": "data/file.txt"} + assert prepared.content == b"hello" + assert prepared.content_type == "application/octet-stream" + assert prepared.response_type is FilesetFileOutput + + +def test_download_file() -> None: + prepared = endpoints.download_file(workspace="default", name="my-fileset", path="data/file.txt") + + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default", "name": "my-fileset", "path": "data/file.txt"} + assert prepared.content is None + assert prepared.response_type is BinaryContent + + +def test_delete_file() -> None: + prepared = endpoints.delete_file(workspace="default", name="my-fileset", path="data/file.txt") + + assert prepared.method == "DELETE" + assert prepared.path_params == {"workspace": "default", "name": "my-fileset", "path": "data/file.txt"} + assert prepared.content is None + assert prepared.response_type is FilesetFileOutput + + +def test_create_fileset_with_project() -> None: + """project field must be preserved in the request body.""" + body = CreateFilesetRequest(name="my-fileset", project="my-project") + prepared = endpoints.create_fileset(workspace="default", body=body) + + import json + + content = json.loads(prepared.content) + assert content["project"] == "my-project" + + +def test_update_fileset_excludes_unset_fields() -> None: + """Only explicitly set fields should be in the request body (exclude_unset).""" + body = UpdateFilesetRequest(description="updated") + prepared = endpoints.update_fileset(workspace="default", name="my-fileset", body=body) + + import json + + content = json.loads(prepared.content) + assert content == {"description": "updated"} + assert "purpose" not in content + assert "metadata" not in content diff --git a/packages/nmp_common/src/nmp/common/api/common.py b/packages/nmp_common/src/nmp/common/api/common.py index 222b22c4a4..bb7777fb25 100644 --- a/packages/nmp_common/src/nmp/common/api/common.py +++ b/packages/nmp_common/src/nmp/common/api/common.py @@ -9,8 +9,9 @@ from nemo_platform_plugin.schema import Page as Page from nemo_platform_plugin.schema import PaginationData as PaginationData +from nemo_platform_plugin.schema import SecretRef as SecretRef from nemo_platform_plugin.schema import Value -from pydantic import BaseModel, Field, RootModel, field_validator, model_validator +from pydantic import BaseModel, Field, field_validator, model_validator from starlette.responses import StreamingResponse T = TypeVar("T") @@ -150,14 +151,3 @@ def validate_date_range(self): if self.start >= self.end: raise ValueError("Start date must be before end date") return self - - -class SecretRef(RootModel): - root: str = Field( - description="Reference to a secret. Format: 'secret_name' (uses request workspace) or 'workspace/secret_name' (explicit workspace).", - pattern=r"^[a-z0-9_-]+(/[a-z0-9_-]+)?$", - examples=[ - "my-secret", - "my-workspace/my-secret", - ], - ) diff --git a/packages/nmp_common/src/nmp/common/files/metadata.py b/packages/nmp_common/src/nmp/common/files/metadata.py index 028beac46e..b3a9482b5b 100644 --- a/packages/nmp_common/src/nmp/common/files/metadata.py +++ b/packages/nmp_common/src/nmp/common/files/metadata.py @@ -3,117 +3,12 @@ """Metadata types for filesets. -The metadata uses a tagged/keyed structure where the key indicates the type: - metadata: {dataset: {schema: {...}}} - -The key in metadata should match the fileset's purpose field. +Re-exported from ``nemo_platform_plugin.files.metadata`` — the canonical +source of truth. This shim keeps existing ``from nmp.common.files.metadata +import …`` statements working without changes. """ -from jsonschema.exceptions import SchemaError -from jsonschema.validators import validator_for -from pydantic import BaseModel, ConfigDict, Field, model_validator - - -class DatasetMetadataContent(BaseModel): - """Content for dataset-type filesets.""" - - # Use `schema_` because `schema` is a BaseModel method. - model_config = ConfigDict(serialize_by_alias=True) - - schema_: dict | str | None = Field( - default=None, - alias="schema", - description="Default row schema for files in this fileset, either inline JSON Schema or a schema_defs key.", - ) - schema_defs: dict[str, dict] = Field( - default_factory=dict, - description="Reusable JSON Schema definitions keyed by name for deduplicating per-file dataset schemas.", - ) - schemas_by_path: dict[str, dict | str] = Field( - default_factory=dict, - description=( - "Optional per-file row schemas keyed by relative path within the fileset. " - "Each value may be inline JSON Schema or a schema_defs key." - ), - ) - - @model_validator(mode="after") - def validate_schema_refs(self) -> "DatasetMetadataContent": - for ref_name, ref_value in [("schema", self.schema_), *self.schemas_by_path.items()]: - if isinstance(ref_value, str) and ref_value not in self.schema_defs: - raise ValueError(f"dataset metadata reference '{ref_name}' points to unknown schema_def '{ref_value}'") - return self - - @model_validator(mode="after") - def validate_json_schemas(self) -> "DatasetMetadataContent": - def _validate_schema_document(schema: dict, ref_name: str) -> None: - validator = validator_for(schema) - try: - validator.check_schema(schema) - except SchemaError as e: - raise ValueError( - f"dataset metadata field '{ref_name}' contains invalid JSON Schema: {e.message}" - ) from e - - if isinstance(self.schema_, dict): - _validate_schema_document(self.schema_, "schema") - - for schema_name, schema in self.schema_defs.items(): - _validate_schema_document(schema, f"schema_defs.{schema_name}") - - for path, schema in self.schemas_by_path.items(): - if isinstance(schema, dict): - _validate_schema_document(schema, f"schemas_by_path.{path}") - - return self - - -class ToolCallingMetadataContent(BaseModel): - """Content for tool-calling configuration on model filesets. - - Stores chat template and tool calling settings that are merged into - the ModelSpec during checkpoint analysis. - """ - - chat_template: str | None = Field( - default=None, - description="Jinja2 chat template for the model.", - ) - tool_call_parser: str | None = Field( - default=None, - description="Name of the tool call parser (e.g., 'openai', 'hermes', 'pythonic', 'llama3_json', 'mistral').", - ) - tool_call_plugin: str | None = Field( - default=None, - description="Reference to a fileset containing a custom tool call plugin Python file. " - "Expected format: '{workspace}/{fileset_name}'.", - ) - auto_tool_choice: bool | None = Field( - default=None, - description="Whether to enable automatic tool choice.", - ) - - -class ModelMetadataContent(BaseModel): - """Content for model-type filesets. - - Contains tool calling configuration that is merged into the ModelSpec - during checkpoint analysis. - """ - - tool_calling: ToolCallingMetadataContent | None = None - - -class FilesetMetadata(BaseModel): - """Tagged metadata container - the key indicates the type. - - Example: - metadata = FilesetMetadata( - dataset=DatasetMetadataContent( - schema={"columns": ["id", "name"]}, - ) - ) - """ - - dataset: DatasetMetadataContent | None = None - model: ModelMetadataContent | None = None +from nemo_platform_plugin.files.metadata import DatasetMetadataContent as DatasetMetadataContent +from nemo_platform_plugin.files.metadata import FilesetMetadata as FilesetMetadata +from nemo_platform_plugin.files.metadata import ModelMetadataContent as ModelMetadataContent +from nemo_platform_plugin.files.metadata import ToolCallingMetadataContent as ToolCallingMetadataContent diff --git a/packages/nmp_common/src/nmp/common/files/storage_config.py b/packages/nmp_common/src/nmp/common/files/storage_config.py index 2be8a0bb5b..67962954d2 100644 --- a/packages/nmp_common/src/nmp/common/files/storage_config.py +++ b/packages/nmp_common/src/nmp/common/files/storage_config.py @@ -3,245 +3,17 @@ """Storage configuration classes for various backends. -These configs can be used by any service that needs to interact with storage backends. +Re-exported from ``nemo_platform_plugin.files.storage_config`` — the canonical +source of truth. This shim keeps existing ``from nmp.common.files.storage_config +import …`` statements working without changes. """ -from __future__ import annotations - -import os -from enum import StrEnum -from pathlib import Path -from typing import ( - Annotated, - Literal, - Self, -) - -from nmp.common.api.common import SecretRef -from pydantic import BaseModel, Field, field_validator, model_validator - - -class StorageConfigType(StrEnum): - LOCAL = "local" - NGC = "ngc" - HUGGINGFACE = "huggingface" - S3 = "s3" - # AZURE_BLOB = "azure_blob" - # GCS = "gcs" - # HTTP = "http" - - -# Default chunk size for reading/streaming files (1MB) -DEFAULT_READ_CHUNK_SIZE = 1 * 1024 * 1024 - - -class BaseStorageConfig(BaseModel): - read_chunk_size: int = Field( - default=DEFAULT_READ_CHUNK_SIZE, - description="Chunk size in bytes for reading/streaming files. " - "Larger chunks reduce async overhead but increase memory per concurrent download. " - "Default: 1MB.", - ) - - def get_secret_references(self) -> dict[str, SecretRef]: - """Get the secret references for the storage config.""" - return {} - - @property - def owns_storage_data(self) -> bool: - """Whether the platform owns the underlying source data for this backend. - - When True, deleting a fileset must also delete the underlying source - data (e.g. local files, S3 objects under our prefix). When False, the - backend points at source data the platform does not own and must not - delete (e.g. read-only external registries like NGC or HuggingFace). - - Defaults to False so external backends are safe by default. - """ - return False - - def copy_config(self, path: str) -> Self: - """ - This method is necessary for when we're using a storage config - as the default storage config. We will create a new fileset that takes - the config-defined storage config and create a fileset within a subpath of - that storage config. - - Only specific backends will be able to support this functionality, - so by default we should raise an error. - """ - raise NotImplementedError() - - -class LocalStorageConfig(BaseStorageConfig): - type: Literal[StorageConfigType.LOCAL] = StorageConfigType.LOCAL - path: str - - # These flags below will likely never be used by end-users, but they're useful - # during iteration to fine-tune performance. - write_buffer_size: int = Field( - default=16 * 1024 * 1024, - description="How many bytes to buffer before flushing to disk", - ) - - @field_validator("path") - @classmethod - def make_path_relative_to_program(cls, v: str) -> str: - """ - This allows the config to pass in absolute paths, ``~``-prefixed - paths (expanded against the running user's home dir), or relative - paths like ``./files_storage`` (joined against cwd). - """ - return str(Path.cwd() / Path(v).expanduser()) - - @property - def owns_storage_data(self) -> bool: - # Deleting a local-backed fileset removes the underlying directory - # (see LocalStorageImpl.delete_all), so we own that data. - return True - - def copy_config(self, path: str) -> Self: - new_subpath = os.path.join(self.path, path) - return self.model_copy(deep=True, update={"path": new_subpath}) - - -class HuggingfaceStorageConfig(BaseStorageConfig): - type: Literal[StorageConfigType.HUGGINGFACE] = StorageConfigType.HUGGINGFACE - repo_id: str = Field(description="Huggingface repository ID (e.g., 'meta-llama/Llama-2-7b')") - repo_type: Literal["model", "dataset", "space"] = Field( - default="model", - description="Type of Huggingface repository: 'model', 'dataset', or 'space'", - ) - revision: str = Field( - default="main", - description="Branch, tag, or commit SHA. Defaults to 'main'", - ) - original_revision: str | None = Field( - default=None, - description="The original revision requested by the user before resolution (e.g., 'main'). " - "The 'revision' field contains the resolved commit SHA.", - ) - - token_secret: SecretRef | None = Field( - default=None, - description="Huggingface API `token` secret name for private repositories", - ) - - endpoint: str = Field( - default="https://huggingface.co", - description="Huggingface Hub endpoint URL. Use for self-hosted instances.", - ) - - def get_secret_references(self) -> dict[str, SecretRef]: - return {"token": self.token_secret} if self.token_secret else {} - - -class NGCStorageConfig(BaseStorageConfig): - type: Literal[StorageConfigType.NGC] = StorageConfigType.NGC - org: str = Field(description="NGC organization name") - team: str = Field(description="NGC team name") - target: str = Field(description="NGC asset name (model or resource)") - target_type: Literal["resource", "model"] = Field( - default="resource", - description="Type of NGC asset: 'resource' or 'model'", - ) - version: str | None = Field( - default=None, - description="NGC asset version. If not provided, defaults to latest version", - ) - original_version: str | None = Field( - default=None, - description="The original version requested by the user before resolution (e.g., 'latest' or None). " - "The 'version' field contains the resolved version ID.", - ) - - api_key_secret: SecretRef = Field(description="NGC API key secret name") - - host: str = Field( - default="https://api.ngc.nvidia.com", - description="NGC API host URL", - ) - - def get_secret_references(self) -> dict[str, SecretRef]: - return {"api_key": self.api_key_secret} - - -class S3StorageConfig(BaseStorageConfig): - type: Literal[StorageConfigType.S3] = StorageConfigType.S3 - bucket: str = Field(description="S3 bucket name") - prefix: str = Field( - default="", - description="Optional prefix (folder path) within the bucket. All operations will be relative to this prefix.", - ) - region: str | None = Field( - default=None, - description="AWS region. If not specified, uses SDK default (env vars, instance metadata, etc.)", - ) - endpoint_url: str | None = Field( - default=None, - description="Custom endpoint URL for S3-compatible storage (e.g., MinIO, Garage, RustFS). " - "If not specified, uses AWS S3.", - ) - use_sdk_auth: bool = Field( - default=False, - description="Use AWS SDK credential chain for authentication (env vars like AWS_ACCESS_KEY_ID, " - "IAM roles, instance profiles, etc.). This option is only available for the platform's default " - "storage backend. User-provided S3 storage must use explicit credentials via " - "access_key_id_secret and secret_access_key_secret.", - ) - access_key_id_secret: SecretRef | None = Field( - default=None, - description="Secret reference for AWS access key ID. Requires use_sdk_auth=False.", - ) - secret_access_key_secret: SecretRef | None = Field( - default=None, - description="Secret reference for AWS secret access key. Requires use_sdk_auth=False.", - ) - signature_version: Literal["s3v4", "s3"] = Field( - default="s3v4", - description="AWS signature version for request signing. " - "Use 's3' for legacy systems that only support signature v2.", - ) - - @model_validator(mode="after") - def validate_auth_config(self) -> Self: - """Validate auth configuration is consistent.""" - has_secrets = self.access_key_id_secret is not None or self.secret_access_key_secret is not None - - if self.use_sdk_auth and has_secrets: - raise ValueError( - "use_sdk_auth=True is mutually exclusive with access_key_id_secret and " - "secret_access_key_secret. Set use_sdk_auth=False to use explicit credentials." - ) - - if not self.use_sdk_auth: - if self.access_key_id_secret is None or self.secret_access_key_secret is None: - raise ValueError( - "Both access_key_id_secret and secret_access_key_secret must be provided when use_sdk_auth=False." - ) - - return self - - def get_secret_references(self) -> dict[str, SecretRef]: - refs: dict[str, SecretRef] = {} - if self.access_key_id_secret: - refs["access_key_id"] = self.access_key_id_secret - if self.secret_access_key_secret: - refs["secret_access_key"] = self.secret_access_key_secret - return refs - - @property - def owns_storage_data(self) -> bool: - # Deleting an S3-backed fileset removes the objects under our prefix - # (see S3StorageImpl.delete_all), so we own that source data. - return True - - def copy_config(self, path: str) -> Self: - """Create a copy with an extended prefix for subpath filesets.""" - new_prefix = f"{self.prefix.rstrip('/')}/{path}" if self.prefix else path - return self.model_copy(deep=True, update={"prefix": new_prefix}) - - -StorageConfig = LocalStorageConfig | NGCStorageConfig | HuggingfaceStorageConfig | S3StorageConfig - -StorageConfigField = Annotated[StorageConfig, Field(discriminator="type")] +from nemo_platform_plugin.files.storage_config import DEFAULT_READ_CHUNK_SIZE as DEFAULT_READ_CHUNK_SIZE +from nemo_platform_plugin.files.storage_config import BaseStorageConfig as BaseStorageConfig +from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig as HuggingfaceStorageConfig +from nemo_platform_plugin.files.storage_config import LocalStorageConfig as LocalStorageConfig +from nemo_platform_plugin.files.storage_config import NGCStorageConfig as NGCStorageConfig +from nemo_platform_plugin.files.storage_config import S3StorageConfig as S3StorageConfig +from nemo_platform_plugin.files.storage_config import StorageConfig as StorageConfig +from nemo_platform_plugin.files.storage_config import StorageConfigField as StorageConfigField +from nemo_platform_plugin.files.storage_config import StorageConfigType as StorageConfigType diff --git a/plugins/example-plugin/tests/test_sdk.py b/plugins/example-plugin/tests/test_sdk.py index d73b5ccdef..7953180822 100644 --- a/plugins/example-plugin/tests/test_sdk.py +++ b/plugins/example-plugin/tests/test_sdk.py @@ -120,7 +120,7 @@ def test_sync_list_items() -> None: ) resp = client.list_items() - page = resp.data() + page = resp.page() assert len(page.items) == 1 assert page.items[0].name == "my-item" @@ -193,7 +193,7 @@ async def test_async_list_items() -> None: ) resp = await client.list_items() - page = resp.data() + page = resp.page() assert len(page.items) == 1 assert page.items[0].name == "my-item" @@ -268,7 +268,7 @@ def test_sync_download_blob_read() -> None: def test_sync_download_blob_stream() -> None: client, mock_http = _sync_client() - raw = httpx.Response(200, content=b"chunk1chunk2", request=httpx.Request("GET", BASE)) + raw = httpx.Response(200, stream=httpx.ByteStream(b"chunk1chunk2"), request=httpx.Request("GET", BASE)) mock_http.stream = _stream_ctx(raw) resp = client.download_blob(name="pic.png") diff --git a/plugins/nemo-data-designer/pyproject.toml b/plugins/nemo-data-designer/pyproject.toml index 96a1c4a413..be318f1cbe 100644 --- a/plugins/nemo-data-designer/pyproject.toml +++ b/plugins/nemo-data-designer/pyproject.toml @@ -57,6 +57,7 @@ data-designer-nemo = [ nemo-platform-plugin = [ "anthropic>=0.88.0", "fastapi>=0.115.4", + "jsonschema>=4.0.0", "lark>=1.1.0", "openai>=1.109.1", "pydantic>=2.10.3", diff --git a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py index 9c5a228389..c388996bbd 100644 --- a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py +++ b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py @@ -9,7 +9,7 @@ from data_designer_nemo.nemotron_personas import WORKSPACE, get_resource_name_for_locale from nemo_data_designer_plugin.cli import personas as personas_module from nemo_platform import NeMoPlatform -from nemo_platform.types.files import NGCStorageConfig +from nemo_platform_plugin.files.storage_config import NGCStorageConfig pytestmark = pytest.mark.integration @@ -76,11 +76,11 @@ def test_make_fileset_creates_requested_locale_with_existing_secret(cli_sdk: NeM assert result.exit_code == 0, result.output filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert [fileset.name for fileset in filesets.data] == [get_resource_name_for_locale("en_US")] + assert [fileset.name for fileset in filesets.items()] == [get_resource_name_for_locale("en_US")] fileset = cli_sdk.files.filesets.retrieve(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE) assert isinstance(fileset.storage, NGCStorageConfig) - assert fileset.storage.api_key_secret == "system/ngc-api-key" + assert fileset.storage.api_key_secret.root == "system/ngc-api-key" def test_make_fileset_creates_secret_from_env_then_fileset( @@ -107,7 +107,7 @@ def test_make_fileset_creates_secret_from_env_then_fileset( fileset = cli_sdk.files.filesets.retrieve(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE) assert isinstance(fileset.storage, NGCStorageConfig) - assert fileset.storage.api_key_secret == "system/my-ngc-key" + assert fileset.storage.api_key_secret.root == "system/my-ngc-key" def test_make_fileset_missing_env_var() -> None: @@ -184,7 +184,7 @@ def test_make_fileset_create_secret_conflict_does_not_create_fileset( assert result.exit_code == 1 assert "already exists" in result.output filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert filesets.data == [] + assert list(filesets.items()) == [] def test_make_fileset_create_secret_internal_error_surfaces_clearly( @@ -213,7 +213,7 @@ def _boom(*args: object, **kwargs: object) -> None: assert "Failed to create secret" in result.output assert "secrets backend exploded" in result.output filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert filesets.data == [] + assert list(filesets.items()) == [] def test_make_fileset_is_idempotent_when_fileset_already_exists(cli_sdk: NeMoPlatform) -> None: @@ -248,7 +248,7 @@ def test_make_fileset_is_idempotent_when_fileset_already_exists(cli_sdk: NeMoPla assert "already exists" in second.output filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert [fileset.name for fileset in filesets.data] == [get_resource_name_for_locale("en_US")] + assert [fileset.name for fileset in filesets.items()] == [get_resource_name_for_locale("en_US")] def test_make_fileset_create_fileset_internal_error_surfaces_clearly(cli_sdk: NeMoPlatform) -> None: @@ -273,4 +273,4 @@ def _boom(*args: object, **kwargs: object) -> None: assert "Failed to create fileset" in result.output assert error_message in result.output filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert filesets.data == [] + assert list(filesets.items()) == [] diff --git a/plugins/nemo-safe-synthesizer/pyproject.toml b/plugins/nemo-safe-synthesizer/pyproject.toml index 1d1a50ff6a..2d0e3eba57 100644 --- a/plugins/nemo-safe-synthesizer/pyproject.toml +++ b/plugins/nemo-safe-synthesizer/pyproject.toml @@ -46,6 +46,7 @@ test = [ nemo-platform-plugin = [ "anthropic>=0.88.0", "fastapi>=0.115.4", + "jsonschema>=4.0.0", "lark>=1.1.0", "openai>=1.109.1", "pydantic>=2.10.3", diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/core/formatters.py b/sdk/python/nemo-platform/src/nemo_platform/cli/core/formatters.py index 37b44a538a..af0a34970c 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/core/formatters.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/core/formatters.py @@ -62,16 +62,18 @@ def display_name(self) -> str: "{namespace}/{name}", # Computed OpenAI path field } -# Common field names for list items in API responses -LIST_ITEM_FIELDS = ["data"] +# Common field names for list items in API responses. +# Checked in order via getattr. Callables (e.g. NemoPaginatedResponse.items) +# are invoked; plain attributes/properties are returned directly. +LIST_ITEM_FIELDS = ["data", "items"] def _extract_items_from_response(data: Any) -> list[Any]: """Extract list items from an API response. Handles various response formats: - - Paginated responses with .data attribute - - Non-paginated responses with .files, .items, etc. + - Paginated responses with .items() method (NemoPaginatedResponse) + - Paginated responses with .data attribute (legacy Page) - Plain lists - Dict responses @@ -81,17 +83,19 @@ def _extract_items_from_response(data: Any) -> list[Any]: Returns: List of items from the response """ - # Try common field names for list items - for field in LIST_ITEM_FIELDS: - if hasattr(data, field): - return getattr(data, field) - if isinstance(data, dict) and field in data: - return data[field] - - # Handle plain lists if isinstance(data, list): return data + for field in LIST_ITEM_FIELDS: + if isinstance(data, dict): + if field in data: + return data[field] + else: + value = getattr(data, field, None) + if value is None: + continue + return list(value()) if callable(value) else value + return [] diff --git a/sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py b/sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py index 536096d311..27f01cc2ce 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py +++ b/sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""FilesetFileSystem - fsspec filesystem for NeMo Platform nemo_platform.filesets.""" +"""FilesetFileSystem - fsspec filesystem for NeMo Platform fileset storage.""" from __future__ import annotations @@ -18,13 +18,8 @@ from fsspec.callbacks import DEFAULT_CALLBACK, Callback from fsspec.spec import AbstractBufferedFile from nemo_platform import AsyncNeMoPlatform, NeMoPlatform -from nemo_platform.types.files import FilesetFile as SDKFilesetFile - -# Conditional import for TestClient detection -try: - from starlette.testclient import TestClient -except ImportError: - TestClient = None +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.types import FilesetFileOutput, ListFilesQueryParams T = TypeVar("T") @@ -105,6 +100,18 @@ async def run_one(coro: Coroutine[Any, Any, T], idx: int) -> None: fsspec.asyn._run_coros_in_chunks = run_coros_in_chunks +def _detect_async_transport(sync_client: Any) -> httpx.AsyncBaseTransport | None: + """Detect if a sync httpx client wraps a TestClient and return ASGITransport.""" + try: + from starlette.testclient import TestClient + + if isinstance(sync_client, TestClient): + return httpx.ASGITransport(app=sync_client.app) + except ImportError: + pass + return None + + class FileInfo(TypedDict): """File or directory info returned by fsspec methods.""" @@ -299,30 +306,31 @@ def build_fileset_ref( class FilesetFileSystem(AsyncFileSystem): """ - fsspec filesystem for NeMo Platform nemo_platform.filesets. + fsspec filesystem for NeMo Platform fileset storage. URL format: fileset://[workspace/]fileset_name[#path] The optional `#` separator distinguishes the fileset name from the file path. If omitted, assumes root of fileset. Workspace is optional - if omitted, - uses the SDK's default workspace. + uses the client's default workspace. Examples: - >>> from nemo_platform import NeMoPlatform - >>> sdk = NeMoPlatform(base_url="http://localhost:8000", workspace="default") - >>> fs = FilesetFileSystem(sdk=sdk) - >>> fs.ls("my-fileset") # root of fileset, workspace from SDK default + >>> from nemo_platform_plugin.files.client import AsyncFilesClient + >>> client = AsyncFilesClient(base_url="http://localhost:8000", workspace="default") + >>> fs = FilesetFileSystem(client=client) + >>> fs.ls("my-fileset") # root of fileset, workspace from client default >>> fs.ls("my-fileset#data/") # specific path within fileset >>> fs.ls("default/my-fileset#data/") # explicit workspace """ protocol = "fileset" + _client: AsyncFilesClient @classmethod def register_fsspec(cls) -> None: """Register the fileset protocol with fsspec. - After calling this, you can use fsspec.filesystem("fileset", sdk=sdk). + After calling this, you can use fsspec.filesystem("fileset", client=client). """ from fsspec import register_implementation @@ -341,60 +349,81 @@ def register_fsspec(cls) -> None: def __init__( self, - sdk: NeMoPlatform | AsyncNeMoPlatform, + *, + client: FilesClient | AsyncFilesClient | None = None, + sdk: NeMoPlatform | AsyncNeMoPlatform | None = None, batch_size: int | None = None, blocksize: int | None = None, **kwargs, ): + if client is None and sdk is None: + raise TypeError("Either 'client' or 'sdk' must be provided") + + # Normalize: convert sdk to a FilesClient so there's one code path. + # AsyncNeMoPlatform → AsyncFilesClient (already async, _ensure_async is a no-op). + # NeMoPlatform → FilesClient (sync, _ensure_async converts to async). + if sdk is not None: + from nemo_platform_plugin.client.adapter import client_from_platform + + if isinstance(sdk, AsyncNeMoPlatform): + client = client_from_platform(sdk, AsyncFilesClient) + else: + client = client_from_platform(sdk, FilesClient) + + async_client = self._ensure_async(client) + is_async = isinstance(client, AsyncFilesClient) + if batch_size is None: batch_size = self.default_batch_size if blocksize is None: blocksize = self.blocksize - # Set asynchronous mode based on SDK type. When asynchronous=False, - # fsspec creates a global daemon event loop (self.loop) that callers - # can use for sync-to-async bridging via fsspec.asyn.sync(). - is_async_sdk = isinstance(sdk, AsyncNeMoPlatform) - super().__init__(asynchronous=is_async_sdk, batch_size=batch_size, blocksize=blocksize, **kwargs) - self._sdk: AsyncNeMoPlatform = self._get_sdk(sdk) + super().__init__(asynchronous=is_async, batch_size=batch_size, blocksize=blocksize, **kwargs) + self._client = async_client - def _get_sdk( - self, - sdk: NeMoPlatform | AsyncNeMoPlatform, - ) -> AsyncNeMoPlatform: - # If already an async SDK, use it as-is to preserve custom transports (e.g., test clients) - if isinstance(sdk, AsyncNeMoPlatform): - return sdk - - # Convert sync SDK to async SDK - transport: httpx.AsyncBaseTransport | None = None - if TestClient is not None and isinstance(sdk._client, TestClient): - # If using a synchronous test client, we should use the ASGITransport - transport = httpx.ASGITransport(app=sdk._client.app) - - return AsyncNeMoPlatform( - workspace=sdk.workspace, - base_url=sdk.base_url, - timeout=sdk.timeout, - max_retries=sdk.max_retries, - default_headers=sdk._custom_headers, - default_query=sdk.default_query, + @staticmethod + def _ensure_async(client: FilesClient | AsyncFilesClient) -> AsyncFilesClient: + """Ensure we have an AsyncFilesClient, converting from sync if needed. + + Preserves subclass behavior: if the sync client has ``_async_cls`` + (e.g. a remapping subclass), that class is used for the async client. + """ + if isinstance(client, AsyncFilesClient): + return client + + import httpx + + # Use _async_cls if the sync client defines one (e.g. _RemappingFilesClient + # → _RemappingAsyncFilesClient), otherwise plain AsyncFilesClient. + async_cls = getattr(client, "_async_cls", None) or AsyncFilesClient + + transport = _detect_async_transport(client._http) + return async_cls( + base_url=client.base_url, + workspace=client.workspace, + auth=client._auth, + default_headers=client._default_headers or None, + retry=client._retry, http_client=httpx.AsyncClient( transport=transport, - base_url=sdk.base_url, - headers=sdk._custom_headers, + base_url=client.base_url, + headers=dict(client._default_headers) if client._default_headers else None, ), ) - def to_fileset_files(self, results: dict[str, Any]) -> list[SDKFilesetFile]: - """Convert fsspec find results to FilesetFile objects. + @property + def _workspace(self) -> str | None: + return self._client.workspace + + def to_fileset_files(self, results: dict[str, Any]) -> list[FilesetFileOutput]: + """Convert fsspec find results to FilesetFileOutput objects. Args: results: Dict from find(detail=True) mapping paths to file info. Returns: - List of FilesetFile objects with path, size, and file_ref. + List of FilesetFileOutput objects with path, size, and file_ref. """ files = [] for name, info in results.items(): @@ -402,7 +431,7 @@ def to_fileset_files(self, results: dict[str, Any]) -> list[SDKFilesetFile]: continue workspace, fileset, file_path = parse_fileset_ref(name, workspace_fallback=None) files.append( - SDKFilesetFile( + FilesetFileOutput( file_ref=f"{workspace}/{fileset}#{file_path}", file_url=f"/apis/files/v2/workspaces/{workspace}/filesets/{fileset}/-/{file_path}", path=file_path, @@ -485,7 +514,7 @@ async def _info(self, path: str, **kwargs) -> FileInfo: Checks dircache first to avoid redundant API calls. For cache misses, uses _ls which populates the cache for all directory levels. """ - _, _, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + _, _, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) path_key = build_fileset_ref(path) parent_path = self._parent(path_key) @@ -527,17 +556,15 @@ async def _info(self, path: str, **kwargs) -> FileInfo: async def _cat_file(self, path: str, start: int | None = None, end: int | None = None, **kwargs) -> bytes: """Fetch file content with optional byte range.""" - workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) if not file_path: raise IsADirectoryError(path) - extra_headers = {} + client = self._client if start is not None or end is not None: - extra_headers["Range"] = f"bytes={start or 0}-{(end - 1) if end else ''}" + client = client.with_headers({"Range": f"bytes={start or 0}-{(end - 1) if end else ''}"}) - response = await self._sdk.files._download_file( - file_path, workspace=workspace, name=fileset, extra_headers=extra_headers or None - ) + response = await client.download_file(workspace=workspace, name=fileset, path=file_path) return await response.read() @classmethod @@ -569,7 +596,7 @@ async def _ls(self, path: str, detail: bool = True, refresh: bool = False, **kwa detail: If True, return list of dicts. If False, return list of paths. refresh: If True, bypass cache and fetch fresh listing. """ - workspace, fileset, prefix = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, prefix = parse_fileset_ref(path, workspace_fallback=self._workspace) prefix = prefix.rstrip("/") path_key = build_fileset_ref(prefix, workspace=workspace, fileset=fileset) @@ -582,7 +609,13 @@ async def _ls(self, path: str, detail: bool = True, refresh: bool = False, **kwa pass # Fetch from backend and populate cache for all directory levels - response = await self._sdk.files._list_files(fileset, workspace=workspace, path=prefix or None) + query_params: ListFilesQueryParams | None = {"path": prefix} if prefix else None + response = await self._client.list_files( + workspace=workspace, + name=fileset, + query_params=query_params, + ) + response = response.data() dir_contents = self._populate_dircache_from_response(response, workspace, fileset, prefix) # Return the listing for the requested path @@ -591,19 +624,19 @@ async def _ls(self, path: str, detail: bool = True, refresh: bool = False, **kwa async def _rm_file(self, path: str, **kwargs) -> None: """Delete a single file.""" - workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) if not file_path: raise ValueError("Cannot delete fileset root via rm") - await self._sdk.files._delete_file(file_path, workspace=workspace, name=fileset) + await self._client.delete_file(workspace=workspace, name=fileset, path=file_path) # Invalidate parent directory's cache since file info is stored there self.invalidate_cache(self._parent(build_fileset_ref(path))) async def _pipe_file(self, path: str, value: bytes, **kwargs) -> None: """Write bytes to a file.""" - workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) if not file_path: raise ValueError("File path required for upload") - await self._sdk.files._upload_file(file_path, body=value, workspace=workspace, name=fileset) + await self._client.upload_file(workspace=workspace, name=fileset, path=file_path, content=value) # Invalidate parent directory's cache since file info is stored there self.invalidate_cache(self._parent(build_fileset_ref(path))) @@ -627,7 +660,7 @@ async def _pipe_stream( content_length: Optional content length for Content-Length header. If not provided, uses chunked transfer encoding. """ - workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace) if not file_path: raise ValueError("File path required for upload") @@ -635,15 +668,11 @@ async def _pipe_stream( if not hasattr(stream, "__anext__"): stream = to_async_iterator(stream) - extra_headers = {"Content-Length": str(content_length)} if content_length is not None else None + client = self._client + if content_length is not None: + client = client.with_headers({"Content-Length": str(content_length)}) - await self._sdk.files._upload_file( - path=file_path, - body=stream, - workspace=workspace, - name=fileset, - extra_headers=extra_headers, - ) + await client.upload_file(workspace=workspace, name=fileset, path=file_path, content=stream) # Invalidate parent directory's cache since file info is stored there self.invalidate_cache(self._parent(build_fileset_ref(path))) @@ -663,7 +692,7 @@ async def _put_file(self, lpath: str, rpath: str, callback: Callback = DEFAULT_C Uses streaming upload to avoid buffering the entire file in memory. Supports per-chunk progress via callback.relative_update(chunk_size). """ - workspace, fileset, file_path = parse_fileset_ref(rpath, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(rpath, workspace_fallback=self._workspace) if not file_path: raise ValueError("File path required for upload") @@ -672,18 +701,17 @@ async def _put_file(self, lpath: str, rpath: str, callback: Callback = DEFAULT_C callback.set_size(file_size) # Create async generator that streams file content with progress - async def stream_file(): + async def stream_file() -> AsyncIterator[bytes]: async with await anyio.open_file(lpath, "rb") as f: while chunk := await f.read(self.blocksize): callback.relative_update(len(chunk)) yield chunk - await self._sdk.files._upload_file( - path=file_path, - body=stream_file(), + await self._client.with_headers({"Content-Length": str(file_size)}).upload_file( workspace=workspace, name=fileset, - extra_headers={"Content-Length": str(file_size)}, + path=file_path, + content=stream_file(), ) # Invalidate parent directory's cache since file info is stored there self.invalidate_cache(self._parent(build_fileset_ref(rpath))) @@ -705,9 +733,15 @@ async def _find( Also populates the dircache so subsequent _ls calls benefit. """ - workspace, fileset, prefix = parse_fileset_ref(path, workspace_fallback=self._sdk.workspace) + workspace, fileset, prefix = parse_fileset_ref(path, workspace_fallback=self._workspace) prefix = prefix.rstrip("/") - response = await self._sdk.files._list_files(fileset, workspace=workspace, path=prefix or None) + query_params: ListFilesQueryParams | None = {"path": prefix} if prefix else None + response = await self._client.list_files( + workspace=workspace, + name=fileset, + query_params=query_params, + ) + response = response.data() # Populate dircache for all directory levels (benefits subsequent _ls calls) self._populate_dircache_from_response(response, workspace, fileset, prefix) @@ -718,7 +752,7 @@ async def _find( # Add root path if withdirs requested if withdirs: - root_path = build_fileset_ref(path, workspace=self._sdk.workspace) + root_path = build_fileset_ref(path, workspace=self._workspace) out[root_path] = {"name": root_path, "size": 0, "type": "directory"} for file_info in response.data: @@ -748,28 +782,27 @@ async def _find( async def _get_file(self, rpath: str, lpath: str, callback: Callback = DEFAULT_CALLBACK, **kwargs) -> None: """Download a file to local path. - Uses with_streaming_response to avoid buffering the entire response in memory. - Uses http_response.aiter_raw() for maximum throughput (bypasses httpx chunking overhead). + Uses streaming response to avoid buffering the entire response in memory. Supports per-chunk progress via callback.relative_update(chunk_size). """ - workspace, fileset, file_path = parse_fileset_ref(rpath, workspace_fallback=self._sdk.workspace) + workspace, fileset, file_path = parse_fileset_ref(rpath, workspace_fallback=self._workspace) if not file_path: return - # Use with_streaming_response to not buffer the data in memory. - async with self._sdk.files.with_streaming_response._download_file( - file_path, workspace=workspace, name=fileset - ) as response: - # Set callback size from Content-Length if available - content_length = response.headers.get("content-length") + response = await self._client.download_file( + workspace=workspace, + name=fileset, + path=file_path, + ) + + async with response.stream() as chunks: + content_length = response.http_response.headers.get("content-length") if content_length: callback.set_size(int(content_length)) - await anyio.Path(lpath).parent.mkdir(parents=True, exist_ok=True) async with await anyio.open_file(lpath, "wb") as f: - # Use aiter_raw() instead of iter_bytes() to bypass httpx chunking overhead. - async for chunk in response.http_response.aiter_raw(self.blocksize): + async for chunk in chunks: await f.write(chunk) callback.relative_update(len(chunk)) @@ -810,7 +843,7 @@ async def _get( return # Normalize rpath to new format for comparison (since _find returns new format paths) - rpath_normalized = build_fileset_ref(rpath, workspace=self._sdk.workspace).rstrip("/") + rpath_normalized = build_fileset_ref(rpath, workspace=self._workspace).rstrip("/") lpath_stripped = lpath.rstrip("/") source_is_file = len(source_files) == 1 and self._strip_protocol(source_files[0]) == rpath_normalized @@ -842,7 +875,7 @@ async def _get( # SPECIAL CASE: Fileset root (workspace/fileset with no file path) always # copies contents directly, matching HuggingFace Hub behavior. Users who want # to preserve the fileset name can include it in local_path. - _, _, file_path = parse_fileset_ref(rpath, workspace_fallback=self._sdk.workspace) + _, _, file_path = parse_fileset_ref(rpath, workspace_fallback=self._workspace) copy_contents_directly = rpath.endswith("/") or not file_path # Extract directory name from the file path portion (e.g., "subdir" from "a/b/subdir") diff --git a/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py b/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py index 44a1d81fdc..f3d42a4014 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py +++ b/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py @@ -1,27 +1,35 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Extended FilesResource classes with FilesetFileSystem support. +"""FilesResource classes with FilesetFileSystem support. -These classes extend the SDK's generated FilesResource classes to add -high-level file operations (upload, download, list, delete) and fsspec -filesystem access. +These classes provide high-level file operations (upload, download, list, delete) +backed by the NemoClient typed HTTP client and fsspec filesystem access. """ import uuid from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass +from functools import cached_property from pathlib import PurePath -from typing import Protocol, runtime_checkable +from typing import Any, Protocol, runtime_checkable +import nemo_platform from fsspec.callbacks import Callback from fsspec.core import has_magic -from nemo_platform import ConflictError -from nemo_platform._compat import cached_property -from nemo_platform.resources.files import AsyncFilesResource as BaseAsyncFilesResource -from nemo_platform.resources.files import FilesResource as BaseFilesResource -from nemo_platform.types.files import CacheStatus, FilesetFile -from nemo_platform.types.files.fileset import Fileset +from nemo_platform_plugin.client.errors import NemoHTTPError +from nemo_platform_plugin.client.response import AsyncNemoPaginatedResponse, NemoPaginatedResponse +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.types import ( + CacheStatus, + CreateFilesetRequest, + FilesetFileOutput, + FilesetMetadata, + FilesetOutput, + FilesetPurpose, + StorageConfig, + UpdateFilesetRequest, +) from nemo_platform.filesets.filesystem.filesystem import ( FilesetFileSystem, @@ -30,6 +38,85 @@ ) +def _build_error_map() -> dict[type[NemoHTTPError], type[nemo_platform.APIStatusError]]: + """Build a mapping from NemoClient errors to Stainless SDK errors. + + Lazy import to avoid hard-coding the Stainless error classes at module level. + This mapping is temporary — remove when all consumers import errors from + nemo_platform_plugin.client.errors instead of nemo_platform (AIRCORE-840). + """ + from nemo_platform_plugin.client import errors + + return { + errors.BadRequestError: nemo_platform.BadRequestError, + errors.AuthenticationError: nemo_platform.AuthenticationError, + errors.PermissionDeniedError: nemo_platform.PermissionDeniedError, + errors.NotFoundError: nemo_platform.NotFoundError, + errors.ConflictError: nemo_platform.ConflictError, + errors.UnprocessableEntityError: nemo_platform.UnprocessableEntityError, + errors.RateLimitError: nemo_platform.RateLimitError, + errors.InternalServerError: nemo_platform.InternalServerError, + } + + +_ERROR_MAP: dict[type[NemoHTTPError], type[nemo_platform.APIStatusError]] | None = None + + +def _get_error_map() -> dict[type[NemoHTTPError], type[nemo_platform.APIStatusError]]: + global _ERROR_MAP + if _ERROR_MAP is None: + _ERROR_MAP = _build_error_map() + return _ERROR_MAP + + +def _raise_as_stainless(e: NemoHTTPError) -> None: + """Re-raise a NemoClient error as its Stainless SDK equivalent. + + Preserves backward compatibility for consumers that catch + ``nemo_platform.NotFoundError`` etc. Remove with AIRCORE-840. + """ + error_map = _get_error_map() + stainless_cls = error_map.get(type(e)) + if stainless_cls is not None: + raise stainless_cls( + message=str(e), + response=e.http_response, + body=e.body, + ) from e + raise + + +class _RemappingFilesClient(FilesClient): + """FilesClient that re-raises NemoClient errors as Stainless SDK errors. + + Wraps ``send()`` so ALL operations through this client (filesets, files, + fsspec) raise Stainless-compatible exceptions. Remove with AIRCORE-840. + """ + + # Used by FilesetFileSystem._ensure_async to create the matching async + # remapping client when converting sync → async. + _async_cls: type[AsyncFilesClient] | None = None + + def send(self, request, *, headers=None, retry=None): # type: ignore[override] + try: + return super().send(request, headers=headers, retry=retry) + except NemoHTTPError as e: + _raise_as_stainless(e) + + +class _RemappingAsyncFilesClient(AsyncFilesClient): + """AsyncFilesClient that re-raises NemoClient errors as Stainless SDK errors.""" + + async def send(self, request, *, headers=None, retry=None): # type: ignore[override] + try: + return await super().send(request, headers=headers, retry=retry) + except NemoHTTPError as e: + _raise_as_stainless(e) + + +_RemappingFilesClient._async_cls = _RemappingAsyncFilesClient + + @dataclass class ListFilesResponse: """Response from listing files in a fileset. @@ -46,7 +133,7 @@ class ListFilesResponse: - None if no cache information is available """ - data: list[FilesetFile] + data: list[FilesetFileOutput] @property def cache_status(self) -> CacheStatus | None: @@ -129,24 +216,250 @@ def _matches_glob(filepath: str, pattern: str) -> bool: return PurePath(filepath).match(pattern) -class FilesResource(BaseFilesResource): - """Extended FilesResource with high-level file operations. +class FilesetsSubResource: + """Fileset CRUD operations (create, retrieve, update, list, delete). + + Wraps ``FilesClient`` methods with higher-level convenience signatures + (unwrapped params, ``exist_ok`` support). + + .. deprecated:: + Temporary shim for the ``sdk.files`` fileset interface. + New code should use ``FilesClient`` directly. + Once all callers are migrated, this class will be removed. + """ + + def __init__(self, client: FilesClient) -> None: + self._client = client + + def create( + self, + *, + name: str, + workspace: str | None = None, + exist_ok: bool = False, + description: str | None = None, + project: str | None = None, + purpose: FilesetPurpose | None = None, + metadata: FilesetMetadata | None = None, + storage: StorageConfig | None = None, + custom_fields: dict[str, Any] | None = None, + cache: bool = False, + ) -> FilesetOutput: + body = CreateFilesetRequest( + name=name, + description=description, + project=project, + purpose=purpose or FilesetPurpose.GENERIC, + metadata=metadata or FilesetMetadata(), + storage=storage, + custom_fields=custom_fields or {}, + cache=cache, + ) + # The server returns an error body on 409, not the entity, so + # exist_ok is handled here with a follow-up GET rather than at + # the endpoint/client level. + try: + return self._client.create_fileset(workspace=workspace, body=body).data() + except nemo_platform.APIStatusError as e: + if e.status_code == 409 and exist_ok: + return self.retrieve(name=name, workspace=workspace) + raise + + def retrieve(self, name: str, *, workspace: str | None = None) -> FilesetOutput: + return self._client.get_fileset(workspace=workspace, name=name).data() + + def update( + self, + name: str, + *, + workspace: str | None = None, + description: str | None = None, + project: str | None = None, + purpose: FilesetPurpose | None = None, + metadata: FilesetMetadata | None = None, + custom_fields: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> FilesetOutput: + # Only include explicitly provided fields so exclude_unset works correctly + kwargs = { + k: v + for k, v in dict( + description=description, + project=project, + purpose=purpose, + metadata=metadata, + custom_fields=custom_fields, + ).items() + if v is not None + } + body = UpdateFilesetRequest(**kwargs) + client = self._client.with_options(timeout=timeout) if timeout is not None else self._client + return client.update_fileset(workspace=workspace, name=name, body=body).data() + + def list( + self, + *, + workspace: str | None = None, + page: int | None = None, + page_size: int | None = None, + sort: str | None = None, + filter: str | dict | None = None, + ) -> NemoPaginatedResponse[FilesetOutput]: + query_params = { + k: v + for k, v in dict( + page=page, + page_size=page_size, + sort=sort, + filter=filter, + ).items() + if v is not None + } + return self._client.list_filesets(workspace=workspace, query_params=query_params or None) + + def delete(self, name: str, *, workspace: str | None = None) -> FilesetOutput: + return self._client.delete_fileset(workspace=workspace, name=name).data() + + +class AsyncFilesetsSubResource: + """Async fileset CRUD operations (create, retrieve, update, list, delete). + + Wraps ``AsyncFilesClient`` methods with higher-level convenience signatures + (unwrapped params, ``exist_ok`` support). + + .. deprecated:: + Temporary shim for the ``sdk.files`` fileset interface. + New code should use ``AsyncFilesClient`` directly. + Once all callers are migrated, this class will be removed. + """ + + def __init__(self, client: AsyncFilesClient) -> None: + self._client = client + + async def create( + self, + *, + name: str, + workspace: str | None = None, + exist_ok: bool = False, + description: str | None = None, + project: str | None = None, + purpose: FilesetPurpose | None = None, + metadata: FilesetMetadata | None = None, + storage: StorageConfig | None = None, + custom_fields: dict[str, Any] | None = None, + cache: bool = False, + ) -> FilesetOutput: + body = CreateFilesetRequest( + name=name, + description=description, + project=project, + purpose=purpose or FilesetPurpose.GENERIC, + metadata=metadata or FilesetMetadata(), + storage=storage, + custom_fields=custom_fields or {}, + cache=cache, + ) + # The server returns an error body on 409, not the entity, so + # exist_ok is handled here with a follow-up GET rather than at + # the endpoint/client level. + try: + return (await self._client.create_fileset(workspace=workspace, body=body)).data() + except nemo_platform.APIStatusError as e: + if e.status_code == 409 and exist_ok: + return await self.retrieve(name=name, workspace=workspace) + raise + + async def retrieve(self, name: str, *, workspace: str | None = None) -> FilesetOutput: + return (await self._client.get_fileset(workspace=workspace, name=name)).data() + + async def update( + self, + name: str, + *, + workspace: str | None = None, + description: str | None = None, + project: str | None = None, + purpose: FilesetPurpose | None = None, + metadata: FilesetMetadata | None = None, + custom_fields: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> FilesetOutput: + kwargs = { + k: v + for k, v in dict( + description=description, + project=project, + purpose=purpose, + metadata=metadata, + custom_fields=custom_fields, + ).items() + if v is not None + } + body = UpdateFilesetRequest(**kwargs) + client = self._client.with_options(timeout=timeout) if timeout is not None else self._client + return (await client.update_fileset(workspace=workspace, name=name, body=body)).data() + + async def list( + self, + *, + workspace: str | None = None, + page: int | None = None, + page_size: int | None = None, + sort: str | None = None, + filter: str | dict | None = None, + ) -> AsyncNemoPaginatedResponse[FilesetOutput]: + query_params = { + k: v + for k, v in dict( + page=page, + page_size=page_size, + sort=sort, + filter=filter, + ).items() + if v is not None + } + return await self._client.list_filesets(workspace=workspace, query_params=query_params or None) + + async def delete(self, name: str, *, workspace: str | None = None) -> FilesetOutput: + return (await self._client.delete_fileset(workspace=workspace, name=name)).data() + + +class FilesResource: + """FilesResource with high-level file operations. Provides convenient methods for uploading, downloading, and listing files. - For fsspec filesystem access, use `sdk.files.fsspec`. + For fsspec filesystem access, use ``resource.fsspec``. """ + def __init__(self, client) -> None: + # _raw_client kept for otlp delegation (Stainless SDK), removed by AIRCORE-840. + self._raw_client = client + + from nemo_platform_plugin.client.adapter import client_from_platform + + self._client = client_from_platform(client, _RemappingFilesClient) + + @cached_property + def filesets(self) -> FilesetsSubResource: + """Access fileset CRUD operations (create, retrieve, update, list, delete).""" + return FilesetsSubResource(self._client) + + @cached_property + def otlp(self): + """Access OTLP log operations (delegated to Stainless SDK resource).""" + from nemo_platform.resources.files.otlp.otlp import OtlpResource + + return OtlpResource(self._raw_client) + @cached_property def fsspec(self) -> FilesetFileSystem: """Access the underlying fsspec filesystem.""" - return FilesetFileSystem(sdk=self._client) + return FilesetFileSystem(client=self._client) def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None: """Create fileset if it doesn't exist (idempotent).""" - try: - self.filesets.create(name=fileset, workspace=workspace) - except ConflictError: - pass # Already exists + self.filesets.create(name=fileset, workspace=workspace, exist_ok=True) def download( self, @@ -170,7 +483,7 @@ def download( local_path: Local destination path (directory). fileset: Fileset name. If not provided, inferred from remote_path (str only). workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. callback: Optional progress callback (e.g., RichProgressCallback). max_workers: Maximum number of concurrent file transfers. @@ -281,7 +594,7 @@ def upload( callback: Callback | None = None, max_workers: int | None = None, fileset_auto_create: bool = False, - ) -> Fileset: + ) -> FilesetOutput: """Upload files from a local path to a fileset. Args: @@ -292,7 +605,7 @@ def upload( Defaults to "" (root of fileset). fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. callback: Optional progress callback (e.g., RichProgressCallback). max_workers: Maximum number of concurrent file transfers. fileset_auto_create: If True, create the fileset if it doesn't exist. @@ -300,7 +613,7 @@ def upload( a unique name is generated (e.g., "fileset-a1b2c3d4"). Returns: - Fileset: The fileset that was uploaded to. Check `fileset.name` to see + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see the generated name when using fileset_auto_create without specifying a fileset. @@ -373,7 +686,7 @@ def upload_content( fileset: str | None = None, workspace: str | None = None, fileset_auto_create: bool = False, - ) -> Fileset: + ) -> FilesetOutput: """Upload in-memory content to a fileset. Args: @@ -384,13 +697,13 @@ def upload_content( - Iterator[bytes]: Generator or iterator yielding byte chunks remote_path: Destination path within the fileset. fileset: Fileset name. If not provided, inferred from remote_path. - workspace: Workspace name. If not provided, uses SDK default. + workspace: Workspace name. If not provided, uses client default. fileset_auto_create: If True, create the fileset if it doesn't exist. When no fileset is specified (neither as param nor in remote_path), a unique name is generated (e.g., "fileset-a1b2c3d4"). Returns: - Fileset: The fileset that was uploaded to. Check `fileset.name` to see + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see the generated name when using fileset_auto_create without specifying a fileset. @@ -478,7 +791,7 @@ def download_content( Args: remote_path: Path of the file within the fileset. fileset: Fileset name. If not provided, inferred from remote_path. - workspace: Workspace name. If not provided, uses SDK default. + workspace: Workspace name. If not provided, uses client default. Returns: bytes: The file content. @@ -532,12 +845,12 @@ def list( Defaults to "" (root of fileset). fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. include_cache_status: Check and return cache status for each file. When False (default), external storage files return None for cache_status. Returns: - ListFilesResponse with data (list of FilesetFile) and cache_status property. + ListFilesResponse with data (list of FilesetFileOutput) and cache_status property. Examples: # List all files in a fileset @@ -585,12 +898,18 @@ def list( # For path prefixes, the API handles filtering server-side api_path = None if has_magic(path) else (path or None) - response = self._list_files( - fileset, + query_params = {} + if api_path is not None: + query_params["path"] = api_path + if include_cache_status: + query_params["include_cache_status"] = True + + response = self._client.list_files( workspace=ws, - include_cache_status=include_cache_status, - path=api_path, + name=fileset, + query_params=query_params or None, ) + response = response.data() files = list(response.data) # Apply glob filtering if needed @@ -613,7 +932,7 @@ def delete( or a relative path (e.g., "data/file.txt") if fileset is provided. fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. Examples: # Delete a file with explicit fileset @@ -638,28 +957,41 @@ def delete( self.fsspec.rm(fileset_ref) -class AsyncFilesResource(BaseAsyncFilesResource): - """Extended AsyncFilesResource with high-level file operations. +class AsyncFilesResource: + """Async FilesResource with high-level file operations. Provides convenient methods for uploading, downloading, and listing files. - For fsspec filesystem access, use `sdk.files.fsspec`. + For fsspec filesystem access, use ``resource.fsspec``. """ + def __init__(self, client) -> None: + # _raw_client kept for otlp delegation (Stainless SDK), removed by AIRCORE-840. + self._raw_client = client + + from nemo_platform_plugin.client.adapter import client_from_platform + + self._client = client_from_platform(client, _RemappingAsyncFilesClient) + @cached_property - def fsspec(self) -> FilesetFileSystem: - """Get a FilesetFileSystem instance pre-configured with this SDK client. + def filesets(self) -> AsyncFilesetsSubResource: + """Access fileset CRUD operations (create, retrieve, update, list, delete).""" + return AsyncFilesetsSubResource(self._client) - This provides fsspec filesystem access. For high-level file - operations, use `sdk.files` instead. - """ - return FilesetFileSystem(sdk=self._client) + @cached_property + def otlp(self): + """Access OTLP log operations (delegated to Stainless SDK resource).""" + from nemo_platform.resources.files.otlp.otlp import AsyncOtlpResource + + return AsyncOtlpResource(self._raw_client) + + @cached_property + def fsspec(self) -> FilesetFileSystem: + """Access the underlying fsspec filesystem.""" + return FilesetFileSystem(client=self._client) async def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None: """Create fileset if it doesn't exist (idempotent).""" - try: - await self.filesets.create(name=fileset, workspace=workspace) - except ConflictError: - pass # Already exists + await self.filesets.create(name=fileset, workspace=workspace, exist_ok=True) async def download( self, @@ -683,7 +1015,7 @@ async def download( local_path: Local destination path (directory). fileset: Fileset name. If not provided, inferred from remote_path (str only). workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. callback: Optional progress callback (e.g., RichProgressCallback). max_workers: Maximum number of concurrent file transfers. @@ -779,7 +1111,7 @@ async def upload( callback: Callback | None = None, max_workers: int | None = None, fileset_auto_create: bool = False, - ) -> Fileset: + ) -> FilesetOutput: """Upload files from a local path to a fileset (async). Args: @@ -790,7 +1122,7 @@ async def upload( Defaults to "" (root of fileset). fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. callback: Optional progress callback (e.g., RichProgressCallback). max_workers: Maximum number of concurrent file transfers. fileset_auto_create: If True, create the fileset if it doesn't exist. @@ -798,7 +1130,7 @@ async def upload( a unique name is generated (e.g., "fileset-a1b2c3d4"). Returns: - Fileset: The fileset that was uploaded to. Check `fileset.name` to see + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see the generated name when using fileset_auto_create without specifying a fileset. @@ -865,7 +1197,7 @@ async def upload_content( fileset: str | None = None, workspace: str | None = None, fileset_auto_create: bool = False, - ) -> Fileset: + ) -> FilesetOutput: """Upload in-memory data to a fileset (async). Args: @@ -876,13 +1208,13 @@ async def upload_content( - AsyncIterator[bytes]: Async iterator yielding byte chunks (streamed) remote_path: Destination path within the fileset. fileset: Fileset name. If not provided, inferred from remote_path. - workspace: Workspace name. If not provided, uses SDK default. + workspace: Workspace name. If not provided, uses client default. fileset_auto_create: If True, create the fileset if it doesn't exist. When no fileset is specified (neither as param nor in remote_path), a unique name is generated (e.g., "fileset-a1b2c3d4"). Returns: - Fileset: The fileset that was uploaded to. Check `fileset.name` to see + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see the generated name when using fileset_auto_create without specifying a fileset. @@ -974,7 +1306,7 @@ async def download_content( Args: remote_path: Path of the file within the fileset. fileset: Fileset name. If not provided, inferred from remote_path. - workspace: Workspace name. If not provided, uses SDK default. + workspace: Workspace name. If not provided, uses client default. Returns: bytes: The file content. @@ -1020,12 +1352,12 @@ async def list( Defaults to "" (root of fileset). fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. include_cache_status: Check and return cache status for each file. When False (default), external storage files return None for cache_status. Returns: - ListFilesResponse with data (list of FilesetFile) and cache_status property. + ListFilesResponse with data (list of FilesetFileOutput) and cache_status property. Examples: # List all files in a fileset @@ -1070,12 +1402,18 @@ async def list( # For path prefixes, the API handles filtering server-side api_path = None if has_magic(path) else (path or None) - response = await self._list_files( - fileset, + query_params = {} + if api_path is not None: + query_params["path"] = api_path + if include_cache_status: + query_params["include_cache_status"] = True + + response = await self._client.list_files( workspace=ws, - include_cache_status=include_cache_status, - path=api_path, + name=fileset, + query_params=query_params or None, ) + response = response.data() files = list(response.data) # Apply glob filtering if needed @@ -1098,7 +1436,7 @@ async def delete( or a relative path (e.g., "data/file.txt") if fileset is provided. fileset: Fileset name. If not provided, inferred from remote_path. workspace: Workspace name. If not provided, inferred from remote_path - or uses the SDK's default workspace. + or uses the client's default workspace. Examples: # Delete a file with explicit fileset diff --git a/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py b/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py index f6fb760786..c0b4a91fd5 100644 --- a/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py +++ b/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py @@ -173,9 +173,8 @@ async def _cleanup_filesets(self, workspace: Workspace) -> None: logger.info(f"Cleaning up filesets for workspace: {workspace.name}") try: filesets_response = await self._nmp_sdk.files.filesets.list(workspace=workspace.name) - filesets = [fileset async for fileset in filesets_response] - for fileset in filesets: + async for fileset in filesets_response.items(): try: logger.info(f"Deleting fileset: {fileset.name}") await self._nmp_sdk.files.filesets.delete( diff --git a/services/core/entities/tests/controllers/test_workspace_cleanup.py b/services/core/entities/tests/controllers/test_workspace_cleanup.py index a4ad959943..d5691bde54 100644 --- a/services/core/entities/tests/controllers/test_workspace_cleanup.py +++ b/services/core/entities/tests/controllers/test_workspace_cleanup.py @@ -37,6 +37,16 @@ async def __anext__(self): raise StopAsyncIteration +class _MockPaginatedResponse: + """Mock for paginated responses that expose .items() for async iteration.""" + + def __init__(self, items): + self._items = items + + def items(self): + return _AsyncIterator(self._items) + + def _make_sdk( jobs: list | None = None, deployments: list | None = None, @@ -49,7 +59,7 @@ def _make_sdk( sdk.jobs.delete = AsyncMock() sdk.inference.deployments.list = AsyncMock(return_value=_AsyncIterator(deployments or [])) sdk.inference.deployments.delete = AsyncMock() - sdk.files.filesets.list = AsyncMock(return_value=_AsyncIterator(filesets or [])) + sdk.files.filesets.list = AsyncMock(return_value=_MockPaginatedResponse(filesets or [])) sdk.files.filesets.delete = AsyncMock() return sdk @@ -148,7 +158,7 @@ async def test_successful_workspace_deletion(self): sdk = MagicMock() sdk.jobs.list = AsyncMock(return_value=_AsyncIterator([])) sdk.inference.deployments.list = AsyncMock(return_value=_AsyncIterator([])) - sdk.files.filesets.list = AsyncMock(return_value=_AsyncIterator([])) + sdk.files.filesets.list = AsyncMock(return_value=_MockPaginatedResponse([])) controller = _make_controller(workspace_repo=repo, nmp_sdk=sdk) @@ -329,7 +339,7 @@ async def test_deletes_filesets(self): fileset.name = "test-fileset" sdk = MagicMock() - sdk.files.filesets.list = AsyncMock(return_value=_AsyncIterator([fileset])) + sdk.files.filesets.list = AsyncMock(return_value=_MockPaginatedResponse([fileset])) sdk.files.filesets.delete = AsyncMock() controller = _make_controller(nmp_sdk=sdk) @@ -349,7 +359,7 @@ async def test_continues_on_individual_fileset_failure(self): fs2.name = "fs2" sdk = MagicMock() - sdk.files.filesets.list = AsyncMock(return_value=_AsyncIterator([fs1, fs2])) + sdk.files.filesets.list = AsyncMock(return_value=_MockPaginatedResponse([fs1, fs2])) sdk.files.filesets.delete = AsyncMock(side_effect=[Exception("fail"), None]) controller = _make_controller(nmp_sdk=sdk) diff --git a/services/core/files/src/nmp/core/files/api/v2/filesets/endpoints.py b/services/core/files/src/nmp/core/files/api/v2/filesets/endpoints.py index f9fb6d6d17..c6faac973d 100644 --- a/services/core/files/src/nmp/core/files/api/v2/filesets/endpoints.py +++ b/services/core/files/src/nmp/core/files/api/v2/filesets/endpoints.py @@ -45,12 +45,17 @@ ) from nmp.core.files.api.v2.filesets.schemas import ( CreateFilesetRequest, + FilesetFileOutput, FilesetFilter, FilesetOutput, FilesetPage, + ListFilesetFilesResponse, UpdateFilesetRequest, + fileset_file_output_from_info, + fileset_output_from_entity, + list_fileset_files_from_infos, ) -from nmp.core.files.app.backends import FileInfo, storage_impl_factory +from nmp.core.files.app.backends import storage_impl_factory from nmp.core.files.app.backends.factory import StorageConfig from nmp.core.files.app.cache import CacheStatus, warm_fileset_cache from nmp.core.files.app.external_hosts import ( @@ -74,7 +79,6 @@ StorageConfigError, StorageUnavailableError, ) -from pydantic import BaseModel from starlette.status import ( HTTP_200_OK, HTTP_400_BAD_REQUEST, @@ -145,54 +149,6 @@ class FilesContext(BaseContext): path: str | None = None -class FilesetFileOutput(BaseModel): - file_ref: str - file_url: str - path: str - size: int - cache_status: CacheStatus | None = None - - @classmethod - def from_file_info( - cls, - workspace: str, - name: str, - file_info: FileInfo, - cache_status: CacheStatus | None = None, - ): - return cls( - file_url=f"/apis/files/v2/workspaces/{workspace}/filesets/{name}/-/{file_info.path}", - file_ref=f"{workspace}/{name}#{file_info.path}", - path=file_info.path, - size=file_info.size, - cache_status=cache_status, - ) - - -class ListFilesetFilesResponse(BaseModel): - data: list[FilesetFileOutput] - - @classmethod - def from_file_infos( - cls, - fileset: Fileset, - file_infos: list[FileInfo], - cache_status_map: dict[str, CacheStatus] | None = None, - ): - cache_status_map = cache_status_map or {} - return cls( - data=[ - FilesetFileOutput.from_file_info( - fileset.workspace, - fileset.name, - fi, - cache_status=cache_status_map.get(fi.path), - ) - for fi in file_infos - ] - ) - - @router.post( "/v2/workspaces/{workspace}/filesets", summary="Create Fileset", @@ -317,7 +273,7 @@ async def create_fileset( ) logger.info(f"Started cache warming for fileset {workspace}/{create_request.name}") - return FilesetOutput.from_entity(created) + return fileset_output_from_entity(created) except EntityConflictError as exc: logger.warning(f"Fileset already exists: {workspace}/{create_request.name}") raise HTTPException( @@ -366,7 +322,7 @@ async def list_filesets( ) return FilesetPage( - data=[FilesetOutput.from_entity(e) for e in res.data], + data=[fileset_output_from_entity(e) for e in res.data], pagination=PaginationData.model_validate(res.pagination.model_dump()), sort=sort, ) @@ -389,7 +345,7 @@ async def retrieve_fileset( """ logger.info(f"GET /filesets/{name} - workspace={workspace}") retrieved = await get_fileset(workspace, name, entity_store) - return FilesetOutput.from_entity(retrieved) + return fileset_output_from_entity(retrieved) @router.delete( @@ -444,7 +400,7 @@ async def delete_fileset( await entity_store.delete(Fileset, fileset.name, workspace=workspace) # Return the fileset data that was captured before deletion - return FilesetOutput.from_entity(fileset) + return fileset_output_from_entity(fileset) @router.patch( @@ -486,7 +442,7 @@ async def update_fileset_metadata( fileset = fileset.model_copy(update=diff) await entity_store.update(fileset) - return FilesetOutput.from_entity(fileset) + return fileset_output_from_entity(fileset) @router.get( @@ -542,7 +498,7 @@ async def list_fileset_files( # External storage without opt-in: null (didn't check) cache_status_map = {} - return ListFilesetFilesResponse.from_file_infos(fileset, files, cache_status_map) + return list_fileset_files_from_infos(fileset, files, cache_status_map) @router.head( @@ -718,7 +674,7 @@ async def upload_file( try: async with streaming_file_upload(request, chunk_processor) as upload: file_info = await storage.upload(path, upload, content_length=content_length) - return FilesetFileOutput.from_file_info( + return fileset_file_output_from_info( workspace=workspace, name=name, file_info=file_info, @@ -786,4 +742,4 @@ async def delete_file( f"File '{path}' not found in fileset '{workspace}/{name}'", ) from e - return FilesetFileOutput.from_file_info(workspace, name, file_info) + return fileset_file_output_from_info(workspace, name, file_info) diff --git a/services/core/files/src/nmp/core/files/api/v2/filesets/schemas.py b/services/core/files/src/nmp/core/files/api/v2/filesets/schemas.py index b2fe6b0473..388cae5425 100644 --- a/services/core/files/src/nmp/core/files/api/v2/filesets/schemas.py +++ b/services/core/files/src/nmp/core/files/api/v2/filesets/schemas.py @@ -1,50 +1,95 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Request and response schemas for filesets API.""" +"""Request and response schemas for filesets API. -from typing import Annotated, Any, Dict, Optional +Response types (FilesetOutput, FilesetFileOutput, ListFilesetFilesResponse) and +request types (CreateFilesetRequest, UpdateFilesetRequest) are imported from +``nemo_platform_plugin.files.types`` — the shared single source of truth. +This module adds server-specific concerns: converter functions that map domain +entities to response DTOs, the FilesetFilter schema, and the FilesetPage alias. +""" + +from typing import Annotated, Optional + +from nemo_platform_plugin.files.types import CacheStatus +from nemo_platform_plugin.files.types import CreateFilesetRequest as CreateFilesetRequest +from nemo_platform_plugin.files.types import FilesetFileOutput as FilesetFileOutput +from nemo_platform_plugin.files.types import FilesetOutput as FilesetOutput +from nemo_platform_plugin.files.types import ListFilesetFilesResponse as ListFilesetFilesResponse +from nemo_platform_plugin.files.types import UpdateFilesetRequest as UpdateFilesetRequest from nmp.common.api.common import Page -from nmp.common.entities import constants from nmp.common.entities.values import DatetimeFilter, Filter, StringFilter, map_entity_field -from nmp.common.files.metadata import FilesetMetadata +from nmp.core.files.app.backends import FileInfo from nmp.core.files.app.backends.base import StorageConfigType -from nmp.core.files.app.backends.factory import StorageConfig from nmp.core.files.entities import Fileset, FilesetPurpose -from pydantic import BaseModel, Field - - -class FilesetOutput(BaseModel): - """Response DTO for fileset operations.""" - - id: str - name: str - workspace: str - description: str - purpose: FilesetPurpose - storage: StorageConfig - metadata: FilesetMetadata - custom_fields: Dict[str, Any] - project: str - created_at: str - updated_at: str - - @classmethod - def from_entity(cls, entity: Fileset) -> "FilesetOutput": - return cls( - id=entity.id, - name=entity.name, - workspace=entity.workspace, - description=entity.description or "", - purpose=entity.purpose, - storage=entity.storage, - metadata=entity.metadata, - custom_fields=entity.custom_fields, - project=entity.project or "", - created_at=entity.created_at.isoformat() if entity.created_at else "", - updated_at=entity.updated_at.isoformat() if entity.updated_at else "", - ) +from pydantic import Field + +FilesetPage = Page[FilesetOutput] + + +# --------------------------------------------------------------------------- +# Entity → DTO converters +# --------------------------------------------------------------------------- + + +def fileset_output_from_entity(entity: Fileset) -> FilesetOutput: + """Convert a Fileset domain entity to a FilesetOutput response DTO.""" + return FilesetOutput( + id=entity.id, + name=entity.name, + workspace=entity.workspace, + description=entity.description or "", + purpose=entity.purpose, + storage=entity.storage, + metadata=entity.metadata, + custom_fields=entity.custom_fields, + project=entity.project or "", + created_at=entity.created_at.isoformat() if entity.created_at else "", + updated_at=entity.updated_at.isoformat() if entity.updated_at else "", + ) + + +def fileset_file_output_from_info( + workspace: str, + name: str, + file_info: FileInfo, + cache_status: CacheStatus | None = None, +) -> FilesetFileOutput: + """Convert a FileInfo to a FilesetFileOutput response DTO.""" + return FilesetFileOutput( + file_url=f"/apis/files/v2/workspaces/{workspace}/filesets/{name}/-/{file_info.path}", + file_ref=f"{workspace}/{name}#{file_info.path}", + path=file_info.path, + size=file_info.size, + cache_status=cache_status, + ) + + +def list_fileset_files_from_infos( + fileset: Fileset, + file_infos: list[FileInfo], + cache_status_map: dict[str, CacheStatus] | None = None, +) -> ListFilesetFilesResponse: + """Convert a list of FileInfos to a ListFilesetFilesResponse.""" + cache_status_map = cache_status_map or {} + return ListFilesetFilesResponse( + data=[ + fileset_file_output_from_info( + fileset.workspace, + fileset.name, + fi, + cache_status=cache_status_map.get(fi.path), + ) + for fi in file_infos + ] + ) + + +# --------------------------------------------------------------------------- +# Filter schema (server-only, not shared with client) +# --------------------------------------------------------------------------- class FilesetFilter(Filter): @@ -68,58 +113,3 @@ class FilesetFilter(Filter): default=None, description="Filter by update date. Supports '$gte' (on or after) and '$lte' (on or before) datetime filters.", ) - - -class CreateFilesetRequest(BaseModel): - name: str = Field( - description=f"The name of the fileset. {constants.REGEX_WORD_CHARACTER_DOT_DASH_DESCRIPTION}", - max_length=constants.MAX_LENGTH_255, - pattern=constants.REGEX_WORD_CHARACTER_DOT_DASH, - examples=["training-data-v1", "llama-checkpoint"], - ) - description: Optional[str] = Field( - default=None, - description="The description of the fileset.", - max_length=constants.MAX_LENGTH_255, - ) - project: Optional[str] = Field( - default=None, - description="The name of the project associated with this fileset.", - ) - storage: StorageConfig | None = Field( - default=None, - description="The storage configuration for the fileset. If not provided, uses default storage.", - ) - - # TODO: Make this a required field eventually - purpose: FilesetPurpose = Field(default=FilesetPurpose.GENERIC, description="The purpose of the fileset.") - metadata: FilesetMetadata = Field( - default_factory=FilesetMetadata, - description="Purpose-specific metadata. Use the purpose as the key (e.g., {dataset: {...}}).", - ) - custom_fields: Dict[str, Any] = Field(default_factory=dict, description="Custom fields for the fileset.") - cache: bool = Field( - default=False, - description="Cache all files after creation. Only applies to external storage.", - ) - - -FilesetPage = Page[FilesetOutput] - - -class UpdateFilesetRequest(BaseModel): - description: str | None = Field( - default=None, - description="The description of the fileset.", - max_length=constants.MAX_LENGTH_255, - ) - project: str | None = Field( - default=None, - description="The name of the project associated with this fileset.", - ) - purpose: FilesetPurpose | None = Field(default=None, description="The purpose of the fileset.") - metadata: FilesetMetadata | None = Field( - default=None, - description="Purpose-specific metadata. Use the purpose as the key (e.g., {dataset: {...}}).", - ) - custom_fields: Dict[str, Any] | None = Field(default=None, description="Custom fields for the fileset.") diff --git a/services/core/files/src/nmp/core/files/app/cache.py b/services/core/files/src/nmp/core/files/app/cache.py index fe0142fe1d..3797d06a89 100644 --- a/services/core/files/src/nmp/core/files/app/cache.py +++ b/services/core/files/src/nmp/core/files/app/cache.py @@ -4,9 +4,9 @@ """Cache utilities for downloading files to cache storage.""" import logging -from enum import StrEnum import anyio +from nemo_platform_plugin.files.types import CacheStatus as CacheStatus from nmp.core.files.app.backends.base import StorageImpl from nmp.core.files.app.file_lock import FileLockManager from nmp.core.files.exceptions import NotFoundError @@ -43,15 +43,6 @@ def reset_background_cache_limiter() -> None: _background_cache_limiter = None -class CacheStatus(StrEnum): - """Cache status for files in external storage backends.""" - - CACHED = "cached" - CACHING = "caching" - NOT_CACHED = "not_cached" - NOT_CACHEABLE = "not_cacheable" - - async def cache_file_directly( source_storage: StorageImpl, cache_storage: StorageImpl, diff --git a/services/core/files/src/nmp/core/files/entities.py b/services/core/files/src/nmp/core/files/entities.py index f69829a8f8..a7b2670975 100644 --- a/services/core/files/src/nmp/core/files/entities.py +++ b/services/core/files/src/nmp/core/files/entities.py @@ -4,9 +4,9 @@ """Domain entities for the Files service.""" from datetime import datetime -from enum import StrEnum from typing import Any, ClassVar, Dict +from nemo_platform_plugin.files.types import FilesetPurpose as FilesetPurpose from nmp.common.entities import constants from nmp.common.entities.client import EntityBase from nmp.common.files.metadata import FilesetMetadata @@ -14,12 +14,6 @@ from pydantic import Field -class FilesetPurpose(StrEnum): - DATASET = "dataset" - GENERIC = "generic" - MODEL = "model" - - class Fileset(EntityBase): """Fileset domain model - represents a fileset entity.""" diff --git a/services/core/files/tests/integration/test_files_basic.py b/services/core/files/tests/integration/test_files_basic.py index 2fd7428006..216b9b8967 100644 --- a/services/core/files/tests/integration/test_files_basic.py +++ b/services/core/files/tests/integration/test_files_basic.py @@ -22,18 +22,15 @@ import pandas as pd import pytest from fastapi.testclient import TestClient -from nemo_platform import ( - APIStatusError, - ConflictError, - NeMoPlatform, - NotFoundError, -) -from nemo_platform.types.files.fileset import Fileset, LocalStorageConfig +from nemo_platform import APIStatusError, ConflictError, NeMoPlatform, NotFoundError +from nemo_platform.types.files.fileset import Fileset +from nemo_platform_plugin.client import errors as nemo_errors from nmp.core.files.testing.utils import ( DEFAULT_WORKSPACE_ID, HTTPXFileSystem, create_fileset, ) +from pydantic import ValidationError class TestFilesBasic: @@ -47,57 +44,67 @@ def test_fileset_list(self, sdk: NeMoPlatform): """Test listing filesets and filtering by workspace.""" with create_fileset(sdk) as fileset1: with create_fileset(sdk) as fileset2: - filesets = sdk.files.filesets.list(workspace=DEFAULT_WORKSPACE_ID) - assert any(fs.id == fileset1.id for fs in filesets.data) - assert any(fs.id == fileset2.id for fs in filesets.data) + filesets = list(sdk.files.filesets.list(workspace=DEFAULT_WORKSPACE_ID).items()) + assert any(fs.id == fileset1.id for fs in filesets) + assert any(fs.id == fileset2.id for fs in filesets) def test_fileset_list_filter_by_name(self, sdk: NeMoPlatform): """Test listing filesets with name filter.""" with create_fileset(sdk) as fileset1: with create_fileset(sdk) as fileset2: # Filter by exact name of fileset1 - filtered = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"name": fileset1.name}, + filtered = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"name": fileset1.name}, + ).items() ) - assert len(filtered.data) == 1 - assert filtered.data[0].id == fileset1.id - assert filtered.data[0].name == fileset1.name + assert len(filtered) == 1 + assert filtered[0].id == fileset1.id + assert filtered[0].name == fileset1.name # Filter by exact name of fileset2 - filtered2 = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"name": fileset2.name}, + filtered2 = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"name": fileset2.name}, + ).items() ) - assert len(filtered2.data) == 1 - assert filtered2.data[0].id == fileset2.id + assert len(filtered2) == 1 + assert filtered2[0].id == fileset2.id # Filter by non-existent name should return empty - filtered_none = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"name": "non-existent-fileset-name"}, + filtered_none = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"name": "non-existent-fileset-name"}, + ).items() ) - assert len(filtered_none.data) == 0 + assert len(filtered_none) == 0 def test_fileset_list_filter_by_purpose(self, sdk: NeMoPlatform): """Test listing filesets with purpose filter.""" with create_fileset(sdk, purpose="dataset") as dataset_fileset: with create_fileset(sdk, purpose="generic") as generic_fileset: # Filter by purpose=dataset - dataset_filesets = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"purpose": "dataset"}, + dataset_filesets = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"purpose": "dataset"}, + ).items() ) - assert any(fs.id == dataset_fileset.id for fs in dataset_filesets.data) - assert not any(fs.id == generic_fileset.id for fs in dataset_filesets.data) + assert any(fs.id == dataset_fileset.id for fs in dataset_filesets) + assert not any(fs.id == generic_fileset.id for fs in dataset_filesets) # Filter by purpose=generic - generic_filesets = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"purpose": "generic"}, + generic_filesets = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"purpose": "generic"}, + ).items() ) - assert any(fs.id == generic_fileset.id for fs in generic_filesets.data) - assert not any(fs.id == dataset_fileset.id for fs in generic_filesets.data) + assert any(fs.id == generic_fileset.id for fs in generic_filesets) + assert not any(fs.id == dataset_fileset.id for fs in generic_filesets) def test_fileset_list_filter_by_storage_type(self, sdk: NeMoPlatform): """Test listing filesets with storage_type filter.""" @@ -105,14 +112,16 @@ def test_fileset_list_filter_by_storage_type(self, sdk: NeMoPlatform): with create_fileset(sdk) as local_fileset1: with create_fileset(sdk) as local_fileset2: # Filter by storage_type=local - local_filesets = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"storage_type": "local"}, + local_filesets = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"storage_type": "local"}, + ).items() ) - assert any(fs.id == local_fileset1.id for fs in local_filesets.data) - assert any(fs.id == local_fileset2.id for fs in local_filesets.data) + assert any(fs.id == local_fileset1.id for fs in local_filesets) + assert any(fs.id == local_fileset2.id for fs in local_filesets) # Verify storage type is local - for fs in local_filesets.data: + for fs in local_filesets: if fs.id in [local_fileset1.id, local_fileset2.id]: assert fs.storage.type == "local" @@ -124,29 +133,29 @@ def test_fileset_list_pagination(self, sdk: NeMoPlatform): stack.enter_context(create_fileset(sdk, purpose="generic")) # Test first page with page_size=2 - page1 = sdk.files.filesets.list( + resp1 = sdk.files.filesets.list( workspace=DEFAULT_WORKSPACE_ID, page=1, page_size=2, ) - assert len(page1.data) == 2 - assert page1.pagination is not None - assert page1.pagination.page == 1 - assert page1.pagination.page_size == 2 + page1 = resp1.page() + assert len(page1.items) == 2 + assert page1.page == 1 + assert page1.page_size == 2 # Test second page - page2 = sdk.files.filesets.list( + resp2 = sdk.files.filesets.list( workspace=DEFAULT_WORKSPACE_ID, page=2, page_size=2, ) - assert len(page2.data) == 2 - assert page2.pagination is not None - assert page2.pagination.page == 2 + page2 = resp2.page() + assert len(page2.items) == 2 + assert page2.page == 2 # Verify pages have different data - page1_ids = {fs.id for fs in page1.data} - page2_ids = {fs.id for fs in page2.data} + page1_ids = {fs.id for fs in page1.items} + page2_ids = {fs.id for fs in page2.items} assert page1_ids.isdisjoint(page2_ids), "Pages should have different filesets" def test_file_upload_download(self, sdk: NeMoPlatform, fileset: Fileset): @@ -340,6 +349,8 @@ def test_error_handling(self, sdk: NeMoPlatform): pass # Expected # Test 3: Try to download non-existent file + # Binary/streaming operations raise errors after send() returns (deferred), + # so they bypass the _RemappingFilesClient.send() override. with create_fileset(sdk) as fileset: try: sdk.files.download_content( @@ -348,7 +359,7 @@ def test_error_handling(self, sdk: NeMoPlatform): workspace=fileset.workspace, ) assert False, "Should have raised NotFoundError for non-existent file" - except NotFoundError: + except (NotFoundError, nemo_errors.NotFoundError): pass # Expected # Test 4: Try to delete non-existent file @@ -360,7 +371,7 @@ def test_error_handling(self, sdk: NeMoPlatform): workspace=fileset.workspace, ) assert False, "Should have raised NotFoundError when deleting non-existent file" - except NotFoundError: + except (NotFoundError, nemo_errors.NotFoundError): pass # Expected # Test 5: List files in non-existent fileset @@ -395,7 +406,7 @@ def test_fileset_create_rejects_user_provided_local_storage(self, sdk: NeMoPlatf name="reject-local-storage", storage={"type": "local", "path": "/etc"}, ) - assert False, "Should have raised APIStatusError for local storage" + assert False, "Should have raised NemoHTTPError for local storage" except APIStatusError as exc: assert exc.status_code == 400 assert "local storage is not allowed" in str(exc.body).lower() @@ -412,7 +423,7 @@ def test_fileset_create_rejects_s3_use_sdk_auth(self, sdk: NeMoPlatform): "use_sdk_auth": True, }, ) - assert False, "Should have raised APIStatusError for S3 with use_sdk_auth=True" + assert False, "Should have raised NemoHTTPError for S3 with use_sdk_auth=True" except APIStatusError as exc: assert exc.status_code == 400 assert "use_sdk_auth=true is not allowed" in str(exc.body).lower() @@ -574,7 +585,10 @@ def test_fileset_create_with_dataset_metadata(self, sdk: NeMoPlatform): def test_fileset_create_rejects_invalid_dataset_schema_metadata(self, sdk: NeMoPlatform): """Test invalid JSON Schema metadata is rejected at fileset create time.""" - with pytest.raises(APIStatusError, match="definitely-not-a-valid-json-schema-type"): + with pytest.raises( + (APIStatusError, ValidationError), + match="definitely-not-a-valid-json-schema-type", + ): with create_fileset( sdk, purpose="dataset", @@ -619,7 +633,7 @@ def test_fileset_delete_removes_storage_data(self, sdk: NeMoPlatform): ) # Verify storage path exists with files - assert isinstance(fileset.storage, LocalStorageConfig) + assert fileset.storage.type == "local" storage_path = Path(fileset.storage.path) assert storage_path.exists() assert (storage_path / "file1.txt").exists() @@ -648,11 +662,13 @@ def test_fileset_list_filter_by_created_at_gte(self, sdk: NeMoPlatform): time.sleep(1) with create_fileset(sdk) as fileset2: # Filter by created_at[gte] should include both new filesets - filtered = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"created_at": {"$gte": before_create.isoformat(timespec="seconds")}}, + filtered = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"created_at": {"$gte": before_create.isoformat(timespec="seconds")}}, + ).items() ) - fileset_ids = {fs.id for fs in filtered.data} + fileset_ids = {fs.id for fs in filtered} assert fileset1.id in fileset_ids assert fileset2.id in fileset_ids @@ -669,11 +685,13 @@ def test_fileset_list_filter_by_created_at_lte(self, sdk: NeMoPlatform): # Filter by created_at[lte] with time after first fileset # should include first fileset but might include second # (depends on timing precision) - filtered = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"created_at": {"$lte": after_first.isoformat()}}, + filtered = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"created_at": {"$lte": after_first.isoformat()}}, + ).items() ) - fileset_ids = {fs.id for fs in filtered.data} + fileset_ids = {fs.id for fs in filtered} assert fileset1.id in fileset_ids def test_fileset_list_filter_by_created_at_range(self, sdk: NeMoPlatform): @@ -684,16 +702,18 @@ def test_fileset_list_filter_by_created_at_range(self, sdk: NeMoPlatform): after_create = datetime.now(timezone.utc) # Filter by date range that includes the fileset - filtered = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={ - "created_at": { - "$gte": before_create.isoformat(), - "$lte": after_create.isoformat(), - } - }, + filtered = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={ + "created_at": { + "$gte": before_create.isoformat(), + "$lte": after_create.isoformat(), + } + }, + ).items() ) - fileset_ids = {fs.id for fs in filtered.data} + fileset_ids = {fs.id for fs in filtered} assert fileset.id in fileset_ids def test_fileset_list_filter_by_created_at_excludes_older(self, sdk: NeMoPlatform): @@ -709,11 +729,13 @@ def test_fileset_list_filter_by_created_at_excludes_older(self, sdk: NeMoPlatfor with create_fileset(sdk) as new_fileset: # Filter by created_at[gte] after old fileset was created - filtered = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"created_at": {"$gte": after_old.isoformat()}}, + filtered = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"created_at": {"$gte": after_old.isoformat()}}, + ).items() ) - fileset_ids = {fs.id for fs in filtered.data} + fileset_ids = {fs.id for fs in filtered} # New fileset should be included assert new_fileset.id in fileset_ids # Old fileset should be excluded @@ -725,11 +747,13 @@ def test_fileset_list_filter_by_updated_at(self, sdk: NeMoPlatform): with create_fileset(sdk) as fileset: # Filter by updated_at[gte] should include the fileset - filtered = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={"updated_at": {"$gte": before_create.isoformat()}}, + filtered = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={"updated_at": {"$gte": before_create.isoformat()}}, + ).items() ) - fileset_ids = {fs.id for fs in filtered.data} + fileset_ids = {fs.id for fs in filtered} assert fileset.id in fileset_ids def test_fileset_list_combined_filters_with_datetime(self, sdk: NeMoPlatform): @@ -739,14 +763,16 @@ def test_fileset_list_combined_filters_with_datetime(self, sdk: NeMoPlatform): with create_fileset(sdk, purpose="dataset") as dataset_fileset: with create_fileset(sdk, purpose="generic") as generic_fileset: # Combine purpose filter with created_at filter - filtered = sdk.files.filesets.list( - workspace=DEFAULT_WORKSPACE_ID, - filter={ - "purpose": "dataset", - "created_at": {"$gte": before_create.isoformat()}, - }, + filtered = list( + sdk.files.filesets.list( + workspace=DEFAULT_WORKSPACE_ID, + filter={ + "purpose": "dataset", + "created_at": {"$gte": before_create.isoformat()}, + }, + ).items() ) - fileset_ids = {fs.id for fs in filtered.data} + fileset_ids = {fs.id for fs in filtered} # Should include dataset fileset assert dataset_fileset.id in fileset_ids # Should exclude generic fileset diff --git a/services/core/files/tests/integration/test_files_sdk.py b/services/core/files/tests/integration/test_files_sdk.py index f1f898dff6..5f159e995e 100644 --- a/services/core/files/tests/integration/test_files_sdk.py +++ b/services/core/files/tests/integration/test_files_sdk.py @@ -8,7 +8,7 @@ - sdk.files.upload_content() - Upload in-memory data - sdk.files.download() - Download files/directories - sdk.files.download_content() - Download file content to memory -- sdk.files.list() - List files with FilesetFile objects +- sdk.files.list() - List files with FilesetFileOutput objects - sdk.files.delete() - Delete files - fileset_auto_create parameter for upload operations @@ -23,16 +23,16 @@ from pathlib import Path import pytest -from nemo_platform import NeMoPlatform, NotFoundError -from nemo_platform.types.files import FilesetFile -from nemo_platform.types.files.fileset import Fileset +from nemo_platform import NeMoPlatform, NotFoundError, PermissionDeniedError +from nemo_platform_plugin.client import errors as nemo_errors +from nemo_platform_plugin.files.types import FilesetFileOutput, FilesetOutput from nmp.core.files.testing.utils import create_fileset, test_fileset_name class TestFilesUpload: """Tests for sdk.files.upload().""" - def test_upload_single_file(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path: Path): + def test_upload_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path: Path): """Test uploading a single file.""" local_file = tmp_path / "upload.txt" local_file.write_text("Hello, World!") @@ -50,7 +50,7 @@ def test_upload_single_file(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path: assert files.data[0].path == "test.txt" assert files.data[0].size == len("Hello, World!") - def test_upload_directory_contents_with_trailing_slash(self, sdk: NeMoPlatform, fileset: Fileset): + def test_upload_directory_contents_with_trailing_slash(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test uploading directory contents (trailing slash on local_path). With trailing slash: `upload("mydir/")` copies the CONTENTS of mydir. @@ -81,7 +81,7 @@ def test_upload_directory_contents_with_trailing_slash(self, sdk: NeMoPlatform, # Should NOT have mydir/ prefix assert not any(p.startswith("mydir/") for p in paths), f"Files should not have 'mydir/' prefix: {paths}" - def test_upload_directory_itself_without_trailing_slash(self, sdk: NeMoPlatform, fileset: Fileset): + def test_upload_directory_itself_without_trailing_slash(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test uploading directory itself (no trailing slash on local_path). Without trailing slash: `upload("mydir")` copies the directory ITSELF. @@ -112,7 +112,7 @@ def test_upload_directory_itself_without_trailing_slash(self, sdk: NeMoPlatform, # Should NOT have files at root assert "file1.txt" not in paths, f"'file1.txt' should not be at root: {paths}" - def test_upload_to_subdirectory(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path: Path): + def test_upload_to_subdirectory(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path: Path): """Test uploading a file to a subdirectory.""" local_file = tmp_path / "nested.txt" local_file.write_text("nested content") @@ -132,7 +132,7 @@ def test_upload_to_subdirectory(self, sdk: NeMoPlatform, fileset: Fileset, tmp_p class TestFilesDownload: """Tests for sdk.files.download().""" - def test_download_single_file(self, sdk: NeMoPlatform, fileset: Fileset): + def test_download_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test downloading a single file.""" # First upload a file test_content = b"Download test content" @@ -154,7 +154,7 @@ def test_download_single_file(self, sdk: NeMoPlatform, fileset: Fileset): downloaded = Path(tmpdir, "downloaded.txt").read_bytes() assert downloaded == test_content - def test_download_directory(self, sdk: NeMoPlatform, fileset: Fileset): + def test_download_directory(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test downloading an entire directory.""" # Upload multiple files sdk.files.upload_content( @@ -189,7 +189,7 @@ def test_download_directory(self, sdk: NeMoPlatform, fileset: Fileset): assert Path(tmpdir, "file2.txt").read_bytes() == b"content2" assert Path(tmpdir, "nested/file3.txt").read_bytes() == b"content3" - def test_download_entire_fileset(self, sdk: NeMoPlatform, fileset: Fileset): + def test_download_entire_fileset(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test downloading all files from a fileset using default remote_path. Downloading a fileset copies contents directly. Users who want a subfolder @@ -227,13 +227,13 @@ def test_download_entire_fileset(self, sdk: NeMoPlatform, fileset: Fileset): class TestFilesList: """Tests for sdk.files.list().""" - def test_list_empty_fileset(self, sdk: NeMoPlatform, fileset: Fileset): + def test_list_empty_fileset(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test listing files in an empty fileset.""" files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) assert files.data == [] - def test_list_returns_fileset_file_objects(self, sdk: NeMoPlatform, fileset: Fileset): - """Test that list returns FilesetFile objects with correct attributes.""" + def test_list_returns_fileset_file_objects(self, sdk: NeMoPlatform, fileset: FilesetOutput): + """Test that list returns FilesetFileOutput objects with correct attributes.""" content = b"test content for size check" sdk.files.upload_content( content=content, @@ -251,7 +251,7 @@ def test_list_returns_fileset_file_objects(self, sdk: NeMoPlatform, fileset: Fil assert file.file_ref == f"{fileset.workspace}/{fileset.name}#test.txt" assert file.file_url == f"/apis/files/v2/workspaces/{fileset.workspace}/filesets/{fileset.name}/-/test.txt" - def test_list_multiple_files(self, sdk: NeMoPlatform, fileset: Fileset): + def test_list_multiple_files(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test listing multiple files.""" sdk.files.upload_content( content=b"a", @@ -283,7 +283,7 @@ def test_list_multiple_files(self, sdk: NeMoPlatform, fileset: Fileset): assert sizes["file2.txt"] == 2 assert sizes["dir/file3.txt"] == 3 - def test_list_subdirectory(self, sdk: NeMoPlatform, fileset: Fileset): + def test_list_subdirectory(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test listing files in a subdirectory.""" sdk.files.upload_content( content=b"root", @@ -320,7 +320,7 @@ def test_list_subdirectory(self, sdk: NeMoPlatform, fileset: Fileset): paths = {f.path for f in files.data} assert paths == {"data/file1.txt", "data/file2.txt"} - def test_list_with_path_format(self, sdk: NeMoPlatform, fileset: Fileset): + def test_list_with_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test listing using full path format instead of explicit fileset param.""" sdk.files.upload_content( content=b"content", @@ -337,7 +337,7 @@ def test_list_with_path_format(self, sdk: NeMoPlatform, fileset: Fileset): assert len(files.data) == 1 assert files.data[0].path == "test.txt" - def test_list_with_glob_pattern(self, sdk: NeMoPlatform, fileset: Fileset): + def test_list_with_glob_pattern(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test listing files matching a glob pattern.""" sdk.files.upload_content( content=b"json", @@ -374,7 +374,7 @@ def test_list_with_glob_pattern(self, sdk: NeMoPlatform, fileset: Fileset): paths = {f.path for f in files.data} assert paths == {"data.json", "config.json"} - def test_list_with_glob_pattern_in_subdirectory(self, sdk: NeMoPlatform, fileset: Fileset): + def test_list_with_glob_pattern_in_subdirectory(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test listing files matching a glob pattern in a subdirectory.""" sdk.files.upload_content( content=b"train", @@ -415,7 +415,7 @@ def test_list_with_glob_pattern_in_subdirectory(self, sdk: NeMoPlatform, fileset class TestFilesGlobDownload: """Tests for sdk.files.download() with glob patterns.""" - def test_download_with_glob_pattern(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path): + def test_download_with_glob_pattern(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path): """Test downloading files matching a glob pattern.""" sdk.files.upload_content( content=b"json content", @@ -453,7 +453,7 @@ def test_download_with_glob_pattern(self, sdk: NeMoPlatform, fileset: Fileset, t assert (tmp_path / "data.json").read_bytes() == b"json content" assert (tmp_path / "config.json").read_bytes() == b"config content" - def test_download_with_glob_pattern_preserves_structure(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path): + def test_download_with_glob_pattern_preserves_structure(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path): """Test that downloading with glob pattern preserves directory structure.""" sdk.files.upload_content( content=b"train data", @@ -495,7 +495,7 @@ def test_download_with_glob_pattern_preserves_structure(self, sdk: NeMoPlatform, class TestFilesDelete: """Tests for sdk.files.delete().""" - def test_delete_single_file(self, sdk: NeMoPlatform, fileset: Fileset): + def test_delete_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test deleting a single file.""" sdk.files.upload_content( content=b"delete me", @@ -519,7 +519,7 @@ def test_delete_single_file(self, sdk: NeMoPlatform, fileset: Fileset): files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 0 - def test_delete_nested_file(self, sdk: NeMoPlatform, fileset: Fileset): + def test_delete_nested_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test deleting a file in a nested directory.""" sdk.files.upload_content( content=b"nested", @@ -537,7 +537,7 @@ def test_delete_nested_file(self, sdk: NeMoPlatform, fileset: Fileset): files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 0 - def test_delete_with_path_format(self, sdk: NeMoPlatform, fileset: Fileset): + def test_delete_with_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test deleting using full path format.""" sdk.files.upload_content( content=b"content", @@ -558,7 +558,7 @@ def test_delete_with_path_format(self, sdk: NeMoPlatform, fileset: Fileset): class TestFilesRoundTrip: """End-to-end tests combining multiple operations.""" - def test_upload_list_download_delete_cycle(self, sdk: NeMoPlatform, fileset: Fileset): + def test_upload_list_download_delete_cycle(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test a complete cycle of file operations.""" with tempfile.TemporaryDirectory() as tmpdir: # Create local files @@ -672,7 +672,7 @@ class TestFilesUploadContent: pytest.param(_chunk_generator(), b"chunk1chunk2chunk3", id="iterator"), ], ) - def test_upload_content(self, sdk: NeMoPlatform, fileset: Fileset, content, expected_bytes: bytes): + def test_upload_content(self, sdk: NeMoPlatform, fileset: FilesetOutput, content, expected_bytes: bytes): """Test uploading different content types.""" result = sdk.files.upload_content( content=content, @@ -681,7 +681,7 @@ def test_upload_content(self, sdk: NeMoPlatform, fileset: Fileset, content, expe workspace=fileset.workspace, ) - assert isinstance(result, Fileset) + assert isinstance(result, FilesetOutput) assert result.name == fileset.name assert result.workspace == fileset.workspace @@ -692,7 +692,7 @@ def test_upload_content(self, sdk: NeMoPlatform, fileset: Fileset, content, expe ) assert downloaded == expected_bytes - def test_upload_content_to_subdirectory(self, sdk: NeMoPlatform, fileset: Fileset): + def test_upload_content_to_subdirectory(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test uploading data to a nested path.""" sdk.files.upload_content( content=b"nested content", @@ -725,7 +725,7 @@ class TestFilesDownloadContent: ), ], ) - def test_download_content(self, sdk: NeMoPlatform, fileset: Fileset, upload_content, expected_bytes: bytes): + def test_download_content(self, sdk: NeMoPlatform, fileset: FilesetOutput, upload_content, expected_bytes: bytes): """Test download_content returns correct bytes for different content types.""" sdk.files.upload_content( content=upload_content, @@ -743,7 +743,7 @@ def test_download_content(self, sdk: NeMoPlatform, fileset: Fileset, upload_cont assert isinstance(result, bytes) assert result == expected_bytes - def test_download_content_with_path_format(self, sdk: NeMoPlatform, fileset: Fileset): + def test_download_content_with_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test download_content using full path format.""" sdk.files.upload_content( content=b"content", @@ -779,8 +779,8 @@ def test_upload_creates_fileset(self, sdk: NeMoPlatform, tmp_path: Path, fileset fileset_auto_create=True, ) - # Verify return type is Fileset with correct info - assert isinstance(result, Fileset) + # Verify return type is FilesetOutput with correct info + assert isinstance(result, FilesetOutput) assert result.name == fileset_name assert result.workspace == workspace @@ -803,8 +803,8 @@ def test_upload_content_creates_fileset(self, sdk: NeMoPlatform, fileset_cleanup fileset_auto_create=True, ) - # Verify return type is Fileset with correct info - assert isinstance(result, Fileset) + # Verify return type is FilesetOutput with correct info + assert isinstance(result, FilesetOutput) assert result.name == fileset_name assert result.workspace == workspace @@ -827,7 +827,7 @@ def test_upload_without_flag_fails_for_nonexistent_fileset(self, sdk: NeMoPlatfo fileset_auto_create=False, ) - def test_existing_fileset_with_flag_succeeds(self, sdk: NeMoPlatform, fileset: Fileset): + def test_existing_fileset_with_flag_succeeds(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test that fileset_auto_create works for existing filesets.""" result = sdk.files.upload_content( content=b"test content", @@ -837,14 +837,14 @@ def test_existing_fileset_with_flag_succeeds(self, sdk: NeMoPlatform, fileset: F fileset_auto_create=True, # Should not fail even though fileset exists ) - assert isinstance(result, Fileset) + assert isinstance(result, FilesetOutput) assert result.name == fileset.name files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 1 - def test_upload_returns_fileset(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path: Path): - """Test that upload() always returns the Fileset entity.""" + def test_upload_returns_fileset(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path: Path): + """Test that upload() always returns the FilesetOutput entity.""" local_file = tmp_path / "test.txt" local_file.write_text("content") @@ -855,8 +855,8 @@ def test_upload_returns_fileset(self, sdk: NeMoPlatform, fileset: Fileset, tmp_p workspace=fileset.workspace, ) - # Even without fileset_auto_create, upload now returns Fileset - assert isinstance(result, Fileset) + # Even without fileset_auto_create, upload now returns FilesetOutput + assert isinstance(result, FilesetOutput) assert result.name == fileset.name assert result.workspace == fileset.workspace @@ -877,8 +877,8 @@ def test_auto_create_generates_name_when_no_fileset_specified( # Register for cleanup now that we know the name fileset_cleanup(result.name) - # Should return a Fileset with a generated name - assert isinstance(result, Fileset) + # Should return a FilesetOutput with a generated name + assert isinstance(result, FilesetOutput) assert result.name.startswith("fileset-") assert len(result.name) == len("fileset-") + 8 # "fileset-" + 8 hex chars @@ -903,7 +903,7 @@ def test_auto_create_uses_fileset_from_path_syntax(self, sdk: NeMoPlatform, file ) # Should use the fileset from the path, not generate a new one - assert isinstance(result, Fileset) + assert isinstance(result, FilesetOutput) assert result.name == fileset_name # Should NOT be "fileset-..." # Verify file was uploaded to correct path @@ -952,7 +952,7 @@ def test_cache_status_aggregation(self, statuses: list, expected: str | None): from nemo_platform.filesets import ListFilesResponse files = [ - FilesetFile( + FilesetFileOutput( path=f"file{i}.txt", size=100 * (i + 1), file_ref=f"ws/fs#file{i}.txt", @@ -968,7 +968,7 @@ def test_cache_status_aggregation(self, statuses: list, expected: str | None): class TestFilesListCacheStatus: """Tests for sdk.files.list() with include_cache_status parameter.""" - def test_list_with_include_cache_status(self, sdk: NeMoPlatform, fileset: Fileset): + def test_list_with_include_cache_status(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test listing files with cache status included.""" sdk.files.upload_content( content=b"test content", @@ -989,7 +989,7 @@ def test_list_with_include_cache_status(self, sdk: NeMoPlatform, fileset: Filese # The important thing is that the parameter is passed through correctly assert files.data[0].path == "test.txt" - def test_list_without_include_cache_status(self, sdk: NeMoPlatform, fileset: Fileset): + def test_list_without_include_cache_status(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test listing files without cache status (default).""" sdk.files.upload_content( content=b"test content", @@ -1011,7 +1011,7 @@ def test_list_without_include_cache_status(self, sdk: NeMoPlatform, fileset: Fil class TestFilesDownloadEdgeCases: """Tests for sdk.files.download() edge cases.""" - def test_download_glob_no_matches(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path): + def test_download_glob_no_matches(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path): """Test downloading with glob pattern that matches no files.""" # Upload a file that won't match the pattern sdk.files.upload_content( @@ -1033,9 +1033,10 @@ def test_download_glob_no_matches(self, sdk: NeMoPlatform, fileset: Fileset, tmp downloaded = list(tmp_path.rglob("*")) assert len([f for f in downloaded if f.is_file()]) == 0 - def test_download_content_non_existent_file(self, sdk: NeMoPlatform, fileset: Fileset): + def test_download_content_non_existent_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test downloading content of a file that doesn't exist raises NotFoundError.""" - with pytest.raises(NotFoundError): + # Binary/streaming errors are deferred (raised after send()), bypassing remapping. + with pytest.raises((NotFoundError, nemo_errors.NotFoundError)): sdk.files.download_content( fileset=fileset.name, workspace=fileset.workspace, @@ -1046,9 +1047,10 @@ def test_download_content_non_existent_file(self, sdk: NeMoPlatform, fileset: Fi class TestFilesDeleteEdgeCases: """Tests for sdk.files.delete() edge cases.""" - def test_delete_non_existent_file(self, sdk: NeMoPlatform, fileset: Fileset): + def test_delete_non_existent_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test deleting a file that doesn't exist raises NotFoundError.""" - with pytest.raises(NotFoundError): + # File delete goes through fsspec rm → deferred error path. + with pytest.raises((NotFoundError, nemo_errors.NotFoundError)): sdk.files.delete( fileset=fileset.name, workspace=fileset.workspace, @@ -1089,22 +1091,22 @@ def test_service_principal_can_set_service_source_and_upload_then_user_cannot_up custom_fields={"service_source": "customizer"}, ) assert created.custom_fields.get("service_source") == "customizer" - sdk_service.files._upload_file( - "data.txt", - b"from service", + sdk_service.files.upload_content( + content=b"from service", + remote_path="data.txt", + fileset=name, workspace=workspace, - name=name, ) files = sdk_service.files.list(fileset=name, workspace=workspace) assert len(files.data) == 1 assert files.data[0].path == "data.txt" # Non-service principal must not be able to upload (fileset is immutable for them). - with pytest.raises(Exception): - sdk_user.files._upload_file( - "user.txt", - b"from user", + with pytest.raises(PermissionDeniedError): + sdk_user.files.upload_content( + content=b"from user", + remote_path="user.txt", + fileset=name, workspace=workspace, - name=name, ) sdk_service.files.filesets.delete(name=name, workspace=workspace) diff --git a/services/core/files/tests/integration/test_fileset_filesystem.py b/services/core/files/tests/integration/test_fileset_filesystem.py index a5dcf6277a..23960fd91c 100644 --- a/services/core/files/tests/integration/test_fileset_filesystem.py +++ b/services/core/files/tests/integration/test_fileset_filesystem.py @@ -1938,39 +1938,6 @@ def test_duckdb_legacy_path_format(self, sdk: NeMoPlatform, fileset: Fileset): assert list(result["id"]) == list(range(1, 11)) -class TestSDKResourceImports: - """Test that SDK resource imports work correctly after vendoring. - - The vendored filesets module exports FilesetsResource and AsyncFilesetsResource. - The *WithRawResponse and *WithStreamingResponse classes are accessible - via sdk.files.filesets.with_raw_response and sdk.files.filesets.with_streaming_response. - """ - - def test_with_raw_response_accessible(self, sdk: NeMoPlatform, fileset: Fileset): - """Test that sdk.files.filesets.with_raw_response is accessible and works.""" - # Access the raw response wrapper - this validates the import works - raw_filesets = sdk.files.filesets.with_raw_response - assert raw_filesets is not None - - # Make an actual API call with raw response - response = raw_filesets.retrieve(fileset.name, workspace=fileset.workspace) - - # Verify we get a raw response wrapper with headers and can parse data - assert hasattr(response, "headers") - assert response.http_response.status_code == 200 - - # Parse the response - this returns the typed object - parsed = response.parse() - assert parsed.name == fileset.name - assert parsed.workspace == fileset.workspace - - def test_with_streaming_response_accessible(self, sdk: NeMoPlatform): - """Test that sdk.files.filesets.with_streaming_response is accessible.""" - # Access the streaming response wrapper - this validates the import works - streaming_filesets = sdk.files.filesets.with_streaming_response - assert streaming_filesets is not None - - class TestDirCache: """Test directory listing caching behavior. diff --git a/uv.lock b/uv.lock index b03f9e45f3..3192a25606 100644 --- a/uv.lock +++ b/uv.lock @@ -4046,6 +4046,7 @@ data-designer-nemo = [ nemo-platform-plugin = [ { name = "anthropic", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "fastapi", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jsonschema", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "lark", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -4075,6 +4076,7 @@ requires-dist = [ { name = "fastapi", marker = "extra == 'test'", specifier = ">=0.115" }, { name = "httpx", specifier = ">=0.27" }, { name = "httpx", marker = "extra == 'test'", specifier = ">=0.27" }, + { name = "jsonschema", marker = "extra == 'nemo-platform-plugin'", specifier = ">=4.0.0" }, { name = "lark", marker = "extra == 'nemo-platform-plugin'", specifier = ">=1.1.0" }, { name = "nemo-platform", editable = "packages/nemo_platform" }, { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, @@ -4642,6 +4644,7 @@ nemo-guardrails-plugin = [ nemo-platform-plugin = [ { name = "anthropic", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "fastapi", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jsonschema", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "lark", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-platform-sdk", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -5040,6 +5043,7 @@ requires-dist = [ { name = "jinja2", marker = "extra == 'nemo-evaluator-sdk'", specifier = ">=3.1.6" }, { name = "jsonpath-ng", marker = "extra == 'nemo-evaluator-sdk'", specifier = ">=1.7.0" }, { name = "jsonschema", marker = "extra == 'nemo-evaluator-sdk'", specifier = ">=4.23.0" }, + { name = "jsonschema", marker = "extra == 'nemo-platform-plugin'", specifier = ">=4.0.0" }, { name = "kubernetes", marker = "extra == 'all'", specifier = ">=30.1.0" }, { name = "kubernetes", marker = "extra == 'all'", specifier = ">=31.0.0" }, { name = "kubernetes", marker = "extra == 'core-service'", specifier = ">=30.1.0" }, @@ -5478,6 +5482,7 @@ source = { editable = "packages/nemo_platform_plugin" } dependencies = [ { name = "anthropic", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "fastapi", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jsonschema", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "lark", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-platform-sdk", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -5517,6 +5522,7 @@ requires-dist = [ { name = "fastapi", specifier = ">=0.115.4" }, { name = "fsspec", marker = "extra == 'nemo-platform-sdk'", specifier = ">=2023.1.0" }, { name = "httpx", marker = "extra == 'nemo-platform-sdk'", specifier = ">=0.23.0,<1" }, + { name = "jsonschema", specifier = ">=4.0.0" }, { name = "lark", specifier = ">=1.1.0" }, { name = "nemo-platform-sdk", editable = "sdk/python/nemo-platform" }, { name = "ngcsdk", marker = "extra == 'nemo-platform-sdk'", specifier = ">=4.8.2" }, @@ -5805,6 +5811,7 @@ dependencies = [ nemo-platform-plugin = [ { name = "anthropic", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "fastapi", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jsonschema", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "lark", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -5830,6 +5837,7 @@ requires-dist = [ { name = "gunicorn", specifier = ">=23.0.0" }, { name = "httpx", specifier = ">=0.27.2" }, { name = "httpx", marker = "extra == 'test'", specifier = ">=0.27" }, + { name = "jsonschema", marker = "extra == 'nemo-platform-plugin'", specifier = ">=4.0.0" }, { name = "lark", marker = "extra == 'nemo-platform-plugin'", specifier = ">=1.1.0" }, { name = "nemo-platform", editable = "packages/nemo_platform" }, { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" },