chore: harden container labeling and e2e tests (#918)

Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>

Author: mckornfield

e2e/services_pool.py

@@ -290,7 +290,12 @@ def _register_config_state(
 
     def _materialize_config_path(self, state: ModuleConfigState) -> ModuleConfigState:
         data_dir = e2e_services_data_dir(self._get_log_dir(), state.key.config_hash)
-        rendered_config_data = _render_e2e_config_for_backend(state.config_data, data_dir, state.harness_config)
+        rendered_config_data = _render_e2e_config_for_backend(
+            state.config_data,
+            data_dir,
+            state.harness_config,
+            state.key.config_hash,
+        )
         rendered_config = yaml.safe_dump(rendered_config_data, default_flow_style=False, sort_keys=True)
         config_path = self._get_generated_config_dir() / f"platform-{state.key.config_hash}.yaml"
         if not config_path.exists():
@@ -546,7 +551,12 @@ def e2e_services_data_dir(log_dir: Path, config_hash: str) -> Path:
     return log_dir / f"data-{config_hash}"
 
 
-def with_e2e_instance_paths(config_data: dict[str, Any], data_dir: Path) -> dict[str, Any]:
+def with_e2e_instance_paths(
+    config_data: dict[str, Any],
+    data_dir: Path,
+    *,
+    resource_scope: str | None = None,
+) -> dict[str, Any]:
     """Return config data with per-instance filesystem paths rooted under ``data_dir``."""
     rendered = deepcopy(config_data)
     subprocess_working_dir = str(data_dir / "subprocess-jobs")
@@ -577,9 +587,27 @@ def with_e2e_instance_paths(config_data: dict[str, Any], data_dir: Path) -> dict
         if isinstance(default_storage_config, dict) and default_storage_config.get("type") == "local":
             default_storage_config["path"] = files_root
 
+    if resource_scope:
+        _stamp_e2e_docker_deployment_scope(rendered, resource_scope)
+
     return rendered
 
 
+def _stamp_e2e_docker_deployment_scope(config_data: dict[str, Any], resource_scope: str) -> None:
+    deployments = config_data.get("deployments")
+    if not isinstance(deployments, dict):
+        return
+    executors = deployments.get("executors")
+    if not isinstance(executors, list):
+        return
+    for executor in executors:
+        if not isinstance(executor, dict) or executor.get("backend") != "docker":
+            continue
+        executor_config = executor.setdefault("config", {})
+        if isinstance(executor_config, dict):
+            executor_config.setdefault("resource_scope", resource_scope)
+
+
 # The authz e2e suite (``e2e/authz_oidc``) editable-installs intentionally-broken
 # fixture plugins (named ``harness-*``) into the shared venv to exercise the authz
 # fail-modes. Their ``nemo.services`` entry points persist for the rest of the pytest
@@ -610,11 +638,11 @@ def _e2e_backend(harness_config: E2EHarnessConfig) -> Literal["subprocess", "doc
 
 
 def _render_e2e_config_for_backend(
-    config_data: dict[str, Any], data_dir: Path, harness_config: E2EHarnessConfig
+    config_data: dict[str, Any], data_dir: Path, harness_config: E2EHarnessConfig, config_hash: str
 ) -> dict[str, Any]:
     if _e2e_backend(harness_config) in {"docker", "docker_compose"}:
         return deepcopy(config_data)
-    return with_e2e_instance_paths(config_data, data_dir)
+    return with_e2e_instance_paths(config_data, data_dir, resource_scope=f"e2e-{config_hash}")
 
 
 class DockerBackendOverrides(TypedDict, total=False):

e2e/test_anonymizer_plugin.py

@@ -14,6 +14,7 @@
 from collections.abc import Iterator
 from contextlib import suppress
 from pathlib import Path
+from typing import cast
 
 import data_designer.config as dd
 import httpx
@@ -28,7 +29,11 @@
     AnonymizerConfigValidationError,
     AnonymizerPreviewError,
 )
-from nemo_anonymizer_plugin.sdk.job_resources import TERMINAL_INCOMPLETE_STATUSES, AnonymizerJobResource
+from nemo_anonymizer_plugin.sdk.job_resources import (
+    MAX_CONSECUTIVE_POLL_ERRORS,
+    TERMINAL_INCOMPLETE_STATUSES,
+    AnonymizerJobResource,
+)
 from nemo_anonymizer_plugin.sdk.resources import AnonymizerPreviewResult
 from nemo_platform import NeMoPlatform
 from nemo_platform_plugin.files.client import FilesClient
@@ -113,13 +118,20 @@ def _workspace_client(sdk: NeMoPlatform, workspace: str) -> NeMoPlatform:
     )
 
 
+def _require_workspace(workspace: str | None) -> str:
+    assert workspace is not None
+    return workspace
+
+
 def _anonymizer_url(sdk: NeMoPlatform, workspace: str, path: str) -> str:
     return f"{str(sdk.base_url).rstrip('/')}/apis/anonymizer/v2/workspaces/{workspace}/{path.lstrip('/')}"
 
 
-def _raw_anonymizer_post(sdk: NeMoPlatform, workspace: str, path: str, payload: dict[str, object]) -> httpx.Response:
+def _raw_anonymizer_post(
+    sdk: NeMoPlatform, workspace: str | None, path: str, payload: dict[str, object]
+) -> httpx.Response:
     return sdk._client.post(
-        _anonymizer_url(sdk, workspace, path),
+        _anonymizer_url(sdk, _require_workspace(workspace), path),
         json=payload,
         headers=_string_headers(sdk),
         timeout=sdk.timeout,
@@ -166,12 +178,12 @@ def _rewrite_config() -> AnonymizerConfig:
     )
 
 
-def _fileset_ref(workspace: str, fileset: str, path: str) -> str:
-    return f"{workspace}/{fileset}#{path}"
+def _fileset_ref(workspace: str | None, fileset: str, path: str) -> str:
+    return f"{_require_workspace(workspace)}/{fileset}#{path}"
 
 
-def _fileset_uri_ref(workspace: str, fileset: str, path: str) -> str:
-    return f"fileset://{workspace}/{fileset}#{path}"
+def _fileset_uri_ref(workspace: str | None, fileset: str, path: str) -> str:
+    return f"fileset://{_require_workspace(workspace)}/{fileset}#{path}"
 
 
 def _input_spec(source: str, *, text_column: str = TEXT_COLUMN) -> AnonymizerInputSpec:
@@ -281,14 +293,32 @@ def _job_name(job: AnonymizerJobResource) -> str:
 
 def _wait_for_anonymizer_job(job: AnonymizerJobResource, *, timeout_seconds: float) -> None:
     deadline = time.monotonic() + timeout_seconds
-    status = job.get_job_status()
+    status = None
+    consecutive_poll_errors = 0
+    last_poll_error: Exception | None = None
     while status not in {"completed", *TERMINAL_INCOMPLETE_STATUSES}:
         if time.monotonic() >= deadline:
-            logs = job.get_logs()
-            tail = logs[-5:] if logs else []
-            raise TimeoutError(f"Anonymizer job {_job_name(job)} timed out with status {status!r}; logs={tail!r}")
-        time.sleep(ANONYMIZER_POLL_INTERVAL_SECONDS)
-        status = job.get_job_status()
+            try:
+                logs = job.get_logs()
+                tail = logs[-5:] if logs else []
+            except Exception as exc:
+                tail = [f"<log retrieval failed: {exc!r}>"]
+            raise TimeoutError(
+                f"Anonymizer job {_job_name(job)} timed out with status {status!r}; "
+                f"last_poll_error={last_poll_error!r}; logs={tail!r}"
+            )
+        try:
+            status = job.get_job_status()
+        except Exception as exc:
+            consecutive_poll_errors += 1
+            last_poll_error = exc
+            if consecutive_poll_errors >= MAX_CONSECUTIVE_POLL_ERRORS:
+                raise
+        else:
+            consecutive_poll_errors = 0
+            last_poll_error = None
+        if status not in {"completed", *TERMINAL_INCOMPLETE_STATUSES}:
+            time.sleep(ANONYMIZER_POLL_INTERVAL_SECONDS)
     assert status == "completed"
 
 
@@ -433,7 +463,11 @@ def test_mock_provider_chat_completion_works_through_minikube_ingress(
         },
     )
 
