refactor: decouple FilesetFileSystem from Starlette test internals and improve upload/download paths

packages/data_designer_nemo/src/data_designer_nemo/fileset_file_seed_reader.py

@@ -8,9 +8,6 @@
 from data_designer_nemo.fileset_file_seed_source import FilesetFileSeedSource
 from data_designer_nemo.sdk_translation import async_to_sync_sdk
 from nemo_platform import AsyncNeMoPlatform, NeMoPlatform
-from nemo_platform.filesets import FilesetFileSystem
-from nemo_platform_plugin.client.adapter import client_from_platform
-from nemo_platform_plugin.files.client import FilesClient
 
 workspace_cvar = ContextVar[str | None]("workspace_cvar", default=None)
 
@@ -28,11 +25,8 @@ def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
         if self._sdk is None:
             raise RuntimeError("FilesetFileSeedReader requires an injected NeMo Platform SDK")
 
-        files_client = client_from_platform(self._sdk, FilesClient)
-        filesystem = FilesetFileSystem(client=files_client)
-
         conn = duckdb.connect()
-        conn.register_filesystem(filesystem)
+        conn.register_filesystem(self._sdk.files.fsspec)
         return conn
 
     def get_dataset_uri(self) -> str:

packages/data_designer_nemo/src/data_designer_nemo/person_reader.py

@@ -8,9 +8,6 @@
 )
 from data_designer_nemo.sdk_translation import async_to_sync_sdk
 from nemo_platform import AsyncNeMoPlatform, NeMoPlatform
-from nemo_platform.filesets import FilesetFileSystem
-from nemo_platform_plugin.client.adapter import client_from_platform
-from nemo_platform_plugin.files.client import FilesClient
 
 
 class FilesetsPersonReader(PersonReader):
@@ -20,15 +17,9 @@ class FilesetsPersonReader(PersonReader):
     top-level) or an :class:`AsyncNeMoPlatform` (API-process path, used
     from a worker thread under :func:`anyio.to_thread.run_sync`).
 
-    DuckDB calls into ``FilesetFileSystem`` synchronously, so the
-    underlying filesystem must be in fsspec's sync mode
-    (``asynchronous=False``) — fsspec then spins up its own daemon event
-    loop for sync→async bridging. ``FilesetFileSystem`` sets
-    ``asynchronous=True`` when given an ``AsyncFilesClient``, which would
-    break DuckDB. So when this reader is constructed with an async SDK we
-    rebuild a sync SDK, derive a sync ``FilesClient``, and hand *that* to
-    ``FilesetFileSystem``. Auth and identity propagate; fsspec stays in
-    sync mode.
+    DuckDB calls into the SDK fileset filesystem synchronously, so when this
+    reader is constructed with an async SDK we rebuild a sync SDK first. Auth
+    and identity propagate; fsspec stays in sync mode.
     """
 
     def __init__(self, sdk: NeMoPlatform | AsyncNeMoPlatform):
@@ -37,10 +28,8 @@ def __init__(self, sdk: NeMoPlatform | AsyncNeMoPlatform):
         self._sdk = sdk
 
     def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
-        files_client = client_from_platform(self._sdk, FilesClient)
-        filesystem = FilesetFileSystem(client=files_client)
         conn = duckdb.connect()
-        conn.register_filesystem(filesystem)
+        conn.register_filesystem(self._sdk.files.fsspec)
         return conn
 
     def get_dataset_uri(self, locale: str) -> str:

packages/filesets/src/filesets/filesystem/filesystem.py

@@ -6,17 +6,20 @@
 from __future__ import annotations
 
 import inspect
+import os
 from collections.abc import AsyncIterator, Coroutine, Iterator, Sequence
 from datetime import datetime, timezone
+from glob import has_magic
 from typing import Any, Literal, TypedDict, TypeVar, overload
 
 import anyio
 import fsspec.asyn
-import httpx
 from anyio import to_thread
 from fsspec.asyn import AbstractAsyncStreamedFile, AsyncFileSystem, _get_batch_size
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
+from fsspec.implementations.local import LocalFileSystem, make_path_posix, trailing_sep
 from fsspec.spec import AbstractBufferedFile
+from fsspec.utils import other_paths
 from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient
 from nemo_platform_plugin.files.types import FilesetFileOutput, ListFilesQueryParams
 
@@ -37,10 +40,8 @@ async def run_coros_in_chunks(
     We admit all work up front and use a CapacityLimiter so a new task can start
     as soon as any slot frees up, which keeps bulk downloads/uploads saturated.
 
-    We still monkey-patch fsspec's helper so inherited AsyncFileSystem bulk ops
-    benefit from AnyIO task-group cancellation. Running tasks are cancelled by
-    the task group; coroutine objects that were never entered are explicitly
-    closed in the outer finally block.
+    Running tasks are cancelled by the task group; coroutine objects that were
+    never entered are explicitly closed in the outer finally block.
     """
 
     if batch_size is None:
@@ -93,24 +94,6 @@ async def run_one(coro: Coroutine[Any, Any, T], idx: int) -> None:
     return results
 
 
