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..8b1079d321 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,9 +231,21 @@ 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]): @@ -266,6 +278,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..0c9a16464f --- /dev/null +++ b/packages/nemo_platform_plugin/tests/entities/test_client_protocols.py @@ -0,0 +1,80 @@ +# 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)) + + +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/plugins/nemo-evaluator/tests/api/service/test_metric_service.py b/plugins/nemo-evaluator/tests/api/service/test_metric_service.py index 7e38f597f2..071452989a 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") 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..0a75b32bd8 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") 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..1da336c026 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") 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..12a56320a9 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") 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..1a95488e35 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") 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..551a4d5acf 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") 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..76d6a87605 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") 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..75709b6e82 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") 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")