diff --git a/e2e/conftest.py b/e2e/conftest.py index 3826cb5bed..976d2c7e5d 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -65,6 +65,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from e2e.services_pool import E2EServicesPool, RunningServices, admin_headers @@ -259,6 +261,12 @@ def sdk(_services: str, _services_instance: RunningServices) -> NeMoPlatform: ) +@pytest.fixture(scope="module") +def files_client(sdk: NeMoPlatform) -> FilesClient: + """Provide a FilesClient derived from the SDK.""" + return client_from_platform(sdk, FilesClient) + + @pytest.fixture(scope="function") def workspace(sdk: NeMoPlatform) -> Iterator[str]: """Create a unique workspace for each test, deleted on teardown.""" diff --git a/e2e/files/test_files.py b/e2e/files/test_files.py index 9db3b2bdc3..d351a37698 100644 --- a/e2e/files/test_files.py +++ b/e2e/files/test_files.py @@ -11,19 +11,21 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nemo_platform_plugin.files.types import FilesetOutput as Fileset @pytest.fixture -def fileset(sdk: NeMoPlatform, workspace: str) -> Iterator[Fileset]: +def fileset(files_client: FilesClient, workspace: str) -> Iterator[Fileset]: """Create a unique fileset for each test with automatic cleanup.""" fileset_name = f"e2e-fileset-{uuid.uuid4().hex[:8]}" - fileset = sdk.files.filesets.create(workspace=workspace, name=fileset_name) + fileset = files_client.create_fileset(body=CreateFilesetRequest(name=fileset_name), workspace=workspace).data() yield fileset try: - sdk.files.filesets.delete(fileset_name, workspace=workspace) + files_client.delete_fileset(name=fileset_name, workspace=workspace) except Exception: - pass # Ignore cleanup errors + pass def test_file_upload_and_download(sdk: NeMoPlatform, workspace: str, fileset: Fileset): diff --git a/e2e/files/test_storage_backends.py b/e2e/files/test_storage_backends.py index 77b8bf5fc3..b09ebc3da2 100644 --- a/e2e/files/test_storage_backends.py +++ b/e2e/files/test_storage_backends.py @@ -15,8 +15,11 @@ from pathlib import Path import pytest -from nemo_platform import BadRequestError, NeMoPlatform -from nemo_platform.types.files import HuggingfaceStorageConfigParam, NGCStorageConfigParam +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.errors import BadRequestError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig, NGCStorageConfig +from nemo_platform_plugin.files.types import CreateFilesetRequest # --------------------------------------------------------------------------- # NGC configuration @@ -40,26 +43,28 @@ @pytest.fixture -def ngc_fileset(sdk: NeMoPlatform, workspace: str, ngc_secret: str) -> Iterator[str]: +def ngc_fileset(files_client: FilesClient, workspace: str, ngc_secret: str) -> Iterator[str]: """Create an NGC-backed fileset, cleaned up after test.""" fileset_name = f"e2e-ngc-fs-{uuid.uuid4().hex[:8]}" - sdk.files.filesets.create( + files_client.create_fileset( workspace=workspace, - name=fileset_name, - description="E2E test NGC-backed fileset", - storage=NGCStorageConfigParam( - api_key_secret=ngc_secret, - org=NGC_ORG, - team=NGC_TEAM, - target=NGC_TARGET, - target_type=NGC_TARGET_TYPE, + body=CreateFilesetRequest( + name=fileset_name, + description="E2E test NGC-backed fileset", + storage=NGCStorageConfig( + api_key_secret=ngc_secret, + org=NGC_ORG, + team=NGC_TEAM, + target=NGC_TARGET, + target_type=NGC_TARGET_TYPE, + ), ), ) yield fileset_name try: - sdk.files.filesets.delete(fileset_name, workspace=workspace) + files_client.delete_fileset(name=fileset_name, workspace=workspace) except Exception: - pass # Best-effort cleanup; the workspace is deleted anyway + pass @pytest.fixture @@ -83,27 +88,29 @@ def hf_secret(sdk: NeMoPlatform, workspace: str, hf_token: str) -> Iterator[str] @pytest.fixture -def hf_fileset(sdk: NeMoPlatform, workspace: str, hf_secret: str) -> Iterator[str]: +def hf_fileset(files_client: FilesClient, workspace: str, hf_secret: str) -> Iterator[str]: """Create a Hugging Face-backed fileset, cleaned up after test.""" fileset_name = f"e2e-hf-fs-{uuid.uuid4().hex[:8]}" - storage = HuggingfaceStorageConfigParam( + storage = HuggingfaceStorageConfig( repo_id=HF_REPO_ID, repo_type=HF_REPO_TYPE, token_secret=hf_secret, ) - sdk.files.filesets.create( + files_client.create_fileset( workspace=workspace, - name=fileset_name, - description="E2E test HF-backed fileset", - storage=storage, + body=CreateFilesetRequest( + name=fileset_name, + description="E2E test HF-backed fileset", + storage=storage, + ), ) yield fileset_name try: - sdk.files.filesets.delete(fileset_name, workspace=workspace) + files_client.delete_fileset(name=fileset_name, workspace=workspace) except Exception: - pass # Best-effort cleanup; the workspace is deleted anyway + pass # =================================================================== @@ -182,6 +189,7 @@ def test_cache_status(self, sdk: NeMoPlatform, workspace: str, ngc_fileset: str) def test_create_error( self, sdk: NeMoPlatform, + files_client: FilesClient, workspace: str, ngc_api_key: str, secret_value: str | None, @@ -193,7 +201,7 @@ def test_create_error( secret_name = f"e2e-ngc-err-{uuid.uuid4().hex[:8]}" sdk.secrets.create(workspace=workspace, name=secret_name, value=value) try: - storage = NGCStorageConfigParam( + storage = NGCStorageConfig( api_key_secret=secret_name, org=storage_overrides.get("org", NGC_ORG), team=storage_overrides.get("team", NGC_TEAM), @@ -201,26 +209,30 @@ def test_create_error( target_type=NGC_TARGET_TYPE, ) with pytest.raises(BadRequestError, match=match): - sdk.files.filesets.create( + files_client.create_fileset( workspace=workspace, - name=f"e2e-ngc-err-{uuid.uuid4().hex[:8]}", - storage=storage, + body=CreateFilesetRequest( + name=f"e2e-ngc-err-{uuid.uuid4().hex[:8]}", + storage=storage, + ), ) finally: sdk.secrets.delete(workspace=workspace, name=secret_name) - def test_create_error_nonexistent_secret(self, sdk: NeMoPlatform, workspace: str): + def test_create_error_nonexistent_secret(self, files_client: FilesClient, workspace: str): """Referencing a secret that doesn't exist is rejected with 400.""" with pytest.raises(BadRequestError, match="Secret not found:"): - sdk.files.filesets.create( + files_client.create_fileset( workspace=workspace, - name=f"e2e-ngc-err-{uuid.uuid4().hex[:8]}", - storage=NGCStorageConfigParam( - api_key_secret="no-such-secret-99999", - org=NGC_ORG, - team=NGC_TEAM, - target=NGC_TARGET, - target_type=NGC_TARGET_TYPE, + body=CreateFilesetRequest( + name=f"e2e-ngc-err-{uuid.uuid4().hex[:8]}", + storage=NGCStorageConfig( + api_key_secret="no-such-secret-99999", + org=NGC_ORG, + team=NGC_TEAM, + target=NGC_TARGET, + target_type=NGC_TARGET_TYPE, + ), ), ) @@ -276,14 +288,16 @@ def test_cache_status(self, sdk: NeMoPlatform, workspace: str, hf_fileset: str): # -- error cases -- - def test_error_nonexistent_repo(self, sdk: NeMoPlatform, workspace: str): + def test_error_nonexistent_repo(self, files_client: FilesClient, workspace: str): """Pointing at a repo that doesn't exist is rejected with 400.""" with pytest.raises(BadRequestError): - sdk.files.filesets.create( + files_client.create_fileset( workspace=workspace, - name=f"e2e-hf-err-{uuid.uuid4().hex[:8]}", - storage=HuggingfaceStorageConfigParam( - repo_id="this-org-does-not-exist/this-repo-does-not-exist-12345", - repo_type="model", + body=CreateFilesetRequest( + name=f"e2e-hf-err-{uuid.uuid4().hex[:8]}", + storage=HuggingfaceStorageConfig( + repo_id="this-org-does-not-exist/this-repo-does-not-exist-12345", + repo_type="model", + ), ), ) diff --git a/e2e/test_data_designer.py b/e2e/test_data_designer.py index 8cb5b03269..2890b19183 100644 --- a/e2e/test_data_designer.py +++ b/e2e/test_data_designer.py @@ -10,8 +10,11 @@ from data_designer_nemo.fileset_file_seed_source import FilesetFileSeedSource from data_designer_nemo.nemotron_personas import WORKSPACE, get_resource_name_for_locale from nemo_data_designer_plugin.sdk.errors import DataDesignerJobError -from nemo_platform import NeMoPlatform, NotFoundError +from nemo_platform import NeMoPlatform from nemo_platform.types.inference import ModelProvider +from nemo_platform_plugin.client.errors import NotFoundError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.testing import MockProviderResponse, add_mock_provider, assert_exit_0, run_nemo_local from nmp.testing.pytest_outcomes import pytest_skip @@ -132,12 +135,12 @@ def test_simple_ndd_config(sdk: NeMoPlatform, workspace: str) -> None: _assert_dataset_equal(job_dataset, expected_job_dataset) -def test_fileset_seed_data(sdk: NeMoPlatform, workspace: str) -> None: +def test_fileset_seed_data(sdk: NeMoPlatform, files_client: FilesClient, workspace: str) -> None: """Tests that the Data Designer *library* plugin that makes Filesets available as seed sources is wired up properly by the Data Designer *platform plugin*. """ fileset_name = "my-fileset" - sdk.files.filesets.create(name=fileset_name, workspace=workspace) + files_client.create_fileset(body=CreateFilesetRequest(name=fileset_name), workspace=workspace) seed_data = pd.DataFrame(data={"seed": ["my-seed"]}) remote_path = "data.parquet" @@ -170,7 +173,9 @@ def test_fileset_seed_data(sdk: NeMoPlatform, workspace: str) -> None: @pytest.fixture -def nemotron_personas_locale(_services: str, sdk: NeMoPlatform, workspace: str, ngc_secret: str) -> Generator[str]: +def nemotron_personas_locale( + _services: str, files_client: FilesClient, workspace: str, ngc_secret: str +) -> Generator[str]: """Invokes the CLI to create a Fileset for Nemotron Personas data. This test does call out to NGC and downloads personas data. Use the smallest locale available @@ -184,7 +189,7 @@ def nemotron_personas_locale(_services: str, sdk: NeMoPlatform, workspace: str, fileset_name = get_resource_name_for_locale(locale) with suppress(NotFoundError): - sdk.files.filesets.delete(fileset_name, workspace=WORKSPACE) + files_client.delete_fileset(name=fileset_name, workspace=WORKSPACE) result = run_nemo_local( "data-designer", @@ -202,7 +207,7 @@ def nemotron_personas_locale(_services: str, sdk: NeMoPlatform, workspace: str, yield locale with suppress(NotFoundError): - sdk.files.filesets.delete(fileset_name, workspace=WORKSPACE) + files_client.delete_fileset(name=fileset_name, workspace=WORKSPACE) def test_nemotron_personas_sampling(sdk: NeMoPlatform, workspace: str, nemotron_personas_locale: str) -> None: diff --git a/e2e/test_jobs_auth.py b/e2e/test_jobs_auth.py index ed83ebc225..463d26d61e 100644 --- a/e2e/test_jobs_auth.py +++ b/e2e/test_jobs_auth.py @@ -12,6 +12,9 @@ import pytest from nemo_platform import NeMoPlatform from nemo_platform_ext.auth.helpers import generate_unsigned_jwt +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nemo_platform_plugin.jobs.api_factory import ( ContainerSpec, CPUExecutionProviderSpec, @@ -121,7 +124,8 @@ def test_job_principal_propagation(sdk: NeMoPlatform): assert completed_job.status == "completed" fileset_name = f"hello-world-{job.name}" - fileset = user_sdk.files.filesets.retrieve(workspace=workspace_name, name=fileset_name) + files = client_from_platform(user_sdk, FilesClient) + fileset = files.get_fileset(workspace=workspace_name, name=fileset_name).data() assert fileset is not None file_content = user_sdk.files.download_content( @@ -150,7 +154,8 @@ def test_job_cannot_access_unauthorized_workspace(sdk: NeMoPlatform): other_sdk = _as_bearer_user(sdk, other_email) fileset_name = "private-data" - owner_sdk.files.filesets.create(workspace=restricted_workspace, name=fileset_name) + files = client_from_platform(owner_sdk, FilesClient) + files.create_fileset(workspace=restricted_workspace, body=CreateFilesetRequest(name=fileset_name)) job = other_sdk.jobs.create( workspace=runner_workspace, 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 239768a91a..a7ff77d61d 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 @@ -9,6 +9,8 @@ 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) @@ -26,7 +28,8 @@ def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection: if self._sdk is None: raise RuntimeError("FilesetFileSeedReader requires an injected NeMo Platform SDK") - filesystem = FilesetFileSystem(sdk=self._sdk) + files_client = client_from_platform(self._sdk, FilesClient) + filesystem = FilesetFileSystem(client=files_client) conn = duckdb.connect() conn.register_filesystem(filesystem) diff --git a/packages/data_designer_nemo/src/data_designer_nemo/nemotron_personas.py b/packages/data_designer_nemo/src/data_designer_nemo/nemotron_personas.py index 5ac44ac0a1..cfd2d6d906 100644 --- a/packages/data_designer_nemo/src/data_designer_nemo/nemotron_personas.py +++ b/packages/data_designer_nemo/src/data_designer_nemo/nemotron_personas.py @@ -6,8 +6,12 @@ from typing import Literal from data_designer.config.utils.constants import NEMOTRON_PERSONAS_DATASET_SIZES -from nemo_platform import ConflictError, NeMoPlatform -from nemo_platform.types.files import NGCStorageConfigParam +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.storage_config import NGCStorageConfig +from nemo_platform_plugin.files.types import CreateFilesetRequest logger = logging.getLogger(__name__) @@ -90,20 +94,23 @@ def sync_nemotron_personas_fileset( def _create_fileset(sdk: NeMoPlatform, locale: str, api_key_secret: str) -> None: - sdk.files.filesets.create( + files = client_from_platform(sdk, FilesClient) + files.create_fileset( workspace=WORKSPACE, - name=get_resource_name_for_locale(locale), - description=f"Nemotron Personas dataset for locale: {locale!r}", - purpose="dataset", - storage=_get_storage_config_for_locale(locale, api_key_secret), - cache=True, + body=CreateFilesetRequest( + name=get_resource_name_for_locale(locale), + description=f"Nemotron Personas dataset for locale: {locale!r}", + purpose="dataset", + storage=_get_storage_config_for_locale(locale, api_key_secret), + cache=True, + ), ) -def _get_storage_config_for_locale(locale: str, api_key_secret: str) -> NGCStorageConfigParam: +def _get_storage_config_for_locale(locale: str, api_key_secret: str) -> NGCStorageConfig: resource_name = get_resource_name_for_locale(locale) - return NGCStorageConfigParam( + return NGCStorageConfig( api_key_secret=api_key_secret, org=NGC_ORG, team=NGC_TEAM, 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 d7ce8d0aab..1cbc9c375a 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 @@ -9,6 +9,8 @@ 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): @@ -21,11 +23,10 @@ class FilesetsPersonReader(PersonReader): 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`` flips into - ``asynchronous=True`` whenever it receives an - :class:`AsyncNeMoPlatform`, which would break DuckDB. So when this - reader is constructed with an async SDK we rebuild a sync SDK from - the async one's base URL / headers / workspace and hand *that* to + 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. """ @@ -36,7 +37,8 @@ def __init__(self, sdk: NeMoPlatform | AsyncNeMoPlatform): self._sdk = sdk def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection: - filesystem = FilesetFileSystem(sdk=self._sdk) + files_client = client_from_platform(self._sdk, FilesClient) + filesystem = FilesetFileSystem(client=files_client) conn = duckdb.connect() conn.register_filesystem(filesystem) return conn diff --git a/packages/data_designer_nemo/src/data_designer_nemo/person_sampling.py b/packages/data_designer_nemo/src/data_designer_nemo/person_sampling.py index 2a826313ed..1fcf2e5499 100644 --- a/packages/data_designer_nemo/src/data_designer_nemo/person_sampling.py +++ b/packages/data_designer_nemo/src/data_designer_nemo/person_sampling.py @@ -6,7 +6,10 @@ import data_designer.config as dd from data_designer_nemo.errors import NDDInternalError from data_designer_nemo.nemotron_personas import get_resource_name_for_locale -from nemo_platform import AsyncNeMoPlatform, NotFoundError, PermissionDeniedError +from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError, PermissionDeniedError +from nemo_platform_plugin.files.client import AsyncFilesClient logger = logging.getLogger(__name__) @@ -18,11 +21,12 @@ async def ensure_nemotron_personas_filesets(config: dd.DataDesignerConfig, sdk: return unreachable_locales = set() + files = client_from_platform(sdk, AsyncFilesClient) for locale in locales: fileset_name = get_resource_name_for_locale(locale) try: - await sdk.files.filesets.retrieve(name=fileset_name, workspace="system") + await files.get_fileset(name=fileset_name, workspace="system") except NotFoundError: logger.error( f"Nemotron personas fileset {fileset_name!r} for locale {locale!r} is missing in workspace 'system'. " diff --git a/packages/data_designer_nemo/src/data_designer_nemo/seed.py b/packages/data_designer_nemo/src/data_designer_nemo/seed.py index 91baf47b70..33429505a3 100644 --- a/packages/data_designer_nemo/src/data_designer_nemo/seed.py +++ b/packages/data_designer_nemo/src/data_designer_nemo/seed.py @@ -8,7 +8,10 @@ from data_designer_nemo.errors import NDDInternalError, NDDInvalidConfigError from data_designer_nemo.fileset_file_seed_source import FilesetFileSeedSource from data_designer_nemo.secret_resolver import validate_secret -from nemo_platform import AsyncNeMoPlatform, NotFoundError, PermissionDeniedError +from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError, PermissionDeniedError +from nemo_platform_plugin.files.client import AsyncFilesClient logger = logging.getLogger(__name__) @@ -27,8 +30,9 @@ async def validate_seed(dd_config: dd.DataDesignerConfig, workspace: str, sdk: A if isinstance(seed_source, FilesetFileSeedSource): workspace, fileset_name = _parse_seed_source_path(seed_source.path, workspace) + files = client_from_platform(sdk, AsyncFilesClient) try: - await sdk.files.filesets.retrieve(name=fileset_name, workspace=workspace) + await files.get_fileset(name=fileset_name, workspace=workspace) except NotFoundError as e: raise NDDInvalidConfigError(f"Could not find fileset {fileset_name!r} in workspace {workspace!r}") from e except PermissionDeniedError as e: diff --git a/packages/data_designer_nemo/tests/unit/test_fileset_file_seed_reader.py b/packages/data_designer_nemo/tests/unit/test_fileset_file_seed_reader.py index 9548bedecc..866a5795a9 100644 --- a/packages/data_designer_nemo/tests/unit/test_fileset_file_seed_reader.py +++ b/packages/data_designer_nemo/tests/unit/test_fileset_file_seed_reader.py @@ -6,6 +6,7 @@ import pytest from data_designer_nemo.fileset_file_seed_reader import FilesetFileSeedReader, workspace_cvar from data_designer_nemo.fileset_file_seed_source import FilesetFileSeedSource +from nemo_platform_plugin.files.client import FilesClient def test_dataset_uri_with_workspace() -> None: @@ -39,12 +40,18 @@ def test_create_duckdb_connection_requires_injected_sdk() -> None: def test_create_duckdb_connection_uses_injected_sdk() -> None: sdk = Mock() conn = Mock() + mock_files_client = Mock() with ( patch("data_designer_nemo.fileset_file_seed_reader.duckdb.connect", return_value=conn), patch("data_designer_nemo.fileset_file_seed_reader.FilesetFileSystem") as fileset_file_system, + patch( + "data_designer_nemo.fileset_file_seed_reader.client_from_platform", + return_value=mock_files_client, + ) as mock_adapter, ): assert FilesetFileSeedReader(sdk).create_duckdb_connection() is conn - fileset_file_system.assert_called_once_with(sdk) + mock_adapter.assert_called_once_with(sdk, FilesClient) + fileset_file_system.assert_called_once_with(client=mock_files_client) conn.register_filesystem.assert_called_once_with(fileset_file_system.return_value) diff --git a/packages/data_designer_nemo/tests/unit/test_person_sampling.py b/packages/data_designer_nemo/tests/unit/test_person_sampling.py index 23556a5a8e..8efb226173 100644 --- a/packages/data_designer_nemo/tests/unit/test_person_sampling.py +++ b/packages/data_designer_nemo/tests/unit/test_person_sampling.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import data_designer.config as dd import pytest @@ -9,7 +9,8 @@ from data_designer_nemo.person_sampling import ( ensure_nemotron_personas_filesets, ) -from nemo_platform import AsyncNeMoPlatform, NotFoundError, PermissionDeniedError +from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.client.errors import NotFoundError, PermissionDeniedError def _make_person_sampler_column(name: str, locale: str) -> dd.SamplerColumnConfig: @@ -27,45 +28,61 @@ def _make_config(*columns: dd.SamplerColumnConfig) -> dd.DataDesignerConfig: return builder.build() +def _mock_http_response(status_code: int) -> MagicMock: + resp = MagicMock() + resp.status_code = status_code + resp.json.return_value = {"detail": "error"} + resp.text = "error" + return resp + + @pytest.mark.asyncio async def test_ensure_nemotron_personas_filesets_checks_each_locale() -> None: sdk = AsyncMock(spec=AsyncNeMoPlatform) - sdk.files.filesets.retrieve = AsyncMock() + mock_files = MagicMock() + mock_files.get_fileset = AsyncMock() config = _make_config( _make_person_sampler_column("person_us", "en_US"), _make_person_sampler_column("person_jp", "ja_JP"), ) - await ensure_nemotron_personas_filesets(config, sdk) + with patch("data_designer_nemo.person_sampling.client_from_platform", return_value=mock_files): + await ensure_nemotron_personas_filesets(config, sdk) - assert sdk.files.filesets.retrieve.await_count == 2 + assert mock_files.get_fileset.await_count == 2 @pytest.mark.asyncio async def test_ensure_nemotron_personas_filesets_raises_error_for_missing_fileset() -> None: sdk = AsyncMock(spec=AsyncNeMoPlatform) - sdk.files.filesets.retrieve.side_effect = NotFoundError("missing", response=MagicMock(), body=None) + mock_files = MagicMock() + mock_files.get_fileset = AsyncMock(side_effect=NotFoundError(_mock_http_response(404))) config = _make_config(_make_person_sampler_column("person", "en_US")) - with pytest.raises(NDDInternalError): - await ensure_nemotron_personas_filesets(config, sdk) + with patch("data_designer_nemo.person_sampling.client_from_platform", return_value=mock_files): + with pytest.raises(NDDInternalError): + await ensure_nemotron_personas_filesets(config, sdk) @pytest.mark.asyncio async def test_ensure_nemotron_personas_filesets_raises_error_for_permission_error() -> None: sdk = AsyncMock(spec=AsyncNeMoPlatform) - sdk.files.filesets.retrieve.side_effect = PermissionDeniedError("denied", response=MagicMock(), body=None) + mock_files = MagicMock() + mock_files.get_fileset = AsyncMock(side_effect=PermissionDeniedError(_mock_http_response(403))) config = _make_config(_make_person_sampler_column("person", "en_US")) - with pytest.raises(NDDInternalError): - await ensure_nemotron_personas_filesets(config, sdk) + with patch("data_designer_nemo.person_sampling.client_from_platform", return_value=mock_files): + with pytest.raises(NDDInternalError): + await ensure_nemotron_personas_filesets(config, sdk) @pytest.mark.asyncio async def test_ensure_nemotron_personas_filesets_raises_internal_error_on_other_errors() -> None: sdk = AsyncMock(spec=AsyncNeMoPlatform) - sdk.files.filesets.retrieve.side_effect = RuntimeError("something went wrong") + mock_files = MagicMock() + mock_files.get_fileset = AsyncMock(side_effect=RuntimeError("something went wrong")) config = _make_config(_make_person_sampler_column("person", "en_US")) - with pytest.raises(NDDInternalError): - await ensure_nemotron_personas_filesets(config, sdk) + with patch("data_designer_nemo.person_sampling.client_from_platform", return_value=mock_files): + with pytest.raises(NDDInternalError): + await ensure_nemotron_personas_filesets(config, sdk) diff --git a/packages/filesets/src/filesets/filesystem/filesystem.py b/packages/filesets/src/filesets/filesystem/filesystem.py index 27f01cc2ce..56230e0e53 100644 --- a/packages/filesets/src/filesets/filesystem/filesystem.py +++ b/packages/filesets/src/filesets/filesystem/filesystem.py @@ -17,7 +17,6 @@ from fsspec.asyn import AbstractAsyncStreamedFile, AsyncFileSystem, _get_batch_size from fsspec.callbacks import DEFAULT_CALLBACK, Callback from fsspec.spec import AbstractBufferedFile -from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient from nemo_platform_plugin.files.types import FilesetFileOutput, ListFilesQueryParams @@ -350,26 +349,11 @@ def register_fsspec(cls) -> None: def __init__( self, *, - client: FilesClient | AsyncFilesClient | None = None, - sdk: NeMoPlatform | AsyncNeMoPlatform | None = None, + client: FilesClient | AsyncFilesClient, 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) @@ -384,22 +368,14 @@ def __init__( @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. - """ + """Ensure we have an AsyncFilesClient, converting from sync if needed.""" 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( + return AsyncFilesClient( base_url=client.base_url, workspace=client.workspace, auth=client._auth, diff --git a/packages/filesets/src/filesets/resources.py b/packages/filesets/src/filesets/resources.py index b61242ec10..198983952c 100644 --- a/packages/filesets/src/filesets/resources.py +++ b/packages/filesets/src/filesets/resources.py @@ -12,23 +12,16 @@ from dataclasses import dataclass from functools import cached_property from pathlib import PurePath -from typing import Any, Protocol, runtime_checkable +from typing import Protocol, runtime_checkable -import nemo_platform from fsspec.callbacks import Callback from fsspec.core import has_magic -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 ( @@ -38,85 +31,6 @@ ) -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. @@ -216,199 +130,6 @@ def _matches_glob(filepath: str, pattern: str) -> bool: return PurePath(filepath).match(pattern) -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, - ) - return self._client.create_fileset(workspace=workspace, body=body, exist_ok=exist_ok).data() - - 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, - ) - return (await self._client.create_fileset(workspace=workspace, body=body, exist_ok=exist_ok)).data() - - 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. @@ -416,15 +137,18 @@ class FilesResource: For fsspec filesystem access, use ``resource.fsspec``. """ - def __init__(self, client) -> None: - from nemo_platform_plugin.client.adapter import client_from_platform + def __init__(self, client, *, files_client: FilesClient | None = None) -> None: + if files_client is not None: + self._client = files_client + else: + from nemo_platform_plugin.client.adapter import client_from_platform - self._client = client_from_platform(client, _RemappingFilesClient) + self._client = client_from_platform(client, FilesClient) @cached_property - def filesets(self) -> FilesetsSubResource: - """Access fileset CRUD operations (create, retrieve, update, list, delete).""" - return FilesetsSubResource(self._client) + def client(self) -> FilesClient: + """Access the underlying FilesClient for direct API calls.""" + return self._client @cached_property def fsspec(self) -> FilesetFileSystem: @@ -433,7 +157,11 @@ def fsspec(self) -> FilesetFileSystem: def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None: """Create fileset if it doesn't exist (idempotent).""" - self.filesets.create(name=fileset, workspace=workspace, exist_ok=True) + self._client.create_fileset( + workspace=workspace, + body=CreateFilesetRequest(name=fileset), + exist_ok=True, + ) def download( self, @@ -650,7 +378,7 @@ def upload( kwargs["callback"] = callback self.fsspec.put(**kwargs) - return self.filesets.retrieve(name=fileset, workspace=ws) + return self._client.get_fileset(name=fileset, workspace=ws).data() def upload_content( self, @@ -751,7 +479,7 @@ def upload_content( case _: raise TypeError(f"Unsupported content type: {type(content)}") - return self.filesets.retrieve(name=fileset, workspace=ws) + return self._client.get_fileset(name=fileset, workspace=ws).data() def download_content( self, @@ -938,15 +666,18 @@ class AsyncFilesResource: For fsspec filesystem access, use ``resource.fsspec``. """ - def __init__(self, client) -> None: - from nemo_platform_plugin.client.adapter import client_from_platform + def __init__(self, client, *, files_client: AsyncFilesClient | None = None) -> None: + if files_client is not None: + self._client = files_client + else: + from nemo_platform_plugin.client.adapter import client_from_platform - self._client = client_from_platform(client, _RemappingAsyncFilesClient) + self._client = client_from_platform(client, AsyncFilesClient) @cached_property - def filesets(self) -> AsyncFilesetsSubResource: - """Access fileset CRUD operations (create, retrieve, update, list, delete).""" - return AsyncFilesetsSubResource(self._client) + def client(self) -> AsyncFilesClient: + """Access the underlying AsyncFilesClient for direct API calls.""" + return self._client @cached_property def fsspec(self) -> FilesetFileSystem: @@ -955,7 +686,11 @@ def fsspec(self) -> FilesetFileSystem: async def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None: """Create fileset if it doesn't exist (idempotent).""" - await self.filesets.create(name=fileset, workspace=workspace, exist_ok=True) + await self._client.create_fileset( + workspace=workspace, + body=CreateFilesetRequest(name=fileset), + exist_ok=True, + ) async def download( self, @@ -1151,7 +886,7 @@ async def upload( kwargs["callback"] = callback await self.fsspec._put(**kwargs) - return await self.filesets.retrieve(name=fileset, workspace=ws) + return (await self._client.get_fileset(name=fileset, workspace=ws)).data() async def upload_content( self, @@ -1256,7 +991,7 @@ async def _read_chunks(f: AsyncReadable, chunk_size: int = 1024 * 1024) -> Async case _: raise TypeError(f"Unsupported content type: {type(content)}") - return await self.filesets.retrieve(name=fileset, workspace=ws) + return (await self._client.get_fileset(name=fileset, workspace=ws)).data() async def download_content( self, diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/__init__.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/__init__.py index d25116490b..fb28186de5 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/__init__.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/__init__.py @@ -8,6 +8,8 @@ from typing import Annotated import typer +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from nemo_platform_ext.cli.commands.api.files import filesets, otlp from nemo_platform_ext.cli.core.context import CLIContext @@ -60,6 +62,7 @@ def upload_files( raw_local_path: str = ctx.params.get("local_path") client = state.get_client() + files = client_from_platform(client, FilesClient) if workspace is None: workspace = client._get_workspace_path_param() @@ -68,7 +71,7 @@ def upload_files( with RichProgressCallback(description="Uploading") as callback: if fileset is not None: # Validate fileset exists before uploading - client.files.filesets.retrieve(fileset, workspace=workspace) + files.get_fileset(name=fileset, workspace=workspace) client.files.upload( local_path=raw_local_path, remote_path=remote_path, diff --git a/packages/nemo_platform_ext/tests/cli/integration/conftest.py b/packages/nemo_platform_ext/tests/cli/integration/conftest.py index ad57200676..6c2d2ac20a 100644 --- a/packages/nemo_platform_ext/tests/cli/integration/conftest.py +++ b/packages/nemo_platform_ext/tests/cli/integration/conftest.py @@ -19,6 +19,8 @@ from click.testing import Result from nemo_platform import NeMoPlatform from nemo_platform_ext.cli.core.context import CLIContext +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from nmp.core.files.service import FilesService from nmp.testing import create_test_client from starlette.testclient import TestClient @@ -40,6 +42,12 @@ def sdk(http_client: TestClient) -> NeMoPlatform: return NeMoPlatform(base_url="http://testserver", http_client=http_client) +@pytest.fixture(scope="module") +def files_client(sdk: NeMoPlatform) -> FilesClient: + """Provide a FilesClient derived from the SDK.""" + return client_from_platform(sdk, FilesClient) + + @pytest.fixture def random_workspace(sdk: NeMoPlatform) -> str: """ diff --git a/packages/nemo_platform_ext/tests/cli/integration/test_filesets.py b/packages/nemo_platform_ext/tests/cli/integration/test_filesets.py index a9603f027a..a351894d9a 100644 --- a/packages/nemo_platform_ext/tests/cli/integration/test_filesets.py +++ b/packages/nemo_platform_ext/tests/cli/integration/test_filesets.py @@ -8,15 +8,20 @@ import pytest from nemo_platform import NeMoPlatform from nemo_platform_ext.cli.app import app +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from ..utils import assert_exit_code from .conftest import NmpCliRunner @pytest.fixture -def test_fileset(sdk: NeMoPlatform, random_workspace: str) -> dict: +def test_fileset(files_client: FilesClient, random_workspace: str) -> dict: """Create a test fileset.""" - fileset = sdk.files.filesets.create(workspace=random_workspace, name="test-fileset") + fileset = files_client.create_fileset( + body=CreateFilesetRequest(name="test-fileset"), workspace=random_workspace + ).data() return {"workspace": random_workspace, "name": fileset.name} @@ -138,7 +143,7 @@ def test_upload_to_nonexistent_fileset_fails( ) assert_exit_code(result, 1) - assert "Not found" in result.stderr + assert "not found" in result.stderr.lower() @pytest.mark.parametrize( ("remote_path", "expected_suffix"), @@ -196,10 +201,8 @@ def test_upload_without_fileset_auto_creates( fileset_name = match.group(1) # Verify fileset exists - fileset = runner.client.files.filesets.retrieve( - name=fileset_name, - workspace=random_workspace, - ) + files = client_from_platform(runner.client, FilesClient) + fileset = files.get_fileset(name=fileset_name, workspace=random_workspace).data() assert fileset.name == fileset_name # Verify file was uploaded @@ -212,7 +215,9 @@ def test_upload_without_fileset_auto_creates( @pytest.fixture -def fileset_with_nested_files(sdk: NeMoPlatform, random_workspace: str, tmp_path: Path) -> dict: +def fileset_with_nested_files( + sdk: NeMoPlatform, files_client: FilesClient, random_workspace: str, tmp_path: Path +) -> dict: """Create a fileset with nested file structure for download tests. Structure: @@ -222,7 +227,9 @@ def fileset_with_nested_files(sdk: NeMoPlatform, random_workspace: str, tmp_path file2.txt file3.txt """ - fileset = sdk.files.filesets.create(workspace=random_workspace, name="download-test-fileset") + fileset = files_client.create_fileset( + body=CreateFilesetRequest(name="download-test-fileset"), workspace=random_workspace + ).data() # Create nested directory structure locally dir_a = tmp_path / "a" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/file_manager.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/file_manager.py index 35db29f0d0..e0904caada 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/file_manager.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/file_manager.py @@ -13,6 +13,9 @@ import fsspec.asyn from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform.filesets import FilesetFileSystem, build_fileset_ref, parse_fileset_ref +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nemo_platform_plugin.jobs.schemas import FileStorageType logger = logging.getLogger(__name__) @@ -139,7 +142,11 @@ class BaseFilesetFileManager: _fs: FilesetFileSystem = field(init=False) def __post_init__(self): - self._fs = FilesetFileSystem(sdk=self.sdk) + if isinstance(self.sdk, AsyncNeMoPlatform): + files_client = client_from_platform(self.sdk, AsyncFilesClient) + else: + files_client = client_from_platform(self.sdk, FilesClient) + self._fs = FilesetFileSystem(client=files_client) def url(self, remote_path: str | None = None) -> str: """Return fileset reference for the given path.""" @@ -163,7 +170,9 @@ async def _validate_storage(self) -> None: except FileNotFoundError: if self.ensure_fileset_exists: logger.info(f"Creating new fileset: [{self.fileset_name}] in workspace [{self.workspace}]") - await self._fs._sdk.files.filesets.create(name=self.fileset_name, workspace=self.workspace) + await self._fs._client.create_fileset( + body=CreateFilesetRequest(name=self.fileset_name), workspace=self.workspace + ) else: raise FileStorageDoesNotExist( f"Fileset [{self.fileset_name}] in workspace [{self.workspace}] does not exist." diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/scheduler.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/scheduler.py index 7fe9c62006..c326d767c9 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/scheduler.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/scheduler.py @@ -103,9 +103,9 @@ def run_local( workspace: Workspace scope for the local context. sdk: Optional :class:`~nemo_platform.NeMoPlatform` handle. Bound to an ``sdk`` kwarg on ``run`` if the signature - declares one. Jobs navigate the SDK tree themselves - (``sdk.files.filesets.create(...)``) — typed per-service - resources land in a follow-up alongside ``NemoSDK``. + declares one. Jobs obtain typed service clients via + ``client_from_platform(sdk, FilesClient)`` for + per-service operations. async_sdk: Optional :class:`~nemo_platform.AsyncNeMoPlatform` handle, bound the same way when an ``async_sdk`` kwarg is declared. Sync ``run`` cannot consume this directly diff --git a/packages/nmp_common/tests/jobs/conftest.py b/packages/nmp_common/tests/jobs/conftest.py index a3f6822bc6..04ca531141 100644 --- a/packages/nmp_common/tests/jobs/conftest.py +++ b/packages/nmp_common/tests/jobs/conftest.py @@ -63,8 +63,6 @@ def mock_fileset_fs(): fs._put_file = AsyncMock() fs._get = AsyncMock() fs._get_file = AsyncMock() - fs._sdk = MagicMock() - fs._sdk.files.filesets.create = AsyncMock() # Provide the fsspec global event loop for sync-to-async bridging fs.loop = fsspec.asyn.get_loop() return fs @@ -84,9 +82,6 @@ def mock_sdk(): sdk._custom_headers = None sdk._client = MagicMock() sdk.files = MagicMock() - sdk.files.filesets = MagicMock() - sdk.files.filesets.create = MagicMock() - sdk.files.filesets.retrieve = MagicMock() sdk.files.upload_content = MagicMock() # Mock list to return ListFilesResponse with empty data by default @@ -105,9 +100,12 @@ def fileset_manager(mock_sdk, mock_fileset_fs) -> FilesetFileManager: """Create FilesetFileManager with mocked FilesetFileSystem. The FilesetFileManager creates FilesetFileSystem in __post_init__, so we must - patch it to inject our mock. + patch both client_from_platform and FilesetFileSystem to inject our mock. """ - with mock.patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem") as mock_fs_class: + with ( + mock.patch("nemo_platform_plugin.jobs.file_manager.client_from_platform"), + mock.patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem") as mock_fs_class, + ): mock_fs_class.return_value = mock_fileset_fs return FilesetFileManager( workspace=DEFAULT_WORKSPACE, diff --git a/packages/nmp_common/tests/jobs/test_file_manager.py b/packages/nmp_common/tests/jobs/test_file_manager.py index 1ec3eb93c1..fe7724c4af 100644 --- a/packages/nmp_common/tests/jobs/test_file_manager.py +++ b/packages/nmp_common/tests/jobs/test_file_manager.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from nemo_platform_plugin.jobs.file_manager import _filter_files_by_patterns from nmp.common.jobs.file_manager import FilesetFileManager, FileStorageType @@ -102,7 +102,10 @@ def test_no_matches(self): def test_fileset_url(): """Test URL generation for fileset storage.""" - with patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem"): + with ( + patch("nemo_platform_plugin.jobs.file_manager.client_from_platform"), + patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem"), + ): mgr = FilesetFileManager( workspace="my-workspace", fileset_name="my-fileset", @@ -114,7 +117,10 @@ def test_fileset_url(): def test_fileset_storage_type(): """Test storage type returns FILESET.""" - with patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem"): + with ( + patch("nemo_platform_plugin.jobs.file_manager.client_from_platform"), + patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem"), + ): mgr = FilesetFileManager( workspace="my-workspace", fileset_name="my-fileset", @@ -129,12 +135,13 @@ def test_fileset_validate_storage_exists(fileset_manager, mock_fileset_fs): mock_fileset_fs._info.assert_called() -def test_fileset_validate_storage_creates(fileset_manager, mock_fileset_fs, mock_sdk): +def test_fileset_validate_storage_creates(fileset_manager, mock_fileset_fs): """Test validate_storage creates fileset when missing.""" + mock_fileset_fs._client = MagicMock() + mock_fileset_fs._client.create_fileset = AsyncMock() mock_fileset_fs._info.side_effect = FileNotFoundError("not found") fileset_manager.validate_storage() - # The async version uses _fs._sdk.files.filesets.create - mock_fileset_fs._sdk.files.filesets.create.assert_called_once() + mock_fileset_fs._client.create_fileset.assert_called_once() def test_fileset_upload_file(tmp_path, fileset_manager, mock_fileset_fs): @@ -261,7 +268,10 @@ async def test_fileset_upload_directory_with_ignore_patterns(tmp_path, mock_sdk, (subdir / "nested.txt").write_text("keep nested") (subdir / "nested.pyc").write_text("skip nested") - with mock.patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem") as mock_fs_class: + with ( + mock.patch("nemo_platform_plugin.jobs.file_manager.client_from_platform"), + mock.patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem") as mock_fs_class, + ): mock_fs_class.return_value = mock_fileset_fs async_manager = AsyncFilesetFileManager( workspace=DEFAULT_WORKSPACE, diff --git a/packages/nmp_customization_common/src/nmp/customization_common/service/platform_client.py b/packages/nmp_customization_common/src/nmp/customization_common/service/platform_client.py index 9a8066bdb2..1398dbd2e9 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/service/platform_client.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/service/platform_client.py @@ -11,6 +11,10 @@ from nemo_platform import AsyncNeMoPlatform from nemo_platform._exceptions import NotFoundError, PermissionDeniedError from nemo_platform.types.models import ModelEntity +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError +from nemo_platform_plugin.client.errors import PermissionDeniedError as ClientPermissionDeniedError +from nemo_platform_plugin.files.client import AsyncFilesClient from nmp.common.entities.utils import parse_entity_ref from nmp.customization_common.schemas.file_io import FileSetRef @@ -24,11 +28,12 @@ async def check_dataset_access(sdk: AsyncNeMoPlatform, dataset_uri: str, default """ ref = FileSetRef.model_validate(dataset_uri) workspace = ref.workspace or default_workspace + files = client_from_platform(sdk, AsyncFilesClient) try: - await sdk.files.filesets.retrieve(workspace=workspace, name=ref.name) - except PermissionDeniedError: + await files.get_fileset(workspace=workspace, name=ref.name) + except ClientPermissionDeniedError: raise PermissionError(f"Access denied to dataset fileset '{workspace}/{ref.name}'") from None - except NotFoundError: + except ClientNotFoundError: raise ValueError( f"Dataset fileset '{ref.name}' not found in workspace '{workspace}'. Verify the dataset exists." ) from None diff --git a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py index 7315380a95..8ccf60597d 100644 --- a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py +++ b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py @@ -29,7 +29,10 @@ from nemo_data_designer_plugin.sdk.resources import DataDesignerResource from nemo_data_designer_plugin.service import DataDesignerService from nemo_platform import AsyncNeMoPlatform, NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.commands import add_function_commands, add_job_commands +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nemo_platform_plugin.job_context import JobContext, StoragePaths from nemo_platform_plugin.job_results import PlatformJobResults from nemo_platform_plugin.jobs.api_factory import PlatformJobSpec @@ -176,8 +179,9 @@ def setup_mock_secret(client_context: ClientContext) -> Generator[None]: @contextmanager def setup_mock_file(client_context: ClientContext) -> Generator[None]: - client_context.sdk.files.filesets.create( - name=FILESET_NAME, + files = client_from_platform(client_context.sdk, FilesClient) + files.create_fileset( + body=CreateFilesetRequest(name=FILESET_NAME), workspace=client_context.sdk.workspace or WORKSPACE_NAME, ) with tempfile.NamedTemporaryFile(suffix=".parquet") as tmpfile: @@ -213,7 +217,8 @@ def setup_mock_nemotron_personas_data( def _create_nemotron_personas_fileset(sdk: NeMoPlatform, persona_data: pd.DataFrame) -> None: fileset_name = get_resource_name_for_locale("en_US") - sdk.files.filesets.create(name=fileset_name, workspace="system") + files = client_from_platform(sdk, FilesClient) + files.create_fileset(body=CreateFilesetRequest(name=fileset_name), workspace="system") with tempfile.NamedTemporaryFile(suffix=".parquet") as tmpfile: persona_data.to_parquet(tmpfile.name, index=False) sdk.files.upload( 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 c388996bbd..c13ac9efaf 100644 --- a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py +++ b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py @@ -9,6 +9,8 @@ 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_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from nemo_platform_plugin.files.storage_config import NGCStorageConfig pytestmark = pytest.mark.integration @@ -75,10 +77,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.items()] == [get_resource_name_for_locale("en_US")] + files = client_from_platform(cli_sdk, FilesClient) + filesets_page = files.list_filesets(workspace=WORKSPACE) + assert [fileset.name for fileset in filesets_page.items()] == [get_resource_name_for_locale("en_US")] - fileset = cli_sdk.files.filesets.retrieve(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE) + fileset = files.get_fileset(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE).data() assert isinstance(fileset.storage, NGCStorageConfig) assert fileset.storage.api_key_secret.root == "system/ngc-api-key" @@ -105,7 +108,8 @@ def test_make_fileset_creates_secret_from_env_then_fileset( secret = cli_sdk.secrets.access("my-ngc-key", workspace="system") assert secret.value == "nvapi-from-env" - fileset = cli_sdk.files.filesets.retrieve(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE) + files = client_from_platform(cli_sdk, FilesClient) + fileset = files.get_fileset(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE).data() assert isinstance(fileset.storage, NGCStorageConfig) assert fileset.storage.api_key_secret.root == "system/my-ngc-key" @@ -183,8 +187,8 @@ 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 list(filesets.items()) == [] + files = client_from_platform(cli_sdk, FilesClient) + assert list(files.list_filesets(workspace=WORKSPACE).items()) == [] def test_make_fileset_create_secret_internal_error_surfaces_clearly( @@ -212,8 +216,8 @@ def _boom(*args: object, **kwargs: object) -> None: assert result.exit_code == 1 assert "Failed to create secret" in result.output assert "secrets backend exploded" in result.output - filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert list(filesets.items()) == [] + files = client_from_platform(cli_sdk, FilesClient) + assert list(files.list_filesets(workspace=WORKSPACE).items()) == [] def test_make_fileset_is_idempotent_when_fileset_already_exists(cli_sdk: NeMoPlatform) -> None: @@ -247,17 +251,19 @@ def test_make_fileset_is_idempotent_when_fileset_already_exists(cli_sdk: NeMoPla assert second.exit_code == 0, second.output assert "already exists" in second.output - filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert [fileset.name for fileset in filesets.items()] == [get_resource_name_for_locale("en_US")] + files = client_from_platform(cli_sdk, FilesClient) + assert [fileset.name for fileset in files.list_filesets(workspace=WORKSPACE).items()] == [ + get_resource_name_for_locale("en_US") + ] def test_make_fileset_create_fileset_internal_error_surfaces_clearly(cli_sdk: NeMoPlatform) -> None: error_message = "kaboom-fileset-error" - def _boom(*args: object, **kwargs: object) -> None: - raise RuntimeError(error_message) + mock_files = Mock() + mock_files.create_fileset.side_effect = RuntimeError(error_message) - with patch.object(cli_sdk.files.filesets, "create", side_effect=_boom): + with patch("data_designer_nemo.nemotron_personas.client_from_platform", return_value=mock_files): result = u.invoke_cli( [ "personas", @@ -272,5 +278,5 @@ def _boom(*args: object, **kwargs: object) -> None: assert result.exit_code == 1 assert "Failed to create fileset" in result.output assert error_message in result.output - filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert list(filesets.items()) == [] + files = client_from_platform(cli_sdk, FilesClient) + assert list(files.list_filesets(workspace=WORKSPACE).items()) == [] diff --git a/plugins/nemo-evaluator/examples/plugin_examples.py b/plugins/nemo-evaluator/examples/plugin_examples.py index f5430757db..adb6e067b5 100644 --- a/plugins/nemo-evaluator/examples/plugin_examples.py +++ b/plugins/nemo-evaluator/examples/plugin_examples.py @@ -40,7 +40,11 @@ ) from nemo_evaluator_sdk.values.results import EvaluationResult from nemo_platform import APIError, AsyncNeMoPlatform, ConflictError, NeMoPlatform, NotFoundError -from nemo_platform.types.files import HuggingfaceStorageConfigParam +from nemo_platform_plugin.client import errors as files_errors +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig +from nemo_platform_plugin.files.types import CreateFilesetRequest if TYPE_CHECKING: import numpy as np @@ -177,28 +181,33 @@ def write_local_helpsteer2_dataset(dataset_path: Path, *, row_count: int) -> Non async def ensure_example_fileset(client: AsyncNeMoPlatform) -> FilesetRef: """Create or reuse the HelpSteer2 fileset, then verify the selected split downloads.""" workspace = client.workspace or DEFAULT_WORKSPACE + files = client_from_platform(client, AsyncFilesClient) try: - fileset = await client.files.filesets.create( - workspace=workspace, - name=DATASET_NAME, - description="NVIDIA HelpSteer2 dataset for quality evaluation", - storage=HuggingfaceStorageConfigParam( - type="huggingface", - repo_id="nvidia/HelpSteer2", - repo_type="dataset", - ), - ) + fileset = ( + await files.create_fileset( + workspace=workspace, + body=CreateFilesetRequest( + name=DATASET_NAME, + description="NVIDIA HelpSteer2 dataset for quality evaluation", + storage=HuggingfaceStorageConfig( + repo_id="nvidia/HelpSteer2", + repo_type="dataset", + ), + ), + ) + ).data() print(f"Registered HelpSteer2 fileset: {fileset}") - except ConflictError: - fileset = await client.files.filesets.retrieve(name=DATASET_NAME, workspace=workspace) + except files_errors.ConflictError: + fileset = (await files.get_fileset(name=DATASET_NAME, workspace=workspace)).data() print(f"{fileset.workspace}/{fileset.name} dataset already registered") - downloaded = await client.files.download_content( - remote_path=HELPSTEER2_REMOTE_PATH, - fileset=fileset.name, + response = await files.download_file( + path=HELPSTEER2_REMOTE_PATH, + name=fileset.name, workspace=fileset.workspace, ) + downloaded = await response.read() rows = _load_helpsteer2_rows(downloaded, limit=2) _validate_helpsteer2_rows(rows) print(f"Verified HelpSteer2 split: {fileset.workspace}/{fileset.name}#{HELPSTEER2_REMOTE_PATH}") @@ -208,28 +217,30 @@ async def ensure_example_fileset(client: AsyncNeMoPlatform) -> FilesetRef: def ensure_example_fileset_sync(client: NeMoPlatform) -> FilesetRef: """Create or reuse the HelpSteer2 fileset with a sync client, then verify the selected split downloads.""" workspace = client.workspace or DEFAULT_WORKSPACE + files = client_from_platform(client, FilesClient) try: - fileset = client.files.filesets.create( + fileset = files.create_fileset( workspace=workspace, - name=DATASET_NAME, - description="NVIDIA HelpSteer2 dataset for quality evaluation", - storage=HuggingfaceStorageConfigParam( - type="huggingface", - repo_id="nvidia/HelpSteer2", - repo_type="dataset", + body=CreateFilesetRequest( + name=DATASET_NAME, + description="NVIDIA HelpSteer2 dataset for quality evaluation", + storage=HuggingfaceStorageConfig( + repo_id="nvidia/HelpSteer2", + repo_type="dataset", + ), ), - ) + ).data() print(f"Registered HelpSteer2 fileset: {fileset}") - except ConflictError: - fileset = client.files.filesets.retrieve(name=DATASET_NAME, workspace=workspace) + except files_errors.ConflictError: + fileset = files.get_fileset(name=DATASET_NAME, workspace=workspace).data() print(f"{fileset.workspace}/{fileset.name} dataset already registered") - downloaded = client.files.download_content( - remote_path=HELPSTEER2_REMOTE_PATH, - fileset=fileset.name, + downloaded = files.download_file( + path=HELPSTEER2_REMOTE_PATH, + name=fileset.name, workspace=fileset.workspace, - ) + ).read() rows = _load_helpsteer2_rows(downloaded, limit=2) _validate_helpsteer2_rows(rows) print(f"Verified HelpSteer2 split: {fileset.workspace}/{fileset.name}#{HELPSTEER2_REMOTE_PATH}") diff --git a/plugins/nemo-evaluator/src/nemo_evaluator/filesets.py b/plugins/nemo-evaluator/src/nemo_evaluator/filesets.py index acc908f345..44ffe32cff 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/filesets.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/filesets.py @@ -11,6 +11,8 @@ import fsspec.asyn 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 AsyncFilesClient, FilesClient from pydantic import Field, RootModel _GLOB_CHARS = {"*", "?", "["} @@ -123,12 +125,16 @@ def _match_path_parts(path_parts: tuple[str, ...], pattern_parts: tuple[str, ... async def _download_fileset_ref( - sdk: AsyncNeMoPlatform, + sdk: AsyncNeMoPlatform | NeMoPlatform, dataset: FilesetRef, destination: str, recursive: bool = True, + *, + fs: FilesetFileSystem | None = None, ) -> Path: - fs = FilesetFileSystem(sdk=sdk) + if fs is None: + files_client = client_from_platform(sdk, AsyncFilesClient) + fs = FilesetFileSystem(client=files_client) ref = dataset.root if "#" in ref: @@ -141,6 +147,7 @@ async def _download_fileset_ref( FilesetRef(root=base_path), destination, recursive=recursive, + fs=fs, ) base_dest = _safe_child_path(Path(destination), base_path) @@ -177,8 +184,9 @@ def _download_fileset_ref_sync( destination: str, recursive: bool = True, ) -> Path: - fs = FilesetFileSystem(sdk=sdk) - result = fsspec.asyn.sync(fs.loop, _download_fileset_ref, fs._sdk, dataset, destination, recursive) + files_client = client_from_platform(sdk, FilesClient) + fs = FilesetFileSystem(client=files_client) + result = fsspec.asyn.sync(fs.loop, _download_fileset_ref, sdk, dataset, destination, recursive, fs=fs) if result is None: raise RuntimeError(f"FilesetRef download returned no path for dataset {dataset.root!r}") return result diff --git a/plugins/nemo-evaluator/src/nemo_evaluator/metric_storage.py b/plugins/nemo-evaluator/src/nemo_evaluator/metric_storage.py index 0a261f470a..2286df230d 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/metric_storage.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/metric_storage.py @@ -22,6 +22,9 @@ import nemo_evaluator.shared.metric_bundles.inline # noqa: F401 from nemo_evaluator.shared.metric_bundles.bundles import MetricBundle from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import AsyncFilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from pydantic import ValidationError #: Filename of the serialized bundle stored within each metric's fileset. @@ -71,26 +74,30 @@ async def store_bundle(sdk: AsyncNeMoPlatform, workspace: str, name: str, bundle """ fileset = _new_fileset_name() body = bundle.model_dump_json().encode("utf-8") + files = client_from_platform(sdk, AsyncFilesClient) try: - await sdk.files.filesets.create( - name=fileset, + description = f"Stored metric bundle for {workspace}/{name}." + await files.create_fileset( + body=CreateFilesetRequest( + name=fileset, + description=description[:255], + ), workspace=workspace, - description=f"Stored metric bundle for {workspace}/{name}.", ) except Exception as exc: raise MetricBundleStorageError(f"failed to create fileset for metric bundle {workspace}/{name}") from exc try: - await sdk.files.upload_content( + await files.upload_file( + path=BUNDLE_FILENAME, content=body, - remote_path=BUNDLE_FILENAME, - fileset=fileset, + name=fileset, workspace=workspace, ) except Exception as exc: # Roll back the just-created (now-empty) fileset so a failed upload # doesn't leak it, then surface a typed storage error. try: - await sdk.files.filesets.delete(name=fileset, workspace=workspace) + await files.delete_fileset(name=fileset, workspace=workspace) except Exception: logger.warning( "Failed to clean up fileset after a failed metric bundle upload; storage may be leaked", @@ -110,8 +117,10 @@ async def load_bundle(sdk: AsyncNeMoPlatform, bundle_ref: str, *, expected_diges verified against it to detect drift or corruption. """ workspace, fileset, path = parse_bundle_ref(bundle_ref) + files = client_from_platform(sdk, AsyncFilesClient) try: - data = await sdk.files.download_content(remote_path=path, fileset=fileset, workspace=workspace) + response = await files.download_file(path=path, workspace=workspace, name=fileset) + data = await response.read() except Exception as exc: raise MetricBundleStorageError(f"failed to download metric bundle from {bundle_ref!r}") from exc try: @@ -130,4 +139,5 @@ async def load_bundle(sdk: AsyncNeMoPlatform, bundle_ref: str, *, expected_diges async def delete_bundle_by_ref(sdk: AsyncNeMoPlatform, bundle_ref: str) -> None: """Delete the specific fileset a bundle reference points at.""" workspace, fileset, _ = parse_bundle_ref(bundle_ref) - await sdk.files.filesets.delete(name=fileset, workspace=workspace) + files = client_from_platform(sdk, AsyncFilesClient) + await files.delete_fileset(name=fileset, workspace=workspace) diff --git a/plugins/nemo-evaluator/tests/api/service/test_metric_service.py b/plugins/nemo-evaluator/tests/api/service/test_metric_service.py index 0a63b8d8eb..edf617f6f7 100644 --- a/plugins/nemo-evaluator/tests/api/service/test_metric_service.py +++ b/plugins/nemo-evaluator/tests/api/service/test_metric_service.py @@ -4,6 +4,7 @@ from __future__ import annotations from datetime import datetime, timezone +from unittest.mock import AsyncMock, patch import pytest from nemo_evaluator.api.schemas import MetricInline @@ -23,36 +24,32 @@ # ---- in-memory fakes ------------------------------------------------------- -class _FakeFilesets: - def __init__(self, store: dict[tuple[str, str], dict[str, bytes]]) -> None: - self._store = store +class _FakeResponse: + def __init__(self, data: bytes) -> None: + self._data = data - async def create(self, *, name, workspace, description=None, exist_ok=False): - self._store.setdefault((workspace, name), {}) - return object() - - async def delete(self, name, *, workspace=None): - self._store.pop((workspace, name), None) - return object() + async def read(self) -> bytes: + return self._data -class _FakeFiles: - def __init__(self, store: dict[tuple[str, str], dict[str, bytes]]) -> None: - self._store = store - self.filesets = _FakeFilesets(store) +class _FakeAsyncFilesClient: + def __init__(self) -> None: + self._store: dict[tuple[str, str], dict[str, bytes]] = {} - async def upload_content(self, *, content, remote_path, fileset, workspace, fileset_auto_create=False): - self._store.setdefault((workspace, fileset), {})[remote_path] = bytes(content) - return object() + async def create_fileset(self, *, body, workspace=None, exist_ok=False): + self._store.setdefault((workspace, body.name), {}) + return AsyncMock(data=lambda: object()) - async def download_content(self, *, remote_path, fileset, workspace): - return self._store[(workspace, fileset)][remote_path] + async def delete_fileset(self, *, name, workspace=None): + self._store.pop((workspace, name), None) + return AsyncMock(data=lambda: object()) + async def upload_file(self, *, path, content, workspace, name): + self._store.setdefault((workspace, name), {})[path] = bytes(content) + return AsyncMock(data=lambda: object()) -class _FakeSDK: - def __init__(self) -> None: - self._store: dict[tuple[str, str], dict[str, bytes]] = {} - self.files = _FakeFiles(self._store) + async def download_file(self, *, path, workspace, name): + return _FakeResponse(self._store[(workspace, name)][path]) class _FakeEntityClient: @@ -94,8 +91,16 @@ async def list(self, entity_cls, *, workspace, filter_operation=None, sort=None, @pytest.fixture -def service() -> MetricService: - return MetricService(_FakeEntityClient(), _FakeSDK()) +def fake_files(): + return _FakeAsyncFilesClient() + + +@pytest.fixture +def service(fake_files): + svc = MetricService(_FakeEntityClient(), object()) + svc._fake_files = fake_files + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_files): + yield svc def _bundle(metric=None) -> MetricInline: @@ -113,7 +118,7 @@ def _fileset_of(service: MetricService, bundle_ref: str) -> tuple[str, str]: # ---- tests ----------------------------------------------------------------- -async def test_create_stores_bundle_and_indexes_entity(service: MetricService) -> None: +async def test_create_stores_bundle_and_indexes_entity(service: MetricService, fake_files) -> None: bundle = _bundle() created = await service.create_metric("exact", bundle, workspace="default") @@ -122,33 +127,30 @@ async def test_create_stores_bundle_and_indexes_entity(service: MetricService) - assert created.payload_kind == "cloudpickle" assert created.payload_digest == bundle.payload.digest assert created.bundle_ref.startswith("default/metric-bundle.") - # Description/labels are sourced from the bundle's metadata. assert created.description == bundle.metadata.description assert created.labels == bundle.metadata.labels - # Bundle bytes live in Files, not in the entity index. - assert _fileset_of(service, created.bundle_ref) in service.sdk._store + assert _fileset_of(service, created.bundle_ref) in fake_files._store -async def test_create_rejects_duplicate_without_clobbering_existing(service: MetricService) -> None: +async def test_create_rejects_duplicate_without_clobbering_existing(service: MetricService, fake_files) -> None: first = await service.create_metric("exact", _bundle(), workspace="default") with pytest.raises(ValueError, match="already exists"): await service.create_metric("exact", _bundle(), workspace="default") - # The original metric's bundle must survive the rejected create's rollback. - assert _fileset_of(service, first.bundle_ref) in service.sdk._store + assert _fileset_of(service, first.bundle_ref) in fake_files._store async def test_get_returns_none_when_missing(service: MetricService) -> None: assert await service.get_metric("default", "nope") is None -async def test_delete_removes_entity_and_bundle(service: MetricService) -> None: +async def test_delete_removes_entity_and_bundle(service: MetricService, fake_files) -> None: created = await service.create_metric("m", _bundle(), workspace="default") assert await service.delete_metric("default", "m") is True assert await service.get_metric("default", "m") is None - assert _fileset_of(service, created.bundle_ref) not in service.sdk._store + assert _fileset_of(service, created.bundle_ref) not in fake_files._store async def test_delete_returns_false_when_missing(service: MetricService) -> None: @@ -161,7 +163,6 @@ async def test_delete_handles_concurrent_delete_race(service: MetricService) -> async def _already_deleted(*_args, **_kwargs): raise EntityNotFoundError("deleted concurrently") - # Simulate another request removing the entity between get and delete. service.entity_client.delete = _already_deleted assert await service.delete_metric("default", "m") is False @@ -180,7 +181,7 @@ async def test_list_returns_workspace_metrics(service: MetricService) -> None: # ---- derived metrics ------------------------------------------------------- -async def test_store_derived_metric_names_by_digest_and_marks_derived(service: MetricService) -> None: +async def test_store_derived_metric_names_by_digest_and_marks_derived(service: MetricService, fake_files) -> None: from nemo_evaluator.api.service.metric_service import _MAX_ENTITY_NAME_LENGTH ref = await service.store_derived_metric(_bundle(), workspace="default") @@ -188,38 +189,33 @@ async def test_store_derived_metric_names_by_digest_and_marks_derived(service: M workspace, _, name = ref.root.partition("/") assert workspace == "default" assert name.startswith("derived.") - # The entity store caps names at 63 chars; the derived name must fit (it 422s otherwise). assert len(name) <= _MAX_ENTITY_NAME_LENGTH - # Stored entity is flagged derived and Files-backed like any metric. entity = service.entity_client.entities[("default", name)] assert entity.derived is True - assert _fileset_of(service, entity.bundle_ref) in service.sdk._store + assert _fileset_of(service, entity.bundle_ref) in fake_files._store async def test_store_derived_metric_distinguishes_full_contract(service: MetricService) -> None: - # Two metrics with an identical payload but a differing bundle-level field (here: metadata) must - # NOT collapse — addressing on payload.digest alone would have silently rebound one onto the other. bundle = _bundle() variant = bundle.model_copy(update={"metadata": bundle.metadata.model_copy(update={"description": "different"})}) - assert bundle.payload.digest == variant.payload.digest # same executable payload... + assert bundle.payload.digest == variant.payload.digest first = await service.store_derived_metric(bundle, workspace="default") second = await service.store_derived_metric(variant, workspace="default") - assert first.root != second.root # ...but distinct derived metrics, not one silently reused + assert first.root != second.root assert len(service.entity_client.entities) == 2 -async def test_store_derived_metric_is_content_addressed_dedup(service: MetricService) -> None: +async def test_store_derived_metric_is_content_addressed_dedup(service: MetricService, fake_files) -> None: bundle = _bundle() first = await service.store_derived_metric(bundle, workspace="default") second = await service.store_derived_metric(bundle, workspace="default") - # Identical content collapses to one stored bundle (same ref, single entity, single fileset). assert first.root == second.root assert len(service.entity_client.entities) == 1 - assert len(service.sdk._store) == 1 + assert len(fake_files._store) == 1 async def test_list_excludes_derived_by_default(service: MetricService) -> None: @@ -235,6 +231,5 @@ async def _spy(entity_cls, *, filter_operation=None, **kwargs): await service.list_metrics("default") await service.list_metrics("default", include_derived=True) - # Default listing injects a filter (NOT derived); include_derived passes none through. assert captured[0] is not None assert captured[1] is None diff --git a/plugins/nemo-evaluator/tests/api/v2/test_metrics_routes.py b/plugins/nemo-evaluator/tests/api/v2/test_metrics_routes.py index 7d3b7e464f..2638f7c9e3 100644 --- a/plugins/nemo-evaluator/tests/api/v2/test_metrics_routes.py +++ b/plugins/nemo-evaluator/tests/api/v2/test_metrics_routes.py @@ -11,6 +11,7 @@ from __future__ import annotations from datetime import datetime, timezone +from unittest.mock import AsyncMock, patch import pytest from fastapi import FastAPI @@ -33,33 +34,30 @@ # ---- in-memory fakes ------------------------------------------------------- -class _FakeFilesets: - def __init__(self, store: dict[tuple[str, str], dict[str, bytes]]) -> None: - self._store = store +class _FakeAsyncFilesClient: + def __init__(self) -> None: + self._store: dict[tuple[str, str], dict[str, bytes]] = {} - async def create(self, *, name, workspace, description=None, exist_ok=False): - self._store.setdefault((workspace, name), {}) - return object() + async def create_fileset(self, *, body, workspace=None, exist_ok=False): + self._store.setdefault((workspace, body.name), {}) + return AsyncMock(data=lambda: object()) - async def delete(self, name, *, workspace=None): + async def delete_fileset(self, *, name, workspace=None): self._store.pop((workspace, name), None) - return object() - + return AsyncMock(data=lambda: object()) -class _FakeFiles: - def __init__(self, store: dict[tuple[str, str], dict[str, bytes]]) -> None: - self._store = store - self.filesets = _FakeFilesets(store) + async def upload_file(self, *, path, content, workspace, name): + self._store.setdefault((workspace, name), {})[path] = bytes(content) + return AsyncMock(data=lambda: object()) - async def upload_content(self, *, content, remote_path, fileset, workspace, fileset_auto_create=False): - self._store.setdefault((workspace, fileset), {})[remote_path] = bytes(content) - return object() + async def download_file(self, *, path, workspace, name): + class _Resp: + async def read(self): + return self._data - -class _FakeSDK: - def __init__(self) -> None: - self._store: dict[tuple[str, str], dict[str, bytes]] = {} - self.files = _FakeFiles(self._store) + resp = _Resp() + resp._data = self._store[(workspace, name)][path] + return resp class _FakeEntityClient: @@ -103,12 +101,12 @@ async def list(self, entity_cls, *, workspace, filter_operation=None, sort=None, @pytest.fixture def client() -> TestClient: app = FastAPI() - # Mirror production mounting: the router's paths are relative; the - # workspace-scoped prefix is applied via RouterSpec in the plugin service. app.include_router(metrics_routes.router, prefix="/v2/workspaces/{workspace}") - service = MetricService(_FakeEntityClient(), _FakeSDK()) + fake_files = _FakeAsyncFilesClient() + service = MetricService(_FakeEntityClient(), object()) app.dependency_overrides[get_metric_service] = lambda: service - return TestClient(app) + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_files): + yield TestClient(app) def _create_body() -> dict: diff --git a/plugins/nemo-evaluator/tests/test_filesets.py b/plugins/nemo-evaluator/tests/test_filesets.py index cce276049d..83e290d769 100644 --- a/plugins/nemo-evaluator/tests/test_filesets.py +++ b/plugins/nemo-evaluator/tests/test_filesets.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import cast +from unittest.mock import patch import pytest from nemo_evaluator.filesets import FilesetRef, download_dataset @@ -15,8 +16,8 @@ class _FakeFilesetFileSystem: - def __init__(self, sdk: object) -> None: - self._sdk = sdk + def __init__(self, *, client: object) -> None: + self._client = client async def _get_file(self, remote_path: str, local_path: str) -> None: raise AssertionError(f"unexpected download to {local_path} from {remote_path}") @@ -30,7 +31,10 @@ async def test_download_dataset_rejects_fragment_path_escape(mocker: MockerFixtu """Fileset fragments should not write outside the requested destination.""" mocker.patch("nemo_evaluator.filesets.FilesetFileSystem", _FakeFilesetFileSystem) - with pytest.raises(ValueError, match="Fileset path escapes destination"): + with ( + patch("nemo_evaluator.filesets.client_from_platform", return_value=object()), + pytest.raises(ValueError, match="Fileset path escapes destination"), + ): await download_dataset( cast(AsyncNeMoPlatform, object()), FilesetRef(root="default/helpsteer2#../../outside.jsonl"), @@ -43,7 +47,10 @@ async def test_download_dataset_rejects_absolute_root_path(mocker: MockerFixture """Fileset roots should not be able to become absolute local paths.""" mocker.patch("nemo_evaluator.filesets.FilesetFileSystem", _FakeFilesetFileSystem) - with pytest.raises(ValueError, match="Fileset path escapes destination"): + with ( + patch("nemo_evaluator.filesets.client_from_platform", return_value=object()), + pytest.raises(ValueError, match="Fileset path escapes destination"), + ): await download_dataset( cast(AsyncNeMoPlatform, object()), FilesetRef(root="/tmp/outside"), diff --git a/plugins/nemo-evaluator/tests/test_metric_refs.py b/plugins/nemo-evaluator/tests/test_metric_refs.py index d30b6d8be4..d5e3616d76 100644 --- a/plugins/nemo-evaluator/tests/test_metric_refs.py +++ b/plugins/nemo-evaluator/tests/test_metric_refs.py @@ -3,6 +3,8 @@ from __future__ import annotations +from unittest.mock import AsyncMock, patch + import pytest from nemo_evaluator.api.schemas import MetricInline from nemo_evaluator.entities import MetricBundleEntity @@ -22,36 +24,32 @@ # ---- in-memory fakes (mirror the storage round-trip) ----------------------- -class _FakeFilesets: - def __init__(self, store) -> None: - self._store = store - - async def create(self, *, name, workspace, description=None, exist_ok=False): - self._store.setdefault((workspace, name), {}) - return object() +class _FakeResponse: + def __init__(self, data: bytes) -> None: + self._data = data - async def delete(self, name, *, workspace=None): - self._store.pop((workspace, name), None) - return object() + async def read(self) -> bytes: + return self._data -class _FakeFiles: - def __init__(self, store) -> None: - self._store = store - self.filesets = _FakeFilesets(store) +class _FakeAsyncFilesClient: + def __init__(self) -> None: + self._store: dict[tuple[str, str], dict[str, bytes]] = {} - async def upload_content(self, *, content, remote_path, fileset, workspace, fileset_auto_create=False): - self._store.setdefault((workspace, fileset), {})[remote_path] = bytes(content) - return object() + async def create_fileset(self, *, body, workspace=None, exist_ok=False): + self._store.setdefault((workspace, body.name), {}) + return AsyncMock(data=lambda: object()) - async def download_content(self, *, remote_path, fileset, workspace): - return self._store[(workspace, fileset)][remote_path] + async def delete_fileset(self, *, name, workspace=None): + self._store.pop((workspace, name), None) + return AsyncMock(data=lambda: object()) + async def upload_file(self, *, path, content, workspace, name): + self._store.setdefault((workspace, name), {})[path] = bytes(content) + return AsyncMock(data=lambda: object()) -class _FakeSDK: - def __init__(self) -> None: - self._store: dict[tuple[str, str], dict[str, bytes]] = {} - self.files = _FakeFiles(self._store) + async def download_file(self, *, path, workspace, name): + return _FakeResponse(self._store[(workspace, name)][path]) class _FakeEntityClient: @@ -75,9 +73,10 @@ def _metric_inline() -> MetricInline: return MetricInline.model_validate_json(_bundle().model_dump_json()) -async def _stored(sdk: _FakeSDK, entity_client: _FakeEntityClient, workspace: str, name: str): +async def _stored(fake_client: _FakeAsyncFilesClient, entity_client: _FakeEntityClient, workspace: str, name: str): bundle = _bundle() - ref = await store_bundle(sdk, workspace, name, bundle) + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client): + ref = await store_bundle(object(), workspace, name, bundle) entity_client.entities[(workspace, name)] = MetricBundleEntity( name=name, workspace=workspace, @@ -103,7 +102,6 @@ def test_parse_metric_ref_bare_name_uses_default_workspace() -> None: @pytest.mark.parametrize("ref", ["", "ws/", "/name", "ws/a/b", "bad name"]) def test_metric_ref_field_rejects_malformed(ref: str) -> None: - # Malformed refs are rejected by the MetricRef field pattern at validation time. with pytest.raises(ValidationError): MetricRef(root=ref) @@ -120,33 +118,35 @@ async def test_resolve_converts_inline_metric_to_runtime_bundle() -> None: async def test_resolve_loads_referenced_bundle() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() entity_client = _FakeEntityClient() - stored = await _stored(sdk, entity_client, "default", "exact") + stored = await _stored(fake_client, entity_client, "default", "exact") - result = await resolve_metric_specs( - [MetricRef(root="default/exact")], - workspace="default", - entity_client=entity_client, - async_sdk=sdk, - ) + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client): + result = await resolve_metric_specs( + [MetricRef(root="default/exact")], + workspace="default", + entity_client=entity_client, + async_sdk=object(), + ) assert len(result) == 1 assert result[0].payload.digest == stored.payload.digest async def test_resolve_mixes_refs_and_inline_preserving_order() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() entity_client = _FakeEntityClient() - await _stored(sdk, entity_client, "default", "exact") + await _stored(fake_client, entity_client, "default", "exact") inline = _metric_inline() - result = await resolve_metric_specs( - [MetricRef(root="exact"), inline], - workspace="default", - entity_client=entity_client, - async_sdk=sdk, - ) + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client): + result = await resolve_metric_specs( + [MetricRef(root="exact"), inline], + workspace="default", + entity_client=entity_client, + async_sdk=object(), + ) assert len(result) == 2 assert result[1].payload.digest == inline.payload.digest @@ -168,7 +168,7 @@ async def test_resolve_missing_metric_raises_clear_error() -> None: [MetricRef(root="default/no-such-metric")], workspace="default", entity_client=_FakeEntityClient(), - async_sdk=_FakeSDK(), + async_sdk=object(), ) @@ -178,7 +178,7 @@ async def test_resolve_ref_without_entity_client_raises() -> None: [MetricRef(root="default/exact")], workspace="default", entity_client=None, - async_sdk=_FakeSDK(), + async_sdk=object(), ) diff --git a/plugins/nemo-evaluator/tests/test_metric_storage.py b/plugins/nemo-evaluator/tests/test_metric_storage.py index ce8f0f41d0..9c06123b88 100644 --- a/plugins/nemo-evaluator/tests/test_metric_storage.py +++ b/plugins/nemo-evaluator/tests/test_metric_storage.py @@ -3,6 +3,8 @@ from __future__ import annotations +from unittest.mock import AsyncMock, patch + import pytest from nemo_evaluator.metric_storage import ( BUNDLE_FILENAME, @@ -18,36 +20,32 @@ from nemo_evaluator_sdk.metrics.exact_match import ExactMatchMetric -class _FakeFilesets: - def __init__(self, store: dict[tuple[str, str], dict[str, bytes]]) -> None: - self._store = store - - async def create(self, *, name, workspace, description=None, exist_ok=False): - self._store.setdefault((workspace, name), {}) - return object() +class _FakeResponse: + def __init__(self, data: bytes) -> None: + self._data = data - async def delete(self, name, *, workspace=None): - self._store.pop((workspace, name), None) - return object() + async def read(self) -> bytes: + return self._data -class _FakeFiles: - def __init__(self, store: dict[tuple[str, str], dict[str, bytes]]) -> None: - self._store = store - self.filesets = _FakeFilesets(store) +class _FakeAsyncFilesClient: + def __init__(self) -> None: + self._store: dict[tuple[str, str], dict[str, bytes]] = {} - async def upload_content(self, *, content, remote_path, fileset, workspace, fileset_auto_create=False): - self._store.setdefault((workspace, fileset), {})[remote_path] = bytes(content) - return object() + async def create_fileset(self, *, body, workspace=None, exist_ok=False): + self._store.setdefault((workspace, body.name), {}) + return AsyncMock(data=lambda: object()) - async def download_content(self, *, remote_path, fileset, workspace): - return self._store[(workspace, fileset)][remote_path] + async def delete_fileset(self, *, name, workspace=None): + self._store.pop((workspace, name), None) + return AsyncMock(data=lambda: object()) + async def upload_file(self, *, path, content, workspace, name): + self._store.setdefault((workspace, name), {})[path] = bytes(content) + return AsyncMock(data=lambda: object()) -class _FakeSDK: - def __init__(self) -> None: - self._store: dict[tuple[str, str], dict[str, bytes]] = {} - self.files = _FakeFiles(self._store) + async def download_file(self, *, path, workspace, name): + return _FakeResponse(self._store[(workspace, name)][path]) def _sample_bundle(): @@ -70,90 +68,103 @@ def test_parse_bundle_ref_rejects_malformed(ref: str) -> None: async def test_store_returns_unique_per_metric_ref() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() bundle = _sample_bundle() - ref1 = await store_bundle(sdk, "default", "my-metric", bundle) - ref2 = await store_bundle(sdk, "default", "my-metric", bundle) + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client): + ref1 = await store_bundle(object(), "default", "my-metric", bundle) + ref2 = await store_bundle(object(), "default", "my-metric", bundle) assert ref1.startswith("default/metric-bundle.") assert ref1.endswith(f"#{BUNDLE_FILENAME}") - # Each upload lands in its own fileset, so a rollback can't clobber another. assert ref1 != ref2 async def test_store_fileset_name_stays_within_limit_for_long_metric_name() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() bundle = _sample_bundle() - long_name = "m" * 255 # MAX_NAME_LENGTH + long_name = "m" * 255 - ref = await store_bundle(sdk, "default", long_name, bundle) + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client): + ref = await store_bundle(object(), "default", long_name, bundle) _, fileset, _ = parse_bundle_ref(ref) - # The Files service caps fileset names at 255 chars. assert len(fileset) <= 255 async def test_store_cleans_up_fileset_on_upload_failure() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() bundle = _sample_bundle() - async def _boom(*args, **kwargs): + async def _boom(*, path, content, workspace, name): raise RuntimeError("network blip during upload") - sdk.files.upload_content = _boom + fake_client.upload_file = _boom - with pytest.raises(MetricBundleStorageError): - await store_bundle(sdk, "default", "my-metric", bundle) - # The fileset created just before the failed upload must not be left orphaned. - assert [key for key in sdk._store if key[1].startswith(FILESET_PREFIX)] == [] + with ( + patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client), + pytest.raises(MetricBundleStorageError), + ): + await store_bundle(object(), "default", "my-metric", bundle) + + assert [key for key in fake_client._store if key[1].startswith(FILESET_PREFIX)] == [] async def test_store_then_load_round_trips_bundle() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() bundle = _sample_bundle() - ref = await store_bundle(sdk, "default", "my-metric", bundle) + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client): + ref = await store_bundle(object(), "default", "my-metric", bundle) + loaded = await load_bundle(object(), ref, expected_digest=bundle.payload.digest) - loaded = await load_bundle(sdk, ref, expected_digest=bundle.payload.digest) assert loaded.metric_type == bundle.metric_type assert loaded.payload.digest == bundle.payload.digest async def test_load_rejects_digest_mismatch() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() bundle = _sample_bundle() - ref = await store_bundle(sdk, "default", "my-metric", bundle) - with pytest.raises(MetricBundleStorageError, match="digest mismatch"): - await load_bundle(sdk, ref, expected_digest="deadbeef") + with ( + patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client), + pytest.raises(MetricBundleStorageError, match="digest mismatch"), + ): + ref = await store_bundle(object(), "default", "my-metric", bundle) + await load_bundle(object(), ref, expected_digest="deadbeef") async def test_load_rejects_corrupt_bundle() -> None: - sdk = _FakeSDK() - # Stored bytes that aren't a valid serialized MetricBundle. - sdk._store[("default", "metric-bundle.deadbeef")] = {"bundle.json": b"not a bundle"} + fake_client = _FakeAsyncFilesClient() + fake_client._store[("default", "metric-bundle.deadbeef")] = {"bundle.json": b"not a bundle"} - with pytest.raises(MetricBundleStorageError, match="corrupt or unreadable"): - await load_bundle(sdk, "default/metric-bundle.deadbeef#bundle.json") + with ( + patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client), + pytest.raises(MetricBundleStorageError, match="corrupt or unreadable"), + ): + await load_bundle(object(), "default/metric-bundle.deadbeef#bundle.json") async def test_load_wraps_download_failure() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() - with pytest.raises(MetricBundleStorageError, match="failed to download metric bundle"): - await load_bundle(sdk, "default/metric-bundle.missing#bundle.json") + with ( + patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client), + pytest.raises(MetricBundleStorageError, match="failed to download metric bundle"), + ): + await load_bundle(object(), "default/metric-bundle.missing#bundle.json") async def test_delete_by_ref_removes_only_that_fileset() -> None: - sdk = _FakeSDK() + fake_client = _FakeAsyncFilesClient() bundle = _sample_bundle() - ref1 = await store_bundle(sdk, "default", "my-metric", bundle) - ref2 = await store_bundle(sdk, "default", "my-metric", bundle) - await delete_bundle_by_ref(sdk, ref1) + with patch("nemo_evaluator.metric_storage.client_from_platform", return_value=fake_client): + ref1 = await store_bundle(object(), "default", "my-metric", bundle) + ref2 = await store_bundle(object(), "default", "my-metric", bundle) + await delete_bundle_by_ref(object(), ref1) _, fileset1, _ = parse_bundle_ref(ref1) _, fileset2, _ = parse_bundle_ref(ref2) - assert ("default", fileset1) not in sdk._store - assert ("default", fileset2) in sdk._store + assert ("default", fileset1) not in fake_client._store + assert ("default", fileset2) in fake_client._store diff --git a/plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py b/plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py index d888862b80..0dc4495df6 100644 --- a/plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py +++ b/plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py @@ -21,7 +21,12 @@ import os import sys -from nemo_platform import APIStatusError, ConflictError, NeMoPlatform, NotFoundError +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError, NemoHTTPError, NotFoundError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig +from nemo_platform_plugin.files.types import CreateFilesetRequest logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") logger = logging.getLogger(__name__) @@ -87,6 +92,7 @@ def create_filesets( Returns: List of created fileset names """ + files = client_from_platform(sdk, FilesClient) created = [] for fileset_config in MODEL_FILESETS: name = fileset_config["name"] @@ -99,7 +105,7 @@ def create_filesets( continue try: - sdk.files.filesets.retrieve(name=name, workspace=workspace) + files.get_fileset(name=name, workspace=workspace) logger.info("Fileset already exists: %s", full_name) created.append(name) continue @@ -107,19 +113,22 @@ def create_filesets( logger.debug("Fileset does not exist yet: %s", full_name) try: - sdk.files.filesets.create( + storage = HuggingfaceStorageConfig.model_validate(fileset_config["storage"]) + files.create_fileset( workspace=workspace, - name=name, - description=fileset_config.get("description", ""), - purpose="generic", - storage=fileset_config["storage"], + body=CreateFilesetRequest( + name=name, + description=fileset_config.get("description", ""), + purpose="generic", + storage=storage, + ), ) logger.info("Created fileset: %s", full_name) created.append(name) except ConflictError: logger.info("Fileset already exists: %s", full_name) created.append(name) - except APIStatusError as e: + except NemoHTTPError as e: logger.error("Failed to create fileset %s: %s", full_name, e) except ValueError as e: logger.error("Invalid fileset config for %s: %s", full_name, e) diff --git a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.py b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.py index d6422961d0..e3e91431f0 100644 --- a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.py +++ b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.py @@ -14,7 +14,11 @@ from nemo_platform import AsyncNeMoPlatform, NotFoundError, PermissionDeniedError from nemo_platform.filesets import FilesetPathError, parse_fileset_ref from nemo_platform_plugin.authz import AuthzScope +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError +from nemo_platform_plugin.client.errors import PermissionDeniedError as ClientPermissionDeniedError from nemo_platform_plugin.entities import EntityClient +from nemo_platform_plugin.files.client import AsyncFilesClient from nemo_platform_plugin.jobs.api_factory import ( ContainerSpec, EnvironmentVariable, @@ -116,13 +120,14 @@ async def job_config_compiler( ds_workspace, fileset_name, _ = parse_fileset_ref(transformed_spec.data_source, workspace_fallback=workspace) except FilesetPathError as e: raise PlatformJobCompilationError(f"Invalid data_source format: {transformed_spec.data_source!r}") from e + files = client_from_platform(sdk, AsyncFilesClient) try: - await sdk.files.filesets.retrieve(name=fileset_name, workspace=ds_workspace) - except NotFoundError as e: + await files.get_fileset(name=fileset_name, workspace=ds_workspace) + except ClientNotFoundError as e: raise PlatformJobCompilationError( f"Could not find fileset {fileset_name!r} in workspace {ds_workspace!r}" ) from e - except PermissionDeniedError as e: + except ClientPermissionDeniedError as e: raise PermissionError(f"Access denied to fileset {fileset_name!r} in workspace {ds_workspace!r}") from e environment = [ diff --git a/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py b/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py index 446aee4637..22311e546a 100644 --- a/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py +++ b/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py @@ -1,10 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from nemo_platform import NotFoundError, PermissionDeniedError +from nemo_platform import NotFoundError +from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError +from nemo_platform_plugin.client.errors import PermissionDeniedError as ClientPermissionDeniedError from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError from nemo_safe_synthesizer.config.replace_pii import ClassifyConfig, Globals, PiiReplacerConfig, StepDefinition from nemo_safe_synthesizer_plugin.api.v2.jobs import endpoints @@ -18,9 +20,15 @@ @pytest.fixture -def mock_sdk(): +def mock_files_client(): + mock_client = MagicMock() + mock_client.get_fileset = AsyncMock() + return mock_client + + +@pytest.fixture +def mock_sdk(mock_files_client): sdk = MagicMock() - sdk.files.filesets.retrieve = AsyncMock() sdk.inference.providers.retrieve = AsyncMock() sdk.models.get_provider_route_openai_url = MagicMock( return_value="http://nmp-host/apis/inference-gateway/v2/workspaces/default/provider/my-nim/-/v1" @@ -28,6 +36,15 @@ def mock_sdk(): return sdk +@pytest.fixture(autouse=True) +def _patch_client_from_platform(mock_files_client): + with patch( + "nemo_safe_synthesizer_plugin.api.v2.jobs.endpoints.client_from_platform", + return_value=mock_files_client, + ): + yield + + @pytest.fixture(autouse=True) def mock_runtime_command(monkeypatch): monkeypatch.setattr(endpoints, "runtime_task_command", lambda _config: ["/runtime/bin/python", "-m", TASK_MODULE]) @@ -58,29 +75,33 @@ async def _compile(spec, mock_sdk): @pytest.mark.asyncio -async def test_job_config_compiler_validates_data_source(mock_sdk): +async def test_job_config_compiler_validates_data_source(mock_sdk, mock_files_client): spec = _make_spec(data_source="my-workspace/my-fileset#data.csv") await _compile(spec, mock_sdk) - mock_sdk.files.filesets.retrieve.assert_awaited_once_with(name="my-fileset", workspace="my-workspace") + mock_files_client.get_fileset.assert_awaited_once_with(name="my-fileset", workspace="my-workspace") @pytest.mark.asyncio -async def test_job_config_compiler_data_source_not_found(mock_sdk): - mock_sdk.files.filesets.retrieve.side_effect = NotFoundError( - message="not found", response=MagicMock(status_code=404), body=None - ) +async def test_job_config_compiler_data_source_not_found(mock_sdk, mock_files_client): + mock_response = MagicMock() + mock_response.status_code = 404 + mock_response.json.return_value = {"detail": "not found"} + mock_response.text = "not found" + mock_files_client.get_fileset.side_effect = ClientNotFoundError(mock_response) with pytest.raises(PlatformJobCompilationError, match="Could not find fileset"): await _compile(_make_spec(), mock_sdk) @pytest.mark.asyncio -async def test_job_config_compiler_data_source_permission_denied(mock_sdk): - mock_sdk.files.filesets.retrieve.side_effect = PermissionDeniedError( - message="denied", response=MagicMock(status_code=403), body=None - ) +async def test_job_config_compiler_data_source_permission_denied(mock_sdk, mock_files_client): + mock_response = MagicMock() + mock_response.status_code = 403 + mock_response.json.return_value = {"detail": "denied"} + mock_response.text = "denied" + mock_files_client.get_fileset.side_effect = ClientPermissionDeniedError(mock_response) with pytest.raises(PermissionError, match="Access denied to fileset"): await _compile(_make_spec(), mock_sdk) diff --git a/plugins/nemo-unsloth/tests/test_jobs.py b/plugins/nemo-unsloth/tests/test_jobs.py index b9158d82c1..8a1fa5c00b 100644 --- a/plugins/nemo-unsloth/tests/test_jobs.py +++ b/plugins/nemo-unsloth/tests/test_jobs.py @@ -56,17 +56,28 @@ def _stub_async_sdk() -> SimpleNamespace: ) +def _mock_files_client() -> AsyncMock: + """Build a mock AsyncFilesClient for check_dataset_access.""" + mock = AsyncMock() + mock.get_fileset.return_value = MagicMock() + return mock + + def _make_canonical(workspace: str = "default", **overrides: Any) -> UnslothJobOutput: spec = UnslothJobInput.model_validate(_input_dict(**overrides)) - return asyncio.run( - UnslothJob.to_spec( - spec, - workspace=workspace, - entity_client=object(), - async_sdk=_stub_async_sdk(), - is_local=False, - ), - ) + with patch( + "nmp.customization_common.service.platform_client.client_from_platform", + return_value=_mock_files_client(), + ): + return asyncio.run( + UnslothJob.to_spec( + spec, + workspace=workspace, + entity_client=object(), + async_sdk=_stub_async_sdk(), + is_local=False, + ), + ) class TestToSpec: diff --git a/plugins/nemo-unsloth/tests/test_schema.py b/plugins/nemo-unsloth/tests/test_schema.py index 5d014fa9bc..bbc4e662db 100644 --- a/plugins/nemo-unsloth/tests/test_schema.py +++ b/plugins/nemo-unsloth/tests/test_schema.py @@ -9,7 +9,7 @@ import json from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from nemo_unsloth_plugin.schema import ( @@ -44,8 +44,19 @@ def _stub_sdk(*, is_embedding: bool = False) -> SimpleNamespace: ) +def _mock_files_client() -> AsyncMock: + """Build a mock AsyncFilesClient for check_dataset_access.""" + mock = AsyncMock() + mock.get_fileset.return_value = MagicMock() + return mock + + def _run_transform(spec: UnslothJobInput) -> UnslothJobOutput: - return asyncio.run(transform_input_to_output(spec, "default", _stub_sdk())) + with patch( + "nmp.customization_common.service.platform_client.client_from_platform", + return_value=_mock_files_client(), + ): + return asyncio.run(transform_input_to_output(spec, "default", _stub_sdk())) class TestCanonicalReexport: @@ -225,7 +236,13 @@ def test_merged_inferred_as_model_type(self) -> None: def test_embedding_model_rejected(self) -> None: sdk = _stub_sdk(is_embedding=True) spec = UnslothJobInput.model_validate(_minimal_payload()) - with pytest.raises(ValueError, match="Embedding-model SFT"): + with ( + patch( + "nmp.customization_common.service.platform_client.client_from_platform", + return_value=_mock_files_client(), + ), + pytest.raises(ValueError, match="Embedding-model SFT"), + ): asyncio.run(transform_input_to_output(spec, "default", sdk)) diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/__init__.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/__init__.py index 1ed357681f..8b5483839c 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/__init__.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/__init__.py @@ -8,6 +8,8 @@ from typing import Annotated import typer +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from nemo_platform.cli.commands.api.files import filesets, otlp from nemo_platform.cli.core.context import CLIContext @@ -60,6 +62,7 @@ def upload_files( raw_local_path: str = ctx.params.get("local_path") client = state.get_client() + files = client_from_platform(client, FilesClient) if workspace is None: workspace = client._get_workspace_path_param() @@ -68,7 +71,7 @@ def upload_files( with RichProgressCallback(description="Uploading") as callback: if fileset is not None: # Validate fileset exists before uploading - client.files.filesets.retrieve(fileset, workspace=workspace) + files.get_fileset(name=fileset, workspace=workspace) client.files.upload( local_path=raw_local_path, remote_path=remote_path, 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 27f01cc2ce..56230e0e53 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 @@ -17,7 +17,6 @@ from fsspec.asyn import AbstractAsyncStreamedFile, AsyncFileSystem, _get_batch_size from fsspec.callbacks import DEFAULT_CALLBACK, Callback from fsspec.spec import AbstractBufferedFile -from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient from nemo_platform_plugin.files.types import FilesetFileOutput, ListFilesQueryParams @@ -350,26 +349,11 @@ def register_fsspec(cls) -> None: def __init__( self, *, - client: FilesClient | AsyncFilesClient | None = None, - sdk: NeMoPlatform | AsyncNeMoPlatform | None = None, + client: FilesClient | AsyncFilesClient, 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) @@ -384,22 +368,14 @@ def __init__( @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. - """ + """Ensure we have an AsyncFilesClient, converting from sync if needed.""" 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( + return AsyncFilesClient( base_url=client.base_url, workspace=client.workspace, auth=client._auth, 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 231b238328..610854e848 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py +++ b/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py @@ -12,23 +12,16 @@ from dataclasses import dataclass from functools import cached_property from pathlib import PurePath -from typing import Any, Protocol, runtime_checkable +from typing import Protocol, runtime_checkable -import nemo_platform from fsspec.callbacks import Callback from fsspec.core import has_magic -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 ( @@ -38,85 +31,6 @@ ) -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. @@ -216,199 +130,6 @@ def _matches_glob(filepath: str, pattern: str) -> bool: return PurePath(filepath).match(pattern) -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, - ) - return self._client.create_fileset(workspace=workspace, body=body, exist_ok=exist_ok).data() - - 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, - ) - return (await self._client.create_fileset(workspace=workspace, body=body, exist_ok=exist_ok)).data() - - 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. @@ -416,15 +137,18 @@ class FilesResource: For fsspec filesystem access, use ``resource.fsspec``. """ - def __init__(self, client) -> None: - from nemo_platform_plugin.client.adapter import client_from_platform + def __init__(self, client, *, files_client: FilesClient | None = None) -> None: + if files_client is not None: + self._client = files_client + else: + from nemo_platform_plugin.client.adapter import client_from_platform - self._client = client_from_platform(client, _RemappingFilesClient) + self._client = client_from_platform(client, FilesClient) @cached_property - def filesets(self) -> FilesetsSubResource: - """Access fileset CRUD operations (create, retrieve, update, list, delete).""" - return FilesetsSubResource(self._client) + def client(self) -> FilesClient: + """Access the underlying FilesClient for direct API calls.""" + return self._client @cached_property def fsspec(self) -> FilesetFileSystem: @@ -433,7 +157,11 @@ def fsspec(self) -> FilesetFileSystem: def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None: """Create fileset if it doesn't exist (idempotent).""" - self.filesets.create(name=fileset, workspace=workspace, exist_ok=True) + self._client.create_fileset( + workspace=workspace, + body=CreateFilesetRequest(name=fileset), + exist_ok=True, + ) def download( self, @@ -650,7 +378,7 @@ def upload( kwargs["callback"] = callback self.fsspec.put(**kwargs) - return self.filesets.retrieve(name=fileset, workspace=ws) + return self._client.get_fileset(name=fileset, workspace=ws).data() def upload_content( self, @@ -751,7 +479,7 @@ def upload_content( case _: raise TypeError(f"Unsupported content type: {type(content)}") - return self.filesets.retrieve(name=fileset, workspace=ws) + return self._client.get_fileset(name=fileset, workspace=ws).data() def download_content( self, @@ -938,15 +666,18 @@ class AsyncFilesResource: For fsspec filesystem access, use ``resource.fsspec``. """ - def __init__(self, client) -> None: - from nemo_platform_plugin.client.adapter import client_from_platform + def __init__(self, client, *, files_client: AsyncFilesClient | None = None) -> None: + if files_client is not None: + self._client = files_client + else: + from nemo_platform_plugin.client.adapter import client_from_platform - self._client = client_from_platform(client, _RemappingAsyncFilesClient) + self._client = client_from_platform(client, AsyncFilesClient) @cached_property - def filesets(self) -> AsyncFilesetsSubResource: - """Access fileset CRUD operations (create, retrieve, update, list, delete).""" - return AsyncFilesetsSubResource(self._client) + def client(self) -> AsyncFilesClient: + """Access the underlying AsyncFilesClient for direct API calls.""" + return self._client @cached_property def fsspec(self) -> FilesetFileSystem: @@ -955,7 +686,11 @@ def fsspec(self) -> FilesetFileSystem: async def _ensure_fileset_exists(self, workspace: str, fileset: str) -> None: """Create fileset if it doesn't exist (idempotent).""" - await self.filesets.create(name=fileset, workspace=workspace, exist_ok=True) + await self._client.create_fileset( + workspace=workspace, + body=CreateFilesetRequest(name=fileset), + exist_ok=True, + ) async def download( self, @@ -1151,7 +886,7 @@ async def upload( kwargs["callback"] = callback await self.fsspec._put(**kwargs) - return await self.filesets.retrieve(name=fileset, workspace=ws) + return (await self._client.get_fileset(name=fileset, workspace=ws)).data() async def upload_content( self, @@ -1256,7 +991,7 @@ async def _read_chunks(f: AsyncReadable, chunk_size: int = 1024 * 1024) -> Async case _: raise TypeError(f"Unsupported content type: {type(content)}") - return await self.filesets.retrieve(name=fileset, workspace=ws) + return (await self._client.get_fileset(name=fileset, workspace=ws)).data() async def download_content( self, diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/conftest.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/conftest.py index b63b7239a3..7b740941c4 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/conftest.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/conftest.py @@ -19,6 +19,8 @@ from click.testing import Result from nemo_platform import NeMoPlatform from nemo_platform.cli.core.context import CLIContext +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from nmp.core.files.service import FilesService from nmp.testing import create_test_client from starlette.testclient import TestClient @@ -40,6 +42,12 @@ def sdk(http_client: TestClient) -> NeMoPlatform: return NeMoPlatform(base_url="http://testserver", http_client=http_client) +@pytest.fixture(scope="module") +def files_client(sdk: NeMoPlatform) -> FilesClient: + """Provide a FilesClient derived from the SDK.""" + return client_from_platform(sdk, FilesClient) + + @pytest.fixture def random_workspace(sdk: NeMoPlatform) -> str: """ diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/test_filesets.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/test_filesets.py index 76c5ef2665..77ba83f8fb 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/test_filesets.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/test_filesets.py @@ -8,15 +8,20 @@ import pytest from nemo_platform import NeMoPlatform from nemo_platform.cli.app import app +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from ..utils import assert_exit_code from .conftest import NmpCliRunner @pytest.fixture -def test_fileset(sdk: NeMoPlatform, random_workspace: str) -> dict: +def test_fileset(files_client: FilesClient, random_workspace: str) -> dict: """Create a test fileset.""" - fileset = sdk.files.filesets.create(workspace=random_workspace, name="test-fileset") + fileset = files_client.create_fileset( + body=CreateFilesetRequest(name="test-fileset"), workspace=random_workspace + ).data() return {"workspace": random_workspace, "name": fileset.name} @@ -138,7 +143,7 @@ def test_upload_to_nonexistent_fileset_fails( ) assert_exit_code(result, 1) - assert "Not found" in result.stderr + assert "not found" in result.stderr.lower() @pytest.mark.parametrize( ("remote_path", "expected_suffix"), @@ -196,10 +201,8 @@ def test_upload_without_fileset_auto_creates( fileset_name = match.group(1) # Verify fileset exists - fileset = runner.client.files.filesets.retrieve( - name=fileset_name, - workspace=random_workspace, - ) + files = client_from_platform(runner.client, FilesClient) + fileset = files.get_fileset(name=fileset_name, workspace=random_workspace).data() assert fileset.name == fileset_name # Verify file was uploaded @@ -212,7 +215,9 @@ def test_upload_without_fileset_auto_creates( @pytest.fixture -def fileset_with_nested_files(sdk: NeMoPlatform, random_workspace: str, tmp_path: Path) -> dict: +def fileset_with_nested_files( + sdk: NeMoPlatform, files_client: FilesClient, random_workspace: str, tmp_path: Path +) -> dict: """Create a fileset with nested file structure for download tests. Structure: @@ -222,7 +227,9 @@ def fileset_with_nested_files(sdk: NeMoPlatform, random_workspace: str, tmp_path file2.txt file3.txt """ - fileset = sdk.files.filesets.create(workspace=random_workspace, name="download-test-fileset") + fileset = files_client.create_fileset( + body=CreateFilesetRequest(name="download-test-fileset"), workspace=random_workspace + ).data() # Create nested directory structure locally dir_a = tmp_path / "a" diff --git a/services/automodel/src/nmp/automodel/tasks/file_io/run.py b/services/automodel/src/nmp/automodel/tasks/file_io/run.py index 031197cdca..a91b27e8e8 100644 --- a/services/automodel/src/nmp/automodel/tasks/file_io/run.py +++ b/services/automodel/src/nmp/automodel/tasks/file_io/run.py @@ -23,12 +23,21 @@ from nemo_platform import ( APIConnectionError, APITimeoutError, - ConflictError, InternalServerError, NeMoPlatform, NotFoundError, ) from nemo_platform.types.files.fileset_file import FilesetFile +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ( + ConflictError, +) +from nemo_platform_plugin.client.errors import ( + InternalServerError as ClientInternalServerError, +) +from nemo_platform_plugin.client.types import RetryPolicy +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest, UpdateFilesetRequest from nmp.automodel.app.constants import SERVICE_NAME from nmp.automodel.tasks.file_io.callbacks import ( CompositeCallback, @@ -63,8 +72,7 @@ logger = logging.getLogger(__name__) -# Timeout configurations for SDK operations (httpx.Timeout for API calls) -CREATE_FILESET_TIMEOUT = httpx.Timeout(10.0, connect=10.0) +CREATE_FILESET_TIMEOUT = 10.0 LIST_FILES_TIMEOUT = httpx.Timeout(10.0, connect=10.0) # Timeout configurations for FilesetFileSystem operations. @@ -387,32 +395,40 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N @retry( stop=stop_after_attempt(MAX_RETRIES), wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type((InternalServerError, APITimeoutError, APIConnectionError)), - reraise=True, # means that the last exception will be re-raised after the last retry attempt + retry=retry_if_exception_type( + ( + InternalServerError, + APITimeoutError, + APIConnectionError, + ClientInternalServerError, + httpx.TimeoutException, + httpx.ConnectError, + ) + ), + reraise=True, ) def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None = None) -> None: """Internal method with retry logic for creating a FileSet.""" + files = client_from_platform(self.sdk, FilesClient).with_options( + timeout=CREATE_FILESET_TIMEOUT, retry=RetryPolicy(max_retries=0) + ) try: - create_kwargs: dict = { - "workspace": fileset.workspace, + body_kwargs: dict = { "name": fileset.name, - "timeout": CREATE_FILESET_TIMEOUT, "custom_fields": {"service_source": "automodel"}, } if metadata is not None: - create_kwargs["metadata"] = metadata - result = self.sdk.with_options(max_retries=0).files.filesets.create(**create_kwargs) + body_kwargs["metadata"] = metadata + result = files.create_fileset(workspace=fileset.workspace, body=CreateFilesetRequest(**body_kwargs)).data() logger.info(f"Created FileSet: {result.workspace}/{result.name}") except ConflictError: - # Fileset already exists - patch metadata so tool_calling etc. are not lost workspace = fileset.workspace or self.job_ctx.workspace if metadata is not None: try: - self.sdk.with_options(max_retries=0).files.filesets.update( - name=fileset.name, + files.update_fileset( workspace=workspace, - metadata=metadata, - timeout=CREATE_FILESET_TIMEOUT, + name=fileset.name, + body=UpdateFilesetRequest(metadata=metadata), ) logger.info(f"Patched existing FileSet metadata: {workspace}/{fileset.name}") except Exception as e: diff --git a/services/automodel/src/nmp/automodel/tasks/model_entity/run.py b/services/automodel/src/nmp/automodel/tasks/model_entity/run.py index b50b604766..246992ff88 100644 --- a/services/automodel/src/nmp/automodel/tasks/model_entity/run.py +++ b/services/automodel/src/nmp/automodel/tasks/model_entity/run.py @@ -36,6 +36,8 @@ ) from nemo_platform.types.models import LoraParam, ModelEntity from nemo_platform.types.shared_params.tool_call_config import ToolCallConfig as ToolCallConfigParam +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from nmp.automodel.app.constants import SERVICE_NAME from nmp.automodel.entities.values import FinetuningType from nmp.common.sdk_factory import get_task_sdk @@ -165,7 +167,9 @@ def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, Mode logger.info(f"Validating fileset exists: {fileset_workspace}/{config.fileset.name}") try: - self.sdk.files.filesets.retrieve(workspace=fileset_workspace, name=config.fileset.name) + client_from_platform(self.sdk, FilesClient).get_fileset( + workspace=fileset_workspace, name=config.fileset.name + ) logger.info(f"Fileset validation successful: {fileset_workspace}/{config.fileset.name}") except Exception as e: logger.error(f"Fileset validation failed: {fileset_workspace}/{config.fileset.name}") diff --git a/services/automodel/tests/test_compiler.py b/services/automodel/tests/test_compiler.py index 7acac971e4..022da249c0 100644 --- a/services/automodel/tests/test_compiler.py +++ b/services/automodel/tests/test_compiler.py @@ -45,8 +45,6 @@ def mock_sdk(): side_effect=lambda name, workspace, verbose=True: _make_mock_model_entity(workspace=workspace, name=name), ) sdk.files = Mock() - sdk.files.filesets = Mock() - sdk.files.filesets.retrieve = AsyncMock(return_value=Mock()) return sdk 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 c0b4a91fd5..565283d600 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 @@ -7,6 +7,8 @@ from nemo_platform import AsyncNeMoPlatform from nemo_platform.types import PlatformJobStatus +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import AsyncFilesClient from nmp.common.api.filter import ComparisonOperation, FilterOperator from nmp.common.controller.controller import Controller from nmp.common.observability import start_span_with_ctx @@ -172,12 +174,13 @@ async def _cleanup_deployments(self, workspace: Workspace) -> None: 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) + files = client_from_platform(self._nmp_sdk, AsyncFilesClient) + filesets_response = await files.list_filesets(workspace=workspace.name) async for fileset in filesets_response.items(): try: logger.info(f"Deleting fileset: {fileset.name}") - await self._nmp_sdk.files.filesets.delete( + await files.delete_fileset( name=fileset.name, workspace=workspace.name, ) diff --git a/services/core/entities/tests/controllers/test_workspace_cleanup.py b/services/core/entities/tests/controllers/test_workspace_cleanup.py index d5691bde54..e156ce6ec5 100644 --- a/services/core/entities/tests/controllers/test_workspace_cleanup.py +++ b/services/core/entities/tests/controllers/test_workspace_cleanup.py @@ -37,31 +37,42 @@ async def __anext__(self): raise StopAsyncIteration -class _MockPaginatedResponse: - """Mock for paginated responses that expose .items() for async iteration.""" +class _MockAsyncPaginatedResponse: + """Mock for AsyncNemoPaginatedResponse that exposes .items() as an async generator.""" def __init__(self, items): self._items = items - def items(self): - return _AsyncIterator(self._items) + async def items(self): + for item in self._items: + yield item + + +def _make_mock_files_client(filesets: list | None = None) -> AsyncMock: + """Build a mock AsyncFilesClient with list_filesets/delete_fileset.""" + mock_files = AsyncMock() + mock_files.list_filesets = AsyncMock(return_value=_MockAsyncPaginatedResponse(filesets or [])) + mock_files.delete_fileset = AsyncMock() + return mock_files def _make_sdk( jobs: list | None = None, deployments: list | None = None, filesets: list | None = None, -) -> MagicMock: - """Build a MagicMock SDK with async mocks wired to the correct paths.""" +) -> tuple[MagicMock, AsyncMock]: + """Build a MagicMock SDK with async mocks wired to the correct paths. + + Returns (sdk, mock_files_client) so tests can assert on files client calls. + """ sdk = MagicMock() sdk.jobs.list = AsyncMock(return_value=_AsyncIterator(jobs or [])) sdk.jobs.cancel = AsyncMock() 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=_MockPaginatedResponse(filesets or [])) - sdk.files.filesets.delete = AsyncMock() - return sdk + mock_files = _make_mock_files_client(filesets) + return sdk, mock_files def _make_job(name: str, status: str = "completed") -> MagicMock: @@ -86,6 +97,9 @@ def _make_controller( ) +_FILES_CLIENT_PATCH = "nmp.core.entities.controllers.workspace_cleanup.client_from_platform" + + class TestWorkspaceCleanupStep: def test_step_skips_when_stop_signal_set(self): import threading @@ -158,11 +172,12 @@ 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=_MockPaginatedResponse([])) + mock_files = _make_mock_files_client([]) controller = _make_controller(workspace_repo=repo, nmp_sdk=sdk) - await controller._async_step() + with patch(_FILES_CLIENT_PATCH, return_value=mock_files): + await controller._async_step() repo.mark_workspace_for_deletion.assert_any_call( name="test-workspace", @@ -338,14 +353,13 @@ async def test_deletes_filesets(self): fileset = MagicMock() fileset.name = "test-fileset" - sdk = MagicMock() - sdk.files.filesets.list = AsyncMock(return_value=_MockPaginatedResponse([fileset])) - sdk.files.filesets.delete = AsyncMock() + mock_files = _make_mock_files_client([fileset]) + controller = _make_controller() - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_filesets(workspace) + with patch(_FILES_CLIENT_PATCH, return_value=mock_files): + await controller._cleanup_filesets(workspace) - sdk.files.filesets.delete.assert_awaited_once_with( + mock_files.delete_fileset.assert_awaited_once_with( name="test-fileset", workspace="test-workspace", ) @@ -358,14 +372,14 @@ async def test_continues_on_individual_fileset_failure(self): fs2 = MagicMock() fs2.name = "fs2" - sdk = MagicMock() - sdk.files.filesets.list = AsyncMock(return_value=_MockPaginatedResponse([fs1, fs2])) - sdk.files.filesets.delete = AsyncMock(side_effect=[Exception("fail"), None]) + mock_files = _make_mock_files_client([fs1, fs2]) + mock_files.delete_fileset = AsyncMock(side_effect=[Exception("fail"), None]) + controller = _make_controller() - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_filesets(workspace) + with patch(_FILES_CLIENT_PATCH, return_value=mock_files): + await controller._cleanup_filesets(workspace) - assert sdk.files.filesets.delete.await_count == 2 + assert mock_files.delete_fileset.await_count == 2 class TestJobCancellationBranches: @@ -374,7 +388,7 @@ class TestJobCancellationBranches: @pytest.mark.asyncio async def test_cancels_pending_jobs(self): workspace = _make_workspace() - sdk = _make_sdk(jobs=[_make_job("pending-job", status="pending")]) + sdk, _ = _make_sdk(jobs=[_make_job("pending-job", status="pending")]) controller = _make_controller(nmp_sdk=sdk) await controller._cleanup_jobs(workspace) @@ -385,7 +399,7 @@ async def test_cancels_pending_jobs(self): @pytest.mark.asyncio async def test_cancels_created_jobs(self): workspace = _make_workspace() - sdk = _make_sdk(jobs=[_make_job("created-job", status="created")]) + sdk, _ = _make_sdk(jobs=[_make_job("created-job", status="created")]) controller = _make_controller(nmp_sdk=sdk) await controller._cleanup_jobs(workspace) @@ -401,7 +415,7 @@ async def test_does_not_cancel_terminal_jobs(self): _make_job("failed", status="error"), _make_job("stopped", status="cancelled"), ] - sdk = _make_sdk(jobs=jobs) + sdk, _ = _make_sdk(jobs=jobs) controller = _make_controller(nmp_sdk=sdk) await controller._cleanup_jobs(workspace) @@ -413,7 +427,7 @@ async def test_does_not_cancel_terminal_jobs(self): async def test_cancel_failure_still_deletes(self): """Regression: cancel() throwing must not prevent delete().""" workspace = _make_workspace() - sdk = _make_sdk(jobs=[_make_job("flaky-job", status="active")]) + sdk, _ = _make_sdk(jobs=[_make_job("flaky-job", status="active")]) sdk.jobs.cancel = AsyncMock(side_effect=Exception("cancel failed")) controller = _make_controller(nmp_sdk=sdk) @@ -430,7 +444,7 @@ async def test_mixed_statuses(self): _make_job("done-job", status="completed"), _make_job("pending-job", status="pending"), ] - sdk = _make_sdk(jobs=jobs) + sdk, _ = _make_sdk(jobs=jobs) controller = _make_controller(nmp_sdk=sdk) await controller._cleanup_jobs(workspace) diff --git a/services/core/files/script/v2_migration.py b/services/core/files/script/v2_migration.py index faa6056cdd..90e192eb5b 100644 --- a/services/core/files/script/v2_migration.py +++ b/services/core/files/script/v2_migration.py @@ -45,6 +45,10 @@ from huggingface_hub import HfApi from nemo_platform import ConflictError, NeMoPlatform, NotFoundError +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest, ListFilesetsQueryParams logger = logging.getLogger(__name__) @@ -337,11 +341,12 @@ def _ensure_fileset( """ if dry_run: return "dry_run" + files = client_from_platform(sdk, FilesClient) try: - sdk.files.filesets.retrieve(name=fileset, workspace=workspace) + files.get_fileset(name=fileset, workspace=workspace) return "exists" - except NotFoundError: - sdk.files.filesets.create(name=fileset, workspace=workspace) + except ClientNotFoundError: + files.create_fileset(body=CreateFilesetRequest(name=fileset), workspace=workspace) return "created" @@ -552,8 +557,9 @@ def run_setup(args: argparse.Namespace) -> int: try: sdk = _get_files_sdk(cfg) + files = client_from_platform(sdk, FilesClient) # Lightweight Files API connectivity check against default workspace. - _ = next(iter(sdk.files.filesets.list(workspace="default", page_size=1)), None) + files.list_filesets(workspace="default", query_params=ListFilesetsQueryParams(page_size=1)) print(f" files service: OK (resolved base_url: {sdk.base_url}, check_workspace=default)") except Exception as exc: print(f" files service: FAIL ({exc})") diff --git a/services/core/files/src/nmp/core/files/testing/utils.py b/services/core/files/src/nmp/core/files/testing/utils.py index 0acce915ee..523b89b26a 100644 --- a/services/core/files/src/nmp/core/files/testing/utils.py +++ b/services/core/files/src/nmp/core/files/testing/utils.py @@ -12,7 +12,9 @@ import httpx from fsspec.spec import AbstractBufferedFile, AbstractFileSystem from nemo_platform import NeMoPlatform -from nemo_platform.types.files.fileset import Fileset +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest, FilesetOutput DEFAULT_WORKSPACE_ID = "default" @@ -109,15 +111,14 @@ def create_fileset( name: str | None = None, workspace: str = DEFAULT_WORKSPACE_ID, **kwargs, -) -> Iterator[Fileset]: +) -> Iterator[FilesetOutput]: if name is None: name = test_fileset_name() - fileset = sdk.files.filesets.create( + files = client_from_platform(sdk, FilesClient) + fileset = files.create_fileset( workspace=workspace, - name=name, - description="Test fileset", - **kwargs, - ) + body=CreateFilesetRequest(name=name, description="Test fileset", **kwargs), + ).data() yield fileset - sdk.files.filesets.delete(name, workspace=workspace) + files.delete_fileset(name=name, workspace=workspace) diff --git a/services/core/files/tests/integration/conftest.py b/services/core/files/tests/integration/conftest.py index 5ab932ad09..cd300f0a44 100644 --- a/services/core/files/tests/integration/conftest.py +++ b/services/core/files/tests/integration/conftest.py @@ -15,7 +15,10 @@ from fastapi import Request from fastapi.testclient import TestClient from nemo_platform import NeMoPlatform -from nemo_platform.types.files.fileset import Fileset +from nemo_platform.filesets.resources import FilesResource +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import FilesetOutput from nmp.common.auth import AuthClient, get_auth_client from nmp.common.auth.models import Principal from nmp.common.config import AuthConfig @@ -106,6 +109,18 @@ def sdk() -> Iterator[NeMoPlatform]: yield sdk +@pytest.fixture +def files_client(sdk: NeMoPlatform) -> FilesClient: + """Provide a FilesClient derived from the SDK.""" + return client_from_platform(sdk, FilesClient) + + +@pytest.fixture +def files_resource(files_client: FilesClient) -> FilesResource: + """Provide a FilesResource backed by the test FilesClient.""" + return FilesResource(None, files_client=files_client) + + @pytest.fixture def sdk_allow_user_local_storage(tmp_path) -> Iterator[NeMoPlatform]: """SDK client with allow_user_local_storage enabled.""" @@ -141,13 +156,13 @@ def files_config() -> FilesConfig: @pytest.fixture -def fileset(sdk: NeMoPlatform) -> Iterator[Fileset]: +def fileset(sdk: NeMoPlatform) -> Iterator[FilesetOutput]: with create_fileset(sdk) as fileset: yield fileset @pytest.fixture -def fileset_cleanup(sdk: NeMoPlatform) -> Iterator[Callable[[str], None]]: +def fileset_cleanup(sdk: NeMoPlatform, files_client: FilesClient) -> Iterator[Callable[[str], None]]: """Fixture that provides a function to register filesets for cleanup. Usage: @@ -164,10 +179,9 @@ def register(name: str, ws: str | None = None) -> None: yield register - # Cleanup all registered filesets for name, ws in to_cleanup: try: - sdk.files.filesets.delete(name=name, workspace=ws) + files_client.delete_fileset(name=name, workspace=ws) except Exception: pass diff --git a/services/core/files/tests/integration/external_storage/test_huggingface_storage.py b/services/core/files/tests/integration/external_storage/test_huggingface_storage.py index c20b861b33..e4336af2d9 100644 --- a/services/core/files/tests/integration/external_storage/test_huggingface_storage.py +++ b/services/core/files/tests/integration/external_storage/test_huggingface_storage.py @@ -19,8 +19,12 @@ import pytest from huggingface_hub import snapshot_download -from nemo_platform import APIStatusError, NeMoPlatform +from nemo_platform import NeMoPlatform from nemo_platform.filesets import FilesetFileSystem +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError as ClientBadRequestError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.core.files.app.backends.base import StorageImpl from nmp.core.files.app.streaming import download_url_streaming from nmp.core.files.testing.utils import create_fileset @@ -49,10 +53,11 @@ def test_fileset_resolves_main_to_commit_sha(self, sdk: NeMoPlatform): }, ) as fileset: # Get the persisted fileset to check resolved values - persisted = sdk.files.filesets.retrieve( + files = client_from_platform(sdk, FilesClient) + persisted = files.get_fileset( name=fileset.name, workspace=fileset.workspace, - ) + ).data() storage = persisted.storage assert storage.type == "huggingface" @@ -87,10 +92,11 @@ def test_fileset_with_explicit_sha_preserves_both(self, sdk: NeMoPlatform): "revision": "main", }, ) as temp_fileset: - temp_persisted = sdk.files.filesets.retrieve( + files = client_from_platform(sdk, FilesClient) + temp_persisted = files.get_fileset( name=temp_fileset.name, workspace=temp_fileset.workspace, - ) + ).data() commit_sha = temp_persisted.storage.revision # Now create a fileset with the explicit SHA @@ -104,10 +110,10 @@ def test_fileset_with_explicit_sha_preserves_both(self, sdk: NeMoPlatform): "revision": commit_sha, # Explicit SHA }, ) as fileset: - persisted = sdk.files.filesets.retrieve( + persisted = files.get_fileset( name=fileset.name, workspace=fileset.workspace, - ) + ).data() storage = persisted.storage @@ -158,20 +164,23 @@ def test_gated_repo_fails_on_fileset_creation(self, sdk: NeMoPlatform): """ name = f"hf-test-{uuid.uuid4().hex[:8]}" - with pytest.raises(APIStatusError) as exc_info: - sdk.files.filesets.create( - name=name, + files = client_from_platform(sdk, FilesClient) + with pytest.raises(ClientBadRequestError) as exc_info: + files.create_fileset( workspace="default", - storage={ - "type": "huggingface", - "repo_id": "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "repo_type": "model", - }, + body=CreateFilesetRequest( + name=name, + storage={ + "type": "huggingface", + "repo_id": "meta-llama/Llama-4-Scout-17B-16E-Instruct", + "repo_type": "model", + }, + ), ) # Should get a 400 error with access denied message assert exc_info.value.status_code == 400 - assert "Access denied" in str(exc_info.value.body) or "gated" in str(exc_info.value.body).lower() + assert "Access denied" in str(exc_info.value) or "gated" in str(exc_info.value).lower() def test_download_file_from_public_dataset(self, sdk: NeMoPlatform): """Test downloading a file from a public Huggingface dataset.""" @@ -251,7 +260,7 @@ def test_file_exists_with_file_path(self, sdk: NeMoPlatform): "repo_type": "model", }, ) as fileset: - fs = FilesetFileSystem(sdk=sdk) + fs = FilesetFileSystem(client=client_from_platform(sdk, FilesClient)) file_path = f"{fileset.workspace}/{fileset.name}#config.json" # This would fail with EntryNotFoundError before the fix @@ -272,7 +281,7 @@ def test_file_exists_with_nonexistent_path_returns_false(self, sdk: NeMoPlatform "repo_type": "model", }, ) as fileset: - fs = FilesetFileSystem(sdk=sdk) + fs = FilesetFileSystem(client=client_from_platform(sdk, FilesClient)) file_path = f"{fileset.workspace}/{fileset.name}#nonexistent/file/path.txt" # Should return False, not raise an error @@ -298,7 +307,7 @@ def test_get_downloads_single_file(self, sdk: NeMoPlatform, tmp_path): "repo_type": "model", }, ) as fileset: - fs = FilesetFileSystem(sdk=sdk) + fs = FilesetFileSystem(client=client_from_platform(sdk, FilesClient)) file_path = f"{fileset.workspace}/{fileset.name}#config.json" # Download single file @@ -330,7 +339,7 @@ def test_get_downloads_directory_with_trailing_slash(self, sdk: NeMoPlatform, tm "repo_type": "model", }, ) as fileset: - fs = FilesetFileSystem(sdk=sdk) + fs = FilesetFileSystem(client=client_from_platform(sdk, FilesClient)) # Trailing slash on source - copy contents directly dir_path = f"{fileset.workspace}/{fileset.name}#/" @@ -358,7 +367,7 @@ def test_get_downloads_directory_without_trailing_slash(self, sdk: NeMoPlatform, "repo_type": "model", }, ) as fileset: - fs = FilesetFileSystem(sdk=sdk) + fs = FilesetFileSystem(client=client_from_platform(sdk, FilesClient)) # No trailing slash on source - for fileset root, copies contents directly dir_path = f"{fileset.workspace}/{fileset.name}#" @@ -393,10 +402,11 @@ def test_cache_path_uses_resolved_sha_not_mutable_ref(self, sdk: NeMoPlatform, c }, ) as fileset: # Get the resolved commit SHA - persisted = sdk.files.filesets.retrieve( + files = client_from_platform(sdk, FilesClient) + persisted = files.get_fileset( name=fileset.name, workspace=fileset.workspace, - ) + ).data() commit_sha = persisted.storage.revision assert commit_sha != "main", "revision should be resolved to SHA" diff --git a/services/core/files/tests/integration/external_storage/test_ngc_storage.py b/services/core/files/tests/integration/external_storage/test_ngc_storage.py index 834c05418b..6e259fc532 100644 --- a/services/core/files/tests/integration/external_storage/test_ngc_storage.py +++ b/services/core/files/tests/integration/external_storage/test_ngc_storage.py @@ -23,6 +23,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from nmp.core.files.app.backends.base import StorageImpl from nmp.core.files.app.streaming import download_url_streaming from nmp.core.files.testing.utils import create_fileset @@ -82,10 +84,11 @@ def test_fileset_resolves_latest_to_version_id(self, sdk: NeMoPlatform, ngc_api_ }, ) as fileset: # Get the persisted fileset to check resolved values - persisted = sdk.files.filesets.retrieve( + files = client_from_platform(sdk, FilesClient) + persisted = files.get_fileset( name=fileset.name, workspace=fileset.workspace, - ) + ).data() storage = persisted.storage assert storage.type == "ngc" @@ -113,10 +116,11 @@ def test_fileset_with_explicit_version_preserves_both(self, sdk: NeMoPlatform, n "api_key_secret": ngc_api_key_secret, }, ) as fileset: - persisted = sdk.files.filesets.retrieve( + files = client_from_platform(sdk, FilesClient) + persisted = files.get_fileset( name=fileset.name, workspace=fileset.workspace, - ) + ).data() storage = persisted.storage @@ -264,10 +268,11 @@ def test_cache_path_uses_resolved_version_not_latest( }, ) as fileset: # Get the resolved version ID - persisted = sdk.files.filesets.retrieve( + files = client_from_platform(sdk, FilesClient) + persisted = files.get_fileset( name=fileset.name, workspace=fileset.workspace, - ) + ).data() version_id = persisted.storage.version assert version_id is not None, "version should be resolved" diff --git a/services/core/files/tests/integration/external_storage/test_s3_storage.py b/services/core/files/tests/integration/external_storage/test_s3_storage.py index 167eee40c7..c2af373848 100644 --- a/services/core/files/tests/integration/external_storage/test_s3_storage.py +++ b/services/core/files/tests/integration/external_storage/test_s3_storage.py @@ -35,9 +35,11 @@ import pytest from aiobotocore.session import get_session from botocore.exceptions import ClientError -from nemo_platform import APIStatusError, NeMoPlatform -from nemo_platform.types.files import S3StorageConfigParam -from nemo_platform.types.files.fileset import Fileset +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError as ClientBadRequestError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest, FilesetOutput from nmp.common.auth import AuthClient, get_auth_client from nmp.common.auth.models import Principal from nmp.common.config import AuthConfig @@ -76,9 +78,9 @@ def s3_storage_config( access_key_secret: str, secret_key_secret: str, prefix: str | None = None, -) -> S3StorageConfigParam: +) -> dict[str, object]: """Helper to create S3 storage config dict with explicit credentials.""" - config: S3StorageConfigParam = { + config: dict[str, object] = { "type": "s3", "bucket": bucket, "endpoint_url": S3_TEST_ENDPOINT, @@ -132,7 +134,7 @@ def s3_credentials(sdk: NeMoPlatform) -> Iterator[tuple[str, str]]: @pytest.fixture -def s3_fileset(sdk: NeMoPlatform, s3_test_bucket: str, s3_credentials: tuple[str, str]) -> Iterator[Fileset]: +def s3_fileset(sdk: NeMoPlatform, s3_test_bucket: str, s3_credentials: tuple[str, str]) -> Iterator[FilesetOutput]: """Create a fileset with S3 storage for testing.""" name = f"s3-test-{uuid.uuid4().hex[:8]}" access_key_secret, secret_key_secret = s3_credentials @@ -160,7 +162,8 @@ def test_fileset_create_with_s3_storage( name, storage=s3_storage_config(s3_test_bucket, access_key_secret, secret_key_secret), ) as fileset: - persisted = sdk.files.filesets.retrieve(name=fileset.name, workspace=fileset.workspace) + files = client_from_platform(sdk, FilesClient) + persisted = files.get_fileset(name=fileset.name, workspace=fileset.workspace).data() assert persisted.storage.type == "s3" assert persisted.storage.bucket == s3_test_bucket @@ -169,19 +172,22 @@ def test_validate_storage_bucket_not_found(self, sdk: NeMoPlatform, s3_credentia name = f"s3-test-{uuid.uuid4().hex[:8]}" access_key_secret, secret_key_secret = s3_credentials - with pytest.raises(APIStatusError) as exc_info: - sdk.files.filesets.create( - name=name, + files = client_from_platform(sdk, FilesClient) + with pytest.raises(ClientBadRequestError) as exc_info: + files.create_fileset( workspace=DEFAULT_WORKSPACE, - storage=s3_storage_config( - f"nonexistent-bucket-{uuid.uuid4().hex[:8]}", - access_key_secret, - secret_key_secret, + body=CreateFilesetRequest( + name=name, + storage=s3_storage_config( + f"nonexistent-bucket-{uuid.uuid4().hex[:8]}", + access_key_secret, + secret_key_secret, + ), ), ) assert exc_info.value.status_code == 400 - assert "Not found" in str(exc_info.value.body) or "bucket" in str(exc_info.value.body).lower() + assert "Not found" in str(exc_info.value) or "bucket" in str(exc_info.value).lower() def test_invalid_credentials(self, sdk: NeMoPlatform, s3_test_bucket: str): """Test that invalid credentials raise an error during fileset creation.""" @@ -193,20 +199,23 @@ def test_invalid_credentials(self, sdk: NeMoPlatform, s3_test_bucket: str): sdk.secrets.create(workspace=DEFAULT_WORKSPACE, name=bad_secret_secret, value="invalid-secret") try: - with pytest.raises(APIStatusError) as exc_info: - sdk.files.filesets.create( - name=name, + files = client_from_platform(sdk, FilesClient) + with pytest.raises(ClientBadRequestError) as exc_info: + files.create_fileset( workspace=DEFAULT_WORKSPACE, - storage=s3_storage_config(s3_test_bucket, bad_access_secret, bad_secret_secret), + body=CreateFilesetRequest( + name=name, + storage=s3_storage_config(s3_test_bucket, bad_access_secret, bad_secret_secret), + ), ) assert exc_info.value.status_code == 400 - assert "Access denied" in str(exc_info.value.body) or "credentials" in str(exc_info.value.body).lower() + assert "Access denied" in str(exc_info.value) or "credentials" in str(exc_info.value).lower() finally: sdk.secrets.delete(workspace=DEFAULT_WORKSPACE, name=bad_access_secret) sdk.secrets.delete(workspace=DEFAULT_WORKSPACE, name=bad_secret_secret) - def test_upload_and_download_roundtrip(self, sdk: NeMoPlatform, s3_fileset: Fileset, tmp_path): + def test_upload_and_download_roundtrip(self, sdk: NeMoPlatform, s3_fileset: FilesetOutput, tmp_path): """Test upload file, download it back, verify content matches.""" test_content = b"Hello, S3 storage backend test!" upload_file = tmp_path / "test-file.txt" @@ -243,7 +252,7 @@ def test_upload_and_download_roundtrip(self, sdk: NeMoPlatform, s3_fileset: File ) assert download_path.read_bytes() == test_content - def test_upload_and_download_empty_file(self, sdk: NeMoPlatform, s3_fileset: Fileset, tmp_path): + def test_upload_and_download_empty_file(self, sdk: NeMoPlatform, s3_fileset: FilesetOutput, tmp_path): """Test upload and download of an empty file. This exercises the edge case where iter_chunked yields no chunks, @@ -274,7 +283,7 @@ def test_upload_and_download_empty_file(self, sdk: NeMoPlatform, s3_fileset: Fil ) assert downloaded_content == test_content - def test_upload_large_file(self, sdk: NeMoPlatform, s3_fileset: Fileset, tmp_path): + def test_upload_large_file(self, sdk: NeMoPlatform, s3_fileset: FilesetOutput, tmp_path): """Test upload of a large file via presigned URL streaming. Validates that large file uploads work correctly through the presigned @@ -303,7 +312,7 @@ def test_upload_large_file(self, sdk: NeMoPlatform, s3_fileset: Fileset, tmp_pat ) assert downloaded_content == test_content - def test_download_with_byte_range(self, sdk: NeMoPlatform, s3_fileset: Fileset, tmp_path): + def test_download_with_byte_range(self, sdk: NeMoPlatform, s3_fileset: FilesetOutput, tmp_path): """Test partial download using HTTP Range header.""" test_content = b"0123456789ABCDEF" upload_file = tmp_path / "range-test.txt" @@ -327,7 +336,7 @@ def test_download_with_byte_range(self, sdk: NeMoPlatform, s3_fileset: Fileset, assert range_response.status_code == 206 # Partial Content assert range_response.read() == b"56789A" - def test_delete_file(self, sdk: NeMoPlatform, s3_fileset: Fileset, tmp_path): + def test_delete_file(self, sdk: NeMoPlatform, s3_fileset: FilesetOutput, tmp_path): """Test upload, delete, verify gone.""" upload_file = tmp_path / "to-delete.txt" upload_file.write_bytes(b"Delete me!") @@ -370,11 +379,14 @@ def test_delete_fileset_with_files( prefix = f"delete-test-{uuid.uuid4().hex[:8]}" # Create fileset with a unique prefix so we can verify cleanup - fileset = sdk.files.filesets.create( - name=name, + files_client = client_from_platform(sdk, FilesClient) + fileset = files_client.create_fileset( workspace=DEFAULT_WORKSPACE, - storage=s3_storage_config(s3_test_bucket, access_key_secret, secret_key_secret, prefix=prefix), - ) + body=CreateFilesetRequest( + name=name, + storage=s3_storage_config(s3_test_bucket, access_key_secret, secret_key_secret, prefix=prefix), + ), + ).data() try: # Upload multiple files to exercise bulk delete @@ -395,26 +407,27 @@ def test_delete_fileset_with_files( ) # Verify files exist - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) - assert len(files.data) == 3 + file_list = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + assert len(file_list.data) == 3 # Delete the fileset - this calls delete_all() on the S3 backend - sdk.files.filesets.delete(name=fileset.name, workspace=fileset.workspace) + files_client.delete_fileset(name=fileset.name, workspace=fileset.workspace) # Verify fileset is gone - with pytest.raises(APIStatusError) as exc_info: - sdk.files.filesets.retrieve(name=name, workspace=DEFAULT_WORKSPACE) - assert exc_info.value.status_code == 404 + from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError + + with pytest.raises(ClientNotFoundError): + files_client.get_fileset(name=name, workspace=DEFAULT_WORKSPACE) except Exception: # Cleanup on failure - delete fileset if it still exists try: - sdk.files.filesets.delete(name=name, workspace=DEFAULT_WORKSPACE) - except APIStatusError: + files_client.delete_fileset(name=name, workspace=DEFAULT_WORKSPACE) + except Exception: pass raise - def test_multiple_files_with_directory_structure(self, sdk: NeMoPlatform, s3_fileset: Fileset, tmp_path): + def test_multiple_files_with_directory_structure(self, sdk: NeMoPlatform, s3_fileset: FilesetOutput, tmp_path): """Test uploading multiple files with directory structure.""" files_to_upload = { "file1.txt": b"content1", @@ -543,11 +556,11 @@ def test_fileset_without_storage_uses_s3_default(self, sdk_with_s3_default: NeMo name = f"default-storage-test-{uuid.uuid4().hex[:8]}" # Create fileset WITHOUT specifying storage - should use S3 default - fileset = sdk_with_s3_default.files.filesets.create( - name=name, + files = client_from_platform(sdk_with_s3_default, FilesClient) + fileset = files.create_fileset( workspace=DEFAULT_WORKSPACE, - # Note: no storage parameter - ) + body=CreateFilesetRequest(name=name), + ).data() try: # Verify it was created with S3 storage @@ -557,7 +570,7 @@ def test_fileset_without_storage_uses_s3_default(self, sdk_with_s3_default: NeMo assert "default-storage" in fileset.storage.prefix assert name in fileset.storage.prefix finally: - sdk_with_s3_default.files.filesets.delete(name=name, workspace=DEFAULT_WORKSPACE) + files.delete_fileset(name=name, workspace=DEFAULT_WORKSPACE) def test_upload_download_with_s3_default(self, sdk_with_s3_default: NeMoPlatform, tmp_path): """Test file upload/download on fileset using S3 default storage.""" @@ -565,10 +578,11 @@ def test_upload_download_with_s3_default(self, sdk_with_s3_default: NeMoPlatform test_content = b"Hello from S3 default storage!" # Create fileset without explicit storage - fileset = sdk_with_s3_default.files.filesets.create( - name=name, + files = client_from_platform(sdk_with_s3_default, FilesClient) + fileset = files.create_fileset( workspace=DEFAULT_WORKSPACE, - ) + body=CreateFilesetRequest(name=name), + ).data() try: # Upload file @@ -598,15 +612,16 @@ def test_upload_download_with_s3_default(self, sdk_with_s3_default: NeMoPlatform ) assert downloaded == test_content finally: - sdk_with_s3_default.files.filesets.delete(name=name, workspace=DEFAULT_WORKSPACE) + files.delete_fileset(name=name, workspace=DEFAULT_WORKSPACE) def test_multiple_filesets_isolated_with_s3_default(self, sdk_with_s3_default: NeMoPlatform, tmp_path): """Test that multiple filesets using S3 default are isolated via prefix.""" name1 = f"default-test-1-{uuid.uuid4().hex[:8]}" name2 = f"default-test-2-{uuid.uuid4().hex[:8]}" - fileset1 = sdk_with_s3_default.files.filesets.create(name=name1, workspace=DEFAULT_WORKSPACE) - fileset2 = sdk_with_s3_default.files.filesets.create(name=name2, workspace=DEFAULT_WORKSPACE) + files = client_from_platform(sdk_with_s3_default, FilesClient) + fileset1 = files.create_fileset(workspace=DEFAULT_WORKSPACE, body=CreateFilesetRequest(name=name1)).data() + fileset2 = files.create_fileset(workspace=DEFAULT_WORKSPACE, body=CreateFilesetRequest(name=name2)).data() try: # Both should have different prefixes (S3 storage has prefix attribute) @@ -651,8 +666,8 @@ def test_multiple_filesets_isolated_with_s3_default(self, sdk_with_s3_default: N assert content1 == b"content for fileset 1" assert content2 == b"content for fileset 2" finally: - sdk_with_s3_default.files.filesets.delete(name=name1, workspace=DEFAULT_WORKSPACE) - sdk_with_s3_default.files.filesets.delete(name=name2, workspace=DEFAULT_WORKSPACE) + files.delete_fileset(name=name1, workspace=DEFAULT_WORKSPACE) + files.delete_fileset(name=name2, workspace=DEFAULT_WORKSPACE) def test_download_from_huggingface_fileset_with_s3_default(self, sdk_with_s3_default: NeMoPlatform): """Test downloading from a HuggingFace fileset works with S3 as default storage. @@ -665,16 +680,19 @@ def test_download_from_huggingface_fileset_with_s3_default(self, sdk_with_s3_def name = f"hf-with-s3-default-{uuid.uuid4().hex[:8]}" # Create a HuggingFace-backed fileset (explicitly specifying storage) - fileset = sdk_with_s3_default.files.filesets.create( - name=name, + files_client = client_from_platform(sdk_with_s3_default, FilesClient) + fileset = files_client.create_fileset( workspace=DEFAULT_WORKSPACE, - storage={ - "type": "huggingface", - "repo_id": "hf-internal-testing/tiny-random-bert", - "repo_type": "model", - "revision": "main", - }, - ) + body=CreateFilesetRequest( + name=name, + storage={ + "type": "huggingface", + "repo_id": "hf-internal-testing/tiny-random-bert", + "repo_type": "model", + "revision": "main", + }, + ), + ).data() try: # List files to verify connection works @@ -707,4 +725,4 @@ def test_download_from_huggingface_fileset_with_s3_default(self, sdk_with_s3_def assert isinstance(parsed, dict), "config.json should be a JSON object" finally: - sdk_with_s3_default.files.filesets.delete(name=name, workspace=DEFAULT_WORKSPACE) + files_client.delete_fileset(name=name, workspace=DEFAULT_WORKSPACE) diff --git a/services/core/files/tests/integration/test_files_basic.py b/services/core/files/tests/integration/test_files_basic.py index 216b9b8967..267361b223 100644 --- a/services/core/files/tests/integration/test_files_basic.py +++ b/services/core/files/tests/integration/test_files_basic.py @@ -22,9 +22,11 @@ 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 -from nemo_platform_plugin.client import errors as nemo_errors +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import BadRequestError, ConflictError, NemoHTTPError, NotFoundError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest, FilesetOutput, UpdateFilesetRequest from nmp.core.files.testing.utils import ( DEFAULT_WORKSPACE_ID, HTTPXFileSystem, @@ -35,28 +37,31 @@ class TestFilesBasic: def test_fileset_get(self, sdk: NeMoPlatform): + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk) as fileset: - fetched = sdk.files.filesets.retrieve(fileset.name, workspace=fileset.workspace) + fetched = files.get_fileset(name=fileset.name, workspace=fileset.workspace).data() assert fetched.id == fileset.id assert fetched.name == fileset.name def test_fileset_list(self, sdk: NeMoPlatform): """Test listing filesets and filtering by workspace.""" + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk) as fileset1: with create_fileset(sdk) as fileset2: - filesets = list(sdk.files.filesets.list(workspace=DEFAULT_WORKSPACE_ID).items()) + filesets = list(files.list_filesets(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.""" + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk) as fileset1: with create_fileset(sdk) as fileset2: # Filter by exact name of fileset1 filtered = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"name": fileset1.name}, + query_params={"filter": {"name": fileset1.name}}, ).items() ) assert len(filtered) == 1 @@ -65,9 +70,9 @@ def test_fileset_list_filter_by_name(self, sdk: NeMoPlatform): # Filter by exact name of fileset2 filtered2 = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"name": fileset2.name}, + query_params={"filter": {"name": fileset2.name}}, ).items() ) assert len(filtered2) == 1 @@ -75,22 +80,23 @@ def test_fileset_list_filter_by_name(self, sdk: NeMoPlatform): # Filter by non-existent name should return empty filtered_none = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"name": "non-existent-fileset-name"}, + query_params={"filter": {"name": "non-existent-fileset-name"}}, ).items() ) assert len(filtered_none) == 0 def test_fileset_list_filter_by_purpose(self, sdk: NeMoPlatform): """Test listing filesets with purpose filter.""" + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk, purpose="dataset") as dataset_fileset: with create_fileset(sdk, purpose="generic") as generic_fileset: # Filter by purpose=dataset dataset_filesets = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"purpose": "dataset"}, + query_params={"filter": {"purpose": "dataset"}}, ).items() ) assert any(fs.id == dataset_fileset.id for fs in dataset_filesets) @@ -98,9 +104,9 @@ def test_fileset_list_filter_by_purpose(self, sdk: NeMoPlatform): # Filter by purpose=generic generic_filesets = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"purpose": "generic"}, + query_params={"filter": {"purpose": "generic"}}, ).items() ) assert any(fs.id == generic_fileset.id for fs in generic_filesets) @@ -108,14 +114,15 @@ def test_fileset_list_filter_by_purpose(self, sdk: NeMoPlatform): def test_fileset_list_filter_by_storage_type(self, sdk: NeMoPlatform): """Test listing filesets with storage_type filter.""" + files = client_from_platform(sdk, FilesClient) # Create filesets with default local storage with create_fileset(sdk) as local_fileset1: with create_fileset(sdk) as local_fileset2: # Filter by storage_type=local local_filesets = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"storage_type": "local"}, + query_params={"filter": {"storage_type": "local"}}, ).items() ) assert any(fs.id == local_fileset1.id for fs in local_filesets) @@ -127,16 +134,16 @@ def test_fileset_list_filter_by_storage_type(self, sdk: NeMoPlatform): def test_fileset_list_pagination(self, sdk: NeMoPlatform): """Test listing filesets with pagination.""" + files = client_from_platform(sdk, FilesClient) with ExitStack() as stack: # Create 5 filesets using ExitStack for automatic cleanup for _ in range(5): stack.enter_context(create_fileset(sdk, purpose="generic")) # Test first page with page_size=2 - resp1 = sdk.files.filesets.list( + resp1 = files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - page=1, - page_size=2, + query_params={"page": 1, "page_size": 2}, ) page1 = resp1.page() assert len(page1.items) == 2 @@ -144,10 +151,9 @@ def test_fileset_list_pagination(self, sdk: NeMoPlatform): assert page1.page_size == 2 # Test second page - resp2 = sdk.files.filesets.list( + resp2 = files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - page=2, - page_size=2, + query_params={"page": 2, "page_size": 2}, ) page2 = resp2.page() assert len(page2.items) == 2 @@ -158,7 +164,7 @@ def test_fileset_list_pagination(self, sdk: NeMoPlatform): 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): + def test_file_upload_download(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test uploading and downloading a file using application/octet-stream.""" test_content = b"Hello, World! This is a test file.\nLine 2\nLine 3" @@ -201,7 +207,7 @@ def test_file_upload_download(self, sdk: NeMoPlatform, fileset: Fileset): ) assert len(files_response.data) == 0, "File should be deleted" - def test_file_upload_nested_paths_and_list(self, sdk: NeMoPlatform, fileset: Fileset): + def test_file_upload_nested_paths_and_list(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test uploading multiple files with nested paths concurrently and listing them.""" # Upload multiple files with nested paths @@ -251,7 +257,7 @@ def upload_file(path_content_tuple): ) assert downloaded == expected_content - def test_file_range_requests_with_duckdb(self, sdk: NeMoPlatform, fileset: Fileset, client: TestClient): + def test_file_range_requests_with_duckdb(self, sdk: NeMoPlatform, fileset: FilesetOutput, client: TestClient): """Test HTTP range requests by querying a parquet file with DuckDB. Uses HTTPXFileSystem to route DuckDB requests through the test client, @@ -306,11 +312,12 @@ def test_file_range_requests_with_duckdb(self, sdk: NeMoPlatform, fileset: Files def test_error_handling(self, sdk: NeMoPlatform): """Test error handling for various 404 scenarios.""" + files = client_from_platform(sdk, FilesClient) # Test 1: Get non-existent fileset try: - sdk.files.filesets.retrieve( - "non-existent-fileset", + files.get_fileset( + name="non-existent-fileset", workspace="non-existent-workspace", ) assert False, "Should have raised NotFoundError" @@ -331,17 +338,17 @@ def test_error_handling(self, sdk: NeMoPlatform): ) # Verify fileset exists - retrieved = sdk.files.filesets.retrieve( - fileset_name_str, + retrieved = files.get_fileset( + name=fileset_name_str, workspace=workspace_str, - ) + ).data() assert retrieved.id == fileset.id # After context manager exits, fileset is deleted # Verify getting the fileset now raises 404 try: - sdk.files.filesets.retrieve( - fileset_name_str, + files.get_fileset( + name=fileset_name_str, workspace=workspace_str, ) assert False, "Should have raised NotFoundError after fileset deletion" @@ -349,8 +356,6 @@ 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( @@ -359,7 +364,7 @@ def test_error_handling(self, sdk: NeMoPlatform): workspace=fileset.workspace, ) assert False, "Should have raised NotFoundError for non-existent file" - except (NotFoundError, nemo_errors.NotFoundError): + except NotFoundError: pass # Expected # Test 4: Try to delete non-existent file @@ -371,7 +376,7 @@ def test_error_handling(self, sdk: NeMoPlatform): workspace=fileset.workspace, ) assert False, "Should have raised NotFoundError when deleting non-existent file" - except (NotFoundError, nemo_errors.NotFoundError): + except NotFoundError: pass # Expected # Test 5: List files in non-existent fileset @@ -386,12 +391,13 @@ def test_error_handling(self, sdk: NeMoPlatform): def test_fileset_create_conflict(self, sdk: NeMoPlatform): """Test that creating a fileset with a duplicate name returns 409 Conflict.""" + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk) as fileset: # Try to create another fileset with the same name and workspace try: - sdk.files.filesets.create( + files.create_fileset( + body=CreateFilesetRequest(name=fileset.name), workspace=fileset.workspace, - name=fileset.name, ) assert False, "Should have raised ConflictError" except ConflictError as e: @@ -400,31 +406,37 @@ def test_fileset_create_conflict(self, sdk: NeMoPlatform): def test_fileset_create_rejects_user_provided_local_storage(self, sdk: NeMoPlatform): """Test that explicitly requesting local storage is rejected.""" + files = client_from_platform(sdk, FilesClient) try: - sdk.files.filesets.create( + files.create_fileset( + body=CreateFilesetRequest( + name="reject-local-storage", + storage={"type": "local", "path": "/etc"}, + ), workspace=DEFAULT_WORKSPACE_ID, - name="reject-local-storage", - storage={"type": "local", "path": "/etc"}, ) assert False, "Should have raised NemoHTTPError for local storage" - except APIStatusError as exc: + except BadRequestError as exc: assert exc.status_code == 400 assert "local storage is not allowed" in str(exc.body).lower() def test_fileset_create_rejects_s3_use_sdk_auth(self, sdk: NeMoPlatform): """Test that S3 storage with use_sdk_auth=True is rejected for user-provided storage.""" + files = client_from_platform(sdk, FilesClient) try: - sdk.files.filesets.create( + files.create_fileset( + body=CreateFilesetRequest( + name="reject-s3-sdk-auth", + storage={ + "type": "s3", + "bucket": "my-bucket", + "use_sdk_auth": True, + }, + ), workspace=DEFAULT_WORKSPACE_ID, - name="reject-s3-sdk-auth", - storage={ - "type": "s3", - "bucket": "my-bucket", - "use_sdk_auth": True, - }, ) assert False, "Should have raised NemoHTTPError for S3 with use_sdk_auth=True" - except APIStatusError as exc: + except BadRequestError as exc: assert exc.status_code == 400 assert "use_sdk_auth=true is not allowed" in str(exc.body).lower() @@ -432,20 +444,24 @@ def test_fileset_create_allows_user_provided_local_storage_when_enabled( self, sdk_allow_user_local_storage: NeMoPlatform, tmp_path: Path ): """Test that explicit local storage is allowed when feature flag is enabled.""" - fileset = sdk_allow_user_local_storage.files.filesets.create( + files = client_from_platform(sdk_allow_user_local_storage, FilesClient) + fileset = files.create_fileset( + body=CreateFilesetRequest( + name="allow-local-storage", + storage={"type": "local", "path": str(tmp_path / "explicit")}, + ), workspace=DEFAULT_WORKSPACE_ID, - name="allow-local-storage", - storage={"type": "local", "path": str(tmp_path / "explicit")}, - ) + ).data() assert fileset.storage.type == "local" assert fileset.storage.path == str(tmp_path / "explicit") # Cleanup because not using create_fileset() helper. - sdk_allow_user_local_storage.files.filesets.delete(fileset.name, workspace=fileset.workspace) + files.delete_fileset(name=fileset.name, workspace=fileset.workspace) def test_fileset_update_partial(self, sdk: NeMoPlatform): """Test that partial updates work - only specified fields are updated.""" + files = client_from_platform(sdk, FilesClient) with create_fileset( sdk, purpose="generic", @@ -455,11 +471,11 @@ def test_fileset_update_partial(self, sdk: NeMoPlatform): original_custom_fields = fileset.custom_fields # Update only description - updated = sdk.files.filesets.update( - fileset.name, + updated = files.update_fileset( + name=fileset.name, workspace=fileset.workspace, - description="Updated description only", - ) + body=UpdateFilesetRequest(description="Updated description only"), + ).data() # Verify description was updated assert updated.description == "Updated description only" @@ -471,15 +487,18 @@ def test_fileset_update_partial(self, sdk: NeMoPlatform): def test_fileset_update_description_purpose_custom_fields(self, sdk: NeMoPlatform): """Test that description, purpose, and custom_fields can all be updated.""" + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk, purpose="generic") as fileset: # Update all three fields - updated = sdk.files.filesets.update( - fileset.name, + updated = files.update_fileset( + name=fileset.name, workspace=fileset.workspace, - description="New description", - purpose="dataset", - custom_fields={"new_key": "new_value", "another": 123}, - ) + body=UpdateFilesetRequest( + description="New description", + purpose="dataset", + custom_fields={"new_key": "new_value", "another": 123}, + ), + ).data() # Verify all fields were updated assert updated.description == "New description" @@ -487,18 +506,19 @@ def test_fileset_update_description_purpose_custom_fields(self, sdk: NeMoPlatfor assert updated.custom_fields == {"new_key": "new_value", "another": 123} # Verify by fetching the fileset again - fetched = sdk.files.filesets.retrieve(fileset.name, workspace=fileset.workspace) + fetched = files.get_fileset(name=fileset.name, workspace=fileset.workspace).data() assert fetched.description == "New description" assert fetched.purpose == "dataset" assert fetched.custom_fields == {"new_key": "new_value", "another": 123} def test_fileset_update_not_found(self, sdk: NeMoPlatform): """Test that updating a non-existent fileset returns 404.""" + files = client_from_platform(sdk, FilesClient) try: - sdk.files.filesets.update( - "non-existent-fileset", + files.update_fileset( + name="non-existent-fileset", workspace=DEFAULT_WORKSPACE_ID, - description="Should fail", + body=UpdateFilesetRequest(description="Should fail"), ) assert False, "Should have raised NotFoundError" except NotFoundError: @@ -506,13 +526,16 @@ def test_fileset_update_not_found(self, sdk: NeMoPlatform): def test_fileset_update_returns_updated_output(self, sdk: NeMoPlatform): """Test that update returns the updated FilesetOutput with correct fields.""" + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk, purpose="generic") as fileset: - updated = sdk.files.filesets.update( - fileset.name, + updated = files.update_fileset( + name=fileset.name, workspace=fileset.workspace, - description="Updated description", - custom_fields={"status": "modified"}, - ) + body=UpdateFilesetRequest( + description="Updated description", + custom_fields={"status": "modified"}, + ), + ).data() # Verify the returned object has all expected fields assert updated.id == fileset.id @@ -586,7 +609,7 @@ 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, ValidationError), + (NemoHTTPError, ValidationError), match="definitely-not-a-valid-json-schema-type", ): with create_fileset( @@ -611,11 +634,12 @@ def test_fileset_default_storage_path(self, sdk: NeMoPlatform): def test_fileset_delete_removes_storage_data(self, sdk: NeMoPlatform): """Test that deleting a fileset also deletes the underlying storage directory.""" + files = client_from_platform(sdk, FilesClient) # Create fileset manually (not using context manager) so we control deletion - fileset = sdk.files.filesets.create( + fileset = files.create_fileset( + body=CreateFilesetRequest(name="delete-storage-test"), workspace=DEFAULT_WORKSPACE_ID, - name="delete-storage-test", - ) + ).data() try: # Upload some files @@ -640,7 +664,7 @@ def test_fileset_delete_removes_storage_data(self, sdk: NeMoPlatform): assert (storage_path / "subdir" / "file2.txt").exists() # Delete the fileset - sdk.files.filesets.delete(fileset.name, workspace=fileset.workspace) + files.delete_fileset(name=fileset.name, workspace=fileset.workspace) # Verify storage directory is gone assert not storage_path.exists() @@ -648,13 +672,14 @@ def test_fileset_delete_removes_storage_data(self, sdk: NeMoPlatform): except Exception: # Clean up on failure try: - sdk.files.filesets.delete(fileset.name, workspace=fileset.workspace) + files.delete_fileset(name=fileset.name, workspace=fileset.workspace) except Exception: pass raise def test_fileset_list_filter_by_created_at_gte(self, sdk: NeMoPlatform): """Test listing filesets with created_at[gte] filter.""" + files = client_from_platform(sdk, FilesClient) # Record time before creating filesets before_create = datetime.now(timezone.utc) - timedelta(seconds=5) @@ -663,9 +688,9 @@ def test_fileset_list_filter_by_created_at_gte(self, sdk: NeMoPlatform): with create_fileset(sdk) as fileset2: # Filter by created_at[gte] should include both new filesets filtered = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"created_at": {"$gte": before_create.isoformat(timespec="seconds")}}, + query_params={"filter": {"created_at": {"$gte": before_create.isoformat(timespec="seconds")}}}, ).items() ) fileset_ids = {fs.id for fs in filtered} @@ -674,6 +699,7 @@ def test_fileset_list_filter_by_created_at_gte(self, sdk: NeMoPlatform): def test_fileset_list_filter_by_created_at_lte(self, sdk: NeMoPlatform): """Test listing filesets with created_at[lte] filter.""" + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk) as fileset1: # Record time after creating first fileset after_first = datetime.now(timezone.utc) @@ -686,9 +712,9 @@ def test_fileset_list_filter_by_created_at_lte(self, sdk: NeMoPlatform): # should include first fileset but might include second # (depends on timing precision) filtered = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"created_at": {"$lte": after_first.isoformat()}}, + query_params={"filter": {"created_at": {"$lte": after_first.isoformat()}}}, ).items() ) fileset_ids = {fs.id for fs in filtered} @@ -696,6 +722,7 @@ def test_fileset_list_filter_by_created_at_lte(self, sdk: NeMoPlatform): def test_fileset_list_filter_by_created_at_range(self, sdk: NeMoPlatform): """Test listing filesets with both created_at[gte] and created_at[lte].""" + files = client_from_platform(sdk, FilesClient) before_create = datetime.now(timezone.utc) with create_fileset(sdk) as fileset: @@ -703,12 +730,14 @@ def test_fileset_list_filter_by_created_at_range(self, sdk: NeMoPlatform): # Filter by date range that includes the fileset filtered = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={ - "created_at": { - "$gte": before_create.isoformat(), - "$lte": after_create.isoformat(), + query_params={ + "filter": { + "created_at": { + "$gte": before_create.isoformat(), + "$lte": after_create.isoformat(), + } } }, ).items() @@ -718,6 +747,7 @@ def test_fileset_list_filter_by_created_at_range(self, sdk: NeMoPlatform): def test_fileset_list_filter_by_created_at_excludes_older(self, sdk: NeMoPlatform): """Test that created_at[gte] filter excludes older filesets.""" + files = client_from_platform(sdk, FilesClient) with create_fileset(sdk) as old_fileset: # Delay to ensure timestamps differ # Using 1s because SQLite doesn't track sub-second precision @@ -730,9 +760,9 @@ 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 = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"created_at": {"$gte": after_old.isoformat()}}, + query_params={"filter": {"created_at": {"$gte": after_old.isoformat()}}}, ).items() ) fileset_ids = {fs.id for fs in filtered} @@ -743,14 +773,15 @@ def test_fileset_list_filter_by_created_at_excludes_older(self, sdk: NeMoPlatfor def test_fileset_list_filter_by_updated_at(self, sdk: NeMoPlatform): """Test listing filesets with updated_at filter.""" + files = client_from_platform(sdk, FilesClient) before_create = datetime.now(timezone.utc) with create_fileset(sdk) as fileset: # Filter by updated_at[gte] should include the fileset filtered = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={"updated_at": {"$gte": before_create.isoformat()}}, + query_params={"filter": {"updated_at": {"$gte": before_create.isoformat()}}}, ).items() ) fileset_ids = {fs.id for fs in filtered} @@ -758,17 +789,20 @@ def test_fileset_list_filter_by_updated_at(self, sdk: NeMoPlatform): def test_fileset_list_combined_filters_with_datetime(self, sdk: NeMoPlatform): """Test combining datetime filters with other filters.""" + files = client_from_platform(sdk, FilesClient) before_create = datetime.now(timezone.utc) 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 = list( - sdk.files.filesets.list( + files.list_filesets( workspace=DEFAULT_WORKSPACE_ID, - filter={ - "purpose": "dataset", - "created_at": {"$gte": before_create.isoformat()}, + query_params={ + "filter": { + "purpose": "dataset", + "created_at": {"$gte": before_create.isoformat()}, + } }, ).items() ) diff --git a/services/core/files/tests/integration/test_files_sdk.py b/services/core/files/tests/integration/test_files_sdk.py index 5f159e995e..35e3b79014 100644 --- a/services/core/files/tests/integration/test_files_sdk.py +++ b/services/core/files/tests/integration/test_files_sdk.py @@ -1,15 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Integration tests for the high-level files SDK (sdk.files.*). +"""Integration tests for the high-level FilesResource. These tests verify: -- sdk.files.upload() - Upload files/directories -- 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 FilesetFileOutput objects -- sdk.files.delete() - Delete files +- files_resource.upload() - Upload files/directories +- files_resource.upload_content() - Upload in-memory data +- files_resource.download() - Download files/directories +- files_resource.download_content() - Download file content to memory +- files_resource.list() - List files with FilesetFileOutput objects +- files_resource.delete() - Delete files - fileset_auto_create parameter for upload operations Uses the create_test_client pattern for fast in-memory testing. @@ -23,21 +23,29 @@ from pathlib import Path import pytest -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 nemo_platform import NeMoPlatform +from nemo_platform.filesets.resources import FilesResource +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError, PermissionDeniedError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import ( + CreateFilesetRequest, + FilesetFileOutput, + FilesetOutput, + UpdateFilesetRequest, +) from nmp.core.files.testing.utils import create_fileset, test_fileset_name class TestFilesUpload: - """Tests for sdk.files.upload().""" + """Tests for files_resource.upload().""" - def test_upload_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path: Path): + def test_upload_single_file(self, files_resource: FilesResource, fileset: FilesetOutput, tmp_path: Path): """Test uploading a single file.""" local_file = tmp_path / "upload.txt" local_file.write_text("Hello, World!") - sdk.files.upload( + files_resource.upload( fileset=fileset.name, workspace=fileset.workspace, local_path=str(local_file), @@ -45,12 +53,12 @@ def test_upload_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp ) # Verify file was uploaded - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 1 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: FilesetOutput): + def test_upload_directory_contents_with_trailing_slash(self, files_resource: FilesResource, fileset: FilesetOutput): """Test uploading directory contents (trailing slash on local_path). With trailing slash: `upload("mydir/")` copies the CONTENTS of mydir. @@ -66,7 +74,7 @@ def test_upload_directory_contents_with_trailing_slash(self, sdk: NeMoPlatform, Path(subdir, "file3.txt").write_text("content3") # Upload with trailing slash - copies CONTENTS - sdk.files.upload( + files_resource.upload( fileset=fileset.name, workspace=fileset.workspace, local_path=f"{mydir}/", @@ -74,14 +82,16 @@ def test_upload_directory_contents_with_trailing_slash(self, sdk: NeMoPlatform, ) # Verify files are at root (not under mydir/) - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) paths = {f.path for f in files.data} assert "file1.txt" in paths, f"Expected 'file1.txt' in {paths}" assert "subdir/file3.txt" in paths, f"Expected 'subdir/file3.txt' in {paths}" # 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: FilesetOutput): + def test_upload_directory_itself_without_trailing_slash( + self, files_resource: FilesResource, fileset: FilesetOutput + ): """Test uploading directory itself (no trailing slash on local_path). Without trailing slash: `upload("mydir")` copies the directory ITSELF. @@ -97,7 +107,7 @@ def test_upload_directory_itself_without_trailing_slash(self, sdk: NeMoPlatform, Path(subdir, "file3.txt").write_text("content3") # Upload WITHOUT trailing slash - copies the directory ITSELF - sdk.files.upload( + files_resource.upload( fileset=fileset.name, workspace=fileset.workspace, local_path=str(mydir), @@ -105,38 +115,38 @@ def test_upload_directory_itself_without_trailing_slash(self, sdk: NeMoPlatform, ) # Verify files are under mydir/ prefix - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) paths = {f.path for f in files.data} assert "mydir/file1.txt" in paths, f"Expected 'mydir/file1.txt' in {paths}" assert "mydir/subdir/file3.txt" in paths, f"Expected 'mydir/subdir/file3.txt' in {paths}" # 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: FilesetOutput, tmp_path: Path): + def test_upload_to_subdirectory(self, files_resource: FilesResource, fileset: FilesetOutput, tmp_path: Path): """Test uploading a file to a subdirectory.""" local_file = tmp_path / "nested.txt" local_file.write_text("nested content") - sdk.files.upload( + files_resource.upload( fileset=fileset.name, workspace=fileset.workspace, local_path=str(local_file), remote_path="a/b/c/nested.txt", ) - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 1 assert files.data[0].path == "a/b/c/nested.txt" class TestFilesDownload: - """Tests for sdk.files.download().""" + """Tests for files_resource.download().""" - def test_download_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_download_single_file(self, files_resource: FilesResource, fileset: FilesetOutput): """Test downloading a single file.""" # First upload a file test_content = b"Download test content" - sdk.files.upload_content( + files_resource.upload_content( content=test_content, remote_path="test.txt", fileset=fileset.name, @@ -144,7 +154,7 @@ def test_download_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): ) with tempfile.TemporaryDirectory() as tmpdir: - sdk.files.download( + files_resource.download( fileset=fileset.name, workspace=fileset.workspace, remote_path="test.txt", @@ -154,22 +164,22 @@ def test_download_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): downloaded = Path(tmpdir, "downloaded.txt").read_bytes() assert downloaded == test_content - def test_download_directory(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_download_directory(self, files_resource: FilesResource, fileset: FilesetOutput): """Test downloading an entire directory.""" # Upload multiple files - sdk.files.upload_content( + files_resource.upload_content( content=b"content1", remote_path="data/file1.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"content2", remote_path="data/file2.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"content3", remote_path="data/nested/file3.txt", fileset=fileset.name, @@ -177,7 +187,7 @@ def test_download_directory(self, sdk: NeMoPlatform, fileset: FilesetOutput): ) with tempfile.TemporaryDirectory() as tmpdir: - sdk.files.download( + files_resource.download( fileset=fileset.name, workspace=fileset.workspace, remote_path="data/", @@ -189,20 +199,20 @@ def test_download_directory(self, sdk: NeMoPlatform, fileset: FilesetOutput): 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: FilesetOutput): + def test_download_entire_fileset(self, files_resource: FilesResource, fileset: FilesetOutput): """Test downloading all files from a fileset using default remote_path. Downloading a fileset copies contents directly. Users who want a subfolder can include the fileset name in local_path. """ # Upload files at root - sdk.files.upload_content( + files_resource.upload_content( content=b"root content", remote_path="root.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"nested content", remote_path="subdir/nested.txt", fileset=fileset.name, @@ -212,7 +222,7 @@ def test_download_entire_fileset(self, sdk: NeMoPlatform, fileset: FilesetOutput with tempfile.TemporaryDirectory() as tmpdir: # Download everything (remote_path defaults to "") # Contents are copied directly to local_path - sdk.files.download( + files_resource.download( fileset=fileset.name, workspace=fileset.workspace, local_path=f"{tmpdir}/", @@ -225,24 +235,24 @@ def test_download_entire_fileset(self, sdk: NeMoPlatform, fileset: FilesetOutput class TestFilesList: - """Tests for sdk.files.list().""" + """Tests for files_resource.list().""" - def test_list_empty_fileset(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_list_empty_fileset(self, files_resource: FilesResource, fileset: FilesetOutput): """Test listing files in an empty fileset.""" - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert files.data == [] - def test_list_returns_fileset_file_objects(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_list_returns_fileset_file_objects(self, files_resource: FilesResource, fileset: FilesetOutput): """Test that list returns FilesetFileOutput objects with correct attributes.""" content = b"test content for size check" - sdk.files.upload_content( + files_resource.upload_content( content=content, remote_path="test.txt", fileset=fileset.name, workspace=fileset.workspace, ) - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 1 file = files.data[0] @@ -251,28 +261,28 @@ 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: FilesetOutput): + def test_list_multiple_files(self, files_resource: FilesResource, fileset: FilesetOutput): """Test listing multiple files.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"a", remote_path="file1.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"bb", remote_path="file2.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"ccc", remote_path="dir/file3.txt", fileset=fileset.name, workspace=fileset.workspace, ) - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 3 paths = {f.path for f in files.data} @@ -283,27 +293,27 @@ def test_list_multiple_files(self, sdk: NeMoPlatform, fileset: FilesetOutput): assert sizes["file2.txt"] == 2 assert sizes["dir/file3.txt"] == 3 - def test_list_subdirectory(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_list_subdirectory(self, files_resource: FilesResource, fileset: FilesetOutput): """Test listing files in a subdirectory.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"root", remote_path="root.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"data1", remote_path="data/file1.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"data2", remote_path="data/file2.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"other", remote_path="other/file.txt", fileset=fileset.name, @@ -311,7 +321,7 @@ def test_list_subdirectory(self, sdk: NeMoPlatform, fileset: FilesetOutput): ) # List only data/ directory - files = sdk.files.list( + files = files_resource.list( fileset=fileset.name, workspace=fileset.workspace, remote_path="data/", @@ -320,9 +330,9 @@ def test_list_subdirectory(self, sdk: NeMoPlatform, fileset: FilesetOutput): 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: FilesetOutput): + def test_list_with_path_format(self, files_resource: FilesResource, fileset: FilesetOutput): """Test listing using full path format instead of explicit fileset param.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"content", remote_path="test.txt", fileset=fileset.name, @@ -330,34 +340,34 @@ def test_list_with_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutput): ) # Use the new path format: workspace/fileset#path - files = sdk.files.list( + files = files_resource.list( remote_path=f"{fileset.workspace}/{fileset.name}#", ) assert len(files.data) == 1 assert files.data[0].path == "test.txt" - def test_list_with_glob_pattern(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_list_with_glob_pattern(self, files_resource: FilesResource, fileset: FilesetOutput): """Test listing files matching a glob pattern.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"json", remote_path="data.json", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"config", remote_path="config.json", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"readme", remote_path="readme.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"nested", remote_path="subdir/nested.json", fileset=fileset.name, @@ -365,7 +375,7 @@ def test_list_with_glob_pattern(self, sdk: NeMoPlatform, fileset: FilesetOutput) ) # List only .json files at root level - files = sdk.files.list( + files = files_resource.list( fileset=fileset.name, workspace=fileset.workspace, remote_path="*.json", @@ -374,27 +384,27 @@ def test_list_with_glob_pattern(self, sdk: NeMoPlatform, fileset: FilesetOutput) 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: FilesetOutput): + def test_list_with_glob_pattern_in_subdirectory(self, files_resource: FilesResource, fileset: FilesetOutput): """Test listing files matching a glob pattern in a subdirectory.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"train", remote_path="data/train.jsonl", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"val", remote_path="data/val.jsonl", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"yaml", remote_path="data/config.yaml", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"other", remote_path="other/file.jsonl", fileset=fileset.name, @@ -402,7 +412,7 @@ def test_list_with_glob_pattern_in_subdirectory(self, sdk: NeMoPlatform, fileset ) # List only .jsonl files in data/ directory - files = sdk.files.list( + files = files_resource.list( fileset=fileset.name, workspace=fileset.workspace, remote_path="data/*.jsonl", @@ -413,23 +423,23 @@ def test_list_with_glob_pattern_in_subdirectory(self, sdk: NeMoPlatform, fileset class TestFilesGlobDownload: - """Tests for sdk.files.download() with glob patterns.""" + """Tests for files_resource.download() with glob patterns.""" - def test_download_with_glob_pattern(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path): + def test_download_with_glob_pattern(self, files_resource: FilesResource, fileset: FilesetOutput, tmp_path): """Test downloading files matching a glob pattern.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"json content", remote_path="data.json", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"config content", remote_path="config.json", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"readme content", remote_path="readme.txt", fileset=fileset.name, @@ -437,7 +447,7 @@ def test_download_with_glob_pattern(self, sdk: NeMoPlatform, fileset: FilesetOut ) # Download only .json files - sdk.files.download( + files_resource.download( fileset=fileset.name, workspace=fileset.workspace, remote_path="*.json", @@ -453,21 +463,23 @@ def test_download_with_glob_pattern(self, sdk: NeMoPlatform, fileset: FilesetOut 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: FilesetOutput, tmp_path): + def test_download_with_glob_pattern_preserves_structure( + self, files_resource: FilesResource, fileset: FilesetOutput, tmp_path + ): """Test that downloading with glob pattern preserves directory structure.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"train data", remote_path="data/train.jsonl", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"val data", remote_path="data/val.jsonl", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.upload_content( + files_resource.upload_content( content=b"yaml", remote_path="data/config.yaml", fileset=fileset.name, @@ -475,7 +487,7 @@ def test_download_with_glob_pattern_preserves_structure(self, sdk: NeMoPlatform, ) # Download only .jsonl files from data/ - sdk.files.download( + files_resource.download( fileset=fileset.name, workspace=fileset.workspace, remote_path="data/*.jsonl", @@ -493,11 +505,11 @@ def test_download_with_glob_pattern_preserves_structure(self, sdk: NeMoPlatform, class TestFilesDelete: - """Tests for sdk.files.delete().""" + """Tests for files_resource.delete().""" - def test_delete_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_delete_single_file(self, files_resource: FilesResource, fileset: FilesetOutput): """Test deleting a single file.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"delete me", remote_path="to_delete.txt", fileset=fileset.name, @@ -505,41 +517,41 @@ def test_delete_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): ) # Verify file exists - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 1 # Delete the file - sdk.files.delete( + files_resource.delete( fileset=fileset.name, workspace=fileset.workspace, remote_path="to_delete.txt", ) # Verify file was deleted - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 0 - def test_delete_nested_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_delete_nested_file(self, files_resource: FilesResource, fileset: FilesetOutput): """Test deleting a file in a nested directory.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"nested", remote_path="a/b/c/nested.txt", fileset=fileset.name, workspace=fileset.workspace, ) - sdk.files.delete( + files_resource.delete( fileset=fileset.name, workspace=fileset.workspace, remote_path="a/b/c/nested.txt", ) - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 0 - def test_delete_with_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_delete_with_path_format(self, files_resource: FilesResource, fileset: FilesetOutput): """Test deleting using full path format.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"content", remote_path="test.txt", fileset=fileset.name, @@ -547,18 +559,18 @@ def test_delete_with_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutput ) # Delete using the new path format - sdk.files.delete( + files_resource.delete( remote_path=f"{fileset.workspace}/{fileset.name}#test.txt", ) - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 0 class TestFilesRoundTrip: """End-to-end tests combining multiple operations.""" - def test_upload_list_download_delete_cycle(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_upload_list_download_delete_cycle(self, files_resource: FilesResource, fileset: FilesetOutput): """Test a complete cycle of file operations.""" with tempfile.TemporaryDirectory() as tmpdir: # Create local files @@ -568,7 +580,7 @@ def test_upload_list_download_delete_cycle(self, sdk: NeMoPlatform, fileset: Fil Path(local_dir, "config.yaml").write_text("setting: true") # Upload - sdk.files.upload( + files_resource.upload( fileset=fileset.name, workspace=fileset.workspace, local_path=f"{local_dir}/", @@ -576,7 +588,7 @@ def test_upload_list_download_delete_cycle(self, sdk: NeMoPlatform, fileset: Fil ) # List and verify - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 2 paths = {f.path for f in files.data} assert paths == {"data.json", "config.yaml"} @@ -585,7 +597,7 @@ def test_upload_list_download_delete_cycle(self, sdk: NeMoPlatform, fileset: Fil # Contents are copied directly download_dir = Path(tmpdir, "download") download_dir.mkdir() - sdk.files.download( + files_resource.download( fileset=fileset.name, workspace=fileset.workspace, local_path=f"{download_dir}/", @@ -597,18 +609,18 @@ def test_upload_list_download_delete_cycle(self, sdk: NeMoPlatform, fileset: Fil assert not (download_dir / fileset.name).exists() # Delete one file - sdk.files.delete( + files_resource.delete( fileset=fileset.name, workspace=fileset.workspace, remote_path="data.json", ) # Verify only one file remains - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 1 assert files.data[0].path == "config.yaml" - def test_large_directory_upload_download(self, sdk: NeMoPlatform): + def test_large_directory_upload_download(self, sdk: NeMoPlatform, files_resource: FilesResource): """Test uploading and downloading a larger directory structure.""" with create_fileset(sdk) as fileset: with tempfile.TemporaryDirectory() as tmpdir: @@ -623,7 +635,7 @@ def test_large_directory_upload_download(self, sdk: NeMoPlatform): (subdir / f"file{i}.txt").write_text(f"content {i}") # Upload - sdk.files.upload( + files_resource.upload( fileset=fileset.name, workspace=fileset.workspace, local_path=f"{upload_dir}/", @@ -631,13 +643,13 @@ def test_large_directory_upload_download(self, sdk: NeMoPlatform): ) # List and verify count - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == file_count # Download (contents copied directly) download_dir = Path(tmpdir, "download") download_dir.mkdir() - sdk.files.download( + files_resource.download( fileset=fileset.name, workspace=fileset.workspace, local_path=f"{download_dir}/", @@ -657,7 +669,7 @@ def _chunk_generator(): class TestFilesUploadContent: - """Tests for sdk.files.upload_content().""" + """Tests for files_resource.upload_content().""" @pytest.mark.parametrize( ("content", "expected_bytes"), @@ -672,9 +684,11 @@ class TestFilesUploadContent: pytest.param(_chunk_generator(), b"chunk1chunk2chunk3", id="iterator"), ], ) - def test_upload_content(self, sdk: NeMoPlatform, fileset: FilesetOutput, content, expected_bytes: bytes): + def test_upload_content( + self, files_resource: FilesResource, fileset: FilesetOutput, content, expected_bytes: bytes + ): """Test uploading different content types.""" - result = sdk.files.upload_content( + result = files_resource.upload_content( content=content, remote_path="test.bin", fileset=fileset.name, @@ -685,29 +699,29 @@ def test_upload_content(self, sdk: NeMoPlatform, fileset: FilesetOutput, content assert result.name == fileset.name assert result.workspace == fileset.workspace - downloaded = sdk.files.download_content( + downloaded = files_resource.download_content( remote_path="test.bin", fileset=fileset.name, workspace=fileset.workspace, ) assert downloaded == expected_bytes - def test_upload_content_to_subdirectory(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_upload_content_to_subdirectory(self, files_resource: FilesResource, fileset: FilesetOutput): """Test uploading data to a nested path.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"nested content", remote_path="a/b/c/nested.txt", fileset=fileset.name, workspace=fileset.workspace, ) - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 1 assert files.data[0].path == "a/b/c/nested.txt" class TestFilesDownloadContent: - """Tests for sdk.files.download_content().""" + """Tests for files_resource.download_content().""" @pytest.mark.parametrize( ("upload_content", "expected_bytes"), @@ -725,16 +739,18 @@ class TestFilesDownloadContent: ), ], ) - def test_download_content(self, sdk: NeMoPlatform, fileset: FilesetOutput, upload_content, expected_bytes: bytes): + def test_download_content( + self, files_resource: FilesResource, fileset: FilesetOutput, upload_content, expected_bytes: bytes + ): """Test download_content returns correct bytes for different content types.""" - sdk.files.upload_content( + files_resource.upload_content( content=upload_content, remote_path="test.bin", fileset=fileset.name, workspace=fileset.workspace, ) - result = sdk.files.download_content( + result = files_resource.download_content( remote_path="test.bin", fileset=fileset.name, workspace=fileset.workspace, @@ -743,9 +759,9 @@ def test_download_content(self, sdk: NeMoPlatform, fileset: FilesetOutput, uploa assert isinstance(result, bytes) assert result == expected_bytes - def test_download_content_with_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_download_content_with_path_format(self, files_resource: FilesResource, fileset: FilesetOutput): """Test download_content using full path format.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"content", remote_path="test.txt", fileset=fileset.name, @@ -753,7 +769,7 @@ def test_download_content_with_path_format(self, sdk: NeMoPlatform, fileset: Fil ) # Use full path format - downloaded = sdk.files.download_content( + downloaded = files_resource.download_content( remote_path=f"{fileset.workspace}/{fileset.name}#test.txt", ) assert downloaded == b"content" @@ -762,7 +778,9 @@ def test_download_content_with_path_format(self, sdk: NeMoPlatform, fileset: Fil class TestFilesUploadAutoCreate: """Tests for fileset_auto_create parameter.""" - def test_upload_creates_fileset(self, sdk: NeMoPlatform, tmp_path: Path, fileset_cleanup: Callable[[str], None]): + def test_upload_creates_fileset( + self, sdk: NeMoPlatform, files_resource: FilesResource, tmp_path: Path, fileset_cleanup: Callable[[str], None] + ): """Test that upload() with fileset_auto_create creates the fileset.""" fileset_name = f"auto-create-upload-{uuid.uuid4().hex[:8]}" workspace = sdk.workspace or "default" @@ -771,7 +789,7 @@ def test_upload_creates_fileset(self, sdk: NeMoPlatform, tmp_path: Path, fileset local_file = tmp_path / "test.txt" local_file.write_text("test content") - result = sdk.files.upload( + result = files_resource.upload( local_path=str(local_file), remote_path="test.txt", fileset=fileset_name, @@ -785,17 +803,19 @@ def test_upload_creates_fileset(self, sdk: NeMoPlatform, tmp_path: Path, fileset assert result.workspace == workspace # Verify fileset was created and file uploaded - files = sdk.files.list(fileset=fileset_name, workspace=workspace) + files = files_resource.list(fileset=fileset_name, workspace=workspace) assert len(files.data) == 1 assert files.data[0].path == "test.txt" - def test_upload_content_creates_fileset(self, sdk: NeMoPlatform, fileset_cleanup: Callable[[str], None]): + def test_upload_content_creates_fileset( + self, sdk: NeMoPlatform, files_resource: FilesResource, fileset_cleanup: Callable[[str], None] + ): """Test that upload_content() with fileset_auto_create creates the fileset.""" fileset_name = f"auto-create-data-{uuid.uuid4().hex[:8]}" workspace = sdk.workspace or "default" fileset_cleanup(fileset_name) - result = sdk.files.upload_content( + result = files_resource.upload_content( content=b"test content", remote_path="test.txt", fileset=fileset_name, @@ -809,17 +829,17 @@ def test_upload_content_creates_fileset(self, sdk: NeMoPlatform, fileset_cleanup assert result.workspace == workspace # Verify fileset was created and file uploaded - files = sdk.files.list(fileset=fileset_name, workspace=workspace) + files = files_resource.list(fileset=fileset_name, workspace=workspace) assert len(files.data) == 1 assert files.data[0].path == "test.txt" - def test_upload_without_flag_fails_for_nonexistent_fileset(self, sdk: NeMoPlatform): + def test_upload_without_flag_fails_for_nonexistent_fileset(self, sdk: NeMoPlatform, files_resource: FilesResource): """Test that upload without flag fails for non-existent fileset.""" fileset_name = f"nonexistent-{uuid.uuid4().hex[:8]}" workspace = sdk.workspace or "default" with pytest.raises(NotFoundError): - sdk.files.upload_content( + files_resource.upload_content( content=b"test", remote_path="test.txt", fileset=fileset_name, @@ -827,9 +847,9 @@ 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: FilesetOutput): + def test_existing_fileset_with_flag_succeeds(self, files_resource: FilesResource, fileset: FilesetOutput): """Test that fileset_auto_create works for existing filesets.""" - result = sdk.files.upload_content( + result = files_resource.upload_content( content=b"test content", remote_path="test.txt", fileset=fileset.name, @@ -840,15 +860,15 @@ def test_existing_fileset_with_flag_succeeds(self, sdk: NeMoPlatform, fileset: F assert isinstance(result, FilesetOutput) assert result.name == fileset.name - files = sdk.files.list(fileset=fileset.name, workspace=fileset.workspace) + files = files_resource.list(fileset=fileset.name, workspace=fileset.workspace) assert len(files.data) == 1 - def test_upload_returns_fileset(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path: Path): + def test_upload_returns_fileset(self, files_resource: FilesResource, fileset: FilesetOutput, tmp_path: Path): """Test that upload() always returns the FilesetOutput entity.""" local_file = tmp_path / "test.txt" local_file.write_text("content") - result = sdk.files.upload( + result = files_resource.upload( local_path=str(local_file), remote_path="test.txt", fileset=fileset.name, @@ -861,12 +881,12 @@ def test_upload_returns_fileset(self, sdk: NeMoPlatform, fileset: FilesetOutput, assert result.workspace == fileset.workspace def test_auto_create_generates_name_when_no_fileset_specified( - self, sdk: NeMoPlatform, fileset_cleanup: Callable[[str], None] + self, sdk: NeMoPlatform, files_resource: FilesResource, fileset_cleanup: Callable[[str], None] ): """Test that fileset_auto_create generates a UUID-based name when no fileset is specified.""" workspace = sdk.workspace or "default" - result = sdk.files.upload_content( + result = files_resource.upload_content( content=b"auto-generated fileset test", remote_path="test.txt", fileset_auto_create=True, @@ -883,18 +903,20 @@ def test_auto_create_generates_name_when_no_fileset_specified( assert len(result.name) == len("fileset-") + 8 # "fileset-" + 8 hex chars # Verify file was uploaded - files = sdk.files.list(fileset=result.name, workspace=workspace) + files = files_resource.list(fileset=result.name, workspace=workspace) assert len(files.data) == 1 assert files.data[0].path == "test.txt" - def test_auto_create_uses_fileset_from_path_syntax(self, sdk: NeMoPlatform, fileset_cleanup: Callable[[str], None]): + def test_auto_create_uses_fileset_from_path_syntax( + self, sdk: NeMoPlatform, files_resource: FilesResource, fileset_cleanup: Callable[[str], None] + ): """Test that fileset_auto_create uses fileset from path when # syntax is used.""" fileset_name = f"path-syntax-{uuid.uuid4().hex[:8]}" workspace = sdk.workspace or "default" fileset_cleanup(fileset_name) # Use the # syntax to embed fileset in path - result = sdk.files.upload_content( + result = files_resource.upload_content( content=b"path syntax test", remote_path=f"{fileset_name}#data/test.txt", fileset_auto_create=True, @@ -907,7 +929,7 @@ def test_auto_create_uses_fileset_from_path_syntax(self, sdk: NeMoPlatform, file assert result.name == fileset_name # Should NOT be "fileset-..." # Verify file was uploaded to correct path - files = sdk.files.list(fileset=fileset_name, workspace=workspace) + files = files_resource.list(fileset=fileset_name, workspace=workspace) assert len(files.data) == 1 assert files.data[0].path == "data/test.txt" @@ -966,11 +988,11 @@ def test_cache_status_aggregation(self, statuses: list, expected: str | None): class TestFilesListCacheStatus: - """Tests for sdk.files.list() with include_cache_status parameter.""" + """Tests for files_resource.list() with include_cache_status parameter.""" - def test_list_with_include_cache_status(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_list_with_include_cache_status(self, files_resource: FilesResource, fileset: FilesetOutput): """Test listing files with cache status included.""" - sdk.files.upload_content( + files_resource.upload_content( content=b"test content", remote_path="test.txt", fileset=fileset.name, @@ -978,7 +1000,7 @@ def test_list_with_include_cache_status(self, sdk: NeMoPlatform, fileset: Filese ) # List with cache status - files = sdk.files.list( + files = files_resource.list( fileset=fileset.name, workspace=fileset.workspace, include_cache_status=True, @@ -989,9 +1011,9 @@ 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: FilesetOutput): + def test_list_without_include_cache_status(self, files_resource: FilesResource, fileset: FilesetOutput): """Test listing files without cache status (default).""" - sdk.files.upload_content( + files_resource.upload_content( content=b"test content", remote_path="test.txt", fileset=fileset.name, @@ -999,7 +1021,7 @@ def test_list_without_include_cache_status(self, sdk: NeMoPlatform, fileset: Fil ) # List without cache status (default) - files = sdk.files.list( + files = files_resource.list( fileset=fileset.name, workspace=fileset.workspace, ) @@ -1009,12 +1031,12 @@ def test_list_without_include_cache_status(self, sdk: NeMoPlatform, fileset: Fil class TestFilesDownloadEdgeCases: - """Tests for sdk.files.download() edge cases.""" + """Tests for files_resource.download() edge cases.""" - def test_download_glob_no_matches(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path): + def test_download_glob_no_matches(self, files_resource: FilesResource, 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( + files_resource.upload_content( content=b"content", remote_path="data.txt", fileset=fileset.name, @@ -1022,7 +1044,7 @@ def test_download_glob_no_matches(self, sdk: NeMoPlatform, fileset: FilesetOutpu ) # Download with glob that matches nothing - sdk.files.download( + files_resource.download( fileset=fileset.name, workspace=fileset.workspace, remote_path="*.json", # No .json files exist @@ -1033,11 +1055,11 @@ def test_download_glob_no_matches(self, sdk: NeMoPlatform, fileset: FilesetOutpu 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: FilesetOutput): + def test_download_content_non_existent_file(self, files_resource: FilesResource, fileset: FilesetOutput): """Test downloading content of a file that doesn't exist raises NotFoundError.""" # Binary/streaming errors are deferred (raised after send()), bypassing remapping. - with pytest.raises((NotFoundError, nemo_errors.NotFoundError)): - sdk.files.download_content( + with pytest.raises(NotFoundError): + files_resource.download_content( fileset=fileset.name, workspace=fileset.workspace, remote_path="non-existent.txt", @@ -1045,13 +1067,13 @@ def test_download_content_non_existent_file(self, sdk: NeMoPlatform, fileset: Fi class TestFilesDeleteEdgeCases: - """Tests for sdk.files.delete() edge cases.""" + """Tests for files_resource.delete() edge cases.""" - def test_delete_non_existent_file(self, sdk: NeMoPlatform, fileset: FilesetOutput): + def test_delete_non_existent_file(self, files_resource: FilesResource, fileset: FilesetOutput): """Test deleting a file that doesn't exist raises NotFoundError.""" # File delete goes through fsspec rm → deferred error path. - with pytest.raises((NotFoundError, nemo_errors.NotFoundError)): - sdk.files.delete( + with pytest.raises(NotFoundError): + files_resource.delete( fileset=fileset.name, workspace=fileset.workspace, remote_path="non-existent.txt", @@ -1063,18 +1085,21 @@ class TestFilesetImmutabilityForNonServicePrincipals: def test_create_fileset_with_service_source_as_default_principal_fails_to_set(self, sdk: NeMoPlatform): """Non-service principal cannot set service_source; it is stripped on create.""" + files = client_from_platform(sdk, FilesClient) workspace = sdk.workspace or "default" name = test_fileset_name() - sdk.files.filesets.create( + files.create_fileset( + body=CreateFilesetRequest( + name=name, + description="Test", + custom_fields={"service_source": "customizer"}, + ), workspace=workspace, - name=name, - description="Test", - custom_fields={"service_source": "customizer"}, ) - created = sdk.files.filesets.retrieve(name=name, workspace=workspace) + created = files.get_fileset(name=name, workspace=workspace).data() # Endpoint strips service_source for non-service principals; fileset must not have it. assert created.custom_fields.get("service_source") is None - sdk.files.filesets.delete(name=name, workspace=workspace) + files.delete_fileset(name=name, workspace=workspace) def test_service_principal_can_set_service_source_and_upload_then_user_cannot_upload( self, sdk_user_and_service: tuple[NeMoPlatform, NeMoPlatform] @@ -1084,11 +1109,17 @@ def test_service_principal_can_set_service_source_and_upload_then_user_cannot_up workspace = sdk_service.workspace or "default" name = test_fileset_name() # Service principal creates fileset with service_source and uploads a file. - created = sdk_service.files.filesets.create( - workspace=workspace, - name=name, - description="Immutability test", - custom_fields={"service_source": "customizer"}, + created = ( + client_from_platform(sdk_service, FilesClient) + .create_fileset( + workspace=workspace, + body=CreateFilesetRequest( + name=name, + description="Immutability test", + custom_fields={"service_source": "customizer"}, + ), + ) + .data() ) assert created.custom_fields.get("service_source") == "customizer" sdk_service.files.upload_content( @@ -1108,7 +1139,7 @@ def test_service_principal_can_set_service_source_and_upload_then_user_cannot_up fileset=name, workspace=workspace, ) - sdk_service.files.filesets.delete(name=name, workspace=workspace) + client_from_platform(sdk_service, FilesClient).delete_fileset(name=name, workspace=workspace) def test_non_service_principal_cannot_overwrite_or_remove_service_source_on_update( self, sdk_user_and_service: tuple[NeMoPlatform, NeMoPlatform] @@ -1118,26 +1149,28 @@ def test_non_service_principal_cannot_overwrite_or_remove_service_source_on_upda workspace = sdk_service.workspace or "default" name = test_fileset_name() # Service principal creates fileset with service_source. - sdk_service.files.filesets.create( + service_files = client_from_platform(sdk_service, FilesClient) + user_files = client_from_platform(sdk_user, FilesClient) + service_files.create_fileset( workspace=workspace, - name=name, - description="Update immutability test", - custom_fields={"service_source": "customizer"}, + body=CreateFilesetRequest( + name=name, + description="Update immutability test", + custom_fields={"service_source": "customizer"}, + ), ) - # User tries to overwrite service_source → must be ignored (preserved). - sdk_user.files.filesets.update( + user_files.update_fileset( name=name, workspace=workspace, - custom_fields={"service_source": "other-service"}, + body=UpdateFilesetRequest(custom_fields={"service_source": "other-service"}), ) - updated = sdk_user.files.filesets.retrieve(name=name, workspace=workspace) + updated = user_files.get_fileset(name=name, workspace=workspace).data() assert updated.custom_fields.get("service_source") == "customizer" - # User tries to remove service_source by sending custom_fields without it → must stay. - sdk_user.files.filesets.update( + user_files.update_fileset( name=name, workspace=workspace, - custom_fields={"other_key": "value"}, + body=UpdateFilesetRequest(custom_fields={"other_key": "value"}), ) - after_remove_attempt = sdk_user.files.filesets.retrieve(name=name, workspace=workspace) + after_remove_attempt = user_files.get_fileset(name=name, workspace=workspace).data() assert after_remove_attempt.custom_fields.get("service_source") == "customizer" - sdk_service.files.filesets.delete(name=name, workspace=workspace) + service_files.delete_fileset(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 23960fd91c..a8a06837fe 100644 --- a/services/core/files/tests/integration/test_fileset_filesystem.py +++ b/services/core/files/tests/integration/test_fileset_filesystem.py @@ -30,7 +30,9 @@ parse_fileset_path, parse_fileset_ref, ) -from nemo_platform.types.files.fileset import Fileset +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import FilesetOutput class TestParseFilesetRef: @@ -298,13 +300,13 @@ def fs(self, sdk: NeMoPlatform) -> FilesetFileSystem: """Create a FilesetFileSystem backed by the test SDK.""" return sdk.files.fsspec - def test_ls_empty_fileset(self, fs: FilesetFileSystem, fileset: Fileset): + def test_ls_empty_fileset(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test listing an empty fileset.""" path = f"{fileset.workspace}/{fileset.name}" result = fs.ls(path) assert result == [] - def test_ls_with_files(self, fs: FilesetFileSystem, fileset: Fileset): + def test_ls_with_files(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test listing a fileset with files.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/file1.txt", b"content1") @@ -317,7 +319,7 @@ def test_ls_with_files(self, fs: FilesetFileSystem, fileset: Fileset): assert f"{base}#file1.txt" in names assert f"{base}#file2.txt" in names - def test_ls_with_directories(self, fs: FilesetFileSystem, fileset: Fileset): + def test_ls_with_directories(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test that nested files show as directories in listing.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/root.txt", b"root") @@ -331,7 +333,7 @@ def test_ls_with_directories(self, fs: FilesetFileSystem, fileset: Fileset): assert types[f"{base}#root.txt"] == "file" assert types[f"{base}#subdir"] == "directory" - def test_ls_subdirectory(self, fs: FilesetFileSystem, fileset: Fileset): + def test_ls_subdirectory(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test listing files in a subdirectory.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/subdir/file1.txt", b"content1") @@ -343,7 +345,7 @@ def test_ls_subdirectory(self, fs: FilesetFileSystem, fileset: Fileset): assert f"{base}#subdir/file1.txt" in result assert f"{base}#subdir/file2.txt" in result - def test_ls_subdirectory_trailing_slash(self, fs: FilesetFileSystem, fileset: Fileset): + def test_ls_subdirectory_trailing_slash(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test that ls with trailing slash returns same result as without. ls("workspace/fileset/subdir/") should return @@ -362,7 +364,7 @@ def test_ls_subdirectory_trailing_slash(self, fs: FilesetFileSystem, fileset: Fi assert len(result_with_slash) == 2 assert set(result_no_slash) == set(result_with_slash) - def test_cat_file(self, fs: FilesetFileSystem, fileset: Fileset): + def test_cat_file(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading file content with cat.""" content = b"Hello, fsspec!" path = f"{fileset.workspace}/{fileset.name}/test.txt" @@ -372,7 +374,7 @@ def test_cat_file(self, fs: FilesetFileSystem, fileset: Fileset): assert result == content - def test_cat_file_with_range(self, fs: FilesetFileSystem, fileset: Fileset): + def test_cat_file_with_range(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading partial file content with byte range.""" content = b"0123456789ABCDEF" path = f"{fileset.workspace}/{fileset.name}/test.txt" @@ -382,7 +384,7 @@ def test_cat_file_with_range(self, fs: FilesetFileSystem, fileset: Fileset): assert result == b"456789" - def test_open_read(self, fs: FilesetFileSystem, fileset: Fileset): + def test_open_read(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test opening a file for reading.""" content = b"File content for reading" path = f"{fileset.workspace}/{fileset.name}/readable.txt" @@ -393,7 +395,7 @@ def test_open_read(self, fs: FilesetFileSystem, fileset: Fileset): assert result == content - def test_open_write(self, fs: FilesetFileSystem, fileset: Fileset): + def test_open_write(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test opening a file for writing.""" path = f"{fileset.workspace}/{fileset.name}/writable.txt" content = b"Written via fsspec" @@ -405,7 +407,7 @@ def test_open_write(self, fs: FilesetFileSystem, fileset: Fileset): result = fs.cat(path) assert result == content - def test_pipe(self, fs: FilesetFileSystem, fileset: Fileset): + def test_pipe(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test writing file content with pipe.""" path = f"{fileset.workspace}/{fileset.name}/piped.txt" content = b"Piped content" @@ -415,7 +417,7 @@ def test_pipe(self, fs: FilesetFileSystem, fileset: Fileset): result = fs.cat(path) assert result == content - def test_put_file(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_put_file(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test uploading a local file with put_file.""" # Create a local file local_file = tmp_path / "upload.txt" @@ -430,7 +432,7 @@ def test_put_file(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path) result = fs.cat(remote_path) assert result == content - def test_put_file_nested_path(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_put_file_nested_path(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test uploading a file to a nested path with put_file.""" # Create a local file local_file = tmp_path / "nested_upload.txt" @@ -449,7 +451,7 @@ def test_put_file_nested_path(self, fs: FilesetFileSystem, fileset: Fileset, tmp parent_info = fs.info(f"{fileset.workspace}/{fileset.name}/subdir/nested") assert parent_info["type"] == "directory" - def test_rm(self, fs: FilesetFileSystem, fileset: Fileset): + def test_rm(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test deleting a file.""" path = f"{fileset.workspace}/{fileset.name}/to_delete.txt" fs.pipe(path, b"delete me") @@ -460,7 +462,7 @@ def test_rm(self, fs: FilesetFileSystem, fileset: Fileset): assert not fs.exists(path) - def test_info(self, fs: FilesetFileSystem, fileset: Fileset): + def test_info(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test getting file info.""" content = b"Content for info test" path = f"{fileset.workspace}/{fileset.name}/info.txt" @@ -474,14 +476,14 @@ def test_info(self, fs: FilesetFileSystem, fileset: Fileset): assert info["size"] == len(content) assert info["type"] == "file" - def test_info_directory(self, fs: FilesetFileSystem, fileset: Fileset): + def test_info_directory(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test getting info for fileset root (directory).""" path = f"{fileset.workspace}/{fileset.name}" info = fs.info(path) assert info["type"] == "directory" - def test_exists(self, fs: FilesetFileSystem, fileset: Fileset): + def test_exists(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test checking file existence.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/exists.txt", b"I exist") @@ -489,7 +491,7 @@ def test_exists(self, fs: FilesetFileSystem, fileset: Fileset): assert fs.exists(f"{base}/exists.txt") assert not fs.exists(f"{base}/does_not_exist.txt") - def test_parquet_read(self, fs: FilesetFileSystem, fileset: Fileset): + def test_parquet_read(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading a parquet file via fsspec. This validates that range requests work correctly for formats @@ -512,7 +514,7 @@ def test_parquet_read(self, fs: FilesetFileSystem, fileset: Fileset): pd.testing.assert_frame_equal(result, df) - def test_protocol_url(self, fs: FilesetFileSystem, fileset: Fileset): + def test_protocol_url(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test that protocol prefix is handled correctly.""" content = b"Protocol test" base = f"{fileset.workspace}/{fileset.name}" @@ -525,12 +527,12 @@ def test_protocol_url(self, fs: FilesetFileSystem, fileset: Fileset): assert fs.cat(path_no_proto) == content assert fs.cat(path_with_proto) == content - def test_fsspec_filesystem_registration(self, sdk: NeMoPlatform, fileset: Fileset): + def test_fsspec_filesystem_registration(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test that FilesetFileSystem can be instantiated via fsspec.filesystem().""" # Protocol is registered by the autouse fixture fs = fsspec.filesystem( "fileset", - sdk=sdk, + client=client_from_platform(sdk, FilesClient), skip_instance_cache=True, ) @@ -541,7 +543,7 @@ def test_fsspec_filesystem_registration(self, sdk: NeMoPlatform, fileset: Filese result = fs.ls(path) assert result == [] - def test_find(self, fs: FilesetFileSystem, fileset: Fileset): + def test_find(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test recursive file discovery with find.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/root.txt", b"root") @@ -558,7 +560,7 @@ def test_find(self, fs: FilesetFileSystem, fileset: Fileset): assert f"{base}#dir1/file1.txt" in result assert f"{base}#dir1/subdir/nested.txt" in result - def test_glob(self, fs: FilesetFileSystem, fileset: Fileset): + def test_glob(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test pattern matching with glob.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/data.csv", b"csv") @@ -579,7 +581,7 @@ def test_glob(self, fs: FilesetFileSystem, fileset: Fileset): # All json files including nested assert len(json_all) == 3 - def test_isdir_isfile(self, fs: FilesetFileSystem, fileset: Fileset): + def test_isdir_isfile(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test isdir and isfile type checking.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/file.txt", b"content") @@ -597,7 +599,7 @@ def test_isdir_isfile(self, fs: FilesetFileSystem, fileset: Fileset): assert fs.isdir(f"{base}/subdir") assert not fs.isfile(f"{base}/subdir") - def test_cat_multiple_files(self, fs: FilesetFileSystem, fileset: Fileset): + def test_cat_multiple_files(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading multiple files at once with cat.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/file1.txt", b"content1") @@ -613,7 +615,7 @@ def test_cat_multiple_files(self, fs: FilesetFileSystem, fileset: Fileset): assert result[f"{base}#file2.txt"] == b"content2" assert result[f"{base}#file3.txt"] == b"content3" - def test_walk(self, fs: FilesetFileSystem, fileset: Fileset): + def test_walk(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test directory traversal with walk.""" base = f"{fileset.workspace}/{fileset.name}" fs.pipe(f"{base}/root.txt", b"root") @@ -630,7 +632,7 @@ def test_walk(self, fs: FilesetFileSystem, fileset: Fileset): # Walk returns files with fileset#path format relative to walked directory assert any("root.txt" in f for f in files) - def test_head_tail(self, fs: FilesetFileSystem, fileset: Fileset): + def test_head_tail(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading first/last bytes of a file.""" content = b"0123456789ABCDEFGHIJ" path = f"{fileset.workspace}/{fileset.name}/test.txt" @@ -644,7 +646,7 @@ def test_head_tail(self, fs: FilesetFileSystem, fileset: Fileset): result = fs.tail(path, size=5) assert result == b"FGHIJ" - def test_get_single_file(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_get_single_file(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test downloading a single file with get(). Tests three behaviors: @@ -676,7 +678,7 @@ def test_get_single_file(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path assert (existing_dir / "test.txt").is_file() assert (existing_dir / "test.txt").read_bytes() == content - def test_get_trailing_slash_semantics(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_get_trailing_slash_semantics(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test trailing slash semantics for get() per fsspec docs. From https://filesystem-spec.readthedocs.io/en/latest/copying.html: @@ -745,7 +747,7 @@ def test_get_trailing_slash_semantics(self, fs: FilesetFileSystem, fileset: File assert (dest_source_slash / "file1.txt").read_bytes() == b"file1" assert not (dest_source_slash / "subdir").exists() - def test_get_fileset_root(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_get_fileset_root(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test downloading from fileset root path (workspace/fileset without subpath). When the source path is just the fileset root (no subpath), the trailing @@ -783,7 +785,7 @@ def test_get_fileset_root(self, fs: FilesetFileSystem, fileset: Fileset, tmp_pat assert (dest_no_slash / "data" / "nested" / "deep.txt").read_bytes() == b"deep" assert not (dest_no_slash / fileset.name).exists() - def test_put_fileset_root(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_put_fileset_root(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test uploading to fileset root path (workspace/fileset without subpath). When the dest path is just the fileset root (no subpath), the trailing @@ -822,7 +824,7 @@ def test_put_fileset_root(self, fs: FilesetFileSystem, fileset: Fileset, tmp_pat assert fs.cat(f"{base}/my_data/file2.txt") == b"file2" assert fs.cat(f"{base}/my_data/subdir/nested.txt") == b"nested" - def test_put_single_file_to_fileset_root(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_put_single_file_to_fileset_root(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test uploading a single file to fileset root path. This tests the case: fs.put("local_file.txt", "workspace/fileset") @@ -861,7 +863,7 @@ def test_put_single_file_to_fileset_root(self, fs: FilesetFileSystem, fileset: F # File should be renamed assert fs.cat(f"{base}/renamed.txt") == b"single file content" - def test_put_file_requires_file_path(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_put_file_requires_file_path(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test that put_file (not put) requires a file path in rpath. The low-level put_file method requires rpath to include a file path. @@ -882,7 +884,7 @@ def test_put_file_requires_file_path(self, fs: FilesetFileSystem, fileset: Files fs.put_file(str(local_file), f"{base}/test.txt") assert fs.cat(f"{base}/test.txt") == b"content" - def test_put_trailing_slash_semantics(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_put_trailing_slash_semantics(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test trailing slash semantics for put() per fsspec docs. Source trailing slash controls whether to preserve source directory name: @@ -918,7 +920,7 @@ def test_put_trailing_slash_semantics(self, fs: FilesetFileSystem, fileset: File def test_download_entire_fileset( self, fs: FilesetFileSystem, - fileset: Fileset, + fileset: FilesetOutput, sample_dataset: Path, tmp_path: Path, ): @@ -974,7 +976,7 @@ def test_download_entire_fileset( assert (download_without_slash / "config" / "settings.json").read_bytes() == b'{"batch_size": 32}' assert not (download_without_slash / fileset.name).exists() - def test_concurrent_download_failure_hang(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_concurrent_download_failure_hang(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test that a failure in one concurrent download doesn't cause a hang (sync).""" from unittest.mock import patch @@ -1004,7 +1006,7 @@ async def failing_get_file(rpath, lpath, **kwargs): # If we get here without hanging, the test passes - def test_concurrent_upload_failure_hang(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_concurrent_upload_failure_hang(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test that a failure in one concurrent upload doesn't cause a hang (sync).""" from unittest.mock import patch @@ -1033,7 +1035,7 @@ async def failing_put_file(lpath, rpath, **kwargs): # If we get here without hanging, the test passes - def test_get_callback_hooks(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + def test_get_callback_hooks(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test that get() properly calls callback hooks for progress tracking (sync).""" from fsspec.callbacks import Callback @@ -1093,15 +1095,15 @@ class TestFilesetFileSystemAsync: @pytest.fixture def fs(self, sdk: NeMoPlatform) -> FilesetFileSystem: """Create a FilesetFileSystem backed by the test SDK.""" - return FilesetFileSystem(sdk=sdk, skip_instance_cache=True) + return FilesetFileSystem(client=client_from_platform(sdk, FilesClient), skip_instance_cache=True) - async def test_ls_empty_fileset(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_ls_empty_fileset(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test listing an empty fileset.""" path = f"{fileset.workspace}/{fileset.name}" result = await fs._ls(path) assert result == [] - async def test_ls_with_files(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_ls_with_files(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test listing a fileset with files.""" base = f"{fileset.workspace}/{fileset.name}" await fs._pipe_file(f"{base}/file1.txt", b"content1") @@ -1114,7 +1116,7 @@ async def test_ls_with_files(self, fs: FilesetFileSystem, fileset: Fileset): assert f"{base}#file1.txt" in names assert f"{base}#file2.txt" in names - async def test_ls_with_directories(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_ls_with_directories(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test that nested files show as directories in listing.""" base = f"{fileset.workspace}/{fileset.name}" await fs._pipe_file(f"{base}/root.txt", b"root") @@ -1128,7 +1130,7 @@ async def test_ls_with_directories(self, fs: FilesetFileSystem, fileset: Fileset assert types[f"{base}#root.txt"] == "file" assert types[f"{base}#subdir"] == "directory" - async def test_ls_subdirectory(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_ls_subdirectory(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test listing files in a subdirectory.""" base = f"{fileset.workspace}/{fileset.name}" await fs._pipe_file(f"{base}/subdir/file1.txt", b"content1") @@ -1140,7 +1142,7 @@ async def test_ls_subdirectory(self, fs: FilesetFileSystem, fileset: Fileset): assert f"{base}#subdir/file1.txt" in result assert f"{base}#subdir/file2.txt" in result - async def test_cat_file(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_cat_file(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading file content with cat.""" content = b"Hello, fsspec!" path = f"{fileset.workspace}/{fileset.name}/test.txt" @@ -1150,7 +1152,7 @@ async def test_cat_file(self, fs: FilesetFileSystem, fileset: Fileset): assert result == content - async def test_cat_file_with_range(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_cat_file_with_range(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading partial file content with byte range.""" content = b"0123456789ABCDEF" path = f"{fileset.workspace}/{fileset.name}/test.txt" @@ -1160,7 +1162,7 @@ async def test_cat_file_with_range(self, fs: FilesetFileSystem, fileset: Fileset assert result == b"456789" - async def test_pipe_file(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_pipe_file(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test writing file content with pipe.""" path = f"{fileset.workspace}/{fileset.name}/piped.txt" content = b"Piped content" @@ -1170,7 +1172,7 @@ async def test_pipe_file(self, fs: FilesetFileSystem, fileset: Fileset): result = await fs._cat_file(path) assert result == content - async def test_put_file(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + async def test_put_file(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test uploading a local file with _put_file.""" # Create a local file local_file = tmp_path / "upload.txt" @@ -1185,7 +1187,7 @@ async def test_put_file(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: result = await fs._cat_file(remote_path) assert result == content - async def test_put_file_nested_path(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + async def test_put_file_nested_path(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test uploading a file to a nested path with _put_file.""" # Create a local file local_file = tmp_path / "nested_upload.txt" @@ -1204,7 +1206,7 @@ async def test_put_file_nested_path(self, fs: FilesetFileSystem, fileset: Filese parent_info = await fs._info(f"{fileset.workspace}/{fileset.name}/subdir/nested") assert parent_info["type"] == "directory" - async def test_rm_file(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_rm_file(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test deleting a file.""" path = f"{fileset.workspace}/{fileset.name}/to_delete.txt" await fs._pipe_file(path, b"delete me") @@ -1217,7 +1219,7 @@ async def test_rm_file(self, fs: FilesetFileSystem, fileset: Fileset): with pytest.raises(FileNotFoundError): await fs._info(path) - async def test_info(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_info(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test getting file info.""" content = b"Content for info test" path = f"{fileset.workspace}/{fileset.name}/info.txt" @@ -1231,14 +1233,14 @@ async def test_info(self, fs: FilesetFileSystem, fileset: Fileset): assert info["size"] == len(content) assert info["type"] == "file" - async def test_info_directory(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_info_directory(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test getting info for fileset root (directory).""" path = f"{fileset.workspace}/{fileset.name}" info = await fs._info(path) assert info["type"] == "directory" - async def test_protocol_url(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_protocol_url(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test that protocol prefix is handled correctly.""" content = b"Protocol test" base = f"{fileset.workspace}/{fileset.name}" @@ -1251,7 +1253,7 @@ async def test_protocol_url(self, fs: FilesetFileSystem, fileset: Fileset): assert await fs._cat_file(path_no_proto) == content assert await fs._cat_file(path_with_proto) == content - async def test_find(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_find(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test recursive file discovery with find (async).""" base = f"{fileset.workspace}/{fileset.name}" await fs._pipe_file(f"{base}/root.txt", b"root") @@ -1268,7 +1270,7 @@ async def test_find(self, fs: FilesetFileSystem, fileset: Fileset): assert f"{base}#dir1/file1.txt" in result assert f"{base}#dir1/subdir/nested.txt" in result - async def test_glob(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_glob(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test pattern matching with glob (async).""" base = f"{fileset.workspace}/{fileset.name}" await fs._pipe_file(f"{base}/data.csv", b"csv") @@ -1289,7 +1291,7 @@ async def test_glob(self, fs: FilesetFileSystem, fileset: Fileset): # All json files including nested assert len(json_all) == 3 - async def test_isdir_isfile(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_isdir_isfile(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test isdir and isfile type checking (async).""" base = f"{fileset.workspace}/{fileset.name}" await fs._pipe_file(f"{base}/file.txt", b"content") @@ -1305,7 +1307,7 @@ async def test_isdir_isfile(self, fs: FilesetFileSystem, fileset: Fileset): # Subdirectory checks assert await fs._isdir(f"{base}/subdir") - async def test_cat_multiple_files(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_cat_multiple_files(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading multiple files at once with cat (async).""" base = f"{fileset.workspace}/{fileset.name}" await fs._pipe_file(f"{base}/file1.txt", b"content1") @@ -1321,7 +1323,7 @@ async def test_cat_multiple_files(self, fs: FilesetFileSystem, fileset: Fileset) assert result[f"{base}#file2.txt"] == b"content2" assert result[f"{base}#file3.txt"] == b"content3" - async def test_head_tail(self, fs: FilesetFileSystem, fileset: Fileset): + async def test_head_tail(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Test reading first/last bytes of a file (async).""" content = b"0123456789ABCDEFGHIJ" path = f"{fileset.workspace}/{fileset.name}/test.txt" @@ -1335,7 +1337,7 @@ async def test_head_tail(self, fs: FilesetFileSystem, fileset: Fileset): result = await fs._cat_file(path, start=15, end=20) assert result == b"FGHIJ" - async def test_get_single_file(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + async def test_get_single_file(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test downloading a single file with _get() (async version). Tests three behaviors: @@ -1370,7 +1372,7 @@ async def test_get_single_file(self, fs: FilesetFileSystem, fileset: Fileset, tm async def test_download_entire_fileset( self, fs: FilesetFileSystem, - fileset: Fileset, + fileset: FilesetOutput, sample_dataset: Path, tmp_path: Path, ): @@ -1426,7 +1428,9 @@ async def test_download_entire_fileset( assert (download_without_slash / "config" / "settings.json").read_bytes() == b'{"batch_size": 32}' assert not (download_without_slash / fileset.name).exists() - async def test_concurrent_download_failure_hang(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + async def test_concurrent_download_failure_hang( + self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path + ): """Test that a failure in one concurrent download doesn't cause a hang. This reproduces an issue where orphaned asyncio tasks from failed concurrent @@ -1461,7 +1465,7 @@ async def failing_get_file(rpath, lpath, **kwargs): # If we get here without hanging, the test passes - async def test_concurrent_upload_failure_hang(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + async def test_concurrent_upload_failure_hang(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test that a failure in one concurrent upload doesn't cause a hang. Similar to test_concurrent_download_failure_hang, this ensures that the @@ -1494,7 +1498,7 @@ async def failing_put_file(lpath, rpath, **kwargs): # If we get here without hanging, the test passes - async def test_batch_size_limits_concurrency(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path: Path): + async def test_batch_size_limits_concurrency(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path: Path): """Test that batch_size properly limits concurrent operations. Creates a filesystem with batch_size=4, then downloads 8 files and verifies @@ -1505,7 +1509,7 @@ async def test_batch_size_limits_concurrency(self, sdk: NeMoPlatform, fileset: F total_files = 8 # Create filesystem with limited concurrency - fs = FilesetFileSystem(sdk=sdk, batch_size=batch_size) + fs = FilesetFileSystem(client=client_from_platform(sdk, FilesClient), batch_size=batch_size) # Upload files base = f"{fileset.workspace}/{fileset.name}" @@ -1560,7 +1564,7 @@ async def run_download(): downloaded_files = list(download_dir.iterdir()) assert len(downloaded_files) == total_files - async def test_get_callback_hooks(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + async def test_get_callback_hooks(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test that _get properly calls callback hooks for progress tracking. This test demonstrates two callback features: @@ -1622,7 +1626,7 @@ def progress_hook(size, value, **kwargs): assert (download_dir / "file2.txt").read_bytes() == b"content2" assert (download_dir / "file3.txt").read_bytes() == b"content3" - async def test_get_per_chunk_callbacks(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + async def test_get_per_chunk_callbacks(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test that _get passes branched callbacks to _get_file for per-chunk progress. This test verifies the full callback hierarchy: @@ -1711,7 +1715,7 @@ def relative_update(self, inc=1): # Verify file was actually downloaded correctly assert (download_dir / "large_file.bin").read_bytes() == large_content - async def test_put_per_chunk_callbacks(self, fs: FilesetFileSystem, fileset: Fileset, tmp_path: Path): + async def test_put_per_chunk_callbacks(self, fs: FilesetFileSystem, fileset: FilesetOutput, tmp_path: Path): """Test that _put passes branched callbacks to _put_file for per-chunk progress. This test verifies the full callback hierarchy for uploads: @@ -1812,7 +1816,7 @@ class TestDuckDBIntegration: protocol registry and creates it automatically. """ - def test_duckdb_parquet_query(self, sdk: NeMoPlatform, fileset: Fileset): + def test_duckdb_parquet_query(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test querying a parquet file with DuckDB via fileset:// protocol.""" # Create test data df = pd.DataFrame( @@ -1835,7 +1839,7 @@ def test_duckdb_parquet_query(self, sdk: NeMoPlatform, fileset: Fileset): ) # Create filesystem via fsspec (how users would configure it) - fs = fsspec.filesystem("fileset", sdk=sdk) + fs = fsspec.filesystem("fileset", client=client_from_platform(sdk, FilesClient)) # Query with DuckDB using fileset:// URL with new # format fileset_url = f"fileset://{fileset.workspace}/{fileset.name}#{file_path}" @@ -1859,7 +1863,7 @@ def test_duckdb_parquet_query(self, sdk: NeMoPlatform, fileset: Fileset): assert len(result) == 2 assert set(result["category"]) == {"A", "B"} - def test_duckdb_parquet_range_read(self, sdk: NeMoPlatform, fileset: Fileset): + def test_duckdb_parquet_range_read(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test that DuckDB performs efficient range reads on parquet files. Parquet files store metadata at the end (footer), so DuckDB reads: @@ -1889,7 +1893,7 @@ def test_duckdb_parquet_range_read(self, sdk: NeMoPlatform, fileset: Fileset): ) # Create filesystem via fsspec - fs = fsspec.filesystem("fileset", sdk=sdk) + fs = fsspec.filesystem("fileset", client=client_from_platform(sdk, FilesClient)) fileset_url = f"fileset://{fileset.workspace}/{fileset.name}#{file_path}" conn = duckdb.connect() conn.register_filesystem(fs) @@ -1901,7 +1905,7 @@ def test_duckdb_parquet_range_read(self, sdk: NeMoPlatform, fileset: Fileset): assert list(result["id"]) == list(range(100, 111)) assert list(result["value"]) == [float(i) for i in range(100, 111)] - def test_duckdb_legacy_path_format(self, sdk: NeMoPlatform, fileset: Fileset): + def test_duckdb_legacy_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test DuckDB queries using legacy workspace/fileset/path format. This validates backwards compatibility with the legacy path format @@ -1926,7 +1930,7 @@ def test_duckdb_legacy_path_format(self, sdk: NeMoPlatform, fileset: Fileset): ) # Create filesystem via fsspec - fs = fsspec.filesystem("fileset", sdk=sdk) + fs = fsspec.filesystem("fileset", client=client_from_platform(sdk, FilesClient)) # Query with DuckDB using LEGACY path format: workspace/fileset/path fileset_url = f"fileset://{fileset.workspace}/{fileset.name}/{file_path}" @@ -1950,7 +1954,7 @@ def fs(self, sdk: NeMoPlatform) -> FilesetFileSystem: """Create a FilesetFileSystem backed by the test SDK.""" return sdk.files.fsspec - def test_ls_populates_cache_for_nested_dirs(self, fs: FilesetFileSystem, fileset: Fileset): + def test_ls_populates_cache_for_nested_dirs(self, fs: FilesetFileSystem, fileset: FilesetOutput): """_ls should populate cache for all directory levels in the response.""" base = f"{fileset.workspace}/{fileset.name}" @@ -1969,7 +1973,7 @@ def test_ls_populates_cache_for_nested_dirs(self, fs: FilesetFileSystem, fileset assert f"{base}#subdir" in fs.dircache assert f"{base}#subdir/nested" in fs.dircache - def test_deeply_nested_tree(self, fs: FilesetFileSystem, fileset: Fileset): + def test_deeply_nested_tree(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Cache should handle deeply nested directory structures.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2003,7 +2007,7 @@ def test_deeply_nested_tree(self, fs: FilesetFileSystem, fileset: Fileset): info = fs.info(f"{base}/a/b/c/mid.txt") assert info["type"] == "file" - def test_nested_ls_uses_cache(self, fs: FilesetFileSystem, fileset: Fileset): + def test_nested_ls_uses_cache(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Subsequent _ls calls for nested paths should use cache.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2028,7 +2032,7 @@ def test_nested_ls_uses_cache(self, fs: FilesetFileSystem, fileset: Fileset): # Nested ls used cache, didn't add new entries assert after_nested_ls == after_root_ls - def test_find_populates_cache(self, fs: FilesetFileSystem, fileset: Fileset): + def test_find_populates_cache(self, fs: FilesetFileSystem, fileset: FilesetOutput): """_find should populate cache for subsequent _ls calls.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2046,7 +2050,7 @@ def test_find_populates_cache(self, fs: FilesetFileSystem, fileset: Fileset): assert base in fs.dircache assert f"{base}#subdir" in fs.dircache - def test_info_uses_cache(self, fs: FilesetFileSystem, fileset: Fileset): + def test_info_uses_cache(self, fs: FilesetFileSystem, fileset: FilesetOutput): """_info should use dircache instead of making API calls.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2069,7 +2073,7 @@ def test_info_uses_cache(self, fs: FilesetFileSystem, fileset: Fileset): assert info["type"] == "file" assert info["size"] == len(b"content") - def test_cache_invalidation_on_write(self, fs: FilesetFileSystem, fileset: Fileset): + def test_cache_invalidation_on_write(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Cache should be invalidated when files are written.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2089,7 +2093,7 @@ def test_cache_invalidation_on_write(self, fs: FilesetFileSystem, fileset: Files assert f"{base}#file1.txt" in result assert f"{base}#file2.txt" in result - def test_cache_invalidation_on_delete(self, fs: FilesetFileSystem, fileset: Fileset): + def test_cache_invalidation_on_delete(self, fs: FilesetFileSystem, fileset: FilesetOutput): """Cache should be invalidated when files are deleted.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2110,7 +2114,7 @@ def test_cache_invalidation_on_delete(self, fs: FilesetFileSystem, fileset: File assert f"{base}#file1.txt" not in result assert f"{base}#file2.txt" in result - def test_refresh_bypasses_cache(self, fs: FilesetFileSystem, fileset: Fileset): + def test_refresh_bypasses_cache(self, fs: FilesetFileSystem, fileset: FilesetOutput): """_ls with refresh=True should bypass cache.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2125,7 +2129,7 @@ def test_refresh_bypasses_cache(self, fs: FilesetFileSystem, fileset: Fileset): result = fs.ls(base, refresh=True) assert len(result) == 1 - def test_info_file_not_found(self, fs: FilesetFileSystem, fileset: Fileset): + def test_info_file_not_found(self, fs: FilesetFileSystem, fileset: FilesetOutput): """_info should raise FileNotFoundError for non-existent paths.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2138,7 +2142,7 @@ def test_info_file_not_found(self, fs: FilesetFileSystem, fileset: Fileset): fs.info(f"{base}/nonexistent.txt") @pytest.mark.parametrize("detail", [True, False]) - def test_ls_detail_parameter(self, fs: FilesetFileSystem, fileset: Fileset, detail: bool): + def test_ls_detail_parameter(self, fs: FilesetFileSystem, fileset: FilesetOutput, detail: bool): """_ls should respect the detail parameter.""" base = f"{fileset.workspace}/{fileset.name}" @@ -2154,10 +2158,10 @@ def test_ls_detail_parameter(self, fs: FilesetFileSystem, fileset: Fileset, deta else: assert all(isinstance(item, str) for item in result) - def test_cache_disabled(self, sdk: NeMoPlatform, fileset: Fileset): + def test_cache_disabled(self, sdk: NeMoPlatform, fileset: FilesetOutput): """When use_listings_cache=False, cache should not be used.""" # Create filesystem with cache disabled - fs = FilesetFileSystem(sdk=sdk) + fs = FilesetFileSystem(client=client_from_platform(sdk, FilesClient)) fs.dircache.use_listings_cache = False base = f"{fileset.workspace}/{fileset.name}" diff --git a/services/core/files/tests/integration/test_filesets_allowed_hosts.py b/services/core/files/tests/integration/test_filesets_allowed_hosts.py index 6d5f2c10db..49e63d6ae3 100644 --- a/services/core/files/tests/integration/test_filesets_allowed_hosts.py +++ b/services/core/files/tests/integration/test_filesets_allowed_hosts.py @@ -7,7 +7,11 @@ from collections.abc import Iterator import pytest -from nemo_platform import APIStatusError, NeMoPlatform +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.common.auth import AuthClient, get_auth_client from nmp.common.auth.models import Principal from nmp.common.config import AuthConfig, Configuration @@ -69,18 +73,20 @@ def test_create_ngc_fileset_with_disallowed_host_rejected( sdk.secrets.create(workspace=DEFAULT_WORKSPACE, name=secret_name, value="nvapi-dummy") try: name = f"ngc-disallowed-{uuid.uuid4().hex[:8]}" - with pytest.raises(APIStatusError) as exc_info: - sdk.files.filesets.create( + with pytest.raises(NemoHTTPError) as exc_info: + client_from_platform(sdk, FilesClient).create_fileset( workspace=DEFAULT_WORKSPACE, - name=name, - storage={ - "type": "ngc", - "host": "https://disallowed.example.com", - "org": "nvidia", - "team": "team", - "target": "some-resource", - "api_key_secret": f"{DEFAULT_WORKSPACE}/{secret_name}", - }, + body=CreateFilesetRequest( + name=name, + storage={ + "type": "ngc", + "host": "https://disallowed.example.com", + "org": "nvidia", + "team": "team", + "target": "some-resource", + "api_key_secret": f"{DEFAULT_WORKSPACE}/{secret_name}", + }, + ), ) assert exc_info.value.status_code == 400 detail = str(exc_info.value).lower() @@ -100,15 +106,17 @@ def test_create_huggingface_fileset_with_disallowed_endpoint_rejected( """Creating a HuggingFace fileset with endpoint outside allowed_external_hosts returns 400.""" sdk = sdk_with_restrictive_hosts name = f"hf-disallowed-{uuid.uuid4().hex[:8]}" - with pytest.raises(APIStatusError) as exc_info: - sdk.files.filesets.create( + with pytest.raises(NemoHTTPError) as exc_info: + client_from_platform(sdk, FilesClient).create_fileset( workspace=DEFAULT_WORKSPACE, - name=name, - storage={ - "type": "huggingface", - "repo_id": "some-org/some-repo", - "endpoint": "https://disallowed.example.com", - }, + body=CreateFilesetRequest( + name=name, + storage={ + "type": "huggingface", + "repo_id": "some-org/some-repo", + "endpoint": "https://disallowed.example.com", + }, + ), ) assert exc_info.value.status_code == 400 detail = str(exc_info.value).lower() diff --git a/services/core/files/tests/integration/test_huggingface_endpoints.py b/services/core/files/tests/integration/test_huggingface_endpoints.py index c2bbb68aa9..8a1c33b387 100644 --- a/services/core/files/tests/integration/test_huggingface_endpoints.py +++ b/services/core/files/tests/integration/test_huggingface_endpoints.py @@ -13,14 +13,14 @@ import httpx from huggingface_hub import HfApi, hf_hub_download, hf_hub_url, snapshot_download from nemo_platform import NeMoPlatform -from nemo_platform.types.files.fileset import Fileset +from nemo_platform_plugin.files.types import FilesetOutput from nmp.core.files.testing.utils import create_fileset class TestHuggingFaceClientLibrary: """Test HuggingFace Hub client library compatibility with the files service.""" - def test_hf_hub_download_nested_files(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path, hf_asgi_client): + def test_hf_hub_download_nested_files(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path, hf_asgi_client): """Test downloading nested files using huggingface_hub client. This test: @@ -68,7 +68,7 @@ def test_hf_hub_download_nested_files(self, sdk: NeMoPlatform, fileset: Fileset, assert downloaded_file.exists(), f"File {path} was not downloaded" assert downloaded_file.read_bytes() == expected_content, f"Content mismatch for {path}" - def test_hf_hub_download_single_file(self, sdk: NeMoPlatform, fileset: Fileset, tmp_path, hf_asgi_client): + def test_hf_hub_download_single_file(self, sdk: NeMoPlatform, fileset: FilesetOutput, tmp_path, hf_asgi_client): """Test downloading a single file using hf_hub_download.""" test_content = b"This is a test file for single download" test_path = "single_file.txt" @@ -95,7 +95,7 @@ def test_hf_hub_download_single_file(self, sdk: NeMoPlatform, fileset: Fileset, with open(local_path, "rb") as f: assert f.read() == test_content - def test_hf_api_list_repo_files(self, sdk: NeMoPlatform, fileset: Fileset, hf_asgi_client): + def test_hf_api_list_repo_files(self, sdk: NeMoPlatform, fileset: FilesetOutput, hf_asgi_client): """Test listing repository files using HfApi.""" test_files = { "file1.txt": b"content1", @@ -122,7 +122,7 @@ def test_hf_api_list_repo_files(self, sdk: NeMoPlatform, fileset: Fileset, hf_as listed_files = {sibling.rfilename for sibling in repo_info.siblings} assert listed_files == set(test_files.keys()) - def test_hf_hub_url_generates_valid_download_url(self, sdk: NeMoPlatform, fileset: Fileset): + def test_hf_hub_url_generates_valid_download_url(self, sdk: NeMoPlatform, fileset: FilesetOutput): """Test that hf_hub_url generates a valid URL for file download. This test verifies that: diff --git a/services/core/files/tests/integration/test_otlp_endpoints.py b/services/core/files/tests/integration/test_otlp_endpoints.py index ef07b427d6..2252bc5934 100644 --- a/services/core/files/tests/integration/test_otlp_endpoints.py +++ b/services/core/files/tests/integration/test_otlp_endpoints.py @@ -14,7 +14,7 @@ import httpx import pytest -from nemo_platform.types.files.fileset import Fileset +from nemo_platform_plugin.files.types import FilesetOutput from opentelemetry.proto.collector.logs.v1 import logs_service_pb2 @@ -155,7 +155,7 @@ def create_request( def test_upload_and_query_logs_roundtrip( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, otlp_request_factory, otlp_format: str, ): @@ -197,7 +197,7 @@ def test_upload_and_query_logs_roundtrip( def test_query_logs_with_filters( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, otlp_request_factory, otlp_format: str, ): @@ -241,7 +241,7 @@ def test_query_logs_with_filters( def test_query_logs_pagination( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, otlp_request_factory, otlp_format: str, ): @@ -308,7 +308,7 @@ def test_query_logs_pagination( def test_multiple_batches_same_partition( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, otlp_request_factory, otlp_format: str, ): @@ -360,7 +360,7 @@ def test_multiple_batches_same_partition( def test_query_empty_fileset( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, ): """Test querying a fileset with no logs returns empty result.""" workspace = fileset.workspace @@ -380,7 +380,7 @@ def test_query_empty_fileset( def test_upload_logs_missing_attributes_partial_success( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, ): """Test that logs with missing required attributes are rejected (JSON format).""" workspace = fileset.workspace @@ -454,7 +454,7 @@ def test_upload_logs_missing_attributes_partial_success( def test_upload_logs_invalid_json( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, ): """Test that invalid JSON returns 400 error.""" workspace = fileset.workspace @@ -471,7 +471,7 @@ def test_upload_logs_invalid_json( def test_upload_logs_invalid_protobuf( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, ): """Test that invalid protobuf returns 400 error.""" workspace = fileset.workspace @@ -522,7 +522,7 @@ def test_upload_logs_nonexistent_fileset( def test_query_logs_invalid_filter_key_returns_400( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, ): """Test invalid filter key is rejected with a 400 response.""" workspace = fileset.workspace @@ -539,7 +539,7 @@ def test_query_logs_invalid_filter_key_returns_400( def test_query_logs_invalid_partition_value_does_not_leak_internal_details( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, ): """Test invalid partition filter input does not leak internal details.""" workspace = fileset.workspace @@ -559,7 +559,7 @@ def test_query_logs_invalid_partition_value_does_not_leak_internal_details( def test_query_logs_invalid_partition_value_returns_400( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, ): """Test unsafe partition filter values are rejected before query execution.""" workspace = fileset.workspace @@ -576,7 +576,7 @@ def test_query_logs_invalid_partition_value_returns_400( def test_query_logs_log_message_allows_apostrophe( client: httpx.Client, - fileset: Fileset, + fileset: FilesetOutput, ): """Test log_message filter remains usable for normal text with apostrophes.""" workspace = fileset.workspace diff --git a/services/core/files/tests/integration/tests_filesets_with_auth_secrets.py b/services/core/files/tests/integration/tests_filesets_with_auth_secrets.py index 0d9739435f..626bb342af 100644 --- a/services/core/files/tests/integration/tests_filesets_with_auth_secrets.py +++ b/services/core/files/tests/integration/tests_filesets_with_auth_secrets.py @@ -8,7 +8,11 @@ from unittest.mock import patch import pytest -from nemo_platform import APIStatusError, NeMoPlatform +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.core.auth.app.bundle import ( build_authorization_data as _real_build_authorization_data, ) @@ -106,17 +110,20 @@ def test_editor_can_create_hf_fileset_with_token_secret( ) editor_sdk = as_user(sdk, editor_email) - created = editor_sdk.files.filesets.create( + files = client_from_platform(editor_sdk, FilesClient) + created = files.create_fileset( workspace=workspace, - name=fileset_name, - description="hf fileset", - storage={ - "type": "huggingface", - "repo_id": "Qwen/Qwen3-0.6B", - "repo_type": "model", - "token_secret": secret_name, - }, - ) + body=CreateFilesetRequest( + name=fileset_name, + description="hf fileset", + storage={ + "type": "huggingface", + "repo_id": "Qwen/Qwen3-0.6B", + "repo_type": "model", + "token_secret": secret_name, + }, + ), + ).data() assert created.name == fileset_name assert created.storage.type == "huggingface" @@ -142,17 +149,20 @@ def test_custom_role_without_secrets_read_denied_with_token_secret( ) user_sdk = as_user(sdk, user_email) - with pytest.raises(APIStatusError) as exc_info: - user_sdk.files.filesets.create( + user_files = client_from_platform(user_sdk, FilesClient) + with pytest.raises(NemoHTTPError) as exc_info: + user_files.create_fileset( workspace=workspace, - name=short_unique_name("fileset"), - description="should fail", - storage={ - "type": "huggingface", - "repo_id": "Qwen/Qwen3-0.6B", - "repo_type": "model", - "token_secret": secret_name, - }, + body=CreateFilesetRequest( + name=short_unique_name("fileset"), + description="should fail", + storage={ + "type": "huggingface", + "repo_id": "Qwen/Qwen3-0.6B", + "repo_type": "model", + "token_secret": secret_name, + }, + ), ) assert exc_info.value.status_code == 400 @@ -176,17 +186,20 @@ def test_missing_secret_returns_secret_not_found_error( ) editor_sdk = as_user(sdk, editor_email) - with pytest.raises(APIStatusError) as exc_info: - editor_sdk.files.filesets.create( + editor_files = client_from_platform(editor_sdk, FilesClient) + with pytest.raises(NemoHTTPError) as exc_info: + editor_files.create_fileset( workspace=workspace, - name=short_unique_name("fileset"), - description="missing secret", - storage={ - "type": "huggingface", - "repo_id": "Qwen/Qwen3-0.6B", - "repo_type": "model", - "token_secret": "does-not-exist", - }, + body=CreateFilesetRequest( + name=short_unique_name("fileset"), + description="missing secret", + storage={ + "type": "huggingface", + "repo_id": "Qwen/Qwen3-0.6B", + "repo_type": "model", + "token_secret": "does-not-exist", + }, + ), ) assert exc_info.value.status_code == 400 @@ -211,17 +224,20 @@ def test_public_hf_without_token_secret_succeeds( ) editor_sdk = as_user(sdk, editor_email) - created = editor_sdk.files.filesets.create( + editor_files = client_from_platform(editor_sdk, FilesClient) + created = editor_files.create_fileset( workspace=workspace, - name=fileset_name, - description="no token", - storage={ - "type": "huggingface", - "repo_id": "Qwen/Qwen3-0.6B", - "repo_type": "model", - # no token_secret - }, - ) + body=CreateFilesetRequest( + name=fileset_name, + description="no token", + storage={ + "type": "huggingface", + "repo_id": "Qwen/Qwen3-0.6B", + "repo_type": "model", + # no token_secret + }, + ), + ).data() assert created.name == fileset_name assert created.storage.type == "huggingface" @@ -253,16 +269,19 @@ async def _list_files_noop(self, path=None): ) editor_sdk = as_user(sdk, editor_email) - editor_sdk.files.filesets.create( + editor_files = client_from_platform(editor_sdk, FilesClient) + editor_files.create_fileset( workspace=workspace, - name=fileset_name, - description="hf fileset read test", - storage={ - "type": "huggingface", - "repo_id": "Qwen/Qwen3-0.6B", - "repo_type": "model", - "token_secret": secret_name, - }, + body=CreateFilesetRequest( + name=fileset_name, + description="hf fileset read test", + storage={ + "type": "huggingface", + "repo_id": "Qwen/Qwen3-0.6B", + "repo_type": "model", + "token_secret": secret_name, + }, + ), ) files = editor_sdk.files.list(fileset=fileset_name, workspace=workspace) diff --git a/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py b/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py index e2b6fd09f0..bdc12a4bf5 100644 --- a/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py +++ b/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py @@ -7,6 +7,10 @@ from typing import Any, Dict, List, Optional, Tuple from nemo_platform import AsyncNeMoPlatform, NotFoundError, PermissionDeniedError +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError +from nemo_platform_plugin.files.client import AsyncFilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.common.api.filter import ComparisonOperation, FilterOperation, FilterOperator, LogicalOperation from nmp.common.api.in_memory_filter import InMemoryFilterRepository from nmp.common.api.parsed_filter import ParsedFilter @@ -258,10 +262,12 @@ async def create_job( job_name = f"{source_prefix}-{short_id}" # Create a fileset to store job artifacts - fileset = await self.sdk.files.filesets.create( + files = client_from_platform(self.sdk, AsyncFilesClient) + fileset_resp = await files.create_fileset( + body=CreateFilesetRequest(name=f"job-fileset-{job_name}"), workspace=workspace, - name=f"job-fileset-{job_name}", ) + fileset = fileset_resp.data() # Create job (ID is assigned by entity store) job = await self.store.create( @@ -451,8 +457,9 @@ async def delete_job(self, job_name: str, workspace: str) -> bool: # Delete job fileset via sdk to properly clean up storage. # Tolerate fileset already being gone (e.g. cleaned up by workspace cleanup). try: - await self.sdk.files.filesets.delete(job_entity.fileset, workspace=workspace) - except NotFoundError: + files = client_from_platform(self.sdk, AsyncFilesClient) + await files.delete_fileset(name=job_entity.fileset, workspace=workspace) + except ClientNotFoundError: logger.warning("Job fileset not found during deletion, may have been cleaned up already", extra=extras) logger.info( diff --git a/services/core/jobs/tests/conftest.py b/services/core/jobs/tests/conftest.py index 03dbfbbbc4..dc36f6da96 100644 --- a/services/core/jobs/tests/conftest.py +++ b/services/core/jobs/tests/conftest.py @@ -5,7 +5,7 @@ import tempfile from pathlib import Path from typing import AsyncGenerator -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest import pytest_asyncio @@ -243,7 +243,20 @@ def sample_job_dict(): @fixture -def mock_nmp_client(): +def _mock_files_client(): + """Create a mock AsyncFilesClient for testing.""" + mock_files = AsyncMock() + mock_fileset = MagicMock() + mock_fileset.name = "test-fileset-id" + mock_response = MagicMock() + mock_response.data.return_value = mock_fileset + mock_files.create_fileset.return_value = mock_response + mock_files.delete_fileset.return_value = None + return mock_files + + +@fixture +def mock_nmp_client(_mock_files_client): """Create a flexible mock of NeMoPlatform for testing.""" mock_client = MagicMock() @@ -257,17 +270,8 @@ def mock_nmp_client(): mock_client.jobs.steps.retrieve = MagicMock() mock_client.jobs.steps.update_status = MagicMock() - # Mock filesets API for log storage - mock_client.files = MagicMock() - mock_client.files.filesets = MagicMock() - mock_client.files.filesets.create = AsyncMock() - mock_client.files.filesets.delete = AsyncMock() - # Return a mock fileset object with an id - mock_fileset = MagicMock() - mock_fileset.name = "test-fileset-id" - mock_client.files.filesets.create.return_value = mock_fileset - - return mock_client + with patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=_mock_files_client): + yield mock_client @fixture diff --git a/services/core/jobs/tests/test_dispatcher.py b/services/core/jobs/tests/test_dispatcher.py index 9dc1511eb2..0e46800685 100644 --- a/services/core/jobs/tests/test_dispatcher.py +++ b/services/core/jobs/tests/test_dispatcher.py @@ -8,7 +8,7 @@ """ import json -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from nmp.common.api.filter import ComparisonOperation, FilterOperator, LogicalOperation, parse_json_filter @@ -192,15 +192,16 @@ async def test_delete_job_missing_fileset_succeeds(mock_dispatcher: JobDispatche """ from unittest.mock import AsyncMock - from nemo_platform import NotFoundError + from nemo_platform_plugin.client.errors import NotFoundError job_id, job_name, _, _, _, _ = await create_test_job_data(mock_store, "delete-missing-fileset-job") - # Simulate the fileset already being gone - mock_dispatcher.sdk.files.filesets.delete = AsyncMock(side_effect=NotFoundError.__new__(NotFoundError)) + # Simulate the fileset already being gone by making the mock files client raise NotFoundError + mock_files = AsyncMock() + mock_files.delete_fileset = AsyncMock(side_effect=NotFoundError.__new__(NotFoundError)) - # delete_job should still succeed and return True - deleted = await mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE) + with patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=mock_files): + deleted = await mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE) assert deleted is True # The job entity itself should be gone @@ -643,58 +644,58 @@ async def test_list_steps_across_multiple_workspaces( # Create entity store with multiple workspaces and projects projects = ["default/test-project", "other-workspace/test-project"] with create_test_client(client_type=EntityClient, projects=projects) as mock_store: - # Create mock SDK + # Create mock SDK with patched files client mock_nmp_client = MagicMock() - mock_nmp_client.files = MagicMock() - mock_nmp_client.files.filesets = MagicMock() - mock_nmp_client.files.filesets.create = AsyncMock() - mock_nmp_client.files.filesets.delete = AsyncMock() - mock_fileset = MagicMock() - mock_fileset.name = "test-fileset-id" - mock_nmp_client.files.filesets.create.return_value = mock_fileset - - # Create dispatcher with the multi-workspace store - mock_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) - - # Create jobs in workspace "default" - job1 = await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) - - job2_request = CreatePlatformJobRequest( - name="test-job-2", - description="Second test job", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job2 = await mock_dispatcher.create_job(job2_request, DEFAULT_WORKSPACE) - - # Create jobs in workspace "other-workspace" - job3_request = CreatePlatformJobRequest( - name="test-job-3", - description="Third test job in other workspace", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job3 = await mock_dispatcher.create_job(job3_request, "other-workspace") - - job4_request = CreatePlatformJobRequest( - name="test-job-4", - description="Fourth test job in other workspace", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job4 = await mock_dispatcher.create_job(job4_request, "other-workspace") + mock_files = AsyncMock() + mock_fileset_obj = MagicMock() + mock_fileset_obj.name = "test-fileset-id" + mock_resp = MagicMock() + mock_resp.data.return_value = mock_fileset_obj + mock_files.create_fileset.return_value = mock_resp + + with patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=mock_files): + # Create dispatcher with the multi-workspace store + mock_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + + # Create jobs in workspace "default" + job1 = await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) + + job2_request = CreatePlatformJobRequest( + name="test-job-2", + description="Second test job", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job2 = await mock_dispatcher.create_job(job2_request, DEFAULT_WORKSPACE) + + # Create jobs in workspace "other-workspace" + job3_request = CreatePlatformJobRequest( + name="test-job-3", + description="Third test job in other workspace", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job3 = await mock_dispatcher.create_job(job3_request, "other-workspace") + + job4_request = CreatePlatformJobRequest( + name="test-job-4", + description="Fourth test job in other workspace", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job4 = await mock_dispatcher.create_job(job4_request, "other-workspace") # List steps in "default" workspace step_filter = PlatformJobStepsListFilter() @@ -774,16 +775,16 @@ async def test_list_steps_with_status_filter( with create_test_client(client_type=EntityClient) as mock_store: mock_nmp_client = MagicMock() - mock_nmp_client.files = MagicMock() - mock_nmp_client.files.filesets = MagicMock() - mock_nmp_client.files.filesets.create = AsyncMock() - mock_nmp_client.files.filesets.delete = AsyncMock() - mock_fileset = MagicMock() - mock_fileset.name = "test-fileset-id" - mock_nmp_client.files.filesets.create.return_value = mock_fileset + mock_files = AsyncMock() + mock_fileset_obj = MagicMock() + mock_fileset_obj.name = "test-fileset-id" + mock_resp = MagicMock() + mock_resp.data.return_value = mock_fileset_obj + mock_files.create_fileset.return_value = mock_resp - mock_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) - await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) + with patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=mock_files): + mock_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) # Filtering by a status that no step has should return nothing — the # important assertion is that the call doesn't raise. @@ -804,7 +805,7 @@ async def test_list_jobs_across_multiple_workspaces( sample_platform_job_request: CreatePlatformJobRequest, ): """Test list_jobs respects workspace filtering.""" - from unittest.mock import AsyncMock, MagicMock + from unittest.mock import MagicMock from nmp.common.entities.client import EntityClient from nmp.core.jobs.api.v2.jobs.schemas import PlatformJobSortField @@ -813,58 +814,58 @@ async def test_list_jobs_across_multiple_workspaces( # Create entity store with multiple workspaces and projects projects = ["default/test-project", "other-workspace/test-project"] with create_test_client(client_type=EntityClient, projects=projects) as mock_store: - # Create mock SDK + # Create mock SDK with patched files client mock_nmp_client = MagicMock() - mock_nmp_client.files = MagicMock() - mock_nmp_client.files.filesets = MagicMock() - mock_nmp_client.files.filesets.create = AsyncMock() - mock_nmp_client.files.filesets.delete = AsyncMock() - mock_fileset = MagicMock() - mock_fileset.name = "test-fileset-id" - mock_nmp_client.files.filesets.create.return_value = mock_fileset - - # Create dispatcher with the multi-workspace store - mock_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) - - # Create jobs in workspace "default" - job1 = await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) - - job2_request = CreatePlatformJobRequest( - name="test-job-2", - description="Second test job", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job2 = await mock_dispatcher.create_job(job2_request, DEFAULT_WORKSPACE) - - # Create jobs in workspace "other-workspace" - job3_request = CreatePlatformJobRequest( - name="test-job-3", - description="Third test job in other workspace", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job3 = await mock_dispatcher.create_job(job3_request, "other-workspace") - - job4_request = CreatePlatformJobRequest( - name="test-job-4", - description="Fourth test job in other workspace", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job4 = await mock_dispatcher.create_job(job4_request, "other-workspace") + mock_files = AsyncMock() + mock_fileset_obj = MagicMock() + mock_fileset_obj.name = "test-fileset-id" + mock_resp = MagicMock() + mock_resp.data.return_value = mock_fileset_obj + mock_files.create_fileset.return_value = mock_resp + + with patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=mock_files): + # Create dispatcher with the multi-workspace store + mock_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + + # Create jobs in workspace "default" + job1 = await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) + + job2_request = CreatePlatformJobRequest( + name="test-job-2", + description="Second test job", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job2 = await mock_dispatcher.create_job(job2_request, DEFAULT_WORKSPACE) + + # Create jobs in workspace "other-workspace" + job3_request = CreatePlatformJobRequest( + name="test-job-3", + description="Third test job in other workspace", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job3 = await mock_dispatcher.create_job(job3_request, "other-workspace") + + job4_request = CreatePlatformJobRequest( + name="test-job-4", + description="Fourth test job in other workspace", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job4 = await mock_dispatcher.create_job(job4_request, "other-workspace") # List jobs in "default" workspace jobs_default, count_default = await mock_dispatcher.list_jobs( diff --git a/services/core/jobs/tests/test_job_logs.py b/services/core/jobs/tests/test_job_logs.py index 04fdbb7405..4d3c752df2 100644 --- a/services/core/jobs/tests/test_job_logs.py +++ b/services/core/jobs/tests/test_job_logs.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from datetime import datetime -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI @@ -58,19 +58,17 @@ def dispatcher(self): """Create a real dispatcher with test entity store and mock SDK.""" projects = ["default/test-project"] with create_test_client(client_type=EntityClient, projects=projects) as mock_store: - # Create mock SDK mock_nmp_client = MagicMock() - mock_nmp_client.files = MagicMock() - mock_nmp_client.files.filesets = MagicMock() - mock_nmp_client.files.filesets.create = AsyncMock() - mock_nmp_client.files.filesets.delete = AsyncMock() - mock_fileset = MagicMock() - mock_fileset.name = "test-fileset-id" - mock_nmp_client.files.filesets.create.return_value = mock_fileset - - # Create dispatcher with the test store - dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) - yield dispatcher + mock_files = AsyncMock() + mock_fileset_obj = MagicMock() + mock_fileset_obj.name = "test-fileset-id" + mock_resp = MagicMock() + mock_resp.data.return_value = mock_fileset_obj + mock_files.create_fileset.return_value = mock_resp + + with patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=mock_files): + dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + yield dispatcher @pytest.fixture def test_client(self, dispatcher, mock_logs_client): diff --git a/services/core/jobs/tests/test_jobs_api.py b/services/core/jobs/tests/test_jobs_api.py index 218c6f6c05..0ed059d7c6 100644 --- a/services/core/jobs/tests/test_jobs_api.py +++ b/services/core/jobs/tests/test_jobs_api.py @@ -1150,7 +1150,7 @@ async def test_get_platform_jobs_steps_list_filter_invalid(): @pytest.mark.asyncio async def test_job_steps_list_global_vs_workspaced(sample_platform_job_request: CreatePlatformJobRequest): """Test that global step listing returns steps from all workspaces while workspaced calls are filtered.""" - from unittest.mock import AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch from nmp.common.entities.client import EntityClient from nmp.testing import create_test_client @@ -1158,58 +1158,58 @@ async def test_job_steps_list_global_vs_workspaced(sample_platform_job_request: # Create entity store with multiple workspaces and projects projects = ["default/test-project", "other-workspace/test-project"] with create_test_client(client_type=EntityClient, projects=projects) as mock_store: - # Create mock SDK + # Create mock SDK with patched files client mock_nmp_client = MagicMock() - mock_nmp_client.files = MagicMock() - mock_nmp_client.files.filesets = MagicMock() - mock_nmp_client.files.filesets.create = AsyncMock() - mock_nmp_client.files.filesets.delete = AsyncMock() - mock_fileset = MagicMock() - mock_fileset.name = "test-fileset-id" - mock_nmp_client.files.filesets.create.return_value = mock_fileset - - # Create dispatcher with the multi-workspace store - mock_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) - - # Create jobs in "default" workspace - job1 = await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) - - job2_request = CreatePlatformJobRequest( - name="test-job-2", - description="Second test job", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job2 = await mock_dispatcher.create_job(job2_request, DEFAULT_WORKSPACE) - - # Create jobs in "other-workspace" - job3_request = CreatePlatformJobRequest( - name="test-job-3", - description="Third test job in other workspace", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job3 = await mock_dispatcher.create_job(job3_request, "other-workspace") - - job4_request = CreatePlatformJobRequest( - name="test-job-4", - description="Fourth test job in other workspace", - project="test-project", - source=TestConstants.SOURCE, - spec=TestConstants.SPEC_BASIC, - platform_spec=TestConstants.PLATFORM_SPEC, - ownership=TestConstants.OWNERSHIP_BASIC, - custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, - ) - job4 = await mock_dispatcher.create_job(job4_request, "other-workspace") + mock_files = AsyncMock() + mock_fileset_obj = MagicMock() + mock_fileset_obj.name = "test-fileset-id" + mock_resp = MagicMock() + mock_resp.data.return_value = mock_fileset_obj + mock_files.create_fileset.return_value = mock_resp + + with patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=mock_files): + # Create dispatcher with the multi-workspace store + mock_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + + # Create jobs in "default" workspace + job1 = await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) + + job2_request = CreatePlatformJobRequest( + name="test-job-2", + description="Second test job", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job2 = await mock_dispatcher.create_job(job2_request, DEFAULT_WORKSPACE) + + # Create jobs in "other-workspace" + job3_request = CreatePlatformJobRequest( + name="test-job-3", + description="Third test job in other workspace", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job3 = await mock_dispatcher.create_job(job3_request, "other-workspace") + + job4_request = CreatePlatformJobRequest( + name="test-job-4", + description="Fourth test job in other workspace", + project="test-project", + source=TestConstants.SOURCE, + spec=TestConstants.SPEC_BASIC, + platform_spec=TestConstants.PLATFORM_SPEC, + ownership=TestConstants.OWNERSHIP_BASIC, + custom_fields=TestConstants.CUSTOM_FIELDS_BASIC, + ) + job4 = await mock_dispatcher.create_job(job4_request, "other-workspace") # Test global listing with wildcard - should return steps from all workspaces step_filter = PlatformJobStepsListFilter() diff --git a/services/core/models/src/nmp/core/models/api/permissions.py b/services/core/models/src/nmp/core/models/api/permissions.py index e9ed8aa8c2..1dfb33d3ad 100644 --- a/services/core/models/src/nmp/core/models/api/permissions.py +++ b/services/core/models/src/nmp/core/models/api/permissions.py @@ -14,7 +14,11 @@ from nemo_platform import AsyncNeMoPlatform from nemo_platform._exceptions import NotFoundError, PermissionDeniedError -from nemo_platform.types.files import Fileset +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError +from nemo_platform_plugin.client.errors import PermissionDeniedError as ClientPermissionDeniedError +from nemo_platform_plugin.files.client import AsyncFilesClient +from nemo_platform_plugin.files.types import FilesetOutput from nmp.common.auth import AuthClient from nmp.common.entities.utils import parse_entity_ref @@ -34,7 +38,7 @@ async def check_secret_access(nmp_sdk: AsyncNeMoPlatform, secret_name: str, work raise ValueError(f"Secret '{secret_name}' not found in workspace '{workspace}'") from None -async def check_fileset_access(nmp_sdk: AsyncNeMoPlatform, fileset: str, workspace: str) -> Fileset: +async def check_fileset_access(nmp_sdk: AsyncNeMoPlatform, fileset: str, workspace: str) -> FilesetOutput: """Check that the current user can access the referenced fileset. Retrieves fileset metadata via the Files API; AuthZ middleware enforces @@ -46,12 +50,13 @@ async def check_fileset_access(nmp_sdk: AsyncNeMoPlatform, fileset: str, workspa """ _fs_ref = parse_entity_ref(fileset, default_workspace=workspace) fs_workspace, fs_name = _fs_ref.workspace, _fs_ref.name + files = client_from_platform(nmp_sdk, AsyncFilesClient) try: - fs = await nmp_sdk.files.filesets.retrieve(workspace=fs_workspace, name=fs_name) + fs = (await files.get_fileset(workspace=fs_workspace, name=fs_name)).data() return fs - except PermissionDeniedError: + except ClientPermissionDeniedError: raise PermissionError(f"Access denied to fileset '{fileset}'") from None - except NotFoundError: + except ClientNotFoundError: raise ValueError(f"Fileset '{fileset}' not found in workspace '{fs_workspace}'") from None diff --git a/services/core/models/src/nmp/core/models/api/service/model_entity_service.py b/services/core/models/src/nmp/core/models/api/service/model_entity_service.py index e2f1ab4d43..c6e1ecb2c8 100644 --- a/services/core/models/src/nmp/core/models/api/service/model_entity_service.py +++ b/services/core/models/src/nmp/core/models/api/service/model_entity_service.py @@ -9,7 +9,8 @@ from collections import defaultdict from nemo_platform import AsyncNeMoPlatform, NotFoundError, PermissionDeniedError -from nemo_platform.types.files import Fileset, FilesetFile, HuggingfaceStorageConfig, NGCStorageConfig +from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig, NGCStorageConfig +from nemo_platform_plugin.files.types import FilesetFileOutput, FilesetOutput from nmp.common.api.common import Page, PaginationData from nmp.common.api.filter import ComparisonOperation, FilterOperation, FilterOperator, LogicalOperation from nmp.common.api.parsed_filter import ParsedFilter @@ -58,7 +59,7 @@ def _repo_id_matches_trusted(repo_id: str, patterns: list[str]) -> bool: async def get_fileset_and_files_list( sdk: AsyncNeMoPlatform, workspace: str, fileset_ref: str | None -) -> tuple[Fileset, list[FilesetFile]]: +) -> tuple[FilesetOutput, list[FilesetFileOutput]]: """Validate that the fileset exists and the user has access.""" if not fileset_ref: raise FilesetValidationError("Fileset reference is required") @@ -200,7 +201,7 @@ def _has_tool_call_plugin(request) -> bool: return False -def fileset_has_tool_call_plugin(fileset: Fileset) -> bool: +def fileset_has_tool_call_plugin(fileset: FilesetOutput) -> bool: """Return True if a fileset's metadata contains a tool_call_plugin value.""" if not fileset.metadata: return False diff --git a/services/core/models/src/nmp/core/models/tasks/model_spec/run.py b/services/core/models/src/nmp/core/models/tasks/model_spec/run.py index 5ae66cc7d8..5cff0d3eba 100644 --- a/services/core/models/src/nmp/core/models/tasks/model_spec/run.py +++ b/services/core/models/src/nmp/core/models/tasks/model_spec/run.py @@ -26,8 +26,11 @@ NeMoPlatformError, NotFoundError, ) -from nemo_platform.types.files import Fileset, HuggingfaceStorageConfig, LocalStorageConfig, NGCStorageConfig from nemo_platform.types.models import ModelEntity +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig, LocalStorageConfig, NGCStorageConfig +from nemo_platform_plugin.files.types import FilesetOutput from nmp.common.entities.utils import parse_entity_ref from nmp.common.model_utils import is_embedding_model from nmp.common.sdk_factory import get_platform_sdk @@ -75,22 +78,25 @@ def __init__(self, sdk: NeMoPlatform, job_ctx: NMPJobContext): self.job_ctx = job_ctx @staticmethod - def _merge_fileset_metadata(fs: Fileset, model_spec: ModelSpec) -> None: + def _merge_fileset_metadata(fs: FilesetOutput, model_spec: ModelSpec) -> None: """Merge tool calling metadata from fileset into model spec. Users can set these values on the fileset at creation time via metadata: - sdk.files.filesets.create( - ..., - metadata={ - "model": { - "tool_calling": { - "chat_template": "", - "tool_call_parser": "llama3_json", - "tool_call_plugin": "default/my-plugin-fileset", - "auto_tool_choice": True, + files = client_from_platform(sdk, FilesClient) + files.create_fileset( + body=CreateFilesetRequest( + ..., + metadata={ + "model": { + "tool_calling": { + "chat_template": "", + "tool_call_parser": "llama3_json", + "tool_call_plugin": "default/my-plugin-fileset", + "auto_tool_choice": True, + }, }, }, - }, + ), ) The model spec task then merges these into the auto-generated ModelSpec so @@ -200,7 +206,8 @@ def analyze_checkpoint(self, config: ModelSpecTaskConfig) -> ModelEntity: # Validate that the fileset exists before creating the model entity logger.info(f"Validating fileset exists: {fileset_workspace}/{fileset_name}") try: - fs = self.sdk.files.filesets.retrieve(workspace=fileset_workspace, name=fileset_name) + files = client_from_platform(self.sdk, FilesClient) + fs = files.get_fileset(workspace=fileset_workspace, name=fileset_name).data() logger.info(f"Fileset validation successful: {fileset_workspace}/{fileset_name}") except Exception as e: logger.error(f"Fileset validation failed: {fileset_workspace}/{fileset_name}") diff --git a/services/core/models/tests/integration/test_chat_template_tool_calling.py b/services/core/models/tests/integration/test_chat_template_tool_calling.py index 878330d419..31d8aa889a 100644 --- a/services/core/models/tests/integration/test_chat_template_tool_calling.py +++ b/services/core/models/tests/integration/test_chat_template_tool_calling.py @@ -175,10 +175,10 @@ def sample_deployment(): def _update_fileset_and_run_task(test_clients, model_name, metadata, tmp_path): """Simulate: user patches fileset → model-spec task runs. - 1. **Update fileset** — mock ``sdk.files.filesets.update()`` to simulate the - user adding ``metadata`` to their fileset. - 2. **Run analyze_checkpoint** — the task calls ``sdk.files.filesets.retrieve()`` - (also mocked to return the now-updated fileset), merges ``metadata.model.tool_calling`` + 1. **Update fileset** — mock the files client to simulate the user adding + ``metadata`` to their fileset. + 2. **Run analyze_checkpoint** — the task calls ``client_from_platform(sdk, FilesClient).get_fileset()`` + (mocked to return the now-updated fileset), merges ``metadata.model.tool_calling`` into a ``ModelSpec``, and calls the *real* ``sdk.models.update(spec=...)``. ``nmp.core.models.parallelism.api`` depends on torch/accelerate (GPU deps @@ -188,24 +188,18 @@ def _update_fileset_and_run_task(test_clients, model_name, metadata, tmp_path): """ sdk = test_clients.sdk - # -- Step 1: User updates fileset with metadata --------------------------- + # -- Step 1: Build the updated fileset representation --------------------- updated_fileset = SimpleNamespace( name=FILESET_NAME, workspace=DEFAULT_WORKSPACE, metadata=metadata, + storage=None, ) - with patch.object(sdk.files.filesets, "update", return_value=updated_fileset) as mock_update: - sdk.files.filesets.update( - FILESET_NAME, - workspace=DEFAULT_WORKSPACE, - metadata=metadata, - ) - mock_update.assert_called_once_with( - FILESET_NAME, - workspace=DEFAULT_WORKSPACE, - metadata=metadata, - ) + mock_files_client = MagicMock() + mock_response = MagicMock() + mock_response.data.return_value = updated_fileset + mock_files_client.get_fileset.return_value = mock_response # -- Step 2: Model-spec background task runs ------------------------------ model_dir = tmp_path / "model" @@ -239,11 +233,11 @@ def _update_fileset_and_run_task(test_clients, model_name, metadata, tmp_path): parent.__path__ = [] modules_patch["nmp.core.models.parallelism"] = parent - # Task calls sdk.files.filesets.retrieve() → gets the updated fileset + # Task calls client_from_platform(sdk, FilesClient).get_fileset() → gets the updated fileset with ( patch.dict(sys.modules, modules_patch), - patch.object(sdk.files.filesets, "retrieve", return_value=updated_fileset), - patch.object(sdk.files, "_list_files", return_value=SimpleNamespace(data=[])), + patch("nmp.core.models.tasks.model_spec.run.client_from_platform", return_value=mock_files_client), + patch.object(sdk.files, "list", return_value=SimpleNamespace(data=[])), patch.object(runner.filesystem_sdk, "get"), patch("nmp.core.models.tasks.model_spec.run.os.listdir", return_value=["config.json"]), ): diff --git a/services/core/models/tests/integration/test_model_entity_service_integration.py b/services/core/models/tests/integration/test_model_entity_service_integration.py index 6611a27e82..bd05a55226 100644 --- a/services/core/models/tests/integration/test_model_entity_service_integration.py +++ b/services/core/models/tests/integration/test_model_entity_service_integration.py @@ -3,13 +3,13 @@ """Integration tests for Model Entity service with in-memory EntityClient.""" -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from nemo_platform import AsyncNeMoPlatform -from nemo_platform.filesets import ListFilesResponse -from nemo_platform.types.files import Fileset, FilesetFile, LocalStorageConfig -from nemo_platform.types.shared import FilesetMetadata +from nemo_platform_plugin.files.metadata import FilesetMetadata +from nemo_platform_plugin.files.storage_config import LocalStorageConfig +from nemo_platform_plugin.files.types import FilesetFileOutput, FilesetOutput, ListFilesetFilesResponse from nmp.common.api.filter import ComparisonOperation, FilterOperator, LogicalOperation from nmp.common.api.parsed_filter import ParsedFilter from nmp.common.entities.client import EntityClient @@ -40,14 +40,37 @@ def entity_client() -> EntityClient: @pytest.fixture -def model_entity_service(entity_client): +def _mock_files_client(): + """Mock the FilesClient returned by client_from_platform in the permissions module.""" + fileset_output = FilesetOutput( + id="fileset-id-123", + name="test-fileset", + workspace="default", + description="Test fileset", + storage=LocalStorageConfig(path="test-path"), + purpose="generic", + project="test-project", + created_at="2026-01-01T00:00:00Z", + updated_at="2026-01-01T00:00:00Z", + custom_fields={"key": "value"}, + metadata=FilesetMetadata(), + ) + mock_fc = AsyncMock() + mock_response = MagicMock() + mock_response.data.return_value = fileset_output + mock_fc.get_fileset.return_value = mock_response + with patch("nmp.core.models.api.permissions.client_from_platform", return_value=mock_fc): + yield mock_fc + + +@pytest.fixture +def model_entity_service(entity_client, _mock_files_client): """Create a ModelEntityService with MockEntityClient for integration testing.""" async_sdk = AsyncMock(spec=AsyncNeMoPlatform) async_sdk.files.list = AsyncMock( - return_value=ListFilesResponse( + return_value=ListFilesetFilesResponse( data=[ - FilesetFile( - id="file-id-123", + FilesetFileOutput( file_ref="file-ref-123", file_url="file-url-123", path="path-123", @@ -57,21 +80,6 @@ def model_entity_service(entity_client): ] ) ) - async_sdk.files.filesets.retrieve = AsyncMock( - return_value=Fileset( - id="fileset-id-123", - name="test-fileset", - workspace="default", - description="Test fileset", - storage=LocalStorageConfig(path="test-path"), - purpose="generic", - project="test-project", - created_at="2026-01-01T00:00:00Z", - updated_at="2026-01-01T00:00:00Z", - custom_fields={"key": "value"}, - metadata=FilesetMetadata(), - ) - ) return ModelEntityService(entity_client, sdk=async_sdk) diff --git a/services/core/models/tests/integration/test_models_with_auth.py b/services/core/models/tests/integration/test_models_with_auth.py index bc63c6458a..9545aa93f4 100644 --- a/services/core/models/tests/integration/test_models_with_auth.py +++ b/services/core/models/tests/integration/test_models_with_auth.py @@ -23,6 +23,9 @@ import pytest from nemo_platform import NeMoPlatform, PermissionDeniedError +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.core.auth.app.bundle import build_authorization_data as _real_build_authorization_data from nmp.core.files.service import FilesService from nmp.core.models.config import config as models_config @@ -1151,7 +1154,9 @@ def test_editor_can_create_model_with_fileset(self, sdk: NeMoPlatform): admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.files.filesets.create(workspace=workspace, name=fileset_name) + client_from_platform(admin_sdk, FilesClient).create_fileset( + workspace=workspace, body=CreateFilesetRequest(name=fileset_name) + ) admin_sdk.files.upload_content( content=b"x", remote_path="placeholder.txt", fileset=fileset_name, workspace=workspace ) @@ -1179,7 +1184,9 @@ def test_editor_can_update_model_with_fileset(self, sdk: NeMoPlatform): admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) admin_sdk.models.create(workspace=workspace, name=model_name) - admin_sdk.files.filesets.create(workspace=workspace, name=fileset_name) + client_from_platform(admin_sdk, FilesClient).create_fileset( + workspace=workspace, body=CreateFilesetRequest(name=fileset_name) + ) admin_sdk.files.upload_content( content=b"x", remote_path="placeholder.txt", fileset=fileset_name, workspace=workspace ) @@ -1207,7 +1214,9 @@ def test_editor_can_create_adapter_with_fileset(self, sdk: NeMoPlatform): admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) admin_sdk.models.create(workspace=workspace, name=model_name) - admin_sdk.files.filesets.create(workspace=workspace, name=fileset_name) + client_from_platform(admin_sdk, FilesClient).create_fileset( + workspace=workspace, body=CreateFilesetRequest(name=fileset_name) + ) admin_sdk.files.upload_content( content=b"x", remote_path="placeholder.txt", fileset=fileset_name, workspace=workspace ) @@ -1345,10 +1354,9 @@ def test_create_model_trust_remote_code_true_has_permission_succeeds(self, sdk: admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.files.filesets.create( + client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, - name=fileset_name, - storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}, + body=CreateFilesetRequest(name=fileset_name, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}), ) grant_workspace_role( admin_sdk, @@ -1376,10 +1384,11 @@ def test_create_model_trust_remote_code_true_without_permission_raises(self, sdk admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.files.filesets.create( + client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, - name=fileset_name, - storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}, + body=CreateFilesetRequest( + name=fileset_name, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"} + ), ) grant_workspace_role( admin_sdk, @@ -1409,10 +1418,9 @@ def test_update_model_trust_remote_code_true_has_permission_succeeds(self, sdk: admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) admin_sdk.models.create(workspace=workspace, name=model_name) - admin_sdk.files.filesets.create( + client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, - name=fileset_name, - storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}, + body=CreateFilesetRequest(name=fileset_name, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}), ) grant_workspace_role( admin_sdk, @@ -1442,10 +1450,11 @@ def test_update_model_trust_remote_code_true_without_permission_raises(self, sdk admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) admin_sdk.models.create(workspace=workspace, name=model_name) - admin_sdk.files.filesets.create( + client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, - name=fileset_name, - storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}, + body=CreateFilesetRequest( + name=fileset_name, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"} + ), ) grant_workspace_role( admin_sdk, @@ -1479,10 +1488,12 @@ def test_update_model_new_fileset_not_trusted_raises_permission_error(self, sdk: admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) # Model created with a trusted fileset (on allow list) so it has trust_remote_code=True. - admin_sdk.files.filesets.create( + client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, - name=trusted_fs, - storage={"type": "huggingface", "repo_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"}, + body=CreateFilesetRequest( + name=trusted_fs, + storage={"type": "huggingface", "repo_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"}, + ), ) admin_sdk.models.create( workspace=workspace, @@ -1491,10 +1502,9 @@ def test_update_model_new_fileset_not_trusted_raises_permission_error(self, sdk: trust_remote_code=True, ) # New fileset resolves to a repo not on the allow list. - admin_sdk.files.filesets.create( + client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, - name=new_fs, - storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}, + body=CreateFilesetRequest(name=new_fs, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}), ) grant_workspace_role( admin_sdk, @@ -1523,10 +1533,12 @@ def test_exact_match_on_allow_list_succeeds(self, sdk: NeMoPlatform): admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.files.filesets.create( + client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, - name=fileset_name, - storage={"type": "huggingface", "repo_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"}, + body=CreateFilesetRequest( + name=fileset_name, + storage={"type": "huggingface", "repo_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"}, + ), ) grant_workspace_role( admin_sdk, diff --git a/services/core/models/tests/integration/test_workspace_iam_models_isolation.py b/services/core/models/tests/integration/test_workspace_iam_models_isolation.py index 37540c22a3..c43c9873cd 100644 --- a/services/core/models/tests/integration/test_workspace_iam_models_isolation.py +++ b/services/core/models/tests/integration/test_workspace_iam_models_isolation.py @@ -25,6 +25,9 @@ import requests from fastapi.testclient import TestClient from nemo_platform import NeMoPlatform, PermissionDeniedError +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.core.files.service import FilesService from nmp.core.models.service import ModelsService from nmp.core.secrets.service import SecretsService @@ -162,8 +165,9 @@ def test_model_and_adapter_iam(self, sdk: NeMoPlatform) -> None: # Filesets: fileset in C (for allow with group); fileset in D (for deny in C) fs_c = short_unique_name("fs-c") fs_d = short_unique_name("fs-d") - admin.files.filesets.create(workspace=ws_c, name=fs_c) - admin.files.filesets.create(workspace=ws_d, name=fs_d) + admin_files = client_from_platform(admin, FilesClient) + admin_files.create_fileset(workspace=ws_c, body=CreateFilesetRequest(name=fs_c)) + admin_files.create_fileset(workspace=ws_d, body=CreateFilesetRequest(name=fs_d)) admin.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_c, workspace=ws_c) admin.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_d, workspace=ws_d) @@ -338,8 +342,12 @@ def h(email: str, groups: list[str] | None = None) -> dict[str, str]: fs_c = short_unique_name("fs-c") fs_d = short_unique_name("fs-d") admin_sdk = as_user(models_auth_context.sdk, TEST_ADMIN_EMAIL) - admin_sdk.files.filesets.create(workspace=ws_c, name=fs_c) - admin_sdk.files.filesets.create(workspace=ws_d, name=fs_d) + client_from_platform(admin_sdk, FilesClient).create_fileset( + workspace=ws_c, body=CreateFilesetRequest(name=fs_c) + ) + client_from_platform(admin_sdk, FilesClient).create_fileset( + workspace=ws_d, body=CreateFilesetRequest(name=fs_d) + ) admin_sdk.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_c, workspace=ws_c) admin_sdk.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_d, workspace=ws_d) diff --git a/services/core/models/tests/unit/api/test_models_api.py b/services/core/models/tests/unit/api/test_models_api.py index b6ecd48dff..4cc12b11f8 100644 --- a/services/core/models/tests/unit/api/test_models_api.py +++ b/services/core/models/tests/unit/api/test_models_api.py @@ -4,7 +4,7 @@ """Tests for Model (ModelEntity) API endpoints.""" from datetime import datetime -from unittest.mock import AsyncMock, MagicMock, Mock +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest from fastapi import FastAPI @@ -53,7 +53,11 @@ def mock_auth_client(): @pytest.fixture def mock_sdk(): """Create a mock SDK for create/update endpoints that depend on get_sdk_client.""" - return AsyncMock() + sdk = AsyncMock() + sdk._custom_headers = {"authorization": "Bearer test"} + sdk.base_url = "http://localhost:8080" + sdk.workspace = "default" + return sdk @pytest.fixture @@ -455,14 +459,19 @@ def test_create_model_adapter_entity_validation_error_returns_422(client, mock_a """Test that entity store validation errors during adapter creation return 422.""" mock_adapter_entity_service.create_adapter.side_effect = EntityValidationError("adapter name invalid") - response = client.post( - "/apis/models/v2/workspaces/nvidia/models/my-model/adapters", - json={ - "name": "my-adapter", - "fileset": "nvidia/my-fileset", - "finetuning_type": "lora", - }, - ) + with patch("nmp.core.models.api.permissions.client_from_platform") as mock_cfp: + mock_files = AsyncMock() + mock_files.get_fileset.return_value = MagicMock() + mock_cfp.return_value = mock_files + + response = client.post( + "/apis/models/v2/workspaces/nvidia/models/my-model/adapters", + json={ + "name": "my-adapter", + "fileset": "nvidia/my-fileset", + "finetuning_type": "lora", + }, + ) assert response.status_code == 422 assert "adapter name invalid" in response.json()["detail"] diff --git a/services/core/models/tests/unit/test_model_entity_service_unit.py b/services/core/models/tests/unit/test_model_entity_service_unit.py index 63b7c264be..405b8099d8 100644 --- a/services/core/models/tests/unit/test_model_entity_service_unit.py +++ b/services/core/models/tests/unit/test_model_entity_service_unit.py @@ -10,15 +10,9 @@ import pytest from nemo_platform import AsyncNeMoPlatform -from nemo_platform.filesets import ListFilesResponse -from nemo_platform.types.files import ( - Fileset, - FilesetFile, - HuggingfaceStorageConfig, - LocalStorageConfig, - NGCStorageConfig, -) -from nemo_platform.types.shared import FilesetMetadata +from nemo_platform_plugin.files.metadata import FilesetMetadata +from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig, LocalStorageConfig, NGCStorageConfig +from nemo_platform_plugin.files.types import FilesetFileOutput, FilesetOutput, ListFilesetFilesResponse from nmp.common.api.common import Page, PaginationData from nmp.common.api.filter import ComparisonOperation, FilterOperator, LogicalOperation from nmp.common.api.parsed_filter import ParsedFilter @@ -189,14 +183,37 @@ def mock_entity_client() -> AsyncMock: @pytest.fixture -def model_entity_service(mock_entity_client): +def _mock_files_client(): + """Mock the FilesClient returned by client_from_platform in the permissions module.""" + fileset_output = FilesetOutput( + id="fileset-id-123", + name="test-fileset", + workspace="default", + description="Test fileset", + storage=LocalStorageConfig(path="test-path"), + purpose="generic", + project="test-project", + created_at="2026-01-01T00:00:00Z", + updated_at="2026-01-01T00:00:00Z", + custom_fields={"key": "value"}, + metadata=FilesetMetadata(), + ) + mock_fc = AsyncMock() + mock_response = MagicMock() + mock_response.data.return_value = fileset_output + mock_fc.get_fileset.return_value = mock_response + with patch("nmp.core.models.api.permissions.client_from_platform", return_value=mock_fc): + yield mock_fc + + +@pytest.fixture +def model_entity_service(mock_entity_client, _mock_files_client): """Create a ModelEntityService with mocked EntityClient.""" async_sdk = AsyncMock(spec=AsyncNeMoPlatform) async_sdk.files.list = AsyncMock( - return_value=ListFilesResponse( + return_value=ListFilesetFilesResponse( data=[ - FilesetFile( - id="file-id-123", + FilesetFileOutput( file_ref="file-ref-123", file_url="file-url-123", path="path-123", @@ -206,21 +223,6 @@ def model_entity_service(mock_entity_client): ] ) ) - async_sdk.files.filesets.retrieve = AsyncMock( - return_value=Fileset( - id="fileset-id-123", - name="test-fileset", - workspace="default", - description="Test fileset", - storage=LocalStorageConfig(path="test-path"), - purpose="generic", - project="test-project", - created_at="2026-01-01T00:00:00Z", - updated_at="2026-01-01T00:00:00Z", - custom_fields={"key": "value"}, - metadata=FilesetMetadata(), - ) - ) return ModelEntityService(mock_entity_client, sdk=async_sdk) @@ -1161,9 +1163,9 @@ async def test_list_adapters_resolves_parent_models_with_canonical_all_workspace assert mock_entity_client.list.call_args_list[1].kwargs["workspace"] == ALL_WORKSPACES -def _hf_fileset(repo_id: str) -> Fileset: - """Create a Fileset with HuggingFace storage for is_trusted_repo_id tests.""" - return Fileset( +def _hf_fileset(repo_id: str) -> FilesetOutput: + """Create a FilesetOutput with HuggingFace storage for is_trusted_repo_id tests.""" + return FilesetOutput( id="fs-1", name="test-fileset", workspace="default", @@ -1178,9 +1180,9 @@ def _hf_fileset(repo_id: str) -> Fileset: ) -def _ngc_fileset(org: str, team: str, target: str) -> Fileset: - """Create a Fileset with NGC storage for is_trusted_repo_id tests (path: org/team/target).""" - return Fileset( +def _ngc_fileset(org: str, team: str, target: str) -> FilesetOutput: + """Create a FilesetOutput with NGC storage for is_trusted_repo_id tests (path: org/team/target).""" + return FilesetOutput( id="fs-1", name="test-fileset", workspace="default", @@ -1280,7 +1282,7 @@ async def test_is_trusted_repo_id_no_match(model_entity_service): async def test_is_trusted_repo_id_non_huggingface_storage(model_entity_service): """When fileset storage is not huggingface or ngc, returns False.""" # Arrange - fileset = Fileset( + fileset = FilesetOutput( id="fs-1", name="test-fileset", workspace="default", diff --git a/services/hello-world/src/nmp/hello_world/tasks/access_fileset/run.py b/services/hello-world/src/nmp/hello_world/tasks/access_fileset/run.py index ba57edb1e6..470fe14b4a 100644 --- a/services/hello-world/src/nmp/hello_world/tasks/access_fileset/run.py +++ b/services/hello-world/src/nmp/hello_world/tasks/access_fileset/run.py @@ -7,7 +7,10 @@ retrieve a fileset and reports whether access was granted or denied. """ -from nemo_platform import APIStatusError, NeMoPlatform +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError +from nemo_platform_plugin.files.client import FilesClient from nmp.common.jobs.config import get_task_config from nmp.common.sdk_factory import get_platform_sdk from pydantic import BaseModel @@ -38,12 +41,13 @@ def run(*, sdk: NeMoPlatform | None = None) -> int: print(f"Attempting to access fileset '{config.fileset}' in workspace '{config.workspace}'") - fileset = sdk.files.filesets.retrieve(workspace=config.workspace, name=config.fileset) + files = client_from_platform(sdk, FilesClient) + fileset = files.get_fileset(workspace=config.workspace, name=config.fileset).data() print(f"Successfully accessed fileset: {fileset.name}") return 0 - except APIStatusError as e: + except NemoHTTPError as e: if e.status_code == 403: print(f"Access denied (403 Forbidden): {e}") elif e.status_code == 404: diff --git a/services/rl/src/nmp/rl/tasks/file_io/run.py b/services/rl/src/nmp/rl/tasks/file_io/run.py index 581567d2c2..fa86215864 100644 --- a/services/rl/src/nmp/rl/tasks/file_io/run.py +++ b/services/rl/src/nmp/rl/tasks/file_io/run.py @@ -23,12 +23,21 @@ from nemo_platform import ( APIConnectionError, APITimeoutError, - ConflictError, InternalServerError, NeMoPlatform, NotFoundError, ) from nemo_platform.types.files.fileset_file import FilesetFile +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ( + ConflictError, +) +from nemo_platform_plugin.client.errors import ( + InternalServerError as ClientInternalServerError, +) +from nemo_platform_plugin.client.types import RetryPolicy +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest, UpdateFilesetRequest from nmp.common.jobs.schemas import PlatformJobStatus from nmp.common.sdk_factory import get_task_sdk from nmp.customization_common.schemas.file_io import ( @@ -67,8 +76,7 @@ # filter filesets by training backend. SERVICE_SOURCE = "rl" -# Timeout configurations for SDK operations (httpx.Timeout for API calls). -CREATE_FILESET_TIMEOUT = httpx.Timeout(10.0, connect=10.0) +CREATE_FILESET_TIMEOUT = 10.0 LIST_FILES_TIMEOUT = httpx.Timeout(10.0, connect=10.0) # Timeout configurations for FilesetFileSystem operations. Passed via @@ -317,34 +325,41 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N @retry( stop=stop_after_attempt(MAX_RETRIES), wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type((InternalServerError, APITimeoutError, APIConnectionError)), + retry=retry_if_exception_type( + ( + InternalServerError, + APITimeoutError, + APIConnectionError, + ClientInternalServerError, + httpx.TimeoutException, + httpx.ConnectError, + ) + ), reraise=True, ) def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None = None) -> None: """Internal method with retry logic for creating a FileSet.""" + files = client_from_platform(self.sdk, FilesClient).with_options( + timeout=CREATE_FILESET_TIMEOUT, retry=RetryPolicy(max_retries=0) + ) try: - create_kwargs: dict = { - "workspace": fileset.workspace, + body_kwargs: dict = { "name": fileset.name, - "timeout": CREATE_FILESET_TIMEOUT, "custom_fields": {"service_source": SERVICE_SOURCE}, } if metadata is not None: - create_kwargs["metadata"] = metadata - result = self.sdk.with_options(max_retries=0).files.filesets.create(**create_kwargs) + body_kwargs["metadata"] = metadata + result = files.create_fileset(workspace=fileset.workspace, body=CreateFilesetRequest(**body_kwargs)).data() logger.info(f"Created FileSet: {result.workspace}/{result.name}") except ConflictError: - # Fileset already exists — patch metadata so tool_calling etc. aren't lost. workspace = fileset.workspace or self.job_ctx.workspace if metadata is not None: - update_kwargs: dict = { - "name": fileset.name, - "workspace": workspace, - "metadata": metadata, - "timeout": CREATE_FILESET_TIMEOUT, - } try: - self.sdk.with_options(max_retries=0).files.filesets.update(**update_kwargs) + files.update_fileset( + workspace=workspace, + name=fileset.name, + body=UpdateFilesetRequest(metadata=metadata), + ) logger.info(f"Patched existing FileSet metadata: {workspace}/{fileset.name}") except Exception as e: logger.warning( diff --git a/services/rl/src/nmp/rl/tasks/model_entity/run.py b/services/rl/src/nmp/rl/tasks/model_entity/run.py index acb3888e0a..192fc36f58 100644 --- a/services/rl/src/nmp/rl/tasks/model_entity/run.py +++ b/services/rl/src/nmp/rl/tasks/model_entity/run.py @@ -20,6 +20,7 @@ import time from pathlib import Path +import httpx from nemo_platform import ( APIConnectionError, APITimeoutError, @@ -37,6 +38,9 @@ ) from nemo_platform.types.models import LoraParam, ModelEntity from nemo_platform.types.shared_params.tool_call_config import ToolCallConfig as ToolCallConfigParam +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import InternalServerError as ClientInternalServerError +from nemo_platform_plugin.files.client import FilesClient from nmp.common.sdk_factory import get_task_sdk from nmp.customization_common.schemas.model_entity import ( DeploymentParameters, @@ -165,12 +169,18 @@ def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, Mode logger.info(f"Validating fileset exists: {fileset_workspace}/{config.fileset.name}") try: - self.sdk.files.filesets.retrieve(workspace=fileset_workspace, name=config.fileset.name) + client_from_platform(self.sdk, FilesClient).get_fileset( + workspace=fileset_workspace, name=config.fileset.name + ) logger.info(f"Fileset validation successful: {fileset_workspace}/{config.fileset.name}") - except (InternalServerError, APITimeoutError, APIConnectionError): - # Transient API failures must propagate so the @retry wrapping - # create_model_entity can retry them, instead of being masked as a - # permanent (non-retryable) ModelEntityCreationError. + except ( + InternalServerError, + APITimeoutError, + APIConnectionError, + ClientInternalServerError, + httpx.TimeoutException, + httpx.ConnectError, + ): raise except Exception as e: logger.error(f"Fileset validation failed: {fileset_workspace}/{config.fileset.name}") diff --git a/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py b/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py index 67bb0d093a..b26bc12842 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py +++ b/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py @@ -23,12 +23,21 @@ from nemo_platform import ( APIConnectionError, APITimeoutError, - ConflictError, InternalServerError, NeMoPlatform, NotFoundError, ) from nemo_platform.types.files.fileset_file import FilesetFile +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ( + ConflictError, +) +from nemo_platform_plugin.client.errors import ( + InternalServerError as ClientInternalServerError, +) +from nemo_platform_plugin.client.types import RetryPolicy +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest, UpdateFilesetRequest from nmp.common.jobs.schemas import PlatformJobStatus from nmp.common.sdk_factory import get_task_sdk from nmp.customization_common.schemas.file_io import ( @@ -67,8 +76,7 @@ # filter filesets by training backend. SERVICE_SOURCE = "unsloth" -# Timeout configurations for SDK operations (httpx.Timeout for API calls). -CREATE_FILESET_TIMEOUT = httpx.Timeout(10.0, connect=10.0) +CREATE_FILESET_TIMEOUT = 10.0 LIST_FILES_TIMEOUT = httpx.Timeout(10.0, connect=10.0) # Timeout configurations for FilesetFileSystem operations. Passed via @@ -291,34 +299,41 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N @retry( stop=stop_after_attempt(MAX_RETRIES), wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type((InternalServerError, APITimeoutError, APIConnectionError)), + retry=retry_if_exception_type( + ( + InternalServerError, + APITimeoutError, + APIConnectionError, + ClientInternalServerError, + httpx.TimeoutException, + httpx.ConnectError, + ) + ), reraise=True, ) def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None = None) -> None: """Internal method with retry logic for creating a FileSet.""" + files = client_from_platform(self.sdk, FilesClient).with_options( + timeout=CREATE_FILESET_TIMEOUT, retry=RetryPolicy(max_retries=0) + ) try: - create_kwargs: dict = { - "workspace": fileset.workspace, + body_kwargs: dict = { "name": fileset.name, - "timeout": CREATE_FILESET_TIMEOUT, "custom_fields": {"service_source": SERVICE_SOURCE}, } if metadata is not None: - create_kwargs["metadata"] = metadata - result = self.sdk.with_options(max_retries=0).files.filesets.create(**create_kwargs) + body_kwargs["metadata"] = metadata + result = files.create_fileset(workspace=fileset.workspace, body=CreateFilesetRequest(**body_kwargs)).data() logger.info(f"Created FileSet: {result.workspace}/{result.name}") except ConflictError: - # Fileset already exists — patch metadata so tool_calling etc. aren't lost. workspace = fileset.workspace or self.job_ctx.workspace if metadata is not None: - update_kwargs: dict = { - "name": fileset.name, - "workspace": workspace, - "metadata": metadata, - "timeout": CREATE_FILESET_TIMEOUT, - } try: - self.sdk.with_options(max_retries=0).files.filesets.update(**update_kwargs) + files.update_fileset( + workspace=workspace, + name=fileset.name, + body=UpdateFilesetRequest(metadata=metadata), + ) logger.info(f"Patched existing FileSet metadata: {workspace}/{fileset.name}") except Exception as e: logger.warning( diff --git a/services/unsloth/src/nmp/unsloth/tasks/model_entity/run.py b/services/unsloth/src/nmp/unsloth/tasks/model_entity/run.py index 355c5cb089..96874ea20c 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/model_entity/run.py +++ b/services/unsloth/src/nmp/unsloth/tasks/model_entity/run.py @@ -37,6 +37,8 @@ ) from nemo_platform.types.models import LoraParam, ModelEntity from nemo_platform.types.shared_params.tool_call_config import ToolCallConfig as ToolCallConfigParam +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from nmp.common.sdk_factory import get_task_sdk from nmp.customization_common.schemas.model_entity import ( DeploymentParameters, @@ -153,7 +155,9 @@ def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, Mode logger.info(f"Validating fileset exists: {fileset_workspace}/{config.fileset.name}") try: - self.sdk.files.filesets.retrieve(workspace=fileset_workspace, name=config.fileset.name) + client_from_platform(self.sdk, FilesClient).get_fileset( + workspace=fileset_workspace, name=config.fileset.name + ) logger.info(f"Fileset validation successful: {fileset_workspace}/{config.fileset.name}") except Exception as e: logger.error(f"Fileset validation failed: {fileset_workspace}/{config.fileset.name}") diff --git a/services/unsloth/tests/test_file_io.py b/services/unsloth/tests/test_file_io.py index c11963afa7..9832038e03 100644 --- a/services/unsloth/tests/test_file_io.py +++ b/services/unsloth/tests/test_file_io.py @@ -20,7 +20,7 @@ import types from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -44,7 +44,7 @@ def _make_runner(sdk, workspace: str = "default", storage_path: Path | None = No return FileIORunner(sdk=sdk, progress_reporter=NoOpProgressReporter(), job_ctx=job_ctx) -def _make_sdk(*, conflict_on_create: bool = False) -> MagicMock: +def _make_sdk() -> MagicMock: """Build a MagicMock SDK with sensible defaults for fluent chaining. ``with_options`` returns the same SDK so chained timeouts don't break @@ -52,15 +52,6 @@ def _make_sdk(*, conflict_on_create: bool = False) -> MagicMock: """ sdk = MagicMock() sdk.with_options.return_value = sdk - - if conflict_on_create: - # Trigger ConflictError on filesets.create by raising the class the - # runner is bound against (see comment in _raise_runner_conflict). - def _raise_conflict(**_kwargs): - _raise_runner_conflict() - - sdk.files.filesets.create.side_effect = _raise_conflict - return sdk @@ -93,62 +84,90 @@ def _make_dir(tmp_path: Path) -> Path: class TestCreateFileset: - def test_creates_fileset_with_service_source_and_metadata(self) -> None: + @patch("nmp.unsloth.tasks.file_io.run.client_from_platform") + def test_creates_fileset_with_service_source_and_metadata(self, mock_cfp) -> None: + from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.customization_common.schemas.file_io import FileSetRef + mock_fc = MagicMock() + mock_cfp.return_value = mock_fc sdk = _make_sdk() runner = _make_runner(sdk) - metadata = {"model": "Qwen/Qwen3-0.6B", "save_method": "lora"} + metadata = {"model": {"tool_calling": {"tool_call_parser": "llama3_json"}}} dest = FileSetRef(workspace="default", name="qwen-test") runner.create_fileset(dest, metadata=metadata) - sdk.files.filesets.create.assert_called_once() - call = sdk.files.filesets.create.call_args + mock_fc.create_fileset.assert_called_once() + call = mock_fc.create_fileset.call_args assert call.kwargs["workspace"] == "default" - assert call.kwargs["name"] == "qwen-test" - assert call.kwargs["custom_fields"] == {"service_source": "unsloth"} - assert call.kwargs["metadata"] == metadata - - def test_conflict_patches_metadata_on_existing(self) -> None: + body = call.kwargs["body"] + assert isinstance(body, CreateFilesetRequest) + assert body.name == "qwen-test" + assert body.custom_fields == {"service_source": "unsloth"} + assert body.metadata is not None + assert body.metadata.model is not None + assert body.metadata.model.tool_calling.tool_call_parser == "llama3_json" + + @patch("nmp.unsloth.tasks.file_io.run.client_from_platform") + def test_conflict_patches_metadata_on_existing(self, mock_cfp) -> None: + from nemo_platform_plugin.files.types import UpdateFilesetRequest from nmp.customization_common.schemas.file_io import FileSetRef - sdk = _make_sdk(conflict_on_create=True) + mock_fc = MagicMock() + mock_fc.create_fileset.side_effect = lambda **_: _raise_runner_conflict() + mock_cfp.return_value = mock_fc + sdk = _make_sdk() runner = _make_runner(sdk) dest = FileSetRef(workspace="default", name="exists") + metadata = {"model": {"tool_calling": {"tool_call_parser": "hermes"}}} - runner.create_fileset(dest, metadata={"model": "x"}) + runner.create_fileset(dest, metadata=metadata) - sdk.files.filesets.update.assert_called_once() - update_call = sdk.files.filesets.update.call_args + mock_fc.update_fileset.assert_called_once() + update_call = mock_fc.update_fileset.call_args assert update_call.kwargs["workspace"] == "default" assert update_call.kwargs["name"] == "exists" - assert update_call.kwargs["metadata"] == {"model": "x"} - - def test_conflict_no_metadata_skips_update(self) -> None: + body = update_call.kwargs["body"] + assert isinstance(body, UpdateFilesetRequest) + assert body.metadata is not None + assert body.metadata.model is not None + assert body.metadata.model.tool_calling.tool_call_parser == "hermes" + + @patch("nmp.unsloth.tasks.file_io.run.client_from_platform") + def test_conflict_no_metadata_skips_update(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef - sdk = _make_sdk(conflict_on_create=True) + mock_fc = MagicMock() + mock_fc.create_fileset.side_effect = lambda **_: _raise_runner_conflict() + mock_cfp.return_value = mock_fc + sdk = _make_sdk() runner = _make_runner(sdk) dest = FileSetRef(workspace="default", name="exists") runner.create_fileset(dest, metadata=None) - sdk.files.filesets.update.assert_not_called() + mock_fc.update_fileset.assert_not_called() + @patch("nmp.unsloth.tasks.file_io.run.client_from_platform") def test_update_failure_is_warning_not_fatal( self, + mock_cfp, caplog: pytest.LogCaptureFixture, ) -> None: from nmp.customization_common.schemas.file_io import FileSetRef - sdk = _make_sdk(conflict_on_create=True) - sdk.files.filesets.update.side_effect = RuntimeError("backend down") + mock_fc = MagicMock() + mock_fc.create_fileset.side_effect = lambda **_: _raise_runner_conflict() + mock_fc.update_fileset.side_effect = RuntimeError("backend down") + mock_cfp.return_value = mock_fc + sdk = _make_sdk() runner = _make_runner(sdk) dest = FileSetRef(workspace="default", name="exists") + metadata = {"model": {"tool_calling": {"tool_call_parser": "hermes"}}} with caplog.at_level("WARNING"): - runner.create_fileset(dest, metadata={"model": "x"}) + runner.create_fileset(dest, metadata=metadata) assert any("Could not patch metadata" in r.getMessage() for r in caplog.records) diff --git a/services/unsloth/tests/test_model_entity.py b/services/unsloth/tests/test_model_entity.py index aedf1115f7..5a817cceec 100644 --- a/services/unsloth/tests/test_model_entity.py +++ b/services/unsloth/tests/test_model_entity.py @@ -16,7 +16,7 @@ import types from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -105,11 +105,14 @@ def test_caps_length_below_60_and_strips_trailing_hyphen(self) -> None: class TestCreateFullEntity: - def test_creates_model_entity_for_full_sft(self) -> None: + @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + def test_creates_model_entity_for_full_sft(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig sdk = _make_sdk() + mock_fc = MagicMock() + mock_cfp.return_value = mock_fc sdk.models.retrieve.return_value = _model_entity(name="base-model") new_me = _model_entity(name="trained-model") sdk.models.create.return_value = new_me @@ -125,18 +128,18 @@ def test_creates_model_entity_for_full_sft(self) -> None: result, deploy_target = runner.create_model_entity(config) - sdk.files.filesets.retrieve.assert_called_once_with(workspace="default", name="trained-model") + mock_fc.get_fileset.assert_called_once_with(workspace="default", name="trained-model") sdk.models.create.assert_called_once() assert deploy_target is new_me - # ``result`` is the output of ``new_me.model_dump()`` — we just assert - # we got *something* back; the actual shape is controlled by the SDK. assert result is not None - def test_conflict_falls_back_to_update(self) -> None: + @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + def test_conflict_falls_back_to_update(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig sdk = _make_sdk() + mock_cfp.return_value = MagicMock() sdk.models.retrieve.return_value = _model_entity(name="base-model") sdk.models.create.side_effect = lambda **_: _raise_runner_conflict() sdk.models.update.return_value = _model_entity(name="trained-model") @@ -157,12 +160,15 @@ def test_conflict_falls_back_to_update(self) -> None: assert update_call.kwargs["name"] == "trained-model" assert update_call.kwargs["workspace"] == "default" - def test_missing_fileset_raises_creation_error(self) -> None: + @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + def test_missing_fileset_raises_creation_error(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityCreationError, ModelEntityTaskConfig sdk = _make_sdk() - sdk.files.filesets.retrieve.side_effect = RuntimeError("fileset missing") + mock_fc = MagicMock() + mock_fc.get_fileset.side_effect = RuntimeError("fileset missing") + mock_cfp.return_value = mock_fc runner = _make_runner(sdk) config = ModelEntityTaskConfig( name="x", @@ -181,12 +187,14 @@ def test_missing_fileset_raises_creation_error(self) -> None: class TestCreateAdapter: - def test_creates_adapter_for_lora(self) -> None: + @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + def test_creates_adapter_for_lora(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig, PEFTConfig from nmp.unsloth.entities.values import FinetuningType sdk = _make_sdk() + mock_cfp.return_value = MagicMock() base_me = _model_entity(name="base-model") sdk.models.retrieve.return_value = base_me sdk.models.adapters.create.return_value = _model_entity(name="adapter-x") @@ -203,15 +211,16 @@ def test_creates_adapter_for_lora(self) -> None: _result, deploy_target = runner.create_model_entity(config) sdk.models.adapters.create.assert_called_once() - # For LoRA, the deploy target is the BASE model, not the adapter. assert deploy_target is base_me - def test_adapter_conflict_falls_back_to_update(self) -> None: + @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + def test_adapter_conflict_falls_back_to_update(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig, PEFTConfig from nmp.unsloth.entities.values import FinetuningType sdk = _make_sdk() + mock_cfp.return_value = MagicMock() sdk.models.retrieve.return_value = _model_entity(name="base-model") sdk.models.adapters.create.side_effect = lambda **_: _raise_runner_conflict() sdk.models.adapters.update.return_value = _model_entity(name="adapter-x") diff --git a/tests/agentic-use/customizer-lora-job-cli/tests/test_outputs.py b/tests/agentic-use/customizer-lora-job-cli/tests/test_outputs.py index 68b858f335..7371113534 100644 --- a/tests/agentic-use/customizer-lora-job-cli/tests/test_outputs.py +++ b/tests/agentic-use/customizer-lora-job-cli/tests/test_outputs.py @@ -14,6 +14,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient WORKSPACE = "lora-training-workspace" FILESET = "sft-training-data" @@ -42,6 +44,11 @@ def client() -> NeMoPlatform: ) +@pytest.fixture +def files_client(client: NeMoPlatform) -> FilesClient: + return client_from_platform(client, FilesClient) + + def _list_automodel_jobs(client: NeMoPlatform) -> list[dict[str, Any]]: """List automodel customization jobs in the eval workspace.""" url = f"{str(client.base_url).rstrip('/')}/apis/customization/v2/workspaces/{WORKSPACE}/automodel/jobs" @@ -59,10 +66,9 @@ def test_workspace_exists(client: NeMoPlatform): assert WORKSPACE in workspace_names, f"Workspace '{WORKSPACE}' not found. Found: {workspace_names}" -def test_fileset_exists(client: NeMoPlatform): +def test_fileset_exists(files_client: FilesClient): """Verify the sft-training-data fileset was created.""" - response = client.files.filesets.list(workspace=WORKSPACE) - fileset_names = [fs.name for fs in response.data] + fileset_names = [fs.name for fs in files_client.list_filesets(workspace=WORKSPACE).page().items] assert FILESET in fileset_names, f"Fileset '{FILESET}' not found. Found: {fileset_names}" diff --git a/tests/agentic-use/evaluator-llm-judge-cli-easy/tests/test_outputs.py b/tests/agentic-use/evaluator-llm-judge-cli-easy/tests/test_outputs.py index 3223c87502..fca01280af 100644 --- a/tests/agentic-use/evaluator-llm-judge-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/evaluator-llm-judge-cli-easy/tests/test_outputs.py @@ -23,6 +23,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient sys.path.insert(0, "/tests/shared") from trace_reader import get_session @@ -52,14 +54,17 @@ def _get_nmp_client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE, access_token=_make_unsigned_jwt()) +def _get_files_client() -> FilesClient: + return client_from_platform(_get_nmp_client(), FilesClient) + + # --- Dataset checks --- def test_fileset_exists() -> None: """Verify the judge-eval-dataset fileset was created.""" - client = _get_nmp_client() - response = client.files.filesets.list() - fileset_names = [fs.name for fs in response.data] + files_client = _get_files_client() + fileset_names = [fs.name for fs in files_client.list_filesets().page().items] assert FILESET in fileset_names, f"Fileset '{FILESET}' not found. Found: {fileset_names}" diff --git a/tests/agentic-use/evaluator-llm-judge-cli/tests/test_outputs.py b/tests/agentic-use/evaluator-llm-judge-cli/tests/test_outputs.py index 3223c87502..fca01280af 100644 --- a/tests/agentic-use/evaluator-llm-judge-cli/tests/test_outputs.py +++ b/tests/agentic-use/evaluator-llm-judge-cli/tests/test_outputs.py @@ -23,6 +23,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient sys.path.insert(0, "/tests/shared") from trace_reader import get_session @@ -52,14 +54,17 @@ def _get_nmp_client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE, access_token=_make_unsigned_jwt()) +def _get_files_client() -> FilesClient: + return client_from_platform(_get_nmp_client(), FilesClient) + + # --- Dataset checks --- def test_fileset_exists() -> None: """Verify the judge-eval-dataset fileset was created.""" - client = _get_nmp_client() - response = client.files.filesets.list() - fileset_names = [fs.name for fs in response.data] + files_client = _get_files_client() + fileset_names = [fs.name for fs in files_client.list_filesets().page().items] assert FILESET in fileset_names, f"Fileset '{FILESET}' not found. Found: {fileset_names}" diff --git a/tests/agentic-use/evaluator-simple-job-cli-easy/tests/test_outputs.py b/tests/agentic-use/evaluator-simple-job-cli-easy/tests/test_outputs.py index c2823ea973..cbd813f6a8 100644 --- a/tests/agentic-use/evaluator-simple-job-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/evaluator-simple-job-cli-easy/tests/test_outputs.py @@ -11,6 +11,8 @@ import os from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient WORKSPACE = "eval-test-workspace" FILESET = "eval-dataset" @@ -21,6 +23,10 @@ def _get_client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url) +def _get_files_client() -> FilesClient: + return client_from_platform(_get_client(), FilesClient) + + def test_workspace_exists(): """Verify the eval-test-workspace was created.""" client = _get_client() @@ -31,9 +37,8 @@ def test_workspace_exists(): def test_fileset_exists(): """Verify the eval-dataset fileset was created.""" - client = _get_client() - response = client.files.filesets.list(workspace=WORKSPACE) - fileset_names = [fs.name for fs in response.data] + files_client = _get_files_client() + fileset_names = [fs.name for fs in files_client.list_filesets(workspace=WORKSPACE).page().items] assert FILESET in fileset_names, f"Fileset '{FILESET}' not found. Found: {fileset_names}" diff --git a/tests/agentic-use/evaluator-simple-job-cli/tests/test_outputs.py b/tests/agentic-use/evaluator-simple-job-cli/tests/test_outputs.py index c2823ea973..cbd813f6a8 100644 --- a/tests/agentic-use/evaluator-simple-job-cli/tests/test_outputs.py +++ b/tests/agentic-use/evaluator-simple-job-cli/tests/test_outputs.py @@ -11,6 +11,8 @@ import os from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient WORKSPACE = "eval-test-workspace" FILESET = "eval-dataset" @@ -21,6 +23,10 @@ def _get_client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url) +def _get_files_client() -> FilesClient: + return client_from_platform(_get_client(), FilesClient) + + def test_workspace_exists(): """Verify the eval-test-workspace was created.""" client = _get_client() @@ -31,9 +37,8 @@ def test_workspace_exists(): def test_fileset_exists(): """Verify the eval-dataset fileset was created.""" - client = _get_client() - response = client.files.filesets.list(workspace=WORKSPACE) - fileset_names = [fs.name for fs in response.data] + files_client = _get_files_client() + fileset_names = [fs.name for fs in files_client.list_filesets(workspace=WORKSPACE).page().items] assert FILESET in fileset_names, f"Fileset '{FILESET}' not found. Found: {fileset_names}" diff --git a/tests/agentic-use/evaluator-tool-calling-cli/tests/test_outputs.py b/tests/agentic-use/evaluator-tool-calling-cli/tests/test_outputs.py index 233dc7459b..2c95c119ea 100644 --- a/tests/agentic-use/evaluator-tool-calling-cli/tests/test_outputs.py +++ b/tests/agentic-use/evaluator-tool-calling-cli/tests/test_outputs.py @@ -17,6 +17,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient sys.path.insert(0, "/tests/shared") from trace_reader import get_session @@ -44,6 +46,10 @@ def _get_client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE, access_token=_make_unsigned_jwt()) +def _get_files_client() -> FilesClient: + return client_from_platform(_get_client(), FilesClient) + + # --- Workspace checks --- @@ -60,9 +66,8 @@ def test_workspace_exists(): def test_fileset_exists(): """Verify the tool-calling-dataset fileset was created.""" - client = _get_client() - response = client.files.filesets.list() - fileset_names = [fs.name for fs in response.data] + files_client = _get_files_client() + fileset_names = [fs.name for fs in files_client.list_filesets().page().items] assert FILESET in fileset_names, f"Fileset '{FILESET}' not found. Found: {fileset_names}" diff --git a/tests/agentic-use/evaluator-zero-config-judge-cli/tests/test_outputs.py b/tests/agentic-use/evaluator-zero-config-judge-cli/tests/test_outputs.py index 59aa52958a..c07eda4270 100644 --- a/tests/agentic-use/evaluator-zero-config-judge-cli/tests/test_outputs.py +++ b/tests/agentic-use/evaluator-zero-config-judge-cli/tests/test_outputs.py @@ -29,6 +29,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient sys.path.insert(0, "/tests/shared") from trace_reader import get_session @@ -58,14 +60,17 @@ def _get_nmp_client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE, access_token=_make_unsigned_jwt()) +def _get_files_client() -> FilesClient: + return client_from_platform(_get_nmp_client(), FilesClient) + + # --- Dataset checks --- def test_fileset_exists() -> None: """Verify the zeroconfig-dataset fileset was created.""" - client = _get_nmp_client() - response = client.files.filesets.list() - fileset_names = [fs.name for fs in response.data] + files_client = _get_files_client() + fileset_names = [fs.name for fs in files_client.list_filesets().page().items] assert FILESET in fileset_names, f"Fileset '{FILESET}' not found. Found: {fileset_names}" diff --git a/tests/agentic-use/files-crud-cli-easy/tests/test_outputs.py b/tests/agentic-use/files-crud-cli-easy/tests/test_outputs.py index e2371acd04..28683cf225 100644 --- a/tests/agentic-use/files-crud-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/files-crud-cli-easy/tests/test_outputs.py @@ -15,6 +15,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from trace_reader import get_session WORKSPACE = "default" @@ -26,18 +28,22 @@ def client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE) -def test_harbor_test_fileset_deleted(client: NeMoPlatform) -> None: +@pytest.fixture +def files_client(client: NeMoPlatform) -> FilesClient: + return client_from_platform(client, FilesClient) + + +def test_harbor_test_fileset_deleted(files_client: FilesClient) -> None: """Test that harbor-test-fileset was deleted after CRUD operations.""" - response = client.files.filesets.list() - fileset_names = [fs.name for fs in response.data] + fileset_names = [fs.name for fs in files_client.list_filesets().page().items] assert "harbor-test-fileset" not in fileset_names, ( f"Fileset 'harbor-test-fileset' should have been deleted but still exists! Found: {fileset_names}" ) -def test_harbor_final_fileset_exists(client: NeMoPlatform) -> None: +def test_harbor_final_fileset_exists(files_client: FilesClient) -> None: """Test that harbor-final-fileset was created and has correct metadata.""" - response = client.files.filesets.retrieve(name="harbor-final-fileset") + response = files_client.get_fileset(name="harbor-final-fileset").data() assert response.name == "harbor-final-fileset", ( f"Expected fileset name 'harbor-final-fileset', got '{response.name}'" ) diff --git a/tests/agentic-use/files-crud-cli/tests/test_outputs.py b/tests/agentic-use/files-crud-cli/tests/test_outputs.py index e2371acd04..28683cf225 100644 --- a/tests/agentic-use/files-crud-cli/tests/test_outputs.py +++ b/tests/agentic-use/files-crud-cli/tests/test_outputs.py @@ -15,6 +15,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from trace_reader import get_session WORKSPACE = "default" @@ -26,18 +28,22 @@ def client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE) -def test_harbor_test_fileset_deleted(client: NeMoPlatform) -> None: +@pytest.fixture +def files_client(client: NeMoPlatform) -> FilesClient: + return client_from_platform(client, FilesClient) + + +def test_harbor_test_fileset_deleted(files_client: FilesClient) -> None: """Test that harbor-test-fileset was deleted after CRUD operations.""" - response = client.files.filesets.list() - fileset_names = [fs.name for fs in response.data] + fileset_names = [fs.name for fs in files_client.list_filesets().page().items] assert "harbor-test-fileset" not in fileset_names, ( f"Fileset 'harbor-test-fileset' should have been deleted but still exists! Found: {fileset_names}" ) -def test_harbor_final_fileset_exists(client: NeMoPlatform) -> None: +def test_harbor_final_fileset_exists(files_client: FilesClient) -> None: """Test that harbor-final-fileset was created and has correct metadata.""" - response = client.files.filesets.retrieve(name="harbor-final-fileset") + response = files_client.get_fileset(name="harbor-final-fileset").data() assert response.name == "harbor-final-fileset", ( f"Expected fileset name 'harbor-final-fileset', got '{response.name}'" ) diff --git a/tests/agentic-use/files-upload-dataset-cli-easy/tests/test_outputs.py b/tests/agentic-use/files-upload-dataset-cli-easy/tests/test_outputs.py index 05270db0cc..c758e60086 100644 --- a/tests/agentic-use/files-upload-dataset-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/files-upload-dataset-cli-easy/tests/test_outputs.py @@ -16,6 +16,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from trace_reader import get_session WORKSPACE = "default" @@ -33,9 +35,14 @@ def client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE) -def test_fileset_exists(client: NeMoPlatform) -> None: +@pytest.fixture +def files_client(client: NeMoPlatform) -> FilesClient: + return client_from_platform(client, FilesClient) + + +def test_fileset_exists(files_client: FilesClient) -> None: """Test that harbor-dataset-fileset was created with correct metadata.""" - response = client.files.filesets.retrieve(name=FILESET_NAME) + response = files_client.get_fileset(name=FILESET_NAME).data() assert response.name == FILESET_NAME, f"Expected fileset name '{FILESET_NAME}', got '{response.name}'" assert response.description == "Dataset fileset for harbor eval", ( f"Expected description 'Dataset fileset for harbor eval', got '{response.description}'" diff --git a/tests/agentic-use/files-upload-dataset-cli/tests/test_outputs.py b/tests/agentic-use/files-upload-dataset-cli/tests/test_outputs.py index 05270db0cc..c758e60086 100644 --- a/tests/agentic-use/files-upload-dataset-cli/tests/test_outputs.py +++ b/tests/agentic-use/files-upload-dataset-cli/tests/test_outputs.py @@ -16,6 +16,8 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient from trace_reader import get_session WORKSPACE = "default" @@ -33,9 +35,14 @@ def client() -> NeMoPlatform: return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE) -def test_fileset_exists(client: NeMoPlatform) -> None: +@pytest.fixture +def files_client(client: NeMoPlatform) -> FilesClient: + return client_from_platform(client, FilesClient) + + +def test_fileset_exists(files_client: FilesClient) -> None: """Test that harbor-dataset-fileset was created with correct metadata.""" - response = client.files.filesets.retrieve(name=FILESET_NAME) + response = files_client.get_fileset(name=FILESET_NAME).data() assert response.name == FILESET_NAME, f"Expected fileset name '{FILESET_NAME}', got '{response.name}'" assert response.description == "Dataset fileset for harbor eval", ( f"Expected description 'Dataset fileset for harbor eval', got '{response.description}'" diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/overrides/files/upload.py b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/overrides/files/upload.py index 10a5b08723..9ec52f7aa1 100644 --- a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/overrides/files/upload.py +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/overrides/files/upload.py @@ -7,6 +7,8 @@ import typer from nemo_platform_ext.cli.core.context import CLIContext from nemo_platform_ext.cli.core.errors import handle_errors +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient app = cast(Any, None) # override-skip: provided by generated file @@ -49,6 +51,7 @@ def upload_files( raw_local_path: str = ctx.params.get("local_path") client = state.get_client() + files = client_from_platform(client, FilesClient) if workspace is None: workspace = client._get_workspace_path_param() @@ -57,7 +60,7 @@ def upload_files( with RichProgressCallback(description="Uploading") as callback: if fileset is not None: # Validate fileset exists before uploading - client.files.filesets.retrieve(fileset, workspace=workspace) + files.get_fileset(name=fileset, workspace=workspace) client.files.upload( local_path=raw_local_path, remote_path=remote_path,