-# Monkey-patch fsspec so inherited AsyncFileSystem bulk ops use the same
-# limiter-based AnyIO runner. This is intentionally semantics-different from
-# upstream chunking because we prefer continuous refill over wave-based batches.
-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."""
 
@@ -350,11 +333,12 @@ def __init__(
         self,
         *,
         client: FilesClient | AsyncFilesClient,
+        async_client: AsyncFilesClient | None = None,
         batch_size: int | None = None,
         blocksize: int | None = None,
         **kwargs,
     ):
-        async_client = self._ensure_async(client)
+        async_client = async_client or self._ensure_async(client)
         is_async = isinstance(client, AsyncFilesClient)
 
         if batch_size is None:
@@ -374,15 +358,13 @@ def _ensure_async(client: FilesClient | AsyncFilesClient) -> AsyncFilesClient:
 
         import httpx
 
-        transport = _detect_async_transport(client._http)
         return AsyncFilesClient(
             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=client.base_url,
                 headers=dict(client._default_headers) if client._default_headers else None,
             ),
@@ -607,7 +589,7 @@ async def _rm_file(self, path: str, **kwargs) -> None:
         # 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:
+    async def _pipe_file(self, path: str, value: bytes, mode: str = "overwrite", **kwargs) -> None:
         """Write bytes to a file."""
         workspace, fileset, file_path = parse_fileset_ref(path, workspace_fallback=self._workspace)
         if not file_path:
@@ -662,7 +644,70 @@ def pipe_stream(
         """Sync wrapper for _pipe_stream. See _pipe_stream for details."""
         return fsspec.asyn.sync(self.loop, self._pipe_stream, path, stream, content_length)
 
-    async def _put_file(self, lpath: str, rpath: str, callback: Callback = DEFAULT_CALLBACK, **kwargs) -> None:
+    async def _put(
+        self,
+        lpath,
+        rpath,
+        recursive=False,
+        callback=DEFAULT_CALLBACK,
+        batch_size=None,
+        maxdepth=None,
+        **kwargs,
+    ):
+        """Copy local file(s) into the fileset."""
+        if isinstance(lpath, list) and isinstance(rpath, list):
+            rpaths = rpath
+            lpaths = lpath
+        else:
+            source_is_str = isinstance(lpath, str)
+            if source_is_str:
+                lpath = make_path_posix(lpath)
+            fs = LocalFileSystem()
+            lpaths = fs.expand_path(lpath, recursive=recursive, maxdepth=maxdepth)
+            if source_is_str and (not recursive or maxdepth is not None):
+                lpaths = [path for path in lpaths if not (trailing_sep(path) or fs.isdir(path))]
+                if not lpaths:
+                    return
+
+            source_is_file = len(lpaths) == 1
+            dest_is_dir = isinstance(rpath, str) and (trailing_sep(rpath) or await self._isdir(rpath))
+
+            rpath = self._strip_protocol(rpath)
+            exists = source_is_str and (
+                (has_magic(lpath) and source_is_file)
+                or (not has_magic(lpath) and dest_is_dir and not trailing_sep(lpath))
+            )
+            rpaths = other_paths(
+                lpaths,
+                rpath,
+                exists=exists,
+                flatten=not source_is_str,
+            )
+
+        is_dir = {path: os.path.isdir(path) for path in lpaths}
+        rdirs = [remote for local, remote in zip(lpaths, rpaths) if is_dir[local]]
+        file_pairs = [(local, remote) for local, remote in zip(lpaths, rpaths) if not is_dir[local]]
+
+        async with anyio.create_task_group() as tg:
+            for directory in rdirs:
+                tg.start_soon(self._makedirs, directory, True)
+
+        callback.set_size(len(file_pairs))
+        put_file = callback.branch_coro(self._put_file)
+        await run_coros_in_chunks(
+            [put_file(local, remote, **kwargs) for local, remote in file_pairs],
+            batch_size=batch_size or self.batch_size,
+            callback=callback,
+        )
+
+    async def _put_file(
+        self,
+        lpath: str,
+        rpath: str,
+        mode: str = "overwrite",
+        callback: Callback = DEFAULT_CALLBACK,
+        **kwargs,
+    ) -> None:
         """Upload a local file to a fileset.
 
         Uses streaming upload to avoid buffering the entire file in memory.
@@ -723,7 +768,7 @@ async def _find(
         self._populate_dircache_from_response(response, workspace, fileset, prefix)
 
         # Build the flat output dict
-        out = {}
+        out: dict[str, FileInfo] = {}
         seen_dirs: set[str] = set()
 
         # Add root path if withdirs requested
@@ -795,7 +840,7 @@ async def _get(
         """Download files using a single _find call for efficiency.
 
         Uses run_coros_in_chunks which provides proper task cancellation
-        via our monkey-patched TaskGroup-based implementation.
+        for the direct list-download path.
 
         When rpath and lpath are both lists, downloads each (remote, local) pair
         directly without path expansion. This is useful for downloading a specific
@@ -808,12 +853,15 @@ async def _get(
             callback.set_size(len(rpath))
             get_file_with_callback = callback.branch_coro(self._get_file)
             await run_coros_in_chunks(
-                [get_file_with_callback(remote, local, **kwargs) for remote, local in zip(rpath, lpath)],
+                [get_file_with_callback(remote, local, **kwargs) for remote, local in zip(rpath, lpath, strict=True)],
                 batch_size=batch_size or self.batch_size,
                 callback=callback,
             )
             return
 
+        if not isinstance(rpath, str) or not isinstance(lpath, str):
+            raise TypeError("rpath and lpath must both be strings or both be lists")
+
         source_files = await self._find(rpath, maxdepth=maxdepth, withdirs=False)
         if not source_files:
             return

packages/filesets/src/filesets/resources.py

@@ -14,7 +14,7 @@
 from pathlib import PurePath
 from typing import Protocol, runtime_checkable
 
-from fsspec.callbacks import Callback
+from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from fsspec.core import has_magic
 from nemo_platform.resources.files.filesets import AsyncFilesetsResource, FilesetsResource
 from nemo_platform.resources.files.otlp.otlp import AsyncOtlpResource, OtlpResource
@@ -24,6 +24,7 @@
     CreateFilesetRequest,
     FilesetFileOutput,
     FilesetOutput,
+    ListFilesQueryParams,
 )
 
 from filesets.filesystem.filesystem import (
@@ -71,16 +72,16 @@ def cache_status(self) -> CacheStatus | None:
 
         # Priority: caching > not_cached > cached > not_cacheable
         if "caching" in statuses:
-            return "caching"
+            return CacheStatus.CACHING
         if "not_cached" in statuses:
-            return "not_cached"
+            return CacheStatus.NOT_CACHED
         if all(s == "cached" for s in statuses):
-            return "cached"
+            return CacheStatus.CACHED
         if all(s == "not_cacheable" for s in statuses):
-            return "not_cacheable"
+            return CacheStatus.NOT_CACHEABLE
 
         # Mixed cached/not_cacheable - return cached since some files are cached
-        return "cached"
+        return CacheStatus.CACHED
 
 
 @runtime_checkable
@@ -139,10 +140,17 @@ class FilesResource:
     For fsspec filesystem access, use ``resource.fsspec``.
     """
 
-    def __init__(self, client, *, files_client: FilesClient | None = None) -> None:
+    def __init__(
+        self,
+        client,
+        *,
+        files_client: FilesClient | None = None,
+        async_files_client: AsyncFilesClient | None = None,
+    ) -> None:
         # Retain the platform client so the generated fileset/otlp sub-resources
         # (which speak to the platform client, not the FilesClient) can be exposed.
         self._platform_client = client
+        self._async_client = async_files_client
         if files_client is not None:
             self._client = files_client
         else:
@@ -168,7 +176,7 @@ def otlp(self) -> OtlpResource:
     @cached_property
     def fsspec(self) -> FilesetFileSystem:
         """Access the underlying fsspec filesystem."""
-        return FilesetFileSystem(client=self._client)
+        return FilesetFileSystem(client=self._client, async_client=self._async_client)
 
     def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None:
         """Create fileset if it doesn't exist (idempotent)."""
@@ -615,7 +623,7 @@ def list(
         # For path prefixes, the API handles filtering server-side
         api_path = None if has_magic(path) else (path or None)
 
-        query_params = {}
+        query_params: ListFilesQueryParams = {}
         if api_path is not None:
             query_params["path"] = api_path
         if include_cache_status:
@@ -794,10 +802,7 @@ async def download(
             # 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]
-            kwargs: dict = {"rpath": rpaths, "lpath": lpaths, "batch_size": max_workers}
-            if callback is not None:
-                kwargs["callback"] = callback
-            await self.fsspec._get(**kwargs)
+            await self.fsspec._get(rpaths, lpaths, batch_size=max_workers, callback=callback or DEFAULT_CALLBACK)
             return
 
         ws, path_fileset, path = parse_fileset_path(
@@ -817,16 +822,16 @@ async def download(
             # 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]
-            kwargs = {"rpath": rpaths, "lpath": lpaths, "batch_size": max_workers}
-            if callback is not None:
-                kwargs["callback"] = callback
-            await self.fsspec._get(**kwargs)
+            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)
-            kwargs = {"rpath": fileset_ref, "lpath": local_path, "recursive": True, "batch_size": max_workers}
-            if callback is not None:
-                kwargs["callback"] = callback
-            await self.fsspec._get(**kwargs)
+            await self.fsspec._get(
+                fileset_ref,
+                local_path,
+                recursive=True,
+                batch_size=max_workers,
+                callback=callback or DEFAULT_CALLBACK,
+            )
 
     async def upload(
         self,
@@ -1129,7 +1134,7 @@ async def list(
         # For path prefixes, the API handles filtering server-side
         api_path = None if has_magic(path) else (path or None)
 
-        query_params = {}
+        query_params: ListFilesQueryParams = {}
         if api_path is not None:
             query_params["path"] = api_path
         if include_cache_status: