feat(client): Create a new entities client#583
Open
matthewgrossman wants to merge 3 commits into
Open
Conversation
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Contributor
📝 WalkthroughWalkthroughThis PR makes retry enabled by default in ChangesRetry, pagination, and entity resource changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant EntityStoreResource
participant AsyncEntitiesClient
participant EntityStoreServer
Caller->>EntityStoreResource: create(entity)
EntityStoreResource->>AsyncEntitiesClient: create_entity(body)
AsyncEntitiesClient->>EntityStoreServer: POST /entities
EntityStoreServer-->>AsyncEntitiesClient: EntityResponse or 409
AsyncEntitiesClient-->>EntityStoreResource: response or ConflictError
EntityStoreResource-->>Caller: EntityT or EntityConflictError
sequenceDiagram
participant NemoClient
participant HTTPServer
participant PaginatedResponse
NemoClient->>HTTPServer: request (DEFAULT_RETRY_POLICY)
HTTPServer-->>NemoClient: 503, retry, then success
NemoClient->>PaginatedResponse: parse page, cache first page
PaginatedResponse->>HTTPServer: fetch next (server-echoed page)
HTTPServer-->>PaginatedResponse: empty items, stop
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Contributor
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (1)
packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/entity_client.py (1)
138-173: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache
TypeAdapterinstances instead of rebuilding per entity.
TypeAdapter(entity_type)andTypeAdapter(attr_type)are constructed fresh on every call — this runs once per item during pagination (_MappedPaginatedResponse._convert/items()), so listing many entities rebuilds adapters repeatedly. Cache by type (e.g.functools.lru_cacheon a small helper, or a module-level dict) to avoid repeated schema-build cost.♻️ Suggested caching approach
+from functools import lru_cache + +@lru_cache(maxsize=None) +def _adapter_for(tp: Any) -> TypeAdapter: + return TypeAdapter(tp) + def _convert_api_entity_to_model(self, entity: EntityResponse, entity_type: EntityTypeLike) -> EntityT: entity_dict = entity.model_dump() entity_dict.update(entity.data) - type_adapter = TypeAdapter(entity_type) + type_adapter = _adapter_for(entity_type) result = type_adapter.validate_python(entity_dict) ... - setattr(result, field_name, TypeAdapter(attr_type).validate_python(raw_value)) + setattr(result, field_name, _adapter_for(attr_type).validate_python(raw_value))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/entity_client.py` around lines 138 - 173, The entity conversion path rebuilds TypeAdapter objects on every item in _convert_api_entity_to_model, which is expensive during paginated listing. Add a cached helper in entity_client.py for both entity_type and attr_type adapter creation (for example via a small lru_cache-backed function or module-level cache), then use that helper instead of calling TypeAdapter(...) directly in _convert_api_entity_to_model. Keep the existing validation logic intact while reusing adapters across repeated calls.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py`:
- Around line 207-211: The default retry policy in the
RetryPolicy/DEFAULT_RETRY_POLICY path is currently applied too broadly and can
replay write request bodies without protection. Update the retry logic so it
only retries safe/idempotent methods by default, and only allows POST/PATCH/PUT
retries when an Idempotency-Key guard is present or equivalent idempotency check
exists. Make the change in the retry-policy setup and any request-classification
helper used by the Nemo platform client types so transient errors like
502/503/504/429 are not blindly retried for non-idempotent writes.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py`:
- Line 17: Remove the TYPE_CHECKING-only handling for NemoEntityClient in
dependencies.py and import it unconditionally from nemo_platform_plugin.entities
alongside EntityClient. Update the dependency provider function that currently
uses the runtime import in the body so its return annotation uses the concrete
NemoEntityClient type instead of a string literal, and keep the symbol names
EntityClient, NemoEntityClient, and the provider function in sync with the
same-package runtime import rule.
- Around line 65-85: `get_nemo_entity_client()` currently creates a fresh async
transport via `get_async_nemo_client()` and wraps it in `NemoEntityClient`
without any cleanup, which can leak sockets over repeated calls. Update the
dependency to manage the client lifecycle explicitly: either reuse a shared
async client instance or make this an async dependency that closes the
`AsyncNemoClient`/`httpx.AsyncClient` after use. Keep the fix localized to
`get_nemo_entity_client`, `get_async_nemo_client`, and the `NemoEntityClient`
construction path so the caller does not retain an unclosed client.
In `@packages/nemo_platform_plugin/tests/client/test_client_options.py`:
- Around line 217-237: The test currently bypasses the constructor default by
passing an explicit RetryPolicy in test_default_policy_retries, so it won’t
catch regressions in NemoClient’s default retry behavior. Update this test to
construct NemoClient without retry= and instead patch the retry backoff/sleep
path to keep it fast, while still asserting the 503 is retried 1 + 3 times via
mock_http.request and the raised NemoHTTPError status_code remains 503.
In `@packages/nemo_platform_plugin/tests/client/test_pagination.py`:
- Around line 430-445: The test is overriding the client’s default retry policy
after construction, so it no longer verifies the constructor-installed behavior.
In test_retries_on_by_default, keep NemoClient untouched after instantiation and
instead patch the retry backoff/sleep behavior used by the retry path so the
default DEFAULT_RETRY_POLICY is what gets exercised. Use the existing
NemoClient, RetryPolicy, and send/LIST_ITEMS setup to confirm 503s are retried
without mutating client._retry.
In `@packages/nmp_common/src/nmp/common/service/base.py`:
- Around line 37-39: Remove the TYPE_CHECKING-only import for NemoEntityClient
in base.py and import it normally at module scope, since the evidence in
client.py shows the plain import is safe. Then update the affected return type
annotations in the Base service methods to use the concrete NemoEntityClient
type instead of quoted forward references, keeping the existing method
signatures and names intact.
---
Nitpick comments:
In
`@packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/entity_client.py`:
- Around line 138-173: The entity conversion path rebuilds TypeAdapter objects
on every item in _convert_api_entity_to_model, which is expensive during
paginated listing. Add a cached helper in entity_client.py for both entity_type
and attr_type adapter creation (for example via a small lru_cache-backed
function or module-level cache), then use that helper instead of calling
TypeAdapter(...) directly in _convert_api_entity_to_model. Keep the existing
validation logic intact while reusing adapters across repeated calls.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3087fe54-f7fb-46d4-b0a6-451b28f497f2
📒 Files selected for processing (19)
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/__init__.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/endpoints.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/entity_client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/legacy.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/types.pypackages/nemo_platform_plugin/tests/client/test_client_options.pypackages/nemo_platform_plugin/tests/client/test_pagination.pypackages/nemo_platform_plugin/tests/entities/test_endpoints.pypackages/nemo_platform_plugin/tests/entities/test_nemo_entity_client.pypackages/nmp_common/src/nmp/common/entities/__init__.pypackages/nmp_common/src/nmp/common/entities/client.pypackages/nmp_common/src/nmp/common/service/__init__.pypackages/nmp_common/src/nmp/common/service/base.pypackages/nmp_common/src/nmp/common/service/dependencies.py
Comment on lines
+207
to
+211
| # Shared default: retries transient failures out of the box. The retryable set | ||
| # is deliberately conservative (pre-processing failures only), so replaying a | ||
| # non-safe method is safe — notably 409 is excluded because it carries the | ||
| # entity store's optimistic-lock conflict, which must never be retried blindly. | ||
| DEFAULT_RETRY_POLICY = RetryPolicy() |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Don’t default-retry write requests without an idempotency guard.
This policy is now the default, and the retry loop replays request bodies. 502/503/504/429 can occur after backend processing, so POST/PATCH/PUT may be duplicated or surface false conflicts. Gate default retries to safe/idempotent methods, or require an Idempotency-Key before retrying writes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py`
around lines 207 - 211, The default retry policy in the
RetryPolicy/DEFAULT_RETRY_POLICY path is currently applied too broadly and can
replay write request bodies without protection. Update the retry logic so it
only retries safe/idempotent methods by default, and only allows POST/PATCH/PUT
retries when an Idempotency-Key guard is present or equivalent idempotency check
exists. Make the change in the retry-policy setup and any request-classification
helper used by the Nemo platform client types so transient errors like
502/503/504/429 are not blindly retried for non-idempotent writes.
| from nemo_platform import AsyncNeMoPlatform | ||
| from nemo_platform_plugin.config import PlatformConfig | ||
| from nemo_platform_plugin.entities import EntityClient | ||
| from nemo_platform_plugin.entities import EntityClient, NemoEntityClient |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Drop TYPE_CHECKING guard for NemoEntityClient — same-package import.
NemoEntityClient lives in nemo_platform_plugin.entities, the same package as this file, and is already imported normally at runtime inside the function body (Line 83). Per path guidelines this class should be imported unconditionally instead of hidden behind TYPE_CHECKING, with a concrete (non-string) return type.
🔧 Proposed fix
-if TYPE_CHECKING:
- from nemo_platform_plugin.entities import EntityClient, NemoEntityClient
+from nemo_platform_plugin.entities import NemoEntityClient
+
+if TYPE_CHECKING:
+ from nemo_platform_plugin.entities import EntityClient-def get_nemo_entity_client() -> "NemoEntityClient":
+def get_nemo_entity_client() -> NemoEntityClient:
...
- from nemo_platform_plugin.client_provider import get_async_nemo_client
- from nemo_platform_plugin.entities import NemoEntityClient
-
+ from nemo_platform_plugin.client_provider import get_async_nemo_client
+
return NemoEntityClient(get_async_nemo_client())As per path instructions for packages/nemo_platform_plugin/src/nemo_*/**/*.py: "do not use TYPE_CHECKING imports for types that are available at runtime in the same package; import them normally instead."
Also applies to: 65-85
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py` at
line 17, Remove the TYPE_CHECKING-only handling for NemoEntityClient in
dependencies.py and import it unconditionally from nemo_platform_plugin.entities
alongside EntityClient. Update the dependency provider function that currently
uses the runtime import in the body so its return annotation uses the concrete
NemoEntityClient type instead of a string literal, and keep the symbol names
EntityClient, NemoEntityClient, and the provider function in sync with the
same-package runtime import rule.
Source: Path instructions
Comment on lines
+65
to
+85
|
|
||
|
|
||
| def get_nemo_entity_client() -> "NemoEntityClient": | ||
| """Return a NemoClient-backed :class:`NemoEntityClient`. | ||
|
|
||
| Unlike the sibling placeholders (``get_sdk_client``, ``get_entity_client``, | ||
| ``get_platform_config``), this has a real default implementation because | ||
| everything it needs lives in ``nemo_platform_plugin``: the client transport, | ||
| the env-based context/header builder (``client_provider``), and the entity | ||
| client itself — no ``nmp-core`` bits required. | ||
|
|
||
| The default builds an env-scoped client (``NMP_BASE_URL`` / ``NMP_PRINCIPAL`` | ||
| via ``client_provider.get_async_nemo_client``), which is the correct context | ||
| for task containers, background jobs, and controllers. For request handling, | ||
| ``Service`` still overrides this via ``app.dependency_overrides`` with a | ||
| request-scoped client (service-principal + on-behalf-of headers). | ||
| """ | ||
| from nemo_platform_plugin.client_provider import get_async_nemo_client | ||
| from nemo_platform_plugin.entities import NemoEntityClient | ||
|
|
||
| return NemoEntityClient(get_async_nemo_client()) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd client_provider.py --exec cat -n {}Repository: NVIDIA-NeMo/nemo-platform
Length of output: 4649
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## files\n'
git ls-files | rg '^(packages/nemo_platform_plugin/src/nemo_platform_plugin/(client|entities|dependencies)\.py|.*client_provider\.py|.*entities\.py)$'
for f in \
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py \
packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py \
packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py \
packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py
do
if [ -f "$f" ]; then
printf '\n## outline: %s\n' "$f"
ast-grep outline "$f" --view expanded || true
printf '\n## head: %s\n' "$f"
wc -l "$f"
fi
doneRepository: NVIDIA-NeMo/nemo-platform
Length of output: 4762
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
paths = [
Path("packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py"),
Path("packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py"),
Path("packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py"),
Path("packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py"),
]
for p in paths:
if p.exists():
print(f"\n### {p} ({sum(1 for _ in p.open())} lines)")
for i, line in enumerate(p.open(), 1):
if any(k in line for k in ["class AsyncNemoClient", "class NemoClient", "def close", "def __aenter__", "def __aexit__", "def __init__", "class NemoEntityClient", "def get_async_nemo_client", "def get_nemo_entity_client", "async def", "httpx", "AsyncClient", "Client"]):
print(f"{i:4d}: {line.rstrip()}")
PYRepository: NVIDIA-NeMo/nemo-platform
Length of output: 3559
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate implementations and lifecycle hooks
rg -n --hidden --glob '!**/.git/**' \
-e 'class AsyncNemoClient|class NemoClient|def close|async def close|__aenter__|__aexit__|httpx\.AsyncClient|httpx\.Client|NemoEntityClient|get_async_nemo_client|get_nemo_entity_client' \
packages/nemo_platform_plugin/src/nemo_platform_pluginRepository: NVIDIA-NeMo/nemo-platform
Length of output: 4751
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file=packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/entity_client.py
wc -l "$file"
sed -n '1,240p' "$file"Repository: NVIDIA-NeMo/nemo-platform
Length of output: 10511
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## dependencies\n'
sed -n '1,140p' packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py
printf '\n## client ctor/close\n'
sed -n '433,490p' packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.pyRepository: NVIDIA-NeMo/nemo-platform
Length of output: 5459
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## entity client\n'
sed -n '1,240p' packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/entity_client.py
printf '\n## dependencies\n'
sed -n '1,140p' packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py
printf '\n## async client lifecycle\n'
sed -n '433,490p' packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.pyRepository: NVIDIA-NeMo/nemo-platform
Length of output: 15745
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## call sites\n'
rg -n --hidden --glob '!**/.git/**' \
-e 'get_nemo_entity_client\(' \
-e 'NemoEntityClient\(' \
-e '\.close\(\)' \
packages/nemo_platform_plugin/src plugins services sdk e2e
printf '\n## dependency overrides / app wiring\n'
rg -n --hidden --glob '!**/.git/**' \
-e 'dependency_overrides' \
-e 'get_nemo_entity_client' \
-e 'yield .*NemoEntityClient|async def .*close' \
packages/nemo_platform_plugin/src plugins services sdk
printf '\n## nearby service lifecycle code\n'
fd -a 'service*.py' packages/nemo_platform_plugin/src services plugins sdk | head -n 50Repository: NVIDIA-NeMo/nemo-platform
Length of output: 48877
Scope or close the async entity client get_nemo_entity_client() creates a new AsyncNemoClient/httpx.AsyncClient each call, but nothing here closes it. If this dependency is used repeatedly in jobs/controllers, add lifecycle management or reuse a shared client to avoid leaking sockets.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py`
around lines 65 - 85, `get_nemo_entity_client()` currently creates a fresh async
transport via `get_async_nemo_client()` and wraps it in `NemoEntityClient`
without any cleanup, which can leak sockets over repeated calls. Update the
dependency to manage the client lifecycle explicitly: either reuse a shared
async client instance or make this an async dependency that closes the
`AsyncNemoClient`/`httpx.AsyncClient` after use. Keep the fix localized to
`get_nemo_entity_client`, `get_async_nemo_client`, and the `NemoEntityClient`
construction path so the caller does not retain an unclosed client.
Comment on lines
+217
to
+237
| def test_default_policy_retries(self) -> None: | ||
| """Without an explicit policy, the client retries transient failures. | ||
|
|
||
| DEFAULT_RETRY_POLICY is applied by the constructor (max_retries=3), so a | ||
| persistent 503 is attempted 1 + 3 times before surfacing. | ||
| """ | ||
| mock_http = MagicMock(spec=httpx.Client) | ||
| mock_http.request.return_value = httpx.Response( | ||
| 503, | ||
| request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), | ||
| json={"detail": "unavailable"}, | ||
| ) | ||
|
|
||
| # Use the default policy's codes but zero backoff to keep the test fast. | ||
| client = NemoClient(base_url=BASE, http_client=mock_http, retry=RetryPolicy(backoff_base=0.0)) | ||
|
|
||
| with pytest.raises(NemoHTTPError) as exc_info: | ||
| client.send(GET_ITEM(name="alice")) | ||
|
|
||
| assert exc_info.value.status_code == 503 | ||
| assert mock_http.request.call_count == 4 # 1 initial + 3 retries |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
This doesn’t test the default policy.
Line 231 passes an explicit RetryPolicy, so the test would pass even if the constructor default regressed to None. Patch sleep instead and omit retry=.
Proposed fix
- def test_default_policy_retries(self) -> None:
+ def test_default_policy_retries(self, monkeypatch: pytest.MonkeyPatch) -> None:
@@
- # Use the default policy's codes but zero backoff to keep the test fast.
- client = NemoClient(base_url=BASE, http_client=mock_http, retry=RetryPolicy(backoff_base=0.0))
+ monkeypatch.setattr("nemo_platform_plugin.client.client.time.sleep", lambda _: None)
+ client = NemoClient(base_url=BASE, http_client=mock_http)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def test_default_policy_retries(self) -> None: | |
| """Without an explicit policy, the client retries transient failures. | |
| DEFAULT_RETRY_POLICY is applied by the constructor (max_retries=3), so a | |
| persistent 503 is attempted 1 + 3 times before surfacing. | |
| """ | |
| mock_http = MagicMock(spec=httpx.Client) | |
| mock_http.request.return_value = httpx.Response( | |
| 503, | |
| request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), | |
| json={"detail": "unavailable"}, | |
| ) | |
| # Use the default policy's codes but zero backoff to keep the test fast. | |
| client = NemoClient(base_url=BASE, http_client=mock_http, retry=RetryPolicy(backoff_base=0.0)) | |
| with pytest.raises(NemoHTTPError) as exc_info: | |
| client.send(GET_ITEM(name="alice")) | |
| assert exc_info.value.status_code == 503 | |
| assert mock_http.request.call_count == 4 # 1 initial + 3 retries | |
| def test_default_policy_retries(self, monkeypatch: pytest.MonkeyPatch) -> None: | |
| """Without an explicit policy, the client retries transient failures. | |
| DEFAULT_RETRY_POLICY is applied by the constructor (max_retries=3), so a | |
| persistent 503 is attempted 1 + 3 times before surfacing. | |
| """ | |
| mock_http = MagicMock(spec=httpx.Client) | |
| mock_http.request.return_value = httpx.Response( | |
| 503, | |
| request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), | |
| json={"detail": "unavailable"}, | |
| ) | |
| monkeypatch.setattr("nemo_platform_plugin.client.client.time.sleep", lambda _: None) | |
| client = NemoClient(base_url=BASE, http_client=mock_http) | |
| with pytest.raises(NemoHTTPError) as exc_info: | |
| client.send(GET_ITEM(name="alice")) | |
| assert exc_info.value.status_code == 503 | |
| assert mock_http.request.call_count == 4 # 1 initial + 3 retries |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/nemo_platform_plugin/tests/client/test_client_options.py` around
lines 217 - 237, The test currently bypasses the constructor default by passing
an explicit RetryPolicy in test_default_policy_retries, so it won’t catch
regressions in NemoClient’s default retry behavior. Update this test to
construct NemoClient without retry= and instead patch the retry backoff/sleep
path to keep it fast, while still asserting the 503 is retried 1 + 3 times via
mock_http.request and the raised NemoHTTPError status_code remains 503.
Comment on lines
+430
to
+445
| def test_retries_on_by_default(self) -> None: | ||
| """A client constructed without an explicit retry policy still retries 503s.""" | ||
| mock_http = MagicMock(spec=httpx.Client) | ||
| mock_http.request.side_effect = [ | ||
| httpx.Response( | ||
| 503, | ||
| request=httpx.Request("GET", f"{BASE}/apis/test/v2/workspaces/default/items"), | ||
| json={"detail": "unavailable"}, | ||
| ), | ||
| _page_response([{"id": 1, "name": "a"}], page=1, total_pages=1), | ||
| ] | ||
|
|
||
| # No retry= passed → DEFAULT_RETRY_POLICY (backoff patched to 0 for speed). | ||
| client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) | ||
| client._retry = RetryPolicy(backoff_base=0.0) # keep default codes, zero sleep | ||
| items = list(client.send(LIST_ITEMS()).items()) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Avoid overriding the default under test.
Line 444 mutates _retry, so this no longer proves the constructor-installed default works. Patch sleep and leave the client untouched.
Proposed fix
- def test_retries_on_by_default(self) -> None:
+ def test_retries_on_by_default(self, monkeypatch: pytest.MonkeyPatch) -> None:
@@
# No retry= passed → DEFAULT_RETRY_POLICY (backoff patched to 0 for speed).
+ monkeypatch.setattr("nemo_platform_plugin.client.client.time.sleep", lambda _: None)
client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http)
- client._retry = RetryPolicy(backoff_base=0.0) # keep default codes, zero sleep
items = list(client.send(LIST_ITEMS()).items())📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def test_retries_on_by_default(self) -> None: | |
| """A client constructed without an explicit retry policy still retries 503s.""" | |
| mock_http = MagicMock(spec=httpx.Client) | |
| mock_http.request.side_effect = [ | |
| httpx.Response( | |
| 503, | |
| request=httpx.Request("GET", f"{BASE}/apis/test/v2/workspaces/default/items"), | |
| json={"detail": "unavailable"}, | |
| ), | |
| _page_response([{"id": 1, "name": "a"}], page=1, total_pages=1), | |
| ] | |
| # No retry= passed → DEFAULT_RETRY_POLICY (backoff patched to 0 for speed). | |
| client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) | |
| client._retry = RetryPolicy(backoff_base=0.0) # keep default codes, zero sleep | |
| items = list(client.send(LIST_ITEMS()).items()) | |
| def test_retries_on_by_default(self, monkeypatch: pytest.MonkeyPatch) -> None: | |
| """A client constructed without an explicit retry policy still retries 503s.""" | |
| mock_http = MagicMock(spec=httpx.Client) | |
| mock_http.request.side_effect = [ | |
| httpx.Response( | |
| 503, | |
| request=httpx.Request("GET", f"{BASE}/apis/test/v2/workspaces/default/items"), | |
| json={"detail": "unavailable"}, | |
| ), | |
| _page_response([{"id": 1, "name": "a"}], page=1, total_pages=1), | |
| ] | |
| # No retry= passed → DEFAULT_RETRY_POLICY (backoff patched to 0 for speed). | |
| monkeypatch.setattr("nemo_platform_plugin.client.client.time.sleep", lambda _: None) | |
| client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) | |
| items = list(client.send(LIST_ITEMS()).items()) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/nemo_platform_plugin/tests/client/test_pagination.py` around lines
430 - 445, The test is overriding the client’s default retry policy after
construction, so it no longer verifies the constructor-installed behavior. In
test_retries_on_by_default, keep NemoClient untouched after instantiation and
instead patch the retry backoff/sleep behavior used by the retry path so the
default DEFAULT_RETRY_POLICY is what gets exercised. Use the existing
NemoClient, RetryPolicy, and send/LIST_ITEMS setup to confirm 503s are retried
without mutating client._retry.
Comment on lines
+37
to
+39
| if TYPE_CHECKING: | ||
| from nemo_platform_plugin.entities import NemoEntityClient | ||
|
|
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Drop TYPE_CHECKING guard for NemoEntityClient — cross-file evidence shows plain import is safe.
packages/nmp_common/src/nmp/common/entities/client.py already does from nemo_platform_plugin.entities import NemoEntityClient as NemoEntityClient as a plain top-level import with no circular-import issue, so gating this import behind TYPE_CHECKING here isn't necessary. Import it normally and use a concrete return type instead of the quoted forward reference.
🔧 Proposed fix
-if TYPE_CHECKING:
- from nemo_platform_plugin.entities import NemoEntityClient
+from nemo_platform_plugin.entities import NemoEntityClient- def get_nemo_entity_client(self, as_service: str | None = None) -> "NemoEntityClient":
+ def get_nemo_entity_client(self, as_service: str | None = None) -> NemoEntityClient:As per coding guidelines: "prefer concrete type hints over string-based type hints, and do not import those types only under TYPE_CHECKING; import them normally when possible."
Also applies to: 158-158
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/nmp_common/src/nmp/common/service/base.py` around lines 37 - 39,
Remove the TYPE_CHECKING-only import for NemoEntityClient in base.py and import
it normally at module scope, since the evidence in client.py shows the plain
import is safe. Then update the affected return type annotations in the Base
service methods to use the concrete NemoEntityClient type instead of quoted
forward references, keeping the existing method signatures and names intact.
Source: Coding guidelines
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/resource.py (1)
365-381: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
delete_by_idnever validates/uses theentity_typeparameter.The method resolves
workspace/entity_type/namefrom the fetchedEntityResponse(entity.entity_type,entity.workspace,entity.name) but the caller-suppliedentity_typeis unused. Compare withdelete(), which correctly validates via_get_entity_type(entity_type). As written,delete_by_id(Widget, some_id)will happily delete an entity of a completely different type ifsome_iddoesn't belong to aWidget.🛡️ Proposed fix: validate fetched entity type matches caller expectation
async def delete_by_id(self, entity_type: EntityTypeLike, entity_id: str) -> DeleteResponse: """Delete an entity by UUID (fetches it first to resolve name/workspace).""" try: entity = (await self._client.get_entity_by_id(id=entity_id)).data() + if entity.entity_type != _get_entity_type(entity_type): + raise EntityNotFoundError(f"Entity with id '{entity_id}' not found") query: ParentQueryParams = {} if entity.parent is not None: query["parent"] = entity.parent🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/resource.py` around lines 365 - 381, `delete_by_id` is ignoring the caller-provided `entity_type`, so it can delete an entity of a different kind than expected. Update `delete_by_id` in `Resource` to validate the requested type with `_get_entity_type(entity_type)` and compare it against the fetched `entity.entity_type` before calling `_client.delete_entity_by_name`; if they differ, raise an appropriate error instead of proceeding. Keep the existing `get_entity_by_id` / `delete_entity_by_name` flow and only use the resolved entity data after the type check passes.
♻️ Duplicate comments (1)
packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py (1)
17-17: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winDrop
TYPE_CHECKINGguard forEntityStoreResource— same issue as priorNemoEntityClientflag.
EntityStoreResourcelives innemo_platform_plugin.entities(same package) and is already imported normally at runtime insideget_entity_store_resource(line 83). Gating it behindTYPE_CHECKINGand using the quoted return type at line 67 repeats the exact pattern previously flagged forNemoEntityClientin this file.🔧 Proposed fix
-if TYPE_CHECKING: - from nemo_platform_plugin.entities import EntityClient, EntityStoreResource +from nemo_platform_plugin.entities import EntityStoreResource + +if TYPE_CHECKING: + from nemo_platform_plugin.entities import EntityClient-def get_entity_store_resource() -> "EntityStoreResource": +def get_entity_store_resource() -> EntityStoreResource: ... from nemo_platform_plugin.client_provider import get_async_nemo_client - from nemo_platform_plugin.entities import EntityStoreResource - return EntityStoreResource.from_client(get_async_nemo_client())As per path instructions for
packages/nemo_platform_plugin/src/nemo_*/**/*.py: "do not useTYPE_CHECKINGimports for types that are available at runtime in the same package; import them normally instead."Also applies to: 67-67, 82-85
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py` at line 17, The `EntityStoreResource` import is incorrectly guarded by `TYPE_CHECKING`, even though it is a same-package runtime symbol already used by `get_entity_store_resource`. Move `EntityStoreResource` to the normal import list in `dependencies.py`, and update the `get_entity_store_resource` return annotation to use the direct type name instead of a quoted string, matching the existing runtime import pattern used for `EntityClient`.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/resource.py`:
- Around line 365-381: `delete_by_id` is ignoring the caller-provided
`entity_type`, so it can delete an entity of a different kind than expected.
Update `delete_by_id` in `Resource` to validate the requested type with
`_get_entity_type(entity_type)` and compare it against the fetched
`entity.entity_type` before calling `_client.delete_entity_by_name`; if they
differ, raise an appropriate error instead of proceeding. Keep the existing
`get_entity_by_id` / `delete_entity_by_name` flow and only use the resolved
entity data after the type check passes.
---
Duplicate comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py`:
- Line 17: The `EntityStoreResource` import is incorrectly guarded by
`TYPE_CHECKING`, even though it is a same-package runtime symbol already used by
`get_entity_store_resource`. Move `EntityStoreResource` to the normal import
list in `dependencies.py`, and update the `get_entity_store_resource` return
annotation to use the direct type name instead of a quoted string, matching the
existing runtime import pattern used for `EntityClient`.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 819b8e00-3ddf-448a-825b-0a520e1d047c
📒 Files selected for processing (10)
packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/__init__.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/entities/resource.pypackages/nemo_platform_plugin/tests/entities/test_resource.pypackages/nmp_common/src/nmp/common/entities/__init__.pypackages/nmp_common/src/nmp/common/entities/client.pypackages/nmp_common/src/nmp/common/service/__init__.pypackages/nmp_common/src/nmp/common/service/base.pypackages/nmp_common/src/nmp/common/service/dependencies.py
✅ Files skipped from review due to trivial changes (1)
- packages/nmp_common/src/nmp/common/service/dependencies.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Signed-off-by: Matthew Grossman mgrossman@nvidia.com
Summary by CodeRabbit
New Features
Bug Fixes
Chores