diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/__init__.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/__init__.py index 09442b20ef..c5abaf6cad 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/__init__.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/__init__.py @@ -68,6 +68,9 @@ from nemo_platform_plugin.entities.base import ( EntityTypeLike as EntityTypeLike, ) +from nemo_platform_plugin.entities.base import ( + EntityUpdateClientProtocol as EntityUpdateClientProtocol, +) from nemo_platform_plugin.entities.base import ( EntityValidationError as EntityValidationError, ) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/base.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/base.py index c485b501ca..a290208c4d 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/base.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities/base.py @@ -231,14 +231,29 @@ class EntityToken(Protocol): class EntityGetterProtocol(Protocol[EntityT]): - """Protocol for entity clients that can fetch entities by workspace/name.""" + """Protocol for entity clients that can fetch entities by workspace/name. - async def get(self, entity_type: Type[EntityT], *, name: str, workspace: str) -> EntityT: ... + ``parent`` addresses a **child** entity, which is unique within + ``(workspace, entity_type, parent, name)`` rather than by name alone. It is optional, so + fetching a root entity is unchanged. + """ + + async def get( + self, + entity_type: Type[EntityT], + *, + name: str, + workspace: str, + parent: Optional[str] = None, + ) -> EntityT: ... class EntityDeleteClientProtocol(EntityGetterProtocol[EntityT], Protocol[EntityT]): """Protocol for entity clients that can list and delete entities.""" + # ``filter_str`` and ``filter_obj`` exist on the client but are deliberately absent here: + # ``filter_operation`` is the sanctioned structured form, and the other two are a JSON-string + # variant and an exact-match shorthand kept for older callers. async def list( self, entity_type: Type[EntityT], @@ -256,6 +271,7 @@ async def delete( name: str, *, workspace: str, + parent: Optional[str] = None, expected_db_version: Optional[int] = None, ) -> object: ... @@ -266,6 +282,20 @@ class EntityClientProtocol(EntityDeleteClientProtocol[EntityT], Protocol[EntityT async def create(self, entity: EntityT) -> EntityT: ... +class EntityUpdateClientProtocol(Protocol[EntityT]): + """Protocol for entity clients that can update an existing entity. + + Separate from :class:`EntityClientProtocol` rather than folded into it: ``update`` is a + read-modify-write against the ``db_version`` optimistic lock, and most services never need it. + Compose it with the CRUD protocol where a service does:: + + class Store(EntityClientProtocol[MyEntity], EntityUpdateClientProtocol[MyEntity], Protocol): + ... + """ + + async def update(self, entity: EntityT, *, original_name: Optional[str] = None) -> EntityT: ... + + class AnyEntityGetterProtocol(Protocol): """Protocol for clients that can fetch any entity model type.""" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entity_client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entity_client.py index bf6c6fc681..770fae0796 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entity_client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entity_client.py @@ -54,6 +54,9 @@ from nemo_platform_plugin.entities import ( EntityNotFoundError as NemoEntityNotFoundError, ) +from nemo_platform_plugin.entities import ( + EntityUpdateClientProtocol as NemoEntityUpdateClientProtocol, +) from nemo_platform_plugin.entities import ( EntityValidationError as NemoEntityValidationError, ) @@ -65,6 +68,7 @@ "NemoEntitiesClient", "NemoEntitiesClientProtocol", "NemoAnyEntityDeleteClientProtocol", + "NemoEntityUpdateClientProtocol", "NemoAnyEntityGetterProtocol", "NemoEntityDeleteClientProtocol", "NemoEntityGetterProtocol", diff --git a/packages/nemo_platform_plugin/tests/entities/test_client_protocols.py b/packages/nemo_platform_plugin/tests/entities/test_client_protocols.py new file mode 100644 index 0000000000..89dd21b7ae --- /dev/null +++ b/packages/nemo_platform_plugin/tests/entities/test_client_protocols.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The entity-client protocols must describe the client they stand in for. + +A protocol that has drifted from its implementation is worse than no protocol: a service typed +against it either fails type-checking on correct code, or type-checks against a method the real +client does not have. These tests pin that relationship so the two cannot silently diverge. +""" + +from __future__ import annotations + +import inspect +from typing import Protocol, TypeVar + +from nemo_platform_plugin.entities import ( + EntityBase, + EntityClient, + EntityClientProtocol, + EntityGetterProtocol, + EntityUpdateClientProtocol, +) + +EntityT = TypeVar("EntityT", bound=EntityBase) + + +class _Entity(EntityBase): + __entity_type__ = "protocol_conformance_probe" + + +class _ReadWriteStore( + EntityClientProtocol[EntityT], + EntityUpdateClientProtocol[EntityT], + Protocol[EntityT], +): + """The shape a service needing the wider surface composes for itself. + + Exists here to prove the pieces *compose*: ``update`` is a separate protocol precisely so a + service can opt into it alongside the CRUD one, rather than declaring a private protocol that + restates the whole surface. + """ + + +def _static_conformance(client: EntityClient) -> _ReadWriteStore[_Entity]: + """Static assertion, checked by ``ty`` rather than at runtime. + + If ``EntityClient`` ever stops satisfying the composed protocols — a renamed method, a changed + signature — this return fails type-checking. The runtime tests below document *which* parts + matter and why; this is what actually catches drift, because structural conformance is a + type-level property no ``hasattr`` check can verify. + """ + return client + + +def _signature(owner: object, method: str) -> inspect.Signature: + return inspect.signature(getattr(owner, method)) + + +#: Parameters the client has that a protocol deliberately does not expose, with the reason. Anything +#: not listed here is treated as accidental under-specification by the signature check below. +_INTENTIONAL_OMISSIONS = { + # ``filter_operation`` is the sanctioned structured form; these two are a JSON-string variant + # and an exact-match shorthand kept for older callers. + ("list", "filter_str"), + ("list", "filter_obj"), +} + + +def test_protocols_expose_every_client_parameter() -> None: + """Catch *under*-specification, which conformance alone cannot. + + Structural conformance is one-directional: a class satisfies a protocol by providing at least + what it declares, so extra parameters on the client pass silently. That is how ``parent`` went + missing from ``delete`` while ``_static_conformance`` reported success — a service typed against + the protocol could not delete a child entity even though its client could. + """ + gaps: dict[str, set[str]] = {} + for method in ("get", "create", "update", "delete", "list"): + protocol = EntityUpdateClientProtocol if method == "update" else EntityClientProtocol + declared = set(_signature(protocol, method).parameters) + available = set(_signature(EntityClient, method).parameters) + missing = {p for p in available - declared if (method, p) not in _INTENTIONAL_OMISSIONS} + if missing: + gaps[method] = missing + assert not gaps, f"protocol is missing client parameters: {gaps}" + + +def test_update_is_its_own_protocol() -> None: + """``update`` is opt-in. Most services never modify an entity in place, and folding ``update`` + into the CRUD protocol would force each of them — and every one of their test doubles — to + satisfy a method they do not use.""" + assert hasattr(EntityUpdateClientProtocol, "update") + assert not hasattr(EntityClientProtocol, "update") + + +def test_update_protocol_matches_the_client() -> None: + protocol = _signature(EntityUpdateClientProtocol, "update").parameters + client = _signature(EntityClient, "update").parameters + assert set(protocol) == set(client) + assert protocol["original_name"].kind is inspect.Parameter.KEYWORD_ONLY + + +def test_getter_accepts_parent_for_child_entities() -> None: + """Child records are unique within ``(workspace, entity_type, parent, name)``, so the parent is + part of their address. It is optional, so fetching a root entity is unchanged.""" + getter = _signature(EntityGetterProtocol, "get").parameters + assert "parent" in getter + assert getter["parent"].default is None + assert "parent" in _signature(EntityClient, "get").parameters diff --git a/packages/nemo_platform_plugin/tests/test_entity_client.py b/packages/nemo_platform_plugin/tests/test_entity_client.py index f09c27dda9..d005c1a728 100644 --- a/packages/nemo_platform_plugin/tests/test_entity_client.py +++ b/packages/nemo_platform_plugin/tests/test_entity_client.py @@ -1,12 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from datetime import datetime, timezone from unittest.mock import AsyncMock, Mock import httpx import pytest from nemo_platform_plugin.client.errors import BadRequestError, NotFoundError -from nemo_platform_plugin.entities import EntityClient, EntityStoreError +from nemo_platform_plugin.entities import EntityBase, EntityClient, EntityStoreError +from nemo_platform_plugin.entities.types import Entity class ExperimentGroup: @@ -115,3 +117,69 @@ async def test_count_by_rejects_non_direct_field() -> None: with pytest.raises(ValueError, match="direct entity data field"): await client.count_by(ExperimentGroup, "data.insight_id") + + +class _Child(EntityBase): + """A child entity — addressed within its parent, not by name alone.""" + + __entity_type__ = "child_probe" + + note: str = "" + + +def _stored_child(parent: str) -> Mock: + """A server response for a child entity, as ``get``/``update`` receive it.""" + entity = Entity( + entity_type="child_probe", + id="child-1", + workspace="default", + parent=parent, + name="child-1", + data={"note": "before"}, + db_version=1, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + updated_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ) + resp = Mock() + resp.data = Mock(return_value=entity) + return resp + + +@pytest.mark.asyncio +async def test_parent_survives_the_get_update_round_trip() -> None: + """``update`` takes no ``parent`` argument — it reads it off the entity. + + That only holds because ``get`` puts it there. If either half breaks, an update to a child + entity silently addresses a root entity of the same name instead, so this pins both halves + together rather than mocking one and asserting the other. + """ + mock_api = Mock() + mock_api.get_entity_by_name = AsyncMock(return_value=_stored_child("parent-1")) + mock_api.update_entity_by_name = AsyncMock(return_value=_stored_child("parent-1")) + client = EntityClient(mock_api) + + fetched = await client.get(_Child, "child-1", workspace="default", parent="parent-1") + get_call = mock_api.get_entity_by_name.await_args + assert get_call is not None + assert get_call.kwargs["query_params"] == {"parent": "parent-1"} + assert fetched.parent == "parent-1" + + fetched.note = "after" + await client.update(fetched) + + update_call = mock_api.update_entity_by_name.await_args + assert update_call is not None + assert update_call.kwargs["query_params"] == {"parent": "parent-1"} + + +@pytest.mark.asyncio +async def test_delete_forwards_parent() -> None: + mock_api = Mock() + mock_api.delete_entity_by_name = AsyncMock(return_value=Mock()) + client = EntityClient(mock_api) + + await client.delete(_Child, "child-1", workspace="default", parent="parent-1") + + call = mock_api.delete_entity_by_name.await_args + assert call is not None + assert call.kwargs["query_params"] == {"parent": "parent-1"} 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 7e38f597f2..24c2a50f14 100644 --- a/plugins/nemo-evaluator/tests/api/service/test_metric_service.py +++ b/plugins/nemo-evaluator/tests/api/service/test_metric_service.py @@ -65,7 +65,9 @@ def __init__(self) -> None: self.delete_error: Exception | None = None self.list_filter_operations: list[FilterOperation | None] = [] - async def get(self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str) -> MetricBundleEntity: + async def get( + self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str, parent: str | None = None + ) -> MetricBundleEntity: key = (workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found") @@ -88,6 +90,7 @@ async def delete( name: str, *, workspace: str, + parent: str | None = None, expected_db_version: int | None = None, ) -> None: if self.delete_error is not None: diff --git a/plugins/nemo-evaluator/tests/api/service/test_result_service.py b/plugins/nemo-evaluator/tests/api/service/test_result_service.py index ec1cecc661..bd0dcb09ac 100644 --- a/plugins/nemo-evaluator/tests/api/service/test_result_service.py +++ b/plugins/nemo-evaluator/tests/api/service/test_result_service.py @@ -37,7 +37,9 @@ def seed(self, entity: _ResultEntityT) -> _ResultEntityT: self.entities[(entity.__entity_type__, entity.workspace, entity.name)] = entity return entity - async def get(self, entity_type: type[_ResultEntityT], *, workspace: str, name: str) -> _ResultEntityT: + async def get( + self, entity_type: type[_ResultEntityT], *, workspace: str, name: str, parent: str | None = None + ) -> _ResultEntityT: key = (entity_type.__entity_type__, workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found") @@ -51,6 +53,7 @@ async def delete( name: str, *, workspace: str, + parent: str | None = None, expected_db_version: int | None = None, ) -> None: # Mirror the real EntityClient: raise NemoEntityNotFoundError when absent. diff --git a/plugins/nemo-evaluator/tests/api/service/test_task_service.py b/plugins/nemo-evaluator/tests/api/service/test_task_service.py index 6b5b333f92..6287817f27 100644 --- a/plugins/nemo-evaluator/tests/api/service/test_task_service.py +++ b/plugins/nemo-evaluator/tests/api/service/test_task_service.py @@ -54,7 +54,9 @@ async def create(self, entity: TaskEntity) -> TaskEntity: self.entities[key] = entity return entity - async def get(self, entity_type: type[TaskEntity], *, workspace: str, name: str) -> TaskEntity: + async def get( + self, entity_type: type[TaskEntity], *, workspace: str, name: str, parent: str | None = None + ) -> TaskEntity: key = (entity_type.__entity_type__, workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found") @@ -66,6 +68,7 @@ async def delete( name: str, *, workspace: str, + parent: str | None = None, expected_db_version: int | None = None, ) -> None: key = (entity_type.__entity_type__, workspace, name) diff --git a/plugins/nemo-evaluator/tests/api/service/test_taskset_service.py b/plugins/nemo-evaluator/tests/api/service/test_taskset_service.py index d1bec5757b..cb9c2c2550 100644 --- a/plugins/nemo-evaluator/tests/api/service/test_taskset_service.py +++ b/plugins/nemo-evaluator/tests/api/service/test_taskset_service.py @@ -43,7 +43,9 @@ async def create(self, entity: TasksetEntity) -> TasksetEntity: self.entities[key] = entity return entity - async def get(self, entity_type: type[TasksetEntity], *, workspace: str, name: str) -> TasksetEntity: + async def get( + self, entity_type: type[TasksetEntity], *, workspace: str, name: str, parent: str | None = None + ) -> TasksetEntity: key = (entity_type.__entity_type__, workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found") @@ -55,6 +57,7 @@ async def delete( name: str, *, workspace: str, + parent: str | None = None, expected_db_version: int | None = None, ) -> None: key = (entity_type.__entity_type__, workspace, name) 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 e9ff7b1591..b5da9c5837 100644 --- a/plugins/nemo-evaluator/tests/api/v2/test_metrics_routes.py +++ b/plugins/nemo-evaluator/tests/api/v2/test_metrics_routes.py @@ -78,7 +78,9 @@ def __init__(self) -> None: self.bump_version_on_next_delete = False self.delete_expected_db_versions: list[int | None] = [] - async def get(self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str) -> MetricBundleEntity: + async def get( + self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str, parent: str | None = None + ) -> MetricBundleEntity: key = (workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found") @@ -103,6 +105,7 @@ async def delete( name: str, *, workspace: str, + parent: str | None = None, expected_db_version: int | None = None, ) -> None: key = (workspace, name) diff --git a/plugins/nemo-evaluator/tests/api/v2/test_results_routes.py b/plugins/nemo-evaluator/tests/api/v2/test_results_routes.py index 24cc48f1e8..2f60248d57 100644 --- a/plugins/nemo-evaluator/tests/api/v2/test_results_routes.py +++ b/plugins/nemo-evaluator/tests/api/v2/test_results_routes.py @@ -39,7 +39,9 @@ def seed(self, entity: _ResultEntityT) -> _ResultEntityT: self.entities[(entity.__entity_type__, entity.workspace, entity.name)] = entity return entity - async def get(self, entity_type: type[_ResultEntityT], *, workspace: str, name: str) -> _ResultEntityT: + async def get( + self, entity_type: type[_ResultEntityT], *, workspace: str, name: str, parent: str | None = None + ) -> _ResultEntityT: key = (entity_type.__entity_type__, workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found") @@ -53,6 +55,7 @@ async def delete( name: str, *, workspace: str, + parent: str | None = None, expected_db_version: int | None = None, ) -> None: # Mirror the real EntityClient: raise NemoEntityNotFoundError when absent. diff --git a/plugins/nemo-evaluator/tests/api/v2/test_tasks_routes.py b/plugins/nemo-evaluator/tests/api/v2/test_tasks_routes.py index e7b09ada9c..f249c52e7b 100644 --- a/plugins/nemo-evaluator/tests/api/v2/test_tasks_routes.py +++ b/plugins/nemo-evaluator/tests/api/v2/test_tasks_routes.py @@ -39,7 +39,9 @@ async def create(self, entity: TaskEntity) -> TaskEntity: self.entities[key] = entity return entity - async def get(self, entity_type: type[TaskEntity], *, workspace: str, name: str) -> TaskEntity: + async def get( + self, entity_type: type[TaskEntity], *, workspace: str, name: str, parent: str | None = None + ) -> TaskEntity: key = (entity_type.__entity_type__, workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found") @@ -51,6 +53,7 @@ async def delete( name: str, *, workspace: str, + parent: str | None = None, expected_db_version: int | None = None, ) -> None: key = (entity_type.__entity_type__, workspace, name) diff --git a/plugins/nemo-evaluator/tests/api/v2/test_tasksets_routes.py b/plugins/nemo-evaluator/tests/api/v2/test_tasksets_routes.py index 277276c686..eb743eea10 100644 --- a/plugins/nemo-evaluator/tests/api/v2/test_tasksets_routes.py +++ b/plugins/nemo-evaluator/tests/api/v2/test_tasksets_routes.py @@ -39,7 +39,9 @@ async def create(self, entity: TasksetEntity) -> TasksetEntity: self.entities[key] = entity return entity - async def get(self, entity_type: type[TasksetEntity], *, workspace: str, name: str) -> TasksetEntity: + async def get( + self, entity_type: type[TasksetEntity], *, workspace: str, name: str, parent: str | None = None + ) -> TasksetEntity: key = (entity_type.__entity_type__, workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found") @@ -51,6 +53,7 @@ async def delete( name: str, *, workspace: str, + parent: str | None = None, expected_db_version: int | None = None, ) -> None: key = (entity_type.__entity_type__, workspace, name) diff --git a/plugins/nemo-evaluator/tests/test_metric_refs.py b/plugins/nemo-evaluator/tests/test_metric_refs.py index 2d41e553c8..94b259a0ca 100644 --- a/plugins/nemo-evaluator/tests/test_metric_refs.py +++ b/plugins/nemo-evaluator/tests/test_metric_refs.py @@ -60,7 +60,9 @@ class _FakeEntityClient: def __init__(self) -> None: self.entities: dict[tuple[str, str], MetricBundleEntity] = {} - async def get(self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str) -> MetricBundleEntity: + async def get( + self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str, parent: str | None = None + ) -> MetricBundleEntity: try: return self.entities[(workspace, name)] except KeyError: diff --git a/plugins/nemo-evaluator/tests/test_task_refs.py b/plugins/nemo-evaluator/tests/test_task_refs.py index 34e8353e10..97b1beb62b 100644 --- a/plugins/nemo-evaluator/tests/test_task_refs.py +++ b/plugins/nemo-evaluator/tests/test_task_refs.py @@ -27,7 +27,9 @@ def __init__(self) -> None: def add(self, entity: EntityBase) -> None: self.entities[(entity.__entity_type__, entity.workspace, entity.name)] = entity - async def get(self, entity_type: type[_EntityT], *, workspace: str, name: str) -> _EntityT: + async def get( + self, entity_type: type[_EntityT], *, workspace: str, name: str, parent: str | None = None + ) -> _EntityT: key = (entity_type.__entity_type__, workspace, name) if key not in self.entities: raise NemoEntityNotFoundError(f"{workspace}/{name} not found")