-    assert SUBSTITUTE_NAME in response["choices"][0]["message"]["content"]
+    choices = cast(list[dict[str, object]], response["choices"])
+    message = cast(dict[str, object], choices[0]["message"])
+    content = message["content"]
+    assert isinstance(content, str)
+    assert SUBSTITUTE_NAME in content
 
 
 def test_file_upload_round_trips_through_minikube_ingress(
@@ -558,7 +592,7 @@ def test_preview_missing_text_column_is_rejected(
 
 
 def test_preview_invalid_strategy_payload_is_rejected(anonymizer_sdk: NeMoPlatform, anonymizer_fileset: str) -> None:
-    payload = {
+    payload: dict[str, object] = {
         "config": {"replace": {"kind": "explode"}, "emit_telemetry": False},
         "data": {
             "source": _fileset_ref(anonymizer_sdk.workspace, anonymizer_fileset, CSV_REMOTE_PATH),

e2e/test_nemo_agents.py

@@ -134,8 +134,15 @@ def _delete_deployment_if_exists(sdk: NeMoPlatform, *, workspace: str, name: str
     try:
         sdk.agents.deployments.delete(name, workspace=workspace)
     except httpx.HTTPStatusError as exc:
-        if exc.response.status_code != 404:
-            raise
+        if exc.response.status_code == 404:
+            return
+        if exc.response.status_code in {409, 500}:
+            try:
+                sdk.agents.deployments.get(name, workspace=workspace)
+            except httpx.HTTPStatusError as get_exc:
+                if get_exc.response.status_code == 404:
+                    return
+        raise
 
 
 def _get_deployment_log_text(sdk: NeMoPlatform, *, workspace: str, name: str) -> str:

packages/nmp_testing/tests/unit/test_e2e_harness.py

@@ -64,3 +64,22 @@ def test_with_e2e_instance_paths_namespaces_local_filesystem_paths(tmp_path):
     jobs_config = cast(dict[str, Any], config_data["jobs"])
     executors = cast(list[dict[str, Any]], jobs_config["executors"])
     assert executors[0]["config"]["working_directory"] == ".tmp/e2e/subprocess-jobs"
+
+
+def test_with_e2e_instance_paths_scopes_docker_deployments_executor(tmp_path):
+    data_dir = tmp_path / "data-abc123def456"
+    config_data: dict[str, Any] = {
+        "deployments": {
+            "executors": [
+                {"name": "local-docker", "backend": "docker", "config": {"pull_images": False}},
+                {"name": "local-k8s", "backend": "k8s", "config": {}},
+            ],
+        },
+    }
+
+    rendered = services_pool.with_e2e_instance_paths(config_data, data_dir, resource_scope="e2e-abc123def456")
+
+    deployments = cast(dict[str, Any], rendered["deployments"])
+    executors = cast(list[dict[str, Any]], deployments["executors"])
+    assert executors[0]["config"]["resource_scope"] == "e2e-abc123def456"
+    assert "resource_scope" not in executors[1]["config"]

plugins/nemo-agents/src/nemo_agents_plugin/api/v2/deployments.py

@@ -45,6 +45,7 @@
 router = APIRouter()
 
 _deployment_filter_dep = make_filter_obj_dep(DeploymentFilter)
+_DELETE_MARK_ATTEMPTS = 3
 
 
 @router.post("/deployments", response_model=AgentDeployment, status_code=201, tags=["Agent Deployments"])
@@ -198,9 +199,42 @@ async def delete_deployment(
     Marks the deployment as ``deleting``.  The controller terminates the
     subprocess and removes the entity on the next reconcile cycle.
     """
+    for attempt in range(_DELETE_MARK_ATTEMPTS):
+        try:
+            await _mark_deployment_deleting_once(
+                entity_client,
+                workspace=workspace,
+                name=name,
+                retrying=attempt > 0,
+            )
+            return
+        except HTTPException:
+            raise
+        except NemoEntityConflictError as exc:
+            if attempt + 1 >= _DELETE_MARK_ATTEMPTS:
+                raise HTTPException(
+                    status_code=409,
+                    detail=f"Deployment '{name}' is being modified concurrently.",
+                ) from exc
+            logger.info("Retrying delete for deployment '%s' after concurrent update", name)
+        except Exception as exc:
+            logger.exception("Failed to mark deployment '%s' as deleting", name)
+            raise HTTPException(status_code=500, detail="Failed to update deployment.") from exc
+
+
+async def _mark_deployment_deleting_once(
+    entity_client: NemoEntitiesClient,
+    *,
+    workspace: str,
+    name: str,
+    retrying: bool,
+) -> None:
     try:
         dep = await entity_client.get(AgentDeployment, name=name, workspace=workspace)
     except NemoEntityNotFoundError as exc:
+        if retrying:
+            logger.info("Deployment '%s' already deleted during delete retry", name)
+            return
         raise HTTPException(
             status_code=404,
             detail=f"Deployment '{name}' not found in workspace '{workspace}'.",
@@ -209,11 +243,12 @@ async def delete_deployment(
         logger.exception("Failed to look up deployment '%s' before delete", name)
         raise HTTPException(status_code=500, detail="Failed to look up deployment.") from exc
 
+    if dep.status == "deleting":
+        return
+
     dep.status = "deleting"
     try:
         await entity_client.update(dep)
     except NemoEntityNotFoundError:
         logger.info("Deployment '%s' already deleted before status update", name)
-    except Exception as exc:
-        logger.exception("Failed to mark deployment '%s' as deleting", name)
-        raise HTTPException(status_code=500, detail="Failed to update deployment.") from exc
+        return

plugins/nemo-agents/tests/unit/test_deployments_api.py

@@ -13,7 +13,8 @@
 from fastapi.testclient import TestClient
 from nemo_agents_plugin.api.v2 import deployments as deployments_router_module
 from nemo_agents_plugin.api.v2.dependencies import get_entity_client
-from nemo_agents_plugin.entities import NEMO_AGENTS_SPEC_CONFIG_FORMAT, Agent, AgentDeployment
+from nemo_agents_plugin.entities import NEMO_AGENTS_SPEC_CONFIG_FORMAT, Agent, AgentDeployment, DeploymentStatus
+from nemo_platform_plugin.entity_client import NemoEntityConflictError, NemoEntityNotFoundError
 
 NOW = datetime.now(timezone.utc)
 
@@ -56,6 +57,19 @@ def _make_agent(
     return agent
 
 
+def _make_deployment(
+    *,
+    name: str = "fabric-dep",
+    workspace: str = "default",
+    agent: str = "fabric-agent",
+    status: DeploymentStatus = "pending",
+) -> AgentDeployment:
+    deployment = AgentDeployment(name=name, workspace=workspace, agent=agent, status=status)
+    deployment._id = f"deployment-{name}-id"
+    deployment._created_at = NOW
+    return deployment
+
+
 def _test_client(mock_entity_client: AsyncMock) -> TestClient:
     app = FastAPI()
     app.include_router(
@@ -106,3 +120,65 @@ def test_create_rejects_invalid_platform_agent_config(self) -> None:
         assert resp.status_code == 400
         assert "Invalid agent config" in resp.json()["detail"]
         mock_entity_client.create.assert_not_called()
+
+
+class TestDeleteDeployment:
+    def test_delete_marks_deployment_deleting(self) -> None:
+        mock_entity_client = AsyncMock()
+        mock_entity_client.get = AsyncMock(return_value=_make_deployment(status="starting"))
+        mock_entity_client.update = AsyncMock(return_value=None)
+        client = _test_client(mock_entity_client)
+
+        resp = client.delete("/apis/agents/v2/workspaces/default/deployments/fabric-dep")
+
+        assert resp.status_code == 204
+        updated: AgentDeployment = mock_entity_client.update.call_args[0][0]
+        assert updated.status == "deleting"
+
+    def test_delete_retries_concurrent_update_conflict(self) -> None:
+        mock_entity_client = AsyncMock()
+        mock_entity_client.get = AsyncMock(
+            side_effect=[
+                _make_deployment(status="pending"),
+                _make_deployment(status="starting"),
+            ]
+        )
+        mock_entity_client.update = AsyncMock(side_effect=[NemoEntityConflictError("conflict"), None])
+        client = _test_client(mock_entity_client)
+
+        resp = client.delete("/apis/agents/v2/workspaces/default/deployments/fabric-dep")
+
+        assert resp.status_code == 204
+        assert mock_entity_client.get.await_count == 2
+        assert mock_entity_client.update.await_count == 2
+
+    def test_delete_returns_success_when_entity_disappears_during_retry(self) -> None:
+        mock_entity_client = AsyncMock()
+        mock_entity_client.get = AsyncMock(
+            side_effect=[
+                _make_deployment(status="pending"),
+                NemoEntityNotFoundError("gone"),
+            ]
+        )
+        mock_entity_client.update = AsyncMock(side_effect=NemoEntityConflictError("conflict"))
+        client = _test_client(mock_entity_client)
+
+        resp = client.delete("/apis/agents/v2/workspaces/default/deployments/fabric-dep")
+
+        assert resp.status_code == 204
+
+    def test_delete_returns_409_when_conflicts_exhausted(self) -> None:
+        mock_entity_client = AsyncMock()
+        mock_entity_client.get = AsyncMock(
+            side_effect=[
+                _make_deployment(status="pending")
+                for _ in range(deployments_router_module._DELETE_MARK_ATTEMPTS)  # noqa: SLF001
+            ]
+        )
+        mock_entity_client.update = AsyncMock(side_effect=NemoEntityConflictError("conflict"))
+        client = _test_client(mock_entity_client)
+
+        resp = client.delete("/apis/agents/v2/workspaces/default/deployments/fabric-dep")
+
+        assert resp.status_code == 409
+        assert mock_entity_client.update.await_count == deployments_router_module._DELETE_MARK_ATTEMPTS  # noqa: SLF001