diff --git a/docs/cli/reference.mdx b/docs/cli/reference.mdx index e3fa8273fd..69875236c5 100644 --- a/docs/cli/reference.mdx +++ b/docs/cli/reference.mdx @@ -8043,7 +8043,7 @@ Some spec fields cannot be represented as CLI flags: `generate.file_extensions`, ### nemo guardrail -Manage guardrails. +Plugin commands for guardrail. **Usage:** @@ -10492,7 +10492,7 @@ Some spec fields cannot be represented as CLI flags: `config.data.max_sequences_ ### nemo experiments -Manage experiments. +Plugin commands for experiments. **Usage:** @@ -10710,7 +10710,7 @@ nemo experiments update [OPTIONS] PATH_NAME ### nemo intake -Intake operations. +Plugin commands for intake. **Usage:** diff --git a/packages/filesets/src/filesets/__init__.py b/packages/filesets/src/filesets/__init__.py index d4fc4d8a0a..288a0e9b49 100644 --- a/packages/filesets/src/filesets/__init__.py +++ b/packages/filesets/src/filesets/__init__.py @@ -7,6 +7,9 @@ fsspec integration for NeMo Platform filesets via the sdk.files.fsspec property. Located at: nemo_platform/filesets/ (after vendoring) + +``ListFilesResponse`` lives in :mod:`filesets.transfer`; it is also re-exported +here for backwards compatibility. """ from .filesystem.callbacks import RichFileProgressCallback as RichFileProgressCallback @@ -17,4 +20,4 @@ from .filesystem.filesystem import build_fileset_ref as build_fileset_ref from .filesystem.filesystem import parse_fileset_path as parse_fileset_path from .filesystem.filesystem import parse_fileset_ref as parse_fileset_ref -from .resources import ListFilesResponse as ListFilesResponse +from .transfer import ListFilesResponse as ListFilesResponse diff --git a/packages/filesets/src/filesets/resources.py b/packages/filesets/src/filesets/resources.py index 23514a1885..4b9249e696 100644 --- a/packages/filesets/src/filesets/resources.py +++ b/packages/filesets/src/filesets/resources.py @@ -7,15 +7,11 @@ 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 Any, Protocol, runtime_checkable -from fsspec.callbacks import DEFAULT_CALLBACK, Callback -from fsspec.core import has_magic +from fsspec.callbacks import Callback from nemo_platform.resources.files.files import ( AsyncFilesResource as GeneratedAsyncFilesResource, ) @@ -25,70 +21,17 @@ from nemo_platform.resources.files.filesets import AsyncFilesetsResource, FilesetsResource from nemo_platform.resources.files.otlp.otlp import AsyncOtlpResource, OtlpResource from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient -from nemo_platform_plugin.files.types import ( - CacheStatus, - CreateFilesetRequest, - FilesetFileOutput, - FilesetOutput, - ListFilesQueryParams, -) +from nemo_platform_plugin.files.types import CreateFilesetRequest, FilesetOutput +from filesets import transfer from filesets.filesystem.filesystem import ( AsyncFilesetFileSystem, FilesetFileSystem, build_fileset_ref, parse_fileset_path, ) - - -@dataclass -class ListFilesResponse: - """Response from listing files in a fileset. - - Attributes: - data: List of files in the fileset. - - Properties: - cache_status: Aggregate cache status of all files. - - "caching" if any file is actively being cached - - "not_cached" if any file is not cached (and none are caching) - - "cached" if all files are fully cached - - "not_cacheable" if all files cannot be cached - - None if no cache information is available - """ - - data: list[FilesetFileOutput] - - @property - def cache_status(self) -> CacheStatus | None: - """Get aggregate cache status of all files. - - Returns the most relevant status based on priority: - - "caching" if any file is actively being cached - - "not_cached" if any file is not cached (and none are caching) - - "cached" if all files are fully cached - - "not_cacheable" if all files cannot be cached - - None if no cache information is available - """ - if not self.data: - return None - - statuses = [f.cache_status for f in self.data if f.cache_status is not None] - if not statuses: - return None - - # Priority: caching > not_cached > cached > not_cacheable - if "caching" in statuses: - return CacheStatus.CACHING - if "not_cached" in statuses: - return CacheStatus.NOT_CACHED - if all(s == "cached" for s in statuses): - return CacheStatus.CACHED - if all(s == "not_cacheable" for s in statuses): - return CacheStatus.NOT_CACHEABLE - - # Mixed cached/not_cacheable - return cached since some files are cached - return CacheStatus.CACHED +from filesets.transfer import ListFilesResponse as ListFilesResponse +from filesets.transfer import generate_fileset_name as _generate_fileset_name @runtime_checkable @@ -109,37 +52,6 @@ async def read(self, size: int = -1) -> bytes: ... AsyncContent = bytes | str | AsyncReadable | AsyncIterator[bytes] -def _generate_fileset_name() -> str: - """Generate a unique fileset name using UUID.""" - return f"fileset-{uuid.uuid4().hex[:8]}" - - -def _matches_glob(filepath: str, pattern: str) -> bool: - """Match filepath against a glob pattern using pathlib. - - Simple patterns (no /) only match top-level files. - Path patterns (with /) match the full relative path from the right. - - Examples: - _matches_glob("train.json", "*.json") -> True - _matches_glob("subdir/nested.json", "*.json") -> False (nested file) - _matches_glob("subdir/nested.json", "subdir/*.json") -> True - _matches_glob("subdir/nested.json", "*/*.json") -> True - - Args: - filepath: The file path to check (relative path within fileset). - pattern: Glob pattern to match against. - - Returns: - True if the filepath matches the pattern. - """ - if "/" not in pattern: - # Simple pattern - only matches top-level files - return "/" not in filepath and PurePath(filepath).match(pattern) - # Path pattern - match from the right - return PurePath(filepath).match(pattern) - - class FilesResource: """FilesResource with high-level file operations. @@ -272,47 +184,16 @@ def download( ... callback=cb ... ) """ - # Handle list of paths - if isinstance(remote_path, list): - if not remote_path: - return - ws = workspace or self._client.workspace - if fileset is None: - raise ValueError("fileset must be provided when remote_path is a list.") - if ws is None: - raise ValueError("workspace must be provided when remote_path is a list.") - # Build list of (remote, local) path pairs preserving directory structure - rpaths = [build_fileset_ref(p, workspace=ws, fileset=fileset) for p in remote_path] - lpaths = [str(PurePath(local_path) / p) for p in remote_path] - self.fsspec.get(rpath=rpaths, lpath=lpaths, callback=callback or DEFAULT_CALLBACK) - return - - ws, path_fileset, path = parse_fileset_path( - remote_path, - workspace_fallback=workspace or self._client.workspace, + transfer.download( + self._client, + remote_path=remote_path, + local_path=local_path, + fileset=fileset, + workspace=workspace, + callback=callback, + max_workers=max_workers, + filesystem=self.fsspec, ) - fileset = fileset or path_fileset - - if fileset is None: - raise ValueError("Fileset must be specified either as a parameter or in the remote_path.") - - # Handle glob patterns by expanding to list of files first - if has_magic(path): - matching_files = self.list(remote_path=path, fileset=fileset, workspace=ws) - if not matching_files.data: - return - # Build list of (remote, local) path pairs preserving directory structure - rpaths = [build_fileset_ref(f.path, workspace=ws, fileset=fileset) for f in matching_files.data] - lpaths = [str(PurePath(local_path) / f.path) for f in matching_files.data] - self.fsspec.get(rpath=rpaths, lpath=lpaths, callback=callback or DEFAULT_CALLBACK) - else: - fileset_ref = build_fileset_ref(path, workspace=ws, fileset=fileset) - self.fsspec.get( - rpath=fileset_ref, - lpath=local_path, - recursive=True, - callback=callback or DEFAULT_CALLBACK, - ) def upload( self, @@ -383,33 +264,18 @@ def upload( ... ) >>> print(f"Uploaded to: {fileset.name}") # e.g., "fileset-a1b2c3d4" """ - ws, path_fileset, path = parse_fileset_path( - remote_path, - workspace_fallback=workspace or self._client.workspace, - ) - fileset = fileset or path_fileset - - if fileset is None: - if fileset_auto_create: - fileset = _generate_fileset_name() - else: - raise ValueError( - "Fileset must be specified either as a parameter or in the remote_path when fileset_auto_create is False." - ) - - fileset_ref = build_fileset_ref(path, workspace=ws, fileset=fileset) - if fileset_auto_create: - self._ensure_fileset_exists(ws, fileset) - - self.fsspec.put( - lpath=local_path, - rpath=fileset_ref, - recursive=True, - callback=callback or DEFAULT_CALLBACK, + return transfer.upload( + self._client, + local_path=local_path, + remote_path=remote_path, + fileset=fileset, + workspace=workspace, + callback=callback, + max_workers=max_workers, + fileset_auto_create=fileset_auto_create, + filesystem=self.fsspec, ) - return self._client.get_fileset(name=fileset, workspace=ws).data() - def upload_content( self, *, @@ -617,37 +483,13 @@ def list( >>> for f in response.data: ... print(f"{f.path}: {f.cache_status}") """ - ws, path_fileset, path = parse_fileset_path( - remote_path, - workspace_fallback=workspace or self._client.workspace, - ) - fileset = fileset or path_fileset - - if fileset is None: - raise ValueError("Fileset must be specified either as a parameter or in the remote_path.") - - # For glob patterns, list all files then filter client-side - # For path prefixes, the API handles filtering server-side - api_path = None if has_magic(path) else (path or None) - - query_params: ListFilesQueryParams = {} - 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, - name=fileset, - query_params=query_params or None, + return transfer.list_files( + self._client, + remote_path=remote_path, + fileset=fileset, + workspace=workspace, + include_cache_status=include_cache_status, ) - response = response.data() - files = list(response.data) - - # Apply glob filtering if needed - if has_magic(path): - files = [f for f in files if _matches_glob(f.path, path)] - return ListFilesResponse(data=files) def delete( self, @@ -676,17 +518,13 @@ def delete( # Delete using full path >>> sdk.files.delete(remote_path="my-fileset#data/old-file.txt") """ - ws, path_fileset, path = parse_fileset_path( - remote_path, - workspace_fallback=workspace or self._client.workspace, + transfer.delete( + self._client, + remote_path=remote_path, + fileset=fileset, + workspace=workspace, + filesystem=self.fsspec, ) - fileset = fileset or path_fileset - - if fileset is None: - raise ValueError("Fileset must be specified either as a parameter or in the remote_path.") - - fileset_ref = build_fileset_ref(path, workspace=ws, fileset=fileset) - self.fsspec.rm(fileset_ref) class AsyncFilesResource: @@ -801,48 +639,16 @@ async def download( ... local_path="./downloads/" ... ) """ - # Handle list of paths - if isinstance(remote_path, list): - if not remote_path: - return - ws = workspace or self._client.workspace - if fileset is None: - raise ValueError("fileset must be provided when remote_path is a list.") - if ws is None: - raise ValueError("workspace must be provided when remote_path is a list.") - # Build list of (remote, local) path pairs preserving directory structure - rpaths = [build_fileset_ref(p, workspace=ws, fileset=fileset) for p in remote_path] - lpaths = [str(PurePath(local_path) / p) for p in remote_path] - await self.fsspec._get(rpaths, lpaths, batch_size=max_workers, callback=callback or DEFAULT_CALLBACK) - return - - ws, path_fileset, path = parse_fileset_path( - remote_path, - workspace_fallback=workspace or self._client.workspace, + await transfer.async_download( + self._client, + remote_path=remote_path, + local_path=local_path, + fileset=fileset, + workspace=workspace, + callback=callback, + max_workers=max_workers, + filesystem=self.fsspec, ) - fileset = fileset or path_fileset - - if fileset is None: - raise ValueError("Fileset must be specified either as a parameter or in the remote_path.") - - # Handle glob patterns by expanding to list of files first - if has_magic(path): - matching_files = await self.list(remote_path=path, fileset=fileset, workspace=ws) - if not matching_files.data: - return - # Build list of (remote, local) path pairs preserving directory structure - rpaths = [build_fileset_ref(f.path, workspace=ws, fileset=fileset) for f in matching_files.data] - lpaths = [str(PurePath(local_path) / f.path) for f in matching_files.data] - await self.fsspec._get(rpaths, lpaths, batch_size=max_workers, callback=callback or DEFAULT_CALLBACK) - else: - fileset_ref = build_fileset_ref(path, workspace=ws, fileset=fileset) - await self.fsspec._get( - fileset_ref, - local_path, - recursive=True, - batch_size=max_workers, - callback=callback or DEFAULT_CALLBACK, - ) async def upload( self, @@ -907,30 +713,17 @@ async def upload( ... ) >>> print(f"Uploaded to: {fileset.name}") # e.g., "fileset-a1b2c3d4" """ - ws, path_fileset, path = parse_fileset_path( - remote_path, - workspace_fallback=workspace or self._client.workspace, + return await transfer.async_upload( + self._client, + local_path=local_path, + remote_path=remote_path, + fileset=fileset, + workspace=workspace, + callback=callback, + max_workers=max_workers, + fileset_auto_create=fileset_auto_create, + filesystem=self.fsspec, ) - fileset = fileset or path_fileset - - if fileset is None: - if fileset_auto_create: - fileset = _generate_fileset_name() - else: - raise ValueError( - "Fileset must be specified either as a parameter or in the remote_path when fileset_auto_create is False." - ) - - fileset_ref = build_fileset_ref(path, workspace=ws, fileset=fileset) - if fileset_auto_create: - await self._ensure_fileset_exists(ws, fileset) - - kwargs: dict = {"lpath": local_path, "rpath": fileset_ref, "recursive": True, "batch_size": max_workers} - if callback is not None: - kwargs["callback"] = callback - await self.fsspec._put(**kwargs) - - return (await self._client.get_fileset(name=fileset, workspace=ws)).data() async def upload_content( self, @@ -1135,34 +928,13 @@ async def list( >>> for f in response.data: ... print(f"{f.path}: {f.cache_status}") """ - ws, path_fileset, path = parse_fileset_path(remote_path, workspace_fallback=workspace or self._client.workspace) - fileset = fileset or path_fileset - - if fileset is None: - raise ValueError("Fileset must be specified either as a parameter or in the remote_path.") - - # For glob patterns, list all files then filter client-side - # For path prefixes, the API handles filtering server-side - api_path = None if has_magic(path) else (path or None) - - query_params: ListFilesQueryParams = {} - 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, - name=fileset, - query_params=query_params or None, + return await transfer.async_list_files( + self._client, + remote_path=remote_path, + fileset=fileset, + workspace=workspace, + include_cache_status=include_cache_status, ) - response = response.data() - files = list(response.data) - - # Apply glob filtering if needed - if has_magic(path): - files = [f for f in files if _matches_glob(f.path, path)] - return ListFilesResponse(data=files) async def delete( self, @@ -1191,11 +963,10 @@ async def delete( # Delete using full path >>> await sdk.files.delete(remote_path="my-fileset#data/old-file.txt") """ - ws, path_fileset, path = parse_fileset_path(remote_path, workspace_fallback=workspace or self._client.workspace) - fileset = fileset or path_fileset - - if fileset is None: - raise ValueError("Fileset must be specified either as a parameter or in the remote_path.") - - fileset_ref = build_fileset_ref(path, workspace=ws, fileset=fileset) - await self.fsspec._rm(fileset_ref) + await transfer.async_delete( + self._client, + remote_path=remote_path, + fileset=fileset, + workspace=workspace, + filesystem=self.fsspec, + ) diff --git a/packages/filesets/src/filesets/transfer.py b/packages/filesets/src/filesets/transfer.py new file mode 100644 index 0000000000..9d19ee0434 --- /dev/null +++ b/packages/filesets/src/filesets/transfer.py @@ -0,0 +1,502 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""High-level fileset transfers on the typed Files client. + +``upload``, ``download``, ``list_files`` and ``delete`` (and their async twins) +drive :class:`~filesets.filesystem.filesystem.FilesetFileSystem` from a +:class:`~nemo_platform_plugin.files.client.FilesClient`. They resolve the +``[workspace/]fileset#path`` reference forms, expand glob patterns, create +filesets on demand, and report progress through fsspec callbacks. The CLI and +the SDK ``FilesResource`` both build on these functions. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from pathlib import PurePath + +from fsspec.callbacks import DEFAULT_CALLBACK, Callback +from fsspec.core import has_magic +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.types import ( + CacheStatus, + CreateFilesetRequest, + FilesetFileOutput, + FilesetOutput, + ListFilesQueryParams, +) + +from filesets.filesystem.filesystem import ( + AsyncFilesetFileSystem, + FilesetFileSystem, + build_fileset_ref, + parse_fileset_path, +) + +FILESET_REQUIRED_MESSAGE = "Fileset must be specified either as a parameter or in the remote_path." +FILESET_REQUIRED_WITHOUT_AUTO_CREATE_MESSAGE = ( + "Fileset must be specified either as a parameter or in the remote_path when fileset_auto_create is False." +) + + +@dataclass +class ListFilesResponse: + """Response from listing files in a fileset. + + Attributes: + data: List of files in the fileset. + + Properties: + cache_status: Aggregate cache status of all files. + - "caching" if any file is actively being cached + - "not_cached" if any file is not cached (and none are caching) + - "cached" if all files are fully cached + - "not_cacheable" if all files cannot be cached + - None if no cache information is available + """ + + data: list[FilesetFileOutput] + + @property + def cache_status(self) -> CacheStatus | None: + """Get aggregate cache status of all files. + + Returns the most relevant status based on priority: + - "caching" if any file is actively being cached + - "not_cached" if any file is not cached (and none are caching) + - "cached" if all files are fully cached + - "not_cacheable" if all files cannot be cached + - None if no cache information is available + """ + if not self.data: + return None + + statuses = [f.cache_status for f in self.data if f.cache_status is not None] + if not statuses: + return None + + # Priority: caching > not_cached > cached > not_cacheable + if "caching" in statuses: + return CacheStatus.CACHING + if "not_cached" in statuses: + return CacheStatus.NOT_CACHED + if all(s == "cached" for s in statuses): + return CacheStatus.CACHED + if all(s == "not_cacheable" for s in statuses): + return CacheStatus.NOT_CACHEABLE + + # Mixed cached/not_cacheable - return cached since some files are cached + return CacheStatus.CACHED + + +def generate_fileset_name() -> str: + """Generate a unique fileset name using UUID.""" + return f"fileset-{uuid.uuid4().hex[:8]}" + + +def matches_glob(filepath: str, pattern: str) -> bool: + """Match filepath against a glob pattern using pathlib. + + Simple patterns (no /) only match top-level files. + Path patterns (with /) match the full relative path from the right. + + Examples: + matches_glob("train.json", "*.json") -> True + matches_glob("subdir/nested.json", "*.json") -> False (nested file) + matches_glob("subdir/nested.json", "subdir/*.json") -> True + matches_glob("subdir/nested.json", "*/*.json") -> True + + Args: + filepath: The file path to check (relative path within fileset). + pattern: Glob pattern to match against. + + Returns: + True if the filepath matches the pattern. + """ + if "/" not in pattern: + # Simple pattern - only matches top-level files + return "/" not in filepath and PurePath(filepath).match(pattern) + # Path pattern - match from the right + return PurePath(filepath).match(pattern) + + +def _resolve_target( + remote_path: str, + *, + fileset: str | None, + workspace: str | None, + client_workspace: str | None, +) -> tuple[str, str | None, str]: + """Return ``(workspace, fileset, path)`` for a remote path that may embed a fileset ref.""" + ws, path_fileset, path = parse_fileset_path(remote_path, workspace_fallback=workspace or client_workspace) + return ws, fileset or path_fileset, path + + +def _resolve_upload_fileset(fileset: str | None, *, fileset_auto_create: bool) -> str: + if fileset is not None: + return fileset + if fileset_auto_create: + return generate_fileset_name() + raise ValueError(FILESET_REQUIRED_WITHOUT_AUTO_CREATE_MESSAGE) + + +def list_query_params(path: str, *, include_cache_status: bool = False) -> ListFilesQueryParams | None: + """Query params ``list_files`` sends for *path*: a prefix goes to the server, a glob is filtered client-side.""" + query_params: ListFilesQueryParams = {} + if not has_magic(path) and path: + query_params["path"] = path + if include_cache_status: + query_params["include_cache_status"] = True + return query_params or None + + +def _filter_listed(files: list[FilesetFileOutput], path: str) -> ListFilesResponse: + if has_magic(path): + files = [f for f in files if matches_glob(f.path, path)] + return ListFilesResponse(data=files) + + +def _pairs_for(paths: list[str], *, workspace: str, fileset: str, local_path: str) -> tuple[list[str], list[str]]: + """Build parallel remote/local path lists that preserve directory structure.""" + rpaths = [build_fileset_ref(p, workspace=workspace, fileset=fileset) for p in paths] + lpaths = [str(PurePath(local_path) / p) for p in paths] + return rpaths, lpaths + + +def _async_transfer_kwargs(callback: Callback | None, max_workers: int | None) -> dict: + kwargs: dict = {"batch_size": max_workers} + if callback is not None: + kwargs["callback"] = callback + return kwargs + + +# --------------------------------------------------------------------------- +# Sync API +# --------------------------------------------------------------------------- + + +def _fs(client: FilesClient, filesystem: FilesetFileSystem | None) -> FilesetFileSystem: + return filesystem if filesystem is not None else FilesetFileSystem(client=client) + + +def list_files( + client: FilesClient, + *, + remote_path: str = "", + fileset: str | None = None, + workspace: str | None = None, + include_cache_status: bool = False, +) -> ListFilesResponse: + """List all files in a fileset path (recursive), with optional glob pattern support. + + Args: + client: Typed Files client. + remote_path: Path within the fileset to list. Can be a full path + (e.g., "workspace/fileset#data/" or "fileset#data/") if fileset is not provided, + or a relative path (e.g., "data/") if fileset is provided. + Supports glob patterns (*, ?, []) for filtering files. + 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 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 FilesetFileOutput) and cache_status property. + """ + ws, fileset, path = _resolve_target( + remote_path, fileset=fileset, workspace=workspace, client_workspace=client.workspace + ) + if fileset is None: + raise ValueError(FILESET_REQUIRED_MESSAGE) + + response = client.list_files( + workspace=ws, + name=fileset, + query_params=list_query_params(path, include_cache_status=include_cache_status), + ).data() + return _filter_listed(list(response.data), path) + + +def download( + client: FilesClient, + *, + remote_path: str | list[str] = "", + local_path: str, + fileset: str | None = None, + workspace: str | None = None, + callback: Callback | None = None, + max_workers: int | None = None, + filesystem: FilesetFileSystem | None = None, +) -> None: + """Download files from a fileset to a local path. + + Args: + client: Typed Files client. + remote_path: Path(s) within the fileset to download. Can be: + - A single path (str): Full path (e.g., "workspace/fileset#data/"), + relative path (e.g., "data/"), or glob pattern (e.g., "*.json"). + - A list of paths (list[str]): Multiple specific file paths to download. + When using a list, fileset and workspace must be provided explicitly. + Defaults to "" (root of fileset). + 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 client's default workspace. + callback: Optional progress callback (e.g., RichProgressCallback). + max_workers: Maximum number of concurrent file transfers. + filesystem: Filesystem to transfer through; defaults to one built on *client*. + """ + fs = _fs(client, filesystem) + callback = callback or DEFAULT_CALLBACK + + if isinstance(remote_path, list): + if not remote_path: + return + ws = workspace or client.workspace + if fileset is None: + raise ValueError("fileset must be provided when remote_path is a list.") + if ws is None: + raise ValueError("workspace must be provided when remote_path is a list.") + rpaths, lpaths = _pairs_for(remote_path, workspace=ws, fileset=fileset, local_path=local_path) + fs.get(rpath=rpaths, lpath=lpaths, callback=callback) + return + + ws, fileset, path = _resolve_target( + remote_path, fileset=fileset, workspace=workspace, client_workspace=client.workspace + ) + if fileset is None: + raise ValueError(FILESET_REQUIRED_MESSAGE) + + if has_magic(path): + matching = list_files(client, remote_path=path, fileset=fileset, workspace=ws) + if not matching.data: + return + rpaths, lpaths = _pairs_for( + [f.path for f in matching.data], workspace=ws, fileset=fileset, local_path=local_path + ) + fs.get(rpath=rpaths, lpath=lpaths, callback=callback) + return + + fs.get( + rpath=build_fileset_ref(path, workspace=ws, fileset=fileset), + lpath=local_path, + recursive=True, + callback=callback, + ) + + +def upload( + client: FilesClient, + *, + local_path: str, + remote_path: str = "", + fileset: str | None = None, + workspace: str | None = None, + callback: Callback | None = None, + max_workers: int | None = None, + fileset_auto_create: bool = False, + filesystem: FilesetFileSystem | None = None, +) -> FilesetOutput: + """Upload files from a local path to a fileset. + + Args: + client: Typed Files client. + local_path: Local source path (file or directory). A trailing slash on a + directory uploads its contents rather than the directory itself. + remote_path: Path within the fileset to upload to. Can be a full path + (e.g., "workspace/fileset#data/" or "fileset#data/") if fileset is not provided, + or a relative path (e.g., "data/") if fileset is provided. + 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 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. + When no fileset is specified (neither as param nor in remote_path), + a unique name is generated (e.g., "fileset-a1b2c3d4"). + filesystem: Filesystem to transfer through; defaults to one built on *client*. + + Returns: + FilesetOutput: The fileset that was uploaded to. Check ``fileset.name`` to see + the generated name when using fileset_auto_create without specifying + a fileset. + """ + ws, fileset, path = _resolve_target( + remote_path, fileset=fileset, workspace=workspace, client_workspace=client.workspace + ) + fileset = _resolve_upload_fileset(fileset, fileset_auto_create=fileset_auto_create) + + fileset_ref = build_fileset_ref(path, workspace=ws, fileset=fileset) + if fileset_auto_create: + client.create_fileset(workspace=ws, body=CreateFilesetRequest(name=fileset), exist_ok=True) + + _fs(client, filesystem).put( + lpath=local_path, rpath=fileset_ref, recursive=True, callback=callback or DEFAULT_CALLBACK + ) + + return client.get_fileset(name=fileset, workspace=ws).data() + + +def delete( + client: FilesClient, + *, + remote_path: str, + fileset: str | None = None, + workspace: str | None = None, + filesystem: FilesetFileSystem | None = None, +) -> None: + """Delete a file from a fileset. + + Args: + client: Typed Files client. + remote_path: Path of the file to delete. Can be a full path + (e.g., "workspace/fileset#data/file.txt") if fileset is not provided, + 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 client's default workspace. + filesystem: Filesystem to delete through; defaults to one built on *client*. + """ + ws, fileset, path = _resolve_target( + remote_path, fileset=fileset, workspace=workspace, client_workspace=client.workspace + ) + if fileset is None: + raise ValueError(FILESET_REQUIRED_MESSAGE) + + _fs(client, filesystem).rm(build_fileset_ref(path, workspace=ws, fileset=fileset)) + + +# --------------------------------------------------------------------------- +# Async API +# --------------------------------------------------------------------------- + + +def _async_fs(client: AsyncFilesClient, filesystem: AsyncFilesetFileSystem | None) -> AsyncFilesetFileSystem: + return filesystem if filesystem is not None else AsyncFilesetFileSystem(client=client) + + +async def async_list_files( + client: AsyncFilesClient, + *, + remote_path: str = "", + fileset: str | None = None, + workspace: str | None = None, + include_cache_status: bool = False, +) -> ListFilesResponse: + """Async twin of :func:`list_files`.""" + ws, fileset, path = _resolve_target( + remote_path, fileset=fileset, workspace=workspace, client_workspace=client.workspace + ) + if fileset is None: + raise ValueError(FILESET_REQUIRED_MESSAGE) + + response = await client.list_files( + workspace=ws, + name=fileset, + query_params=list_query_params(path, include_cache_status=include_cache_status), + ) + return _filter_listed(list(response.data().data), path) + + +async def async_download( + client: AsyncFilesClient, + *, + remote_path: str | list[str] = "", + local_path: str, + fileset: str | None = None, + workspace: str | None = None, + callback: Callback | None = None, + max_workers: int | None = None, + filesystem: AsyncFilesetFileSystem | None = None, +) -> None: + """Async twin of :func:`download`.""" + fs = _async_fs(client, filesystem) + cb = callback or DEFAULT_CALLBACK + + if isinstance(remote_path, list): + if not remote_path: + return + ws = workspace or client.workspace + if fileset is None: + raise ValueError("fileset must be provided when remote_path is a list.") + if ws is None: + raise ValueError("workspace must be provided when remote_path is a list.") + rpaths, lpaths = _pairs_for(remote_path, workspace=ws, fileset=fileset, local_path=local_path) + await fs._get(rpaths, lpaths, batch_size=max_workers, callback=cb) + return + + ws, fileset, path = _resolve_target( + remote_path, fileset=fileset, workspace=workspace, client_workspace=client.workspace + ) + if fileset is None: + raise ValueError(FILESET_REQUIRED_MESSAGE) + + if has_magic(path): + matching = await async_list_files(client, remote_path=path, fileset=fileset, workspace=ws) + if not matching.data: + return + rpaths, lpaths = _pairs_for( + [f.path for f in matching.data], workspace=ws, fileset=fileset, local_path=local_path + ) + await fs._get(rpaths, lpaths, batch_size=max_workers, callback=cb) + return + + await fs._get( + build_fileset_ref(path, workspace=ws, fileset=fileset), + local_path, + recursive=True, + batch_size=max_workers, + callback=cb, + ) + + +async def async_upload( + client: AsyncFilesClient, + *, + local_path: str, + remote_path: str = "", + fileset: str | None = None, + workspace: str | None = None, + callback: Callback | None = None, + max_workers: int | None = None, + fileset_auto_create: bool = False, + filesystem: AsyncFilesetFileSystem | None = None, +) -> FilesetOutput: + """Async twin of :func:`upload`.""" + ws, fileset, path = _resolve_target( + remote_path, fileset=fileset, workspace=workspace, client_workspace=client.workspace + ) + fileset = _resolve_upload_fileset(fileset, fileset_auto_create=fileset_auto_create) + + fileset_ref = build_fileset_ref(path, workspace=ws, fileset=fileset) + if fileset_auto_create: + await client.create_fileset(workspace=ws, body=CreateFilesetRequest(name=fileset), exist_ok=True) + + await _async_fs(client, filesystem)._put( + lpath=local_path, rpath=fileset_ref, recursive=True, **_async_transfer_kwargs(callback, max_workers) + ) + + return (await client.get_fileset(name=fileset, workspace=ws)).data() + + +async def async_delete( + client: AsyncFilesClient, + *, + remote_path: str, + fileset: str | None = None, + workspace: str | None = None, + filesystem: AsyncFilesetFileSystem | None = None, +) -> None: + """Async twin of :func:`delete`.""" + ws, fileset, path = _resolve_target( + remote_path, fileset=fileset, workspace=workspace, client_workspace=client.workspace + ) + if fileset is None: + raise ValueError(FILESET_REQUIRED_MESSAGE) + + await _async_fs(client, filesystem)._rm(build_fileset_ref(path, workspace=ws, fileset=fileset)) diff --git a/packages/filesets/tests/test_transfer.py b/packages/filesets/tests/test_transfer.py new file mode 100644 index 0000000000..d7726d0f2a --- /dev/null +++ b/packages/filesets/tests/test_transfer.py @@ -0,0 +1,544 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the Stainless-free transfer helpers in ``filesets.transfer``. + +A small in-memory fileset server answers both the sync client and the async +client the filesystem builds from it, so the tests pin the request sequence, +paths, query params, and bodies that each helper produces. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path + +import httpx +import pytest +from filesets import transfer +from filesets.filesystem.filesystem import FilesetFileSystem +from filesets.transfer import ListFilesResponse +from nemo_platform_plugin.client.errors import NotFoundError +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.types import CacheStatus, FilesetFileOutput +from starlette.requests import Request +from starlette.responses import Response +from starlette.types import Receive, Scope, Send + +BASE = "http://test" +WORKSPACE = "default" + + +def _fileset_json(workspace: str, name: str) -> dict: + return { + "id": f"id-{name}", + "name": name, + "workspace": workspace, + "description": "", + "purpose": "generic", + "storage": {"type": "local", "path": f"/data/{name}"}, + "metadata": {}, + "custom_fields": {}, + "project": "", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + } + + +@dataclass +class FakeFilesServer: + """In-memory fileset store keyed by ``(workspace, fileset)``; records every request.""" + + filesets: dict[tuple[str, str], dict[str, bytes]] = field(default_factory=dict) + requests: list[httpx.Request] = field(default_factory=list) + + def calls(self) -> list[tuple[str, str]]: + return [(request.method, request.url.path) for request in self.requests] + + def __call__(self, request: httpx.Request) -> httpx.Response: + request.read() + self.requests.append(request) + parts = request.url.path.lstrip("/").split("/") + # apis/files/v2/workspaces/{ws}/filesets[/{name}[/files | /-/{path}]] + workspace = parts[4] + if len(parts) == 6: + if request.method == "POST": + name = json.loads(request.content)["name"] + if (workspace, name) in self.filesets: + return httpx.Response(409, json={"detail": "exists"}) + self.filesets[(workspace, name)] = {} + return httpx.Response(201, json=_fileset_json(workspace, name)) + return httpx.Response(405) + name = parts[6] + files = self.filesets.get((workspace, name)) + if files is None: + return httpx.Response(404, json={"detail": f"Fileset '{name}' not found"}) + if len(parts) == 7: + return httpx.Response(200, json=_fileset_json(workspace, name)) + if parts[7] == "files": + prefix = request.url.params.get("path", "") + data = [ + { + "file_ref": f"{workspace}/{name}#{path}", + "file_url": f"/apis/files/v2/workspaces/{workspace}/filesets/{name}/-/{path}", + "path": path, + "size": len(content), + } + for path, content in sorted(files.items()) + if path.startswith(prefix) + ] + return httpx.Response(200, json={"data": data}) + path = "/".join(parts[8:]) + file_json = { + "file_ref": f"{workspace}/{name}#{path}", + "file_url": request.url.path, + "path": path, + "size": len(files.get(path, b"")), + } + if request.method == "PUT": + files[path] = request.content + return httpx.Response(200, json={**file_json, "size": len(request.content)}) + if path not in files: + return httpx.Response(404, json={"detail": "File not found"}) + if request.method == "GET": + return httpx.Response(200, content=files[path], headers={"content-length": str(len(files[path]))}) + if request.method == "DELETE": + del files[path] + return httpx.Response(200, json=file_json) + return httpx.Response(405) + + async def asgi(self, scope: Scope, receive: Receive, send: Send) -> None: + incoming = Request(scope, receive) + body = await incoming.body() + response = self(httpx.Request(incoming.method, str(incoming.url), headers=incoming.headers.raw, content=body)) + await Response(content=response.content, status_code=response.status_code, headers=dict(response.headers))( + scope, receive, send + ) + + +class _HttpClient(httpx.Client): + def __init__(self, server: FakeFilesServer) -> None: + super().__init__(transport=httpx.MockTransport(server), base_url=BASE) + self._server = server + + @property + def asgi_app(self): + return self._server.asgi + + +@pytest.fixture +def server() -> FakeFilesServer: + return FakeFilesServer() + + +@pytest.fixture +def client(server: FakeFilesServer) -> FilesClient: + return FilesClient(base_url=BASE, workspace=WORKSPACE, http_client=_HttpClient(server)) + + +@pytest.fixture +def async_client(server: FakeFilesServer) -> AsyncFilesClient: + return AsyncFilesClient( + base_url=BASE, workspace=WORKSPACE, http_client=httpx.AsyncClient(transport=httpx.ASGITransport(server.asgi)) + ) + + +def _put(request: httpx.Request) -> tuple[str, str]: + return request.method, request.url.path + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("filepath", "pattern", "expected"), + [ + ("train.json", "*.json", True), + ("subdir/nested.json", "*.json", False), + ("subdir/nested.json", "subdir/*.json", True), + ("subdir/nested.json", "*/*.json", True), + ("a/b/c.txt", "b/*.txt", True), + ("a/b/c.txt", "*.md", False), + ], +) +def test_matches_glob(filepath: str, pattern: str, expected: bool) -> None: + assert transfer.matches_glob(filepath, pattern) is expected + + +def test_generate_fileset_name_is_unique_and_prefixed() -> None: + names = {transfer.generate_fileset_name() for _ in range(5)} + assert len(names) == 5 + assert all(name.startswith("fileset-") and len(name) == len("fileset-") + 8 for name in names) + + +def _file(path: str, cache_status: CacheStatus | None) -> FilesetFileOutput: + return FilesetFileOutput(file_ref=f"ws/fs#{path}", file_url="/x", path=path, size=1, cache_status=cache_status) + + +@pytest.mark.parametrize( + ("statuses", "expected"), + [ + ([], None), + ([None, None], None), + ([CacheStatus.CACHED, CacheStatus.CACHING], CacheStatus.CACHING), + ([CacheStatus.CACHED, CacheStatus.NOT_CACHED], CacheStatus.NOT_CACHED), + ([CacheStatus.CACHED, CacheStatus.CACHED], CacheStatus.CACHED), + ([CacheStatus.NOT_CACHEABLE, CacheStatus.NOT_CACHEABLE], CacheStatus.NOT_CACHEABLE), + ([CacheStatus.CACHED, CacheStatus.NOT_CACHEABLE], CacheStatus.CACHED), + ], +) +def test_list_files_response_cache_status(statuses: list[CacheStatus | None], expected: CacheStatus | None) -> None: + response = ListFilesResponse(data=[_file(f"f{i}", status) for i, status in enumerate(statuses)]) + assert response.cache_status == expected + + +# --------------------------------------------------------------------------- +# list_files +# --------------------------------------------------------------------------- + + +def test_list_files_root(server: FakeFilesServer, client: FilesClient) -> None: + server.filesets[(WORKSPACE, "fs")] = {"a.txt": b"a", "d/b.txt": b"bb"} + + response = transfer.list_files(client, fileset="fs") + + assert [(f.path, f.size) for f in response.data] == [("a.txt", 1), ("d/b.txt", 2)] + assert server.calls() == [("GET", "/apis/files/v2/workspaces/default/filesets/fs/files")] + assert dict(server.requests[0].url.params) == {} + + +def test_list_files_prefix_is_sent_as_path_param(server: FakeFilesServer, client: FilesClient) -> None: + server.filesets[(WORKSPACE, "fs")] = {"a.txt": b"a", "d/b.txt": b"bb"} + + response = transfer.list_files(client, fileset="fs", remote_path="d/", include_cache_status=True) + + assert [f.path for f in response.data] == ["d/b.txt"] + assert dict(server.requests[0].url.params) == {"path": "d/", "include_cache_status": "true"} + + +def test_list_files_glob_filters_client_side(server: FakeFilesServer, client: FilesClient) -> None: + server.filesets[(WORKSPACE, "fs")] = {"a.json": b"a", "b.txt": b"b", "d/c.json": b"c"} + + response = transfer.list_files(client, fileset="fs", remote_path="*.json") + + assert [f.path for f in response.data] == ["a.json"] + assert dict(server.requests[0].url.params) == {} + + +def test_list_files_parses_fileset_ref_in_remote_path(server: FakeFilesServer, client: FilesClient) -> None: + server.filesets[("other", "fs")] = {"d/x.txt": b"x"} + + response = transfer.list_files(client, remote_path="other/fs#d/") + + assert [f.path for f in response.data] == ["d/x.txt"] + assert server.calls() == [("GET", "/apis/files/v2/workspaces/other/filesets/fs/files")] + + +def test_list_files_requires_fileset(client: FilesClient) -> None: + with pytest.raises(ValueError, match="Fileset must be specified"): + transfer.list_files(client, remote_path="d/") + + +def test_list_files_missing_fileset_raises_not_found(client: FilesClient) -> None: + with pytest.raises(NotFoundError): + transfer.list_files(client, fileset="nope") + + +# --------------------------------------------------------------------------- +# upload +# --------------------------------------------------------------------------- + + +def test_upload_single_file(server: FakeFilesServer, client: FilesClient, tmp_path: Path) -> None: + server.filesets[(WORKSPACE, "fs")] = {} + local = tmp_path / "a.txt" + local.write_bytes(b"hello") + + result = transfer.upload(client, local_path=str(local), fileset="fs", remote_path="data/") + + assert result.name == "fs" + assert server.filesets[(WORKSPACE, "fs")] == {"data/a.txt": b"hello"} + put = next(r for r in server.requests if r.method == "PUT") + assert put.url.path == "/apis/files/v2/workspaces/default/filesets/fs/-/data/a.txt" + assert put.headers["content-length"] == "5" + assert server.calls()[-1] == ("GET", "/apis/files/v2/workspaces/default/filesets/fs") + + +def test_upload_directory_keeps_name_without_trailing_slash( + server: FakeFilesServer, client: FilesClient, tmp_path: Path +) -> None: + server.filesets[(WORKSPACE, "fs")] = {} + src = tmp_path / "src" + (src / "n").mkdir(parents=True) + (src / "a.txt").write_bytes(b"a") + (src / "n" / "b.txt").write_bytes(b"b") + + transfer.upload(client, local_path=str(src), fileset="fs") + + assert server.filesets[(WORKSPACE, "fs")] == {"src/a.txt": b"a", "src/n/b.txt": b"b"} + + +def test_upload_directory_contents_with_trailing_slash( + server: FakeFilesServer, client: FilesClient, tmp_path: Path +) -> None: + server.filesets[(WORKSPACE, "fs")] = {} + src = tmp_path / "src" + (src / "n").mkdir(parents=True) + (src / "a.txt").write_bytes(b"a") + (src / "n" / "b.txt").write_bytes(b"b") + + transfer.upload(client, local_path=f"{src}/", fileset="fs", remote_path="up/") + + assert server.filesets[(WORKSPACE, "fs")] == {"up/a.txt": b"a", "up/n/b.txt": b"b"} + + +def test_upload_fileset_from_remote_path_ref(server: FakeFilesServer, client: FilesClient, tmp_path: Path) -> None: + server.filesets[("other", "fs")] = {} + local = tmp_path / "a.txt" + local.write_bytes(b"x") + + result = transfer.upload(client, local_path=str(local), remote_path="other/fs#dir/") + + assert result.workspace == "other" + assert server.filesets[("other", "fs")] == {"dir/a.txt": b"x"} + + +def test_upload_requires_fileset_without_auto_create(client: FilesClient, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="fileset_auto_create is False"): + transfer.upload(client, local_path=str(tmp_path)) + + +def test_upload_auto_create_named_fileset_is_idempotent( + server: FakeFilesServer, client: FilesClient, tmp_path: Path +) -> None: + local = tmp_path / "a.txt" + local.write_bytes(b"x") + + first = transfer.upload(client, local_path=str(local), fileset="new", fileset_auto_create=True) + second = transfer.upload(client, local_path=str(local), fileset="new", fileset_auto_create=True) + + assert first.name == second.name == "new" + posts = [r for r in server.requests if r.method == "POST"] + assert [json.loads(r.content) for r in posts] == [{"name": "new"}, {"name": "new"}] + # The second create 409s and is resolved by re-fetching the fileset (exist_ok). + assert server.filesets[(WORKSPACE, "new")] == {"a.txt": b"x"} + + +def test_upload_auto_create_generates_name(server: FakeFilesServer, client: FilesClient, tmp_path: Path) -> None: + local = tmp_path / "a.txt" + local.write_bytes(b"x") + + result = transfer.upload(client, local_path=str(local), fileset_auto_create=True) + + assert result.name.startswith("fileset-") + assert server.filesets[(WORKSPACE, result.name)] == {"a.txt": b"x"} + assert server.calls()[0] == ("POST", "/apis/files/v2/workspaces/default/filesets") + + +def test_upload_uses_supplied_filesystem(server: FakeFilesServer, client: FilesClient, tmp_path: Path) -> None: + server.filesets[(WORKSPACE, "fs")] = {} + local = tmp_path / "a.txt" + local.write_bytes(b"x") + fs = FilesetFileSystem(client=client) + + transfer.upload(client, local_path=str(local), fileset="fs", filesystem=fs) + + assert server.filesets[(WORKSPACE, "fs")] == {"a.txt": b"x"} + + +# --------------------------------------------------------------------------- +# download +# --------------------------------------------------------------------------- + + +@pytest.fixture +def populated(server: FakeFilesServer) -> dict[str, bytes]: + files = {"a/file1.txt": b"content1", "a/b/file2.txt": b"content2", "a/b/file3.txt": b"content3", "r.json": b"{}"} + server.filesets[(WORKSPACE, "fs")] = dict(files) + return files + + +def test_download_single_file_into_existing_directory( + populated: dict[str, bytes], client: FilesClient, tmp_path: Path +) -> None: + out = tmp_path / "out" + out.mkdir() + + transfer.download(client, fileset="fs", remote_path="a/b/file2.txt", local_path=str(out)) + + assert (out / "file2.txt").read_bytes() == b"content2" + + +def test_download_single_file_to_exact_path(populated: dict[str, bytes], client: FilesClient, tmp_path: Path) -> None: + dest = tmp_path / "renamed.txt" + + transfer.download(client, fileset="fs", remote_path="a/b/file2.txt", local_path=str(dest)) + + assert dest.read_bytes() == b"content2" + + +def test_download_directory_contents_with_trailing_slash( + populated: dict[str, bytes], client: FilesClient, tmp_path: Path +) -> None: + out = tmp_path / "out" + + transfer.download(client, fileset="fs", remote_path="a/", local_path=f"{out}/") + + assert (out / "file1.txt").read_bytes() == b"content1" + assert (out / "b" / "file2.txt").read_bytes() == b"content2" + assert (out / "b" / "file3.txt").read_bytes() == b"content3" + + +def test_download_directory_keeps_name_without_trailing_slash( + populated: dict[str, bytes], client: FilesClient, tmp_path: Path +) -> None: + out = tmp_path / "out" + + transfer.download(client, fileset="fs", remote_path="a/b", local_path=str(out)) + + assert (out / "b" / "file2.txt").read_bytes() == b"content2" + assert (out / "b" / "file3.txt").read_bytes() == b"content3" + assert not (out / "file1.txt").exists() + + +def test_download_fileset_root_copies_contents( + populated: dict[str, bytes], client: FilesClient, tmp_path: Path +) -> None: + out = tmp_path / "out" + + transfer.download(client, fileset="fs", local_path=str(out)) + + assert sorted(p.relative_to(out).as_posix() for p in out.rglob("*") if p.is_file()) == sorted(populated) + + +def test_download_glob_preserves_relative_paths( + populated: dict[str, bytes], server: FakeFilesServer, client: FilesClient, tmp_path: Path +) -> None: + out = tmp_path / "out" + + transfer.download(client, fileset="fs", remote_path="a/b/*.txt", local_path=str(out)) + + assert (out / "a" / "b" / "file2.txt").read_bytes() == b"content2" + assert (out / "a" / "b" / "file3.txt").read_bytes() == b"content3" + assert not (out / "a" / "file1.txt").exists() + downloads = sorted(r.url.path for r in server.requests if r.method == "GET" and "/-/" in r.url.path) + assert downloads == [ + "/apis/files/v2/workspaces/default/filesets/fs/-/a/b/file2.txt", + "/apis/files/v2/workspaces/default/filesets/fs/-/a/b/file3.txt", + ] + + +def test_download_glob_without_matches_is_noop( + populated: dict[str, bytes], server: FakeFilesServer, client: FilesClient, tmp_path: Path +) -> None: + transfer.download(client, fileset="fs", remote_path="*.parquet", local_path=str(tmp_path)) + + assert server.calls() == [("GET", "/apis/files/v2/workspaces/default/filesets/fs/files")] + + +def test_download_list_of_paths(populated: dict[str, bytes], client: FilesClient, tmp_path: Path) -> None: + out = tmp_path / "out" + + transfer.download(client, fileset="fs", remote_path=["a/file1.txt", "r.json"], local_path=str(out)) + + assert (out / "a" / "file1.txt").read_bytes() == b"content1" + assert (out / "r.json").read_bytes() == b"{}" + + +def test_download_list_requires_fileset_and_workspace(server: FakeFilesServer, tmp_path: Path) -> None: + no_workspace = FilesClient(base_url=BASE, http_client=_HttpClient(server)) + + with pytest.raises(ValueError, match="fileset must be provided"): + transfer.download(no_workspace, remote_path=["a"], local_path=str(tmp_path)) + with pytest.raises(ValueError, match="workspace must be provided"): + transfer.download(no_workspace, fileset="fs", remote_path=["a"], local_path=str(tmp_path)) + + +def test_download_empty_list_is_noop(server: FakeFilesServer, client: FilesClient, tmp_path: Path) -> None: + transfer.download(client, fileset="fs", remote_path=[], local_path=str(tmp_path)) + + assert server.requests == [] + + +def test_download_requires_fileset(client: FilesClient, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Fileset must be specified"): + transfer.download(client, remote_path="a/", local_path=str(tmp_path)) + + +def test_download_missing_fileset_raises_not_found(client: FilesClient, tmp_path: Path) -> None: + with pytest.raises(NotFoundError): + transfer.download(client, fileset="nope", local_path=str(tmp_path)) + + +# --------------------------------------------------------------------------- +# delete +# --------------------------------------------------------------------------- + + +def test_delete_file(populated: dict[str, bytes], server: FakeFilesServer, client: FilesClient) -> None: + transfer.delete(client, fileset="fs", remote_path="a/b/file2.txt") + + assert server.calls() == [("DELETE", "/apis/files/v2/workspaces/default/filesets/fs/-/a/b/file2.txt")] + assert "a/b/file2.txt" not in server.filesets[(WORKSPACE, "fs")] + + +def test_delete_with_fileset_ref(populated: dict[str, bytes], server: FakeFilesServer, client: FilesClient) -> None: + transfer.delete(client, remote_path="default/fs#r.json") + + assert server.calls() == [("DELETE", "/apis/files/v2/workspaces/default/filesets/fs/-/r.json")] + + +def test_delete_requires_fileset(client: FilesClient) -> None: + with pytest.raises(ValueError, match="Fileset must be specified"): + transfer.delete(client, remote_path="a.txt") + + +def test_delete_missing_file_raises_not_found(populated: dict[str, bytes], client: FilesClient) -> None: + with pytest.raises(NotFoundError): + transfer.delete(client, fileset="fs", remote_path="nope.txt") + + +# --------------------------------------------------------------------------- +# async twins +# --------------------------------------------------------------------------- + + +async def test_async_roundtrip(server: FakeFilesServer, async_client: AsyncFilesClient, tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + (src / "a.txt").write_bytes(b"a") + (src / "b.json").write_bytes(b"{}") + + result = await transfer.async_upload( + async_client, local_path=f"{src}/", fileset="fs", remote_path="up/", fileset_auto_create=True + ) + assert result.name == "fs" + assert server.filesets[(WORKSPACE, "fs")] == {"up/a.txt": b"a", "up/b.json": b"{}"} + + listed = await transfer.async_list_files(async_client, fileset="fs", remote_path="up/*.json") + assert [f.path for f in listed.data] == ["up/b.json"] + + out = tmp_path / "out" + await transfer.async_download(async_client, fileset="fs", remote_path="up/", local_path=f"{out}/") + assert (out / "a.txt").read_bytes() == b"a" + assert (out / "b.json").read_bytes() == b"{}" + + await transfer.async_download(async_client, fileset="fs", remote_path=["up/a.txt"], local_path=str(out / "l")) + assert (out / "l" / "up" / "a.txt").read_bytes() == b"a" + + await transfer.async_delete(async_client, fileset="fs", remote_path="up/a.txt") + assert server.filesets[(WORKSPACE, "fs")] == {"up/b.json": b"{}"} + + +async def test_async_requires_fileset(async_client: AsyncFilesClient, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Fileset must be specified"): + await transfer.async_list_files(async_client, remote_path="x/") + with pytest.raises(ValueError, match="Fileset must be specified"): + await transfer.async_download(async_client, remote_path="x/", local_path=str(tmp_path)) + with pytest.raises(ValueError, match="Fileset must be specified"): + await transfer.async_delete(async_client, remote_path="x") + with pytest.raises(ValueError, match="fileset_auto_create is False"): + await transfer.async_upload(async_client, local_path=str(tmp_path)) diff --git a/packages/models/src/models/resources.py b/packages/models/src/models/resources.py index 7ce619680b..e6c950319e 100644 --- a/packages/models/src/models/resources.py +++ b/packages/models/src/models/resources.py @@ -7,7 +7,6 @@ import logging import time from collections.abc import Awaitable, Callable -from dataclasses import dataclass from datetime import datetime from typing import TypeVar @@ -17,75 +16,20 @@ from nemo_platform.types.inference import ModelDeployment, ModelProvider from nemo_platform.types.inference.gateway.openai.v1 import OpenAIModelResp from nemo_platform.types.models import ModelEntity +from nemo_platform_plugin.models.refs import ( + ResolvedModelReference, + first_provider_ref, + model_entity_route_openai_url, + parse_workspace_name_ref, + resolved_model_reference, + warn_provider_host_url_resolution_failure, +) _T = TypeVar("_T") _TRANSIENT_GATEWAY_STATUS_CODES = {429, 502, 503, 504} _logger = logging.getLogger(__name__) -@dataclass(frozen=True, slots=True) -class ResolvedModelReference: - """Inference route details for a workspace-qualified model reference.""" - - url: str - name: str - host_url: str | None - - -def parse_workspace_name_ref(ref: str, *, label: str, expected_format: str = "workspace/name") -> tuple[str, str]: - """Parse a strict workspace-qualified reference.""" - workspace, separator, name = ref.partition("/") - if separator != "/" or not workspace or not name or "/" in name: - raise ValueError(f"{label} must be in format '{expected_format}'") - return workspace, name - - -def first_provider_ref(model_providers: list[str] | None) -> tuple[str, str, str] | None: - if not model_providers: - return None - - provider_ref = model_providers[0] - try: - provider_workspace, provider_name = parse_workspace_name_ref(provider_ref, label="Provider reference") - except ValueError: - _logger.warning("Invalid provider reference format", extra={"provider_ref": provider_ref}) - return None - return provider_ref, provider_workspace, provider_name - - -def model_entity_route_openai_url(*, base_url: str, workspace: str, name: str) -> str: - """OpenAI SDK-compatible URL for a model-entity proxy route.""" - return f"{base_url.rstrip('/')}/apis/inference-gateway/v2/workspaces/{workspace}/model/{name}/-/v1" - - -def resolved_model_reference( - *, - base_url: str, - name: str, - route_workspace: str, - route_model_name: str, - host_url: str | None, -) -> ResolvedModelReference: - """Build route details for a resolved model entity.""" - return ResolvedModelReference( - url=model_entity_route_openai_url(base_url=base_url, workspace=route_workspace, name=route_model_name), - name=name, - host_url=host_url, - ) - - -def warn_provider_host_url_resolution_failure( - provider_ref: str, - exc: Exception, - *, - not_found_error_type: type[Exception], -) -> None: - if isinstance(exc, not_found_error_type): - _logger.warning("Provider not found during host_url resolution", extra={"provider_ref": provider_ref}) - return - _logger.warning("Failed to resolve provider host_url", extra={"provider_ref": provider_ref}, exc_info=True) - - def _seconds_since_creation(entry_timestamp: datetime | str | None, created_at: datetime | None) -> int | None: """Seconds from deployment creation to the entry timestamp. Returns None if either is missing or not comparable.""" if created_at is None or entry_timestamp is None: diff --git a/packages/nemo_platform/pyproject.toml b/packages/nemo_platform/pyproject.toml index 03b8ad68cd..bf6f8c913b 100644 --- a/packages/nemo_platform/pyproject.toml +++ b/packages/nemo_platform/pyproject.toml @@ -201,6 +201,8 @@ intake-service = [ "pydantic>=2.9.2, <3.0.0", "pydantic-settings>=2.6.1, <3.0.0", "nmp-common", + "nemo-platform-plugin", + "nemo-platform-ext", "clickhouse-connect>=0.7,<1.0", "docker>=7.1.0", "opentelemetry-proto>=1.27.0", @@ -670,6 +672,9 @@ data-designer = "nemo_data_designer_plugin.cli.main:DataDesignerCLI" evaluator = "nemo_evaluator.cli:EvaluatorPluginCLI" insights = "nemo_insights_plugin.cli:InsightsCLI" customization = "nemo_customizer.cli:CustomizationCLI" +guardrail = "nemo_guardrails_plugin.cli:GuardrailCLI" +intake = "nmp.intake.cli:IntakeCLI" +experiments = "nmp.intake.cli:ExperimentsCLI" # Generated from [tool.bundle-package]; do not edit this table by hand. [project.entry-points."nemo.cli.agents"] @@ -890,7 +895,7 @@ nmp-inference-gateway = { source = "../../services/core/inference-gateway/src/nm nmp-guardrails = { source = "../../services/guardrails/src/nmp/guardrails", module = "nmp/guardrails", deps_group = "guardrails-service" } nmp-platform-seed = { source = "../../services/platform-seed/src/nmp/platform_seed", module = "nmp/platform_seed", deps_group = "platform-seed-service" } nmp-hello-world = { source = "../../services/hello-world/src/nmp/hello_world", module = "nmp/hello_world", deps_group = "hello-world-service" } -nmp-intake = { source = "../../services/intake/src/nmp/intake", module = "nmp/intake", deps_group = "intake-service" } +nmp-intake = { source = "../../services/intake/src/nmp/intake", module = "nmp/intake", deps_group = "intake-service", inherit = { "entry-points" = ["nemo.*"] } } # Customization task packages: compile glue and schemas only. Their container # entrypoint scripts are deliberately not re-exported, and the GPU stacks they # drive (torch, unsloth, nemo-rl) live in the training images, not this wheel. diff --git a/packages/nemo_platform_ext/scripts/docs_generator.py b/packages/nemo_platform_ext/scripts/docs_generator.py index d88b7583e1..a80ca2d4e3 100644 --- a/packages/nemo_platform_ext/scripts/docs_generator.py +++ b/packages/nemo_platform_ext/scripts/docs_generator.py @@ -708,7 +708,10 @@ def _escape_mdx_line(line: str) -> str: "customization", "data-designer", "evaluator", + "experiments", + "guardrail", "insights", + "intake", "safe-synthesizer", ) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/app.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/app.py index 38b47571a3..b89bb6e278 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/app.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/app.py @@ -330,9 +330,9 @@ def main( def _version_callback(value: bool) -> None: """Print version information and exit.""" if value: - import nemo_platform + from nemo_platform_ext.cli.version import client_version - typer.echo(f"nemo version {nemo_platform.__version__}") + typer.echo(f"nemo version {client_version()}") raise typer.Exit() diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/__init__.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/__init__.py index 6fad1fd8c3..2c9677db40 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/__init__.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/__init__.py @@ -15,14 +15,6 @@ kind="group", hidden=True, ), - TopLevelEntry( - import_path=f"{__package__}.experiments:app", - name="experiments", - help="Manage experiments.", - panel="Functional plugins", - kind="group", - hidden=False, - ), TopLevelEntry( import_path=f"{__package__}.files:app", name="files", @@ -31,14 +23,6 @@ kind="group", hidden=False, ), - TopLevelEntry( - import_path=f"{__package__}.guardrail:app", - name="guardrail", - help="Manage guardrails.", - panel="Functional plugins", - kind="group", - hidden=False, - ), TopLevelEntry( import_path=f"{__package__}.inference:app", name="inference", @@ -47,14 +31,6 @@ kind="group", hidden=False, ), - TopLevelEntry( - import_path=f"{__package__}.intake:app", - name="intake", - help="Intake operations.", - panel="Functional plugins", - kind="group", - hidden=False, - ), TopLevelEntry( import_path=f"{__package__}.models:app", name="models", diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/intake/__init__.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/intake/__init__.py deleted file mode 100644 index 04767cd639..0000000000 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/intake/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# NOTE: This file is auto-generated -from __future__ import annotations - -from importlib import import_module as _importlib_import_module - -from nemo_platform_ext.cli.core.help_formatter import create_typer_app - -_cli_child_annotations = _importlib_import_module("nemo_platform_ext.cli.commands.api.intake.annotations") -_cli_child_evaluator_results = _importlib_import_module("nemo_platform_ext.cli.commands.api.intake.evaluator_results") -_cli_child_ingest = _importlib_import_module("nemo_platform_ext.cli.commands.api.intake.ingest") -_cli_child_sessions = _importlib_import_module("nemo_platform_ext.cli.commands.api.intake.sessions") -_cli_child_spans = _importlib_import_module("nemo_platform_ext.cli.commands.api.intake.spans") -_cli_child_traces = _importlib_import_module("nemo_platform_ext.cli.commands.api.intake.traces") - -app = create_typer_app(name="intake", help="Intake operations") - -app.add_typer(_cli_child_annotations.app, name="annotations") -app.add_typer(_cli_child_evaluator_results.app, name="evaluator-results") -app.add_typer(_cli_child_ingest.app, name="ingest") -app.add_typer(_cli_child_sessions.app, name="sessions") -app.add_typer(_cli_child_spans.app, name="spans") -app.add_typer(_cli_child_traces.app, name="traces") diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/intake/annotations.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/intake/annotations.py deleted file mode 100644 index 1abd943cb7..0000000000 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/intake/annotations.py +++ /dev/null @@ -1,273 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# NOTE: This file is auto-generated -from __future__ import annotations - -from typing import Annotated, Literal - -import typer - -from nemo_platform_ext.cli.core.api import build_kwargs, merge_filter_dict -from nemo_platform_ext.cli.core.code_generator import handle_code_generation -from nemo_platform_ext.cli.core.context import CLIContext -from nemo_platform_ext.cli.core.errors import handle_errors -from nemo_platform_ext.cli.core.formatters import ( - Column, - check_output_columns_with_format, - format_output, - validate_stream_output_format, -) -from nemo_platform_ext.cli.core.help_formatter import collect_warnings, create_typer_app -from nemo_platform_ext.cli.core.pagination import PaginationType, fetch_all_pages, warn_if_more_pages -from nemo_platform_ext.cli.core.stdin_utils import read_data_input_with_flags, read_payload, validate_required_fields -from nemo_platform_ext.cli.core.types import ( - EntityOutputFormatOption, - ListOutputFormatOption, - NoTruncateOption, - OutputColumnsOption, - StreamOutputOption, -) - -app = create_typer_app(name="annotations", help="Manage annotations") - - -@app.command("create") -@collect_warnings -@handle_errors -def create_annotations( - ctx: typer.Context, - name: Annotated[str | None, typer.Argument()] = None, - workspace: Annotated[str | None, typer.Option("--workspace")] = None, - kind: Annotated[ - Literal["feedback", "note", "metadata", "label"] | None, typer.Option("--kind", help="(required)") - ] = None, - session_id: Annotated[str | None, typer.Option("--session-id", help="(required)")] = None, - value: Annotated[str | None, typer.Option("--value")] = None, - span_id: Annotated[str | None, typer.Option("--span-id")] = None, - text: Annotated[str | None, typer.Option("--text")] = None, - metadata: Annotated[str | None, typer.Option("--metadata", help="JSON string")] = None, - value_type: Annotated[Literal["text", "numeric"] | None, typer.Option("--value-type")] = None, - exist_ok: Annotated[bool | None, typer.Option("--exist-ok")] = None, - input_file: Annotated[ - str | None, - typer.Option("--input-file", help="Path to JSON file (use '-' for stdin)", rich_help_panel="Input Options"), - ] = None, - input_data: Annotated[ - str | None, - typer.Option("--input-data", help="Input data for the request (JSON or YAML)", rich_help_panel="Input Options"), - ] = None, - output_format: EntityOutputFormatOption = None, -) -> None: - """Create annotations. - - [bold red]Required fields:[/] kind, session_id - - [green]Examples:[/] - nemo intake annotations create --input-file config.json - nemo intake annotations create --input-data '{"kind": "value", "session_id": "value"}' - echo '{"json": "data"}' | nemo intake annotations create --input-file - - nemo intake annotations create --