From 575ea2e5c2d514c63bac04e38925bfb71b58ae66 Mon Sep 17 00:00:00 2001 From: Ben McCown Date: Mon, 27 Jul 2026 11:33:43 -0600 Subject: [PATCH 1/5] feat(agents): authenticate deployed agents via service-principal sidecar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When platform auth is enabled, a deployed agent's Inference Gateway calls were rejected with 401 because the agent (a NAT runtime whose OpenAI HTTP client we do not control) carries no platform credential — its LLM api_key is the placeholder "not-used". This injects a loopback auth-proxy sidecar (the nmp-api image running `nemo services run --sidecars auth-proxy`) into k8s/docker agent deployments when auth is enabled. The agent targets the sidecar on localhost; the sidecar forwards to the platform stamping X-NMP-Principal-Id: service:agents, which the OPA policy authorizes via the ServiceSystem role. This is the same static service identity the platform's own SDK clients use (get_platform_sdk as_service=...); the proxy exists only for workloads that cannot set the header themselves. The LoRA adapters sidecar already self-injects service:models via the SDK and is unaffected. - New workload_proxy sidecar: loopback FastAPI forwarder that strips inbound auth/principal headers and stamps the service principal, streaming responses. Registered as "auth-proxy" in AVAILABLE_SIDECARS. - Agents backend injects the sidecar (as a native init-container sidecar with restartPolicy=Always) and points the agent's llms.*.base_url at it when auth is enabled; unchanged behavior when auth is off. - Fix k8s compiler exec-probe: V1Probe needs the `_exec` kwarg (`exec` is a Python keyword); the sidecar uses an exec readiness probe because it binds loopback only and a pod-IP httpGet probe would be refused. - Unit tests for the sidecar forwarder, agents injection/base_url wiring, and the exec-probe compiler path. Verified end-to-end on kind with auth enabled: agent deployment reaches 2/2 Running (sidecar ready via exec probe) and the agent's IGW call authenticates (reaches model resolution instead of 401). Signed-off-by: Ben McCown --- .../nmp/common/auth/workload_proxy/main.py | 159 ++++++++++++++++++ .../tests/auth/test_workload_proxy.py | 65 +++++++ .../src/nmp/platform_runner/registry.py | 1 + .../src/nemo_agents_plugin/config.py | 19 +++ .../runner/deployments_backend.py | 99 ++++++++++- .../tests/unit/test_runner_deployments.py | 78 +++++++++ .../backends/k8s/compiler.py | 3 +- .../tests/unit/backends/k8s/test_compiler.py | 19 +++ 8 files changed, 441 insertions(+), 2 deletions(-) create mode 100644 packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py create mode 100644 packages/nmp_common/tests/auth/test_workload_proxy.py diff --git a/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py b/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py new file mode 100644 index 0000000000..3e072745c2 --- /dev/null +++ b/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Service-principal auth-proxy sidecar. + +Runs inside a deployed workload's pod as a loopback forwarder. A co-located +workload whose HTTP client we do not control (e.g. a NAT agent calling the +Inference Gateway) points its platform base URL at this proxy +(``http://127.0.0.1:``) and sends no credentials of its own. The proxy +stamps a service-principal identity header (``X-NMP-Principal-Id: service:``) +on every forwarded request, which the platform authorizes via the ServiceSystem +role. This is the same static service-identity the platform's own SDK clients +use (``get_platform_sdk(as_service=...)``); the proxy exists only for workloads +that cannot set the header themselves. + +Started via ``nemo services run --sidecars auth-proxy``. +""" + +from __future__ import annotations + +import logging +import os +import threading +from collections.abc import AsyncIterator + +import httpx +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import StreamingResponse +from starlette.background import BackgroundTask + +logger = logging.getLogger(__name__) + +# Loopback host + port the proxy listens on. The workload targets this address. +AUTH_PROXY_HOST_ENVVAR = "NMP_AUTH_PROXY_HOST" +AUTH_PROXY_PORT_ENVVAR = "NMP_AUTH_PROXY_PORT" +# Service-principal name stamped on forwarded requests (e.g. "agents"). +AUTH_PROXY_PRINCIPAL_ENVVAR = "NMP_AUTH_PROXY_PRINCIPAL" +DEFAULT_AUTH_PROXY_HOST = "127.0.0.1" +DEFAULT_AUTH_PROXY_PORT = 8090 +DEFAULT_AUTH_PROXY_PRINCIPAL = "agents" + +_READ_TIMEOUT_ENVVAR = "NMP_AUTH_PROXY_READ_TIMEOUT" +_PRINCIPAL_ID_HEADER = "x-nmp-principal-id" + +# Hop-by-hop headers must not be forwarded (RFC 7230). We also drop the +# workload's own authorization and principal headers (we set the principal), and +# host/content-length (rewritten by httpx / chunked transfer). +_STRIP_REQUEST_HEADERS = frozenset( + { + "host", + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", + "authorization", + "content-length", + _PRINCIPAL_ID_HEADER, + "x-nmp-principal-on-behalf-of", + } +) +_STRIP_RESPONSE_HEADERS = frozenset( + { + "connection", + "keep-alive", + "transfer-encoding", + "content-length", + } +) + + +def _upstream_base_url() -> str: + """Return the platform base URL to forward to (env override or platform config).""" + from nemo_platform_plugin.config import get_platform_config + + return (os.environ.get("NEMO_BASE_URL") or os.environ.get("NMP_BASE_URL") or get_platform_config().base_url).rstrip( + "/" + ) + + +def build_app(*, base_url: str, principal: str) -> FastAPI: + """Build the forwarding FastAPI app for the given upstream and service principal.""" + app = FastAPI(title="nmp-auth-proxy") + principal_id = principal if principal.startswith("service:") else f"service:{principal}" + read_timeout = float(os.environ.get(_READ_TIMEOUT_ENVVAR, "300")) + timeout = httpx.Timeout(connect=10.0, read=read_timeout, write=60.0, pool=10.0) + client = httpx.AsyncClient(base_url=base_url, timeout=timeout, follow_redirects=False) + + @app.get("/healthz") + async def healthz() -> dict[str, str]: + return {"status": "ok"} + + @app.api_route( + "/{path:path}", + methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"], + ) + async def forward(request: Request, path: str) -> StreamingResponse: + headers = {k: v for k, v in request.headers.items() if k.lower() not in _STRIP_REQUEST_HEADERS} + headers[_PRINCIPAL_ID_HEADER] = principal_id + url = httpx.URL(path="/" + path, query=request.url.query.encode("utf-8")) + body = await request.body() + upstream = client.build_request(request.method, url, headers=headers, content=body) + response = await client.send(upstream, stream=True) + + async def _body() -> AsyncIterator[bytes]: + try: + async for chunk in response.aiter_raw(): + yield chunk + finally: + await response.aclose() + + resp_headers = {k: v for k, v in response.headers.items() if k.lower() not in _STRIP_RESPONSE_HEADERS} + return StreamingResponse( + _body(), + status_code=response.status_code, + headers=resp_headers, + background=BackgroundTask(response.aclose), + ) + + return app + + +def run(parent_stop_signal: threading.Event | None = None) -> None: + """Sidecar entrypoint. Serves the loopback auth-proxy until stopped.""" + base_url = _upstream_base_url() + principal = os.environ.get(AUTH_PROXY_PRINCIPAL_ENVVAR, DEFAULT_AUTH_PROXY_PRINCIPAL) + host = os.environ.get(AUTH_PROXY_HOST_ENVVAR, DEFAULT_AUTH_PROXY_HOST) + port = int(os.environ.get(AUTH_PROXY_PORT_ENVVAR, str(DEFAULT_AUTH_PROXY_PORT))) + app = build_app(base_url=base_url, principal=principal) + + config = uvicorn.Config(app, host=host, port=port, log_level="info", access_log=False) + server = uvicorn.Server(config) + + logger.info("Starting auth-proxy sidecar on %s:%s -> %s (principal=service:%s)", host, port, base_url, principal) + if parent_stop_signal is None: + server.run() + return + + thread = threading.Thread(target=server.run, name="auth-proxy-uvicorn", daemon=True) + thread.start() + try: + while not parent_stop_signal.is_set(): + parent_stop_signal.wait(timeout=1) + finally: + server.should_exit = True + thread.join(timeout=10) + logger.info("auth-proxy sidecar stopped") + + +if __name__ == "__main__": + logging.basicConfig( + level=os.getenv("LOG_LEVEL", "INFO").upper(), + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + run() diff --git a/packages/nmp_common/tests/auth/test_workload_proxy.py b/packages/nmp_common/tests/auth/test_workload_proxy.py new file mode 100644 index 0000000000..346bac8083 --- /dev/null +++ b/packages/nmp_common/tests/auth/test_workload_proxy.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the service-principal auth-proxy sidecar forwarder.""" + +from __future__ import annotations + +import httpx +import respx +from fastapi.testclient import TestClient +from nmp.common.auth.workload_proxy.main import build_app + + +@respx.mock +def test_forward_stamps_service_principal_and_preserves_path() -> None: + upstream = "http://nemo-platform-api:8080" + route = respx.post(f"{upstream}/apis/inference-gateway/v2/workspaces/default/openai/-/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + app = build_app(base_url=upstream, principal="agents") + client = TestClient(app) + + resp = client.post( + "/apis/inference-gateway/v2/workspaces/default/openai/-/v1/chat/completions", + json={"model": "m", "messages": []}, + headers={"authorization": "Bearer not-used"}, + ) + + assert resp.status_code == 200 + assert resp.json() == {"ok": True} + assert route.called + sent = route.calls.last.request + # The proxy sets the service-principal identity and drops the placeholder auth. + assert sent.headers["x-nmp-principal-id"] == "service:agents" + assert "authorization" not in {k.lower() for k in sent.headers} + + +@respx.mock +def test_forward_normalizes_bare_principal_name() -> None: + upstream = "http://nemo-platform-api:8080" + route = respx.get(f"{upstream}/apis/entities/v2/workspaces").mock(return_value=httpx.Response(200, json={})) + # Already-prefixed principal is passed through unchanged. + app = build_app(base_url=upstream, principal="service:models") + client = TestClient(app) + client.get("/apis/entities/v2/workspaces") + assert route.calls.last.request.headers["x-nmp-principal-id"] == "service:models" + + +@respx.mock +def test_forward_passes_through_upstream_status() -> None: + upstream = "http://nemo-platform-api:8080" + respx.get(f"{upstream}/apis/entities/v2/workspaces").mock(return_value=httpx.Response(403, json={"detail": "no"})) + app = build_app(base_url=upstream, principal="agents") + client = TestClient(app) + + resp = client.get("/apis/entities/v2/workspaces") + assert resp.status_code == 403 + + +def test_healthz_does_not_require_upstream() -> None: + app = build_app(base_url="http://nemo-platform-api:8080", principal="agents") + client = TestClient(app) + resp = client.get("/healthz") + assert resp.status_code == 200 + assert resp.json() == {"status": "ok"} diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py b/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py index ab8bcefa45..79693fcbc6 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py @@ -37,6 +37,7 @@ AVAILABLE_SIDECARS: dict[str, str] = { "adapters": "nmp.core.models.sidecars.adapters.main:run", + "auth-proxy": "nmp.common.auth.workload_proxy.main:run", } SERVICE_SIDECAR_DEPENDENCIES: dict[str, set[str]] = { diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/config.py b/plugins/nemo-agents/src/nemo_agents_plugin/config.py index d18ad6e4f9..cea05d1f1c 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/config.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/config.py @@ -106,6 +106,25 @@ class DeploymentsRunnerConfig(BaseModel): "via a ConfigMap subPath. Matches the image's NAT_CONFIG_FILE convention." ), ) + auth_proxy_image_name: str = Field( + default="nmp-api", + description=( + "Image name for the auth-proxy sidecar injected when platform auth is enabled. Qualified " + "with the platform image registry/tag. Must be an nmp-api image (runs " + "`nemo services run --sidecars auth-proxy`)." + ), + ) + auth_proxy_image: str = Field( + default="", + description=( + "Optional fully-qualified image override for the auth-proxy sidecar. When set, used " + "verbatim instead of qualifying auth_proxy_image_name with the platform registry/tag." + ), + ) + auth_proxy_port: int = Field( + default=8090, + description="Loopback port the auth-proxy sidecar listens on; deployed agents target it as their inference base URL.", + ) class AgentsConfig(NemoConfig): diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/runner/deployments_backend.py b/plugins/nemo-agents/src/nemo_agents_plugin/runner/deployments_backend.py index 6a685fe28c..f364211b57 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/runner/deployments_backend.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/runner/deployments_backend.py @@ -37,6 +37,7 @@ Deployment, DeploymentConfig, EnvVar, + ExecAction, HTTPGetAction, Probe, VolumeMount, @@ -44,6 +45,7 @@ from nemo_platform.resources.entities import AsyncEntitiesResource from nemo_platform_plugin.config import LOOPBACK_ADDRESSES from nemo_platform_plugin.entity_client import NemoEntitiesClient, NemoEntityNotFoundError +from nemo_platform_plugin.jobs.image import get_qualified_image from nemo_platform_plugin.sdk_provider import get_async_platform_sdk logger = logging.getLogger(__name__) @@ -53,6 +55,24 @@ _PLUGIN_WHEELS_VOLUME = "plugin-wheels" _PLUGIN_WHEELS_MOUNT = "/opt/nemo/plugin-wheels" _NAT_CONFIG_ENV = "NAT_CONFIG_PATH" +_AUTH_PROXY_CONTAINER_NAME = "auth-proxy" +_AUTH_PROXY_HOST_ENVVAR = "NMP_AUTH_PROXY_HOST" +_AUTH_PROXY_PORT_ENVVAR = "NMP_AUTH_PROXY_PORT" +_AUTH_PROXY_PRINCIPAL_ENVVAR = "NMP_AUTH_PROXY_PRINCIPAL" +_AUTH_PROXY_PRINCIPAL = "agents" +_NATIVE_SIDECAR_RESTART_POLICY = "Always" + + +def _platform_auth_enabled() -> bool: + """Return whether platform auth is enabled (deployed agents then need a principal).""" + try: + from nmp.common.config import get_auth_config + + return bool(get_auth_config().enabled) + except Exception: + logger.debug("Could not resolve auth config; assuming auth disabled", exc_info=True) + return False + # On delete, wait up to this long for the deployments controller to tear down the # container and remove the Deployment entity before we drop the DeploymentConfig. @@ -188,6 +208,51 @@ def _info_from_deployment(deployment: Deployment) -> DeploymentInfo: return info +class AuthProxySpec: + """Parameters for injecting the service-principal auth-proxy sidecar.""" + + def __init__(self, *, image: str, port: int, upstream_base_url: str, principal: str) -> None: + self.image = image + self.port = port + self.upstream_base_url = upstream_base_url + self.principal = principal + + +def _build_auth_proxy_sidecar(spec: AuthProxySpec) -> Container: + """Build the native-sidecar Container that runs the service-principal auth proxy. + + The sidecar forwards the agent's loopback platform calls to *upstream_base_url* + with an ``X-NMP-Principal-Id: service:`` header so they authorize + under the ServiceSystem role. The deployed agent targets it on localhost. + """ + return Container( + name=_AUTH_PROXY_CONTAINER_NAME, + image=spec.image, + command=["nemo", "services", "run", "--sidecars", "auth-proxy"], + env=[ + EnvVar(name="NMP_BASE_URL", value=spec.upstream_base_url), + EnvVar(name=_AUTH_PROXY_PRINCIPAL_ENVVAR, value=spec.principal), + EnvVar(name=_AUTH_PROXY_HOST_ENVVAR, value="127.0.0.1"), + EnvVar(name=_AUTH_PROXY_PORT_ENVVAR, value=str(spec.port)), + ], + ).model_copy( + update={ + "restart_policy": _NATIVE_SIDECAR_RESTART_POLICY, + # The proxy binds loopback only (reachable solely by the co-located + # agent), so an httpGet probe against the pod IP would be refused. Use + # an exec probe that curls localhost from inside the container netns. + "readiness_probe": Probe( + exec=ExecAction( + command=["sh", "-c", f"curl -sf http://127.0.0.1:{spec.port}/healthz"], + ), + initialDelaySeconds=1, + periodSeconds=5, + failureThreshold=12, + ), + } + ) + + def build_deployment_config( *, name: str, @@ -199,6 +264,7 @@ def build_deployment_config( mode: DeploymentMode, plugin_wheels_init_image: str | None = None, labels: dict[str, str] | None = None, + auth_proxy: AuthProxySpec | None = None, ) -> DeploymentConfig: """Compile an agent into a long-running ``DeploymentConfig`` (Always). @@ -295,6 +361,14 @@ def build_deployment_config( } ) + if auth_proxy is not None: + # Native sidecar (init container with restartPolicy=Always): starts before + # the agent and forwards its loopback platform calls with a service-principal + # identity header so they authorize when auth is on. Modeled as an init + # container because the deployments plugin only allows per-container + # restart_policy there (matches the LoRA adapters sidecar pattern). + init_containers.append(_build_auth_proxy_sidecar(auth_proxy)) + return DeploymentConfig( name=name, workspace=workspace, @@ -366,7 +440,29 @@ async def create_deployment( except UnreachableGatewayURLError as exc: logger.error("Refusing to deploy agent %r: %s", name, exc) return DeploymentInfo(name=name, status="failed", error=str(exc)) - config = rewrite_config_base_urls(config, gateway) + + # When platform auth is enabled, the agent carries no platform credential, + # so route its inference calls through a loopback auth-proxy sidecar that + # stamps a service-principal identity header. The agent targets the sidecar + # on localhost; the sidecar forwards to *gateway* (the reachable platform). + auth_proxy: AuthProxySpec | None = None + if _platform_auth_enabled(): + proxy_port = self._config.auth_proxy_port + proxy_base_url = f"http://127.0.0.1:{proxy_port}" + # The sidecar runs the nmp-api image (it needs the `nemo` CLI), NOT the + # agent runtime image. Qualify with the platform registry/tag unless an + # explicit override is configured. + proxy_image = self._config.auth_proxy_image or get_qualified_image(self._config.auth_proxy_image_name) + auth_proxy = AuthProxySpec( + image=proxy_image, + port=proxy_port, + upstream_base_url=gateway, + principal=_AUTH_PROXY_PRINCIPAL, + ) + config = rewrite_config_base_urls(config, proxy_base_url) + else: + config = rewrite_config_base_urls(config, gateway) + deployment_config = build_deployment_config( name=name, workspace=workspace, @@ -380,6 +476,7 @@ async def create_deployment( "nemo.agents/deployment": name, "nemo.agents/mode": deployment_mode, }, + auth_proxy=auth_proxy, ) await entities.create(deployment_config) try: diff --git a/plugins/nemo-agents/tests/unit/test_runner_deployments.py b/plugins/nemo-agents/tests/unit/test_runner_deployments.py index bf19993ce4..abc57f6763 100644 --- a/plugins/nemo-agents/tests/unit/test_runner_deployments.py +++ b/plugins/nemo-agents/tests/unit/test_runner_deployments.py @@ -302,6 +302,84 @@ async def test_create_deployment_k8s_without_internal_url_fails() -> None: entities.create.assert_not_awaited() +@pytest.mark.asyncio +async def test_create_deployment_k8s_auth_on_injects_auth_proxy_sidecar() -> None: + backend = _backend( + default_image="nmp-api:latest", + default_executor="k8s", + k8s_internal_base_url="http://nmp-api:8080", + auth_proxy_port=8090, + ) + entities = AsyncMock() + backend._entities = entities + config = { + "llms": { + "llm": { + "_type": "openai", + "base_url": "http://localhost:8080/apis/inference-gateway/v2/workspaces/default/openai/-/v1", + } + } + } + with ( + patch("nemo_agents_plugin.runner.deployments_backend.get_base_url", return_value="http://localhost:8080"), + patch("nemo_agents_plugin.runner.deployments_backend._platform_auth_enabled", return_value=True), + ): + info = await backend.create_deployment( + workspace="default", name="hello-dep", config=config, port=0, deployment_mode="k8s" + ) + assert info.status == "starting" + created_config = entities.create.await_args_list[0].args[0] + + # Agent's inference base_url now points at the loopback auth-proxy sidecar. + baked = yaml.safe_load(created_config.config_files[0].content) + assert baked["llms"]["llm"]["base_url"] == ( + "http://127.0.0.1:8090/apis/inference-gateway/v2/workspaces/default/openai/-/v1" + ) + + # Main container list is just the agent; the auth-proxy is a native sidecar + # (init container with restartPolicy=Always). + assert [c.name for c in created_config.containers] == ["agent"] + proxy = next(c for c in created_config.init_containers if c.name == "auth-proxy") + assert proxy.restart_policy == "Always" + proxy_env = {e.name: e.value for e in proxy.env} + assert proxy_env["NMP_BASE_URL"] == "http://nmp-api:8080" + assert proxy_env["NMP_AUTH_PROXY_PRINCIPAL"] == "agents" + + +@pytest.mark.asyncio +async def test_create_deployment_k8s_auth_off_no_sidecar() -> None: + backend = _backend( + default_image="nmp-api:latest", + default_executor="k8s", + k8s_internal_base_url="http://nmp-api:8080", + ) + entities = AsyncMock() + backend._entities = entities + config = { + "llms": { + "llm": { + "_type": "openai", + "base_url": "http://localhost:8080/apis/inference-gateway/v2/workspaces/default/openai/-/v1", + } + } + } + with ( + patch("nemo_agents_plugin.runner.deployments_backend.get_base_url", return_value="http://localhost:8080"), + patch("nemo_agents_plugin.runner.deployments_backend._platform_auth_enabled", return_value=False), + ): + info = await backend.create_deployment( + workspace="default", name="hello-dep", config=config, port=0, deployment_mode="k8s" + ) + assert info.status == "starting" + created_config = entities.create.await_args_list[0].args[0] + assert [c.name for c in created_config.containers] == ["agent"] + baked = yaml.safe_load(created_config.config_files[0].content) + # Auth off: agent talks directly to the internal Service DNS (PR #899 behavior). + assert baked["llms"]["llm"]["base_url"] == ( + "http://nmp-api:8080/apis/inference-gateway/v2/workspaces/default/openai/-/v1" + ) + + @pytest.mark.asyncio async def test_create_deployment_missing_image_fails() -> None: backend = _backend(default_image="") diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/compiler.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/compiler.py index d3a43be341..9d7205d59a 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/compiler.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/compiler.py @@ -199,7 +199,8 @@ def _build_probe(probe: Probe | None) -> Any | None: "failure_threshold": probe.failure_threshold, } if probe.exec_action is not None: - kwargs["exec"] = k8s.client.V1ExecAction(command=list(probe.exec_action.command)) + # V1Probe uses `_exec` (Python keyword `exec` cannot be a kwarg name). + kwargs["_exec"] = k8s.client.V1ExecAction(command=list(probe.exec_action.command)) elif probe.http_get is not None: kwargs["http_get"] = k8s.client.V1HTTPGetAction( path=probe.http_get.path, diff --git a/plugins/nemo-deployments/tests/unit/backends/k8s/test_compiler.py b/plugins/nemo-deployments/tests/unit/backends/k8s/test_compiler.py index 402db73cb1..ae79541679 100644 --- a/plugins/nemo-deployments/tests/unit/backends/k8s/test_compiler.py +++ b/plugins/nemo-deployments/tests/unit/backends/k8s/test_compiler.py @@ -11,6 +11,7 @@ from kubernetes.client import ApiClient from nemo_deployments_plugin.backends.k8s.compiler import ( DeploymentConfigError, + _build_probe, build_configmap_body, compile_workload, configmap_data_key, @@ -23,7 +24,10 @@ ConfigFile, Container, ContainerPort, + ExecAction, + HTTPGetAction, K8sDeploymentConfig, + Probe, ) from nemo_platform_plugin.config import ImagePullSecret @@ -32,6 +36,21 @@ def _serialized(obj: object) -> dict: return ApiClient().sanitize_for_serialization(obj) +def test_build_probe_exec_action() -> None: + # V1Probe requires the `_exec` kwarg (exec is a Python keyword); regression + # for the exec-probe path used by loopback-only sidecars. + probe = _build_probe(Probe(exec=ExecAction(command=["sh", "-c", "curl -sf http://127.0.0.1:8090/healthz"]))) + serialized = _serialized(probe) + assert serialized["exec"]["command"] == ["sh", "-c", "curl -sf http://127.0.0.1:8090/healthz"] + + +def test_build_probe_http_get_action() -> None: + probe = _build_probe(Probe(httpGet=HTTPGetAction(path="/health", port=8000))) + serialized = _serialized(probe) + assert serialized["httpGet"]["path"] == "/health" + assert serialized["httpGet"]["port"] == 8000 + + def test_configmap_data_key_sanitizes_paths() -> None: assert configmap_data_key("/etc/app/config.yaml") == "etc__app__config.yaml" From 5af39855da25d1a4bb7fe64b56ff13bbdf3edba7 Mon Sep 17 00:00:00 2001 From: Ben McCown Date: Mon, 27 Jul 2026 12:01:52 -0600 Subject: [PATCH 2/5] refactor(deployments): own auth-proxy sidecar compilation via config flags Addresses PR feedback: the auth-proxy sidecar belongs to nemo-deployments (which compiles containers into deployments), not nemo-agents. Also removes a direct nmp_common reference from the agents plugin. - DeploymentConfig gains `auth_proxy_sidecar: bool` and `auth_proxy_sidecar_identity: str | None`. The deployments plugin compiles and injects the sidecar (nemo_deployments_plugin.auth_proxy) from these flags, interpolating the identity into the service-principal header. Injection is a no-op when platform auth is disabled. - nemo-agents now just sets the two flags (identity="agents") and points the agent's llms.*.base_url at the proxy port when auth is enabled; it no longer builds the sidecar container or reads deployments image/port config. - Plugins no longer import nmp_common: added nemo_platform_plugin.auth .platform_auth_enabled() wrapper; both agents and deployments consult it. - Sidecar image/port config moved from AgentsConfig to DeploymentsConfig. - Tests: deployments-side build_auth_proxy_container (requested/auth-off no-op/ identity default) and agents-side flag-setting + base_url wiring. Signed-off-by: Ben McCown --- .../src/nemo_platform_plugin/auth.py | 28 +++++ .../src/nemo_agents_plugin/config.py | 19 --- .../runner/deployments_backend.py | 109 +++--------------- .../tests/unit/test_runner_deployments.py | 31 ++--- .../src/nemo_deployments_plugin/auth_proxy.py | 93 +++++++++++++++ .../backends/k8s/compiler.py | 10 +- .../src/nemo_deployments_plugin/config.py | 15 +++ .../src/nemo_deployments_plugin/entities.py | 17 +++ .../tests/unit/test_auth_proxy.py | 63 ++++++++++ 9 files changed, 258 insertions(+), 127 deletions(-) create mode 100644 packages/nemo_platform_plugin/src/nemo_platform_plugin/auth.py create mode 100644 plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py create mode 100644 plugins/nemo-deployments/tests/unit/test_auth_proxy.py diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/auth.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/auth.py new file mode 100644 index 0000000000..9a057ddf84 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/auth.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Auth helpers exposed to plugins. + +Plugins must not import ``nmp_common`` directly. This module wraps the pieces of +the platform auth configuration that plugins need. The underlying auth config +lives in ``nmp_common`` (only present in the platform process image), so it is +imported lazily and failures degrade to "disabled" rather than raising in +environments without it. +""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + + +def platform_auth_enabled() -> bool: + """Return whether platform authentication is enabled.""" + try: + from nmp.common.config import get_auth_config + + return bool(get_auth_config().enabled) + except Exception: + logger.debug("Could not resolve auth config; assuming auth disabled", exc_info=True) + return False diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/config.py b/plugins/nemo-agents/src/nemo_agents_plugin/config.py index cea05d1f1c..d18ad6e4f9 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/config.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/config.py @@ -106,25 +106,6 @@ class DeploymentsRunnerConfig(BaseModel): "via a ConfigMap subPath. Matches the image's NAT_CONFIG_FILE convention." ), ) - auth_proxy_image_name: str = Field( - default="nmp-api", - description=( - "Image name for the auth-proxy sidecar injected when platform auth is enabled. Qualified " - "with the platform image registry/tag. Must be an nmp-api image (runs " - "`nemo services run --sidecars auth-proxy`)." - ), - ) - auth_proxy_image: str = Field( - default="", - description=( - "Optional fully-qualified image override for the auth-proxy sidecar. When set, used " - "verbatim instead of qualifying auth_proxy_image_name with the platform registry/tag." - ), - ) - auth_proxy_port: int = Field( - default=8090, - description="Loopback port the auth-proxy sidecar listens on; deployed agents target it as their inference base URL.", - ) class AgentsConfig(NemoConfig): diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/runner/deployments_backend.py b/plugins/nemo-agents/src/nemo_agents_plugin/runner/deployments_backend.py index f364211b57..bc4aac021c 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/runner/deployments_backend.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/runner/deployments_backend.py @@ -30,6 +30,7 @@ ) from nemo_agents_plugin.runner.backend import DeploymentInfo, ExternalLog, LogLocation, RunnerBackend from nemo_agents_plugin.utils import get_base_url, get_internal_base_url +from nemo_deployments_plugin.auth_proxy import auth_proxy_port from nemo_deployments_plugin.entities import ( ConfigFile, Container, @@ -37,15 +38,14 @@ Deployment, DeploymentConfig, EnvVar, - ExecAction, HTTPGetAction, Probe, VolumeMount, ) from nemo_platform.resources.entities import AsyncEntitiesResource +from nemo_platform_plugin.auth import platform_auth_enabled from nemo_platform_plugin.config import LOOPBACK_ADDRESSES from nemo_platform_plugin.entity_client import NemoEntitiesClient, NemoEntityNotFoundError -from nemo_platform_plugin.jobs.image import get_qualified_image from nemo_platform_plugin.sdk_provider import get_async_platform_sdk logger = logging.getLogger(__name__) @@ -55,23 +55,7 @@ _PLUGIN_WHEELS_VOLUME = "plugin-wheels" _PLUGIN_WHEELS_MOUNT = "/opt/nemo/plugin-wheels" _NAT_CONFIG_ENV = "NAT_CONFIG_PATH" -_AUTH_PROXY_CONTAINER_NAME = "auth-proxy" -_AUTH_PROXY_HOST_ENVVAR = "NMP_AUTH_PROXY_HOST" -_AUTH_PROXY_PORT_ENVVAR = "NMP_AUTH_PROXY_PORT" -_AUTH_PROXY_PRINCIPAL_ENVVAR = "NMP_AUTH_PROXY_PRINCIPAL" -_AUTH_PROXY_PRINCIPAL = "agents" -_NATIVE_SIDECAR_RESTART_POLICY = "Always" - - -def _platform_auth_enabled() -> bool: - """Return whether platform auth is enabled (deployed agents then need a principal).""" - try: - from nmp.common.config import get_auth_config - - return bool(get_auth_config().enabled) - except Exception: - logger.debug("Could not resolve auth config; assuming auth disabled", exc_info=True) - return False +_AUTH_PROXY_IDENTITY = "agents" # On delete, wait up to this long for the deployments controller to tear down the @@ -208,51 +192,6 @@ def _info_from_deployment(deployment: Deployment) -> DeploymentInfo: return info -class AuthProxySpec: - """Parameters for injecting the service-principal auth-proxy sidecar.""" - - def __init__(self, *, image: str, port: int, upstream_base_url: str, principal: str) -> None: - self.image = image - self.port = port - self.upstream_base_url = upstream_base_url - self.principal = principal - - -def _build_auth_proxy_sidecar(spec: AuthProxySpec) -> Container: - """Build the native-sidecar Container that runs the service-principal auth proxy. - - The sidecar forwards the agent's loopback platform calls to *upstream_base_url* - with an ``X-NMP-Principal-Id: service:`` header so they authorize - under the ServiceSystem role. The deployed agent targets it on localhost. - """ - return Container( - name=_AUTH_PROXY_CONTAINER_NAME, - image=spec.image, - command=["nemo", "services", "run", "--sidecars", "auth-proxy"], - env=[ - EnvVar(name="NMP_BASE_URL", value=spec.upstream_base_url), - EnvVar(name=_AUTH_PROXY_PRINCIPAL_ENVVAR, value=spec.principal), - EnvVar(name=_AUTH_PROXY_HOST_ENVVAR, value="127.0.0.1"), - EnvVar(name=_AUTH_PROXY_PORT_ENVVAR, value=str(spec.port)), - ], - ).model_copy( - update={ - "restart_policy": _NATIVE_SIDECAR_RESTART_POLICY, - # The proxy binds loopback only (reachable solely by the co-located - # agent), so an httpGet probe against the pod IP would be refused. Use - # an exec probe that curls localhost from inside the container netns. - "readiness_probe": Probe( - exec=ExecAction( - command=["sh", "-c", f"curl -sf http://127.0.0.1:{spec.port}/healthz"], - ), - initialDelaySeconds=1, - periodSeconds=5, - failureThreshold=12, - ), - } - ) - - def build_deployment_config( *, name: str, @@ -264,7 +203,7 @@ def build_deployment_config( mode: DeploymentMode, plugin_wheels_init_image: str | None = None, labels: dict[str, str] | None = None, - auth_proxy: AuthProxySpec | None = None, + auth_proxy_identity: str | None = None, ) -> DeploymentConfig: """Compile an agent into a long-running ``DeploymentConfig`` (Always). @@ -361,14 +300,8 @@ def build_deployment_config( } ) - if auth_proxy is not None: - # Native sidecar (init container with restartPolicy=Always): starts before - # the agent and forwards its loopback platform calls with a service-principal - # identity header so they authorize when auth is on. Modeled as an init - # container because the deployments plugin only allows per-container - # restart_policy there (matches the LoRA adapters sidecar pattern). - init_containers.append(_build_auth_proxy_sidecar(auth_proxy)) - + # Request the auth-proxy sidecar via the DeploymentConfig flags; the + # deployments plugin compiles and injects it (and no-ops when auth is off). return DeploymentConfig( name=name, workspace=workspace, @@ -381,6 +314,8 @@ def build_deployment_config( ConfigFile(path=config_mount_path, content=nat_yaml), ], "restart_policy": "Always", + "auth_proxy_sidecar": auth_proxy_identity is not None, + "auth_proxy_sidecar_identity": auth_proxy_identity, } ) @@ -442,24 +377,14 @@ async def create_deployment( return DeploymentInfo(name=name, status="failed", error=str(exc)) # When platform auth is enabled, the agent carries no platform credential, - # so route its inference calls through a loopback auth-proxy sidecar that - # stamps a service-principal identity header. The agent targets the sidecar - # on localhost; the sidecar forwards to *gateway* (the reachable platform). - auth_proxy: AuthProxySpec | None = None - if _platform_auth_enabled(): - proxy_port = self._config.auth_proxy_port - proxy_base_url = f"http://127.0.0.1:{proxy_port}" - # The sidecar runs the nmp-api image (it needs the `nemo` CLI), NOT the - # agent runtime image. Qualify with the platform registry/tag unless an - # explicit override is configured. - proxy_image = self._config.auth_proxy_image or get_qualified_image(self._config.auth_proxy_image_name) - auth_proxy = AuthProxySpec( - image=proxy_image, - port=proxy_port, - upstream_base_url=gateway, - principal=_AUTH_PROXY_PRINCIPAL, - ) - config = rewrite_config_base_urls(config, proxy_base_url) + # so route its inference calls through a loopback auth-proxy sidecar (the + # deployments plugin compiles the sidecar from the auth_proxy flags). The + # agent targets the sidecar on localhost; the sidecar forwards to the + # platform with a service-principal identity header. + auth_proxy_identity: str | None = None + if platform_auth_enabled(): + auth_proxy_identity = _AUTH_PROXY_IDENTITY + config = rewrite_config_base_urls(config, f"http://127.0.0.1:{auth_proxy_port()}") else: config = rewrite_config_base_urls(config, gateway) @@ -476,7 +401,7 @@ async def create_deployment( "nemo.agents/deployment": name, "nemo.agents/mode": deployment_mode, }, - auth_proxy=auth_proxy, + auth_proxy_identity=auth_proxy_identity, ) await entities.create(deployment_config) try: diff --git a/plugins/nemo-agents/tests/unit/test_runner_deployments.py b/plugins/nemo-agents/tests/unit/test_runner_deployments.py index abc57f6763..21e3e15d19 100644 --- a/plugins/nemo-agents/tests/unit/test_runner_deployments.py +++ b/plugins/nemo-agents/tests/unit/test_runner_deployments.py @@ -303,12 +303,13 @@ async def test_create_deployment_k8s_without_internal_url_fails() -> None: @pytest.mark.asyncio -async def test_create_deployment_k8s_auth_on_injects_auth_proxy_sidecar() -> None: +async def test_create_deployment_k8s_auth_on_requests_auth_proxy_sidecar() -> None: + # Agents layer only sets the DeploymentConfig flags + points the agent at the + # proxy port; the deployments plugin compiles the actual sidecar container. backend = _backend( default_image="nmp-api:latest", default_executor="k8s", k8s_internal_base_url="http://nmp-api:8080", - auth_proxy_port=8090, ) entities = AsyncMock() backend._entities = entities @@ -322,7 +323,8 @@ async def test_create_deployment_k8s_auth_on_injects_auth_proxy_sidecar() -> Non } with ( patch("nemo_agents_plugin.runner.deployments_backend.get_base_url", return_value="http://localhost:8080"), - patch("nemo_agents_plugin.runner.deployments_backend._platform_auth_enabled", return_value=True), + patch("nemo_agents_plugin.runner.deployments_backend.platform_auth_enabled", return_value=True), + patch("nemo_agents_plugin.runner.deployments_backend.auth_proxy_port", return_value=8090), ): info = await backend.create_deployment( workspace="default", name="hello-dep", config=config, port=0, deployment_mode="k8s" @@ -330,21 +332,19 @@ async def test_create_deployment_k8s_auth_on_injects_auth_proxy_sidecar() -> Non assert info.status == "starting" created_config = entities.create.await_args_list[0].args[0] - # Agent's inference base_url now points at the loopback auth-proxy sidecar. + # The DeploymentConfig requests the sidecar with the agents identity; the + # agents layer does not build the container itself. + assert created_config.auth_proxy_sidecar is True + assert created_config.auth_proxy_sidecar_identity == "agents" + assert [c.name for c in created_config.containers] == ["agent"] + assert [c.name for c in created_config.init_containers] == [] + + # Agent's inference base_url points at the loopback proxy port. baked = yaml.safe_load(created_config.config_files[0].content) assert baked["llms"]["llm"]["base_url"] == ( "http://127.0.0.1:8090/apis/inference-gateway/v2/workspaces/default/openai/-/v1" ) - # Main container list is just the agent; the auth-proxy is a native sidecar - # (init container with restartPolicy=Always). - assert [c.name for c in created_config.containers] == ["agent"] - proxy = next(c for c in created_config.init_containers if c.name == "auth-proxy") - assert proxy.restart_policy == "Always" - proxy_env = {e.name: e.value for e in proxy.env} - assert proxy_env["NMP_BASE_URL"] == "http://nmp-api:8080" - assert proxy_env["NMP_AUTH_PROXY_PRINCIPAL"] == "agents" - @pytest.mark.asyncio async def test_create_deployment_k8s_auth_off_no_sidecar() -> None: @@ -365,14 +365,15 @@ async def test_create_deployment_k8s_auth_off_no_sidecar() -> None: } with ( patch("nemo_agents_plugin.runner.deployments_backend.get_base_url", return_value="http://localhost:8080"), - patch("nemo_agents_plugin.runner.deployments_backend._platform_auth_enabled", return_value=False), + patch("nemo_agents_plugin.runner.deployments_backend.platform_auth_enabled", return_value=False), ): info = await backend.create_deployment( workspace="default", name="hello-dep", config=config, port=0, deployment_mode="k8s" ) assert info.status == "starting" created_config = entities.create.await_args_list[0].args[0] - assert [c.name for c in created_config.containers] == ["agent"] + assert created_config.auth_proxy_sidecar is False + assert created_config.auth_proxy_sidecar_identity is None baked = yaml.safe_load(created_config.config_files[0].content) # Auth off: agent talks directly to the internal Service DNS (PR #899 behavior). assert baked["llms"]["llm"]["base_url"] == ( diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py new file mode 100644 index 0000000000..a24f82187a --- /dev/null +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Auth-proxy sidecar compilation. + +A DeploymentConfig with ``auth_proxy_sidecar=True`` gets a loopback auth-proxy +sidecar injected: the nmp-api image running ``nemo services run --sidecars +auth-proxy``, which stamps ``X-NMP-Principal-Id: service:`` on the +workload's platform calls. The workload targets the proxy on localhost. + +Injection is a no-op when platform auth is disabled — the workload's calls are +already trusted on the internal network, so no identity header is needed. +""" + +from __future__ import annotations + +import logging + +from nemo_deployments_plugin.config import DeploymentsConfig +from nemo_deployments_plugin.entities import ( + Container, + DeploymentConfig, + EnvVar, + ExecAction, + Probe, + RestartPolicy, +) +from nemo_platform_plugin.auth import platform_auth_enabled +from nemo_platform_plugin.config import get_nemo_config +from nemo_platform_plugin.jobs.image import get_qualified_image + +logger = logging.getLogger(__name__) + +AUTH_PROXY_CONTAINER_NAME = "auth-proxy" +_NATIVE_SIDECAR_RESTART_POLICY: RestartPolicy = "Always" +_AUTH_PROXY_PRINCIPAL_ENVVAR = "NMP_AUTH_PROXY_PRINCIPAL" +_AUTH_PROXY_HOST_ENVVAR = "NMP_AUTH_PROXY_HOST" +_AUTH_PROXY_PORT_ENVVAR = "NMP_AUTH_PROXY_PORT" +_DEFAULT_IDENTITY = "agents" + + +def auth_proxy_port() -> int: + """Return the loopback port the auth-proxy sidecar listens on.""" + return get_nemo_config(DeploymentsConfig).auth_proxy_port + + +def _upstream_base_url() -> str: + from nemo_platform_plugin.config import get_platform_config + + return get_platform_config().base_url.rstrip("/") + + +def build_auth_proxy_container(config: DeploymentConfig) -> Container | None: + """Return the auth-proxy sidecar Container for *config*, or None. + + Returns None when the config does not request the sidecar, or when platform + auth is disabled (the sidecar would be pointless — internal calls are already + trusted). + """ + if not config.auth_proxy_sidecar: + return None + if not platform_auth_enabled(): + logger.debug("auth_proxy_sidecar requested but platform auth is disabled; skipping sidecar injection") + return None + + deployments_config = get_nemo_config(DeploymentsConfig) + identity = config.auth_proxy_sidecar_identity or _DEFAULT_IDENTITY + port = deployments_config.auth_proxy_port + image = deployments_config.auth_proxy_image or get_qualified_image(deployments_config.auth_proxy_image_name) + + return Container( + name=AUTH_PROXY_CONTAINER_NAME, + image=image, + command=["nemo", "services", "run", "--sidecars", "auth-proxy"], + env=[ + EnvVar(name="NMP_BASE_URL", value=_upstream_base_url()), + EnvVar(name=_AUTH_PROXY_PRINCIPAL_ENVVAR, value=identity), + EnvVar(name=_AUTH_PROXY_HOST_ENVVAR, value="127.0.0.1"), + EnvVar(name=_AUTH_PROXY_PORT_ENVVAR, value=str(port)), + ], + ).model_copy( + update={ + "restart_policy": _NATIVE_SIDECAR_RESTART_POLICY, + # The proxy binds loopback only, so a pod-IP httpGet probe would be + # refused; exec-probe curls localhost inside the container netns. + "readiness_probe": Probe( + exec=ExecAction(command=["sh", "-c", f"curl -sf http://127.0.0.1:{port}/healthz"]), + initialDelaySeconds=1, + periodSeconds=5, + failureThreshold=12, + ), + } + ) diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/compiler.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/compiler.py index 9d7205d59a..bb8340c1f6 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/compiler.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/compiler.py @@ -16,6 +16,7 @@ from typing import Any from kubernetes.client.rest import ApiException +from nemo_deployments_plugin.auth_proxy import build_auth_proxy_container from nemo_deployments_plugin.backends.k8s.client import k8s_client_module from nemo_deployments_plugin.backends.k8s.status import resource_labels_match from nemo_deployments_plugin.backends.labels import k8s_deployment_configmap_name, k8s_volume_resource_name @@ -383,13 +384,20 @@ def compile_workload( if configmap_name is not None: volumes = [*volumes, _build_config_file_volume(configmap_name, config.config_files)] + ordered_init = list(_ordered_init_containers(config)) + # Auth-proxy sidecar (native sidecar with restartPolicy=Always) is appended so + # it starts before the main workload and keeps running. No-op when the config + # does not request it or platform auth is disabled. + auth_proxy = build_auth_proxy_container(config) + if auth_proxy is not None: + ordered_init.append(auth_proxy) init_containers = [ build_container( container, config=config, include_probes=container.restart_policy == NATIVE_SIDECAR_RESTART_POLICY, ) - for container in _ordered_init_containers(config) + for container in ordered_init ] main_containers = [ build_container(container, config=config, include_probes=True) for container in config.containers diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/config.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/config.py index b789ac5d60..3d50fcc8e0 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/config.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/config.py @@ -72,3 +72,18 @@ class DeploymentsConfig(NemoConfig): description="Default executor when Deployment.executor is unset.", ) controller: ControllerConfig = Field(default_factory=ControllerConfig) + auth_proxy_image_name: str = Field( + default="nmp-api", + description=( + "Image name for the auth-proxy sidecar (qualified with the platform image registry/tag). " + "Must be an nmp-api image (runs `nemo services run --sidecars auth-proxy`)." + ), + ) + auth_proxy_image: str = Field( + default="", + description="Optional fully-qualified image override for the auth-proxy sidecar.", + ) + auth_proxy_port: int = Field( + default=8090, + description="Loopback port the auth-proxy sidecar listens on; workloads target it as their platform base URL.", + ) diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py index b0d653a56c..4f810b4bfb 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py @@ -275,6 +275,23 @@ class DeploymentConfig(NemoEntity, entity_type=ENTITY_TYPE_DEPLOYMENT_CONFIG): drift_recovery: DriftRecoveryPolicy = Field(default_factory=DriftRecoveryPolicy, alias="driftRecovery") labels: dict[str, str] = Field(default_factory=dict) backend_config: DeploymentBackendConfig = Field(default_factory=DeploymentBackendConfig, alias="backendConfig") + auth_proxy_sidecar: bool = Field( + default=False, + alias="authProxySidecar", + description=( + "Inject a loopback auth-proxy sidecar that stamps a service-principal identity header on the " + "workload's platform calls. No-op when platform auth is disabled. The workload must target the " + "proxy on localhost (see auth_proxy_sidecar_port)." + ), + ) + auth_proxy_sidecar_identity: str | None = Field( + default=None, + alias="authProxySidecarIdentity", + description=( + "Service-principal name the auth-proxy sidecar stamps (interpolated into " + "'X-NMP-Principal-Id: service:'). Required when auth_proxy_sidecar is True." + ), + ) model_config = {"populate_by_name": True} diff --git a/plugins/nemo-deployments/tests/unit/test_auth_proxy.py b/plugins/nemo-deployments/tests/unit/test_auth_proxy.py new file mode 100644 index 0000000000..d94e5ccb8c --- /dev/null +++ b/plugins/nemo-deployments/tests/unit/test_auth_proxy.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for auth-proxy sidecar compilation in the deployments plugin.""" + +from __future__ import annotations + +from unittest.mock import patch + +from nemo_deployments_plugin.auth_proxy import AUTH_PROXY_CONTAINER_NAME, build_auth_proxy_container +from nemo_deployments_plugin.entities import DeploymentConfig + +_MOD = "nemo_deployments_plugin.auth_proxy" + + +def _config(**kwargs) -> DeploymentConfig: + return DeploymentConfig(name="dep", workspace="default", **kwargs) + + +def test_no_sidecar_when_not_requested() -> None: + with patch(f"{_MOD}.platform_auth_enabled", return_value=True): + assert build_auth_proxy_container(_config(auth_proxy_sidecar=False)) is None + + +def test_no_sidecar_when_auth_disabled() -> None: + # Requested but auth off -> no-op. + with patch(f"{_MOD}.platform_auth_enabled", return_value=False): + assert ( + build_auth_proxy_container(_config(auth_proxy_sidecar=True, auth_proxy_sidecar_identity="agents")) is None + ) + + +def test_builds_sidecar_when_requested_and_auth_on() -> None: + with ( + patch(f"{_MOD}.platform_auth_enabled", return_value=True), + patch(f"{_MOD}.get_qualified_image", return_value="my-registry/nmp-api:local"), + patch(f"{_MOD}._upstream_base_url", return_value="http://nemo-platform-api:8080"), + ): + container = build_auth_proxy_container(_config(auth_proxy_sidecar=True, auth_proxy_sidecar_identity="agents")) + assert container is not None + assert container.name == AUTH_PROXY_CONTAINER_NAME + assert container.image == "my-registry/nmp-api:local" + assert container.command == ["nemo", "services", "run", "--sidecars", "auth-proxy"] + assert container.restart_policy == "Always" + env = {e.name: e.value for e in container.env} + assert env["NMP_AUTH_PROXY_PRINCIPAL"] == "agents" + assert env["NMP_BASE_URL"] == "http://nemo-platform-api:8080" + # Loopback exec probe (proxy binds 127.0.0.1, so pod-IP httpGet would be refused). + assert container.readiness_probe is not None + assert container.readiness_probe.exec_action is not None + assert "127.0.0.1" in " ".join(container.readiness_probe.exec_action.command) + + +def test_identity_defaults_to_agents_when_unset() -> None: + with ( + patch(f"{_MOD}.platform_auth_enabled", return_value=True), + patch(f"{_MOD}.get_qualified_image", return_value="img"), + patch(f"{_MOD}._upstream_base_url", return_value="http://x:8080"), + ): + container = build_auth_proxy_container(_config(auth_proxy_sidecar=True)) + assert container is not None + env = {e.name: e.value for e in container.env} + assert env["NMP_AUTH_PROXY_PRINCIPAL"] == "agents" From 3bbed59e87bdffd9901c65dca4d1a49a60da71fd Mon Sep 17 00:00:00 2001 From: Ben McCown Date: Mon, 27 Jul 2026 12:24:19 -0600 Subject: [PATCH 3/5] feat(deployments): inject auth-proxy sidecar for docker; document auth fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses PR feedback. - Docker backend now injects the auth-proxy sidecar too (build_docker_plan), reusing build_auth_proxy_container. The sidecar shares the primary container's network namespace, so the agent reaches it on localhost:8090 — the same loopback address used in k8s. Previously the sidecar was injected only in the k8s compiler, so auth-on docker agent deployments got no identity and would 401. - The docker sidecar's upstream (NMP_BASE_URL) is rewritten to a docker-reachable host via determine_loopback_override() (e.g. host.docker.internal on macOS) when the platform base URL is a loopback, mirroring the agent-side rewrite from #899. K8s uses the Service DNS verbatim. - Documented platform_auth_enabled()'s fail-to-False behavior: the realistic failure is ImportError (package used outside the platform image); other failures are effectively unreachable in the controller (missing config -> defaults; malformed config would have crashed startup; the read is cached). - Tests: docker auth-proxy injection, auth-off no-op, and loopback upstream rewrite. Signed-off-by: Ben McCown --- .../src/nemo_platform_plugin/auth.py | 17 ++++++- .../src/nemo_deployments_plugin/auth_proxy.py | 37 +++++++++++---- .../backends/docker/containers.py | 8 ++++ .../unit/backends/docker/test_containers.py | 45 +++++++++++++++++++ 4 files changed, 97 insertions(+), 10 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/auth.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/auth.py index 9a057ddf84..d8b035b5ea 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/auth.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/auth.py @@ -18,7 +18,22 @@ def platform_auth_enabled() -> bool: - """Return whether platform authentication is enabled.""" + """Return whether platform authentication is enabled. + + Returns ``False`` on any failure to resolve the auth config. The realistic + failure is ``ImportError``: ``nmp_common`` ships only in the platform process + image, so when this package is used standalone (outside the platform) there + is no auth config and "disabled" is the correct answer. + + Other failures are effectively unreachable in the context that matters here + (the deployment controller, which runs *inside* the platform image): a + missing config file resolves to defaults (``enabled=False``) rather than + raising, and a malformed/invalid config file would have already crashed the + platform service at startup before any deployment is reconciled. The config + read is cached from that successful startup load. We therefore accept the + narrow, largely theoretical fail-open window rather than propagate and block + deployments on a transient/unexpected error. + """ try: from nmp.common.config import get_auth_config diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py index a24f82187a..d15febdad7 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py @@ -15,6 +15,7 @@ from __future__ import annotations import logging +from urllib.parse import urlsplit from nemo_deployments_plugin.config import DeploymentsConfig from nemo_deployments_plugin.entities import ( @@ -26,7 +27,7 @@ RestartPolicy, ) from nemo_platform_plugin.auth import platform_auth_enabled -from nemo_platform_plugin.config import get_nemo_config +from nemo_platform_plugin.config import LOOPBACK_ADDRESSES, get_nemo_config from nemo_platform_plugin.jobs.image import get_qualified_image logger = logging.getLogger(__name__) @@ -44,18 +45,36 @@ def auth_proxy_port() -> int: return get_nemo_config(DeploymentsConfig).auth_proxy_port -def _upstream_base_url() -> str: - from nemo_platform_plugin.config import get_platform_config +def _upstream_base_url(*, docker: bool) -> str: + """Return the platform base URL the sidecar forwards to, reachable from its container. - return get_platform_config().base_url.rstrip("/") - - -def build_auth_proxy_container(config: DeploymentConfig) -> Container | None: + In docker mode the platform base URL is often a host loopback the container + cannot reach; substitute the docker-reachable host (e.g. host.docker.internal) + the same way jobs do. In k8s the base URL is the in-cluster Service DNS and is + used verbatim. + """ + from nemo_platform_plugin.config import determine_loopback_override, get_platform_config + + base_url = get_platform_config().base_url.rstrip("/") + if not docker: + return base_url + override = determine_loopback_override() + if not override: + return base_url + parts = urlsplit(base_url) + if (parts.hostname or "").lower() not in LOOPBACK_ADDRESSES: + return base_url + netloc = override if parts.port is None else f"{override}:{parts.port}" + return parts._replace(netloc=netloc).geturl() + + +def build_auth_proxy_container(config: DeploymentConfig, *, docker: bool = False) -> Container | None: """Return the auth-proxy sidecar Container for *config*, or None. Returns None when the config does not request the sidecar, or when platform auth is disabled (the sidecar would be pointless — internal calls are already - trusted). + trusted). Pass ``docker=True`` so the sidecar's upstream is rewritten to a + docker-reachable host. """ if not config.auth_proxy_sidecar: return None @@ -73,7 +92,7 @@ def build_auth_proxy_container(config: DeploymentConfig) -> Container | None: image=image, command=["nemo", "services", "run", "--sidecars", "auth-proxy"], env=[ - EnvVar(name="NMP_BASE_URL", value=_upstream_base_url()), + EnvVar(name="NMP_BASE_URL", value=_upstream_base_url(docker=docker)), EnvVar(name=_AUTH_PROXY_PRINCIPAL_ENVVAR, value=identity), EnvVar(name=_AUTH_PROXY_HOST_ENVVAR, value="127.0.0.1"), EnvVar(name=_AUTH_PROXY_PORT_ENVVAR, value=str(port)), diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/containers.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/containers.py index b5575423e5..8406fee809 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/containers.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/containers.py @@ -8,6 +8,7 @@ from dataclasses import dataclass, field from typing import Any +from nemo_deployments_plugin.auth_proxy import build_auth_proxy_container from nemo_deployments_plugin.backends.labels import docker_volume_name from nemo_deployments_plugin.entities import Container, DeploymentConfig, DockerDeploymentConfig, VolumeMount from nemo_deployments_plugin.types import RestartPolicy @@ -58,6 +59,13 @@ def build_docker_plan(config: DeploymentConfig) -> DockerDeploymentPlan: primary = config.containers[0] sidecars = list(config.containers[1:]) + # Auth-proxy sidecar (no-op unless requested and platform auth is enabled). + # It shares the primary's netns, so the primary reaches it on localhost — + # the same loopback address the workload targets in k8s. Declares no ports. + auth_proxy = build_auth_proxy_container(config, docker=True) + if auth_proxy is not None: + sidecars.append(auth_proxy) + # Sidecars share the primary's netns, so they cannot publish their own host # ports. (The primary owns the published ports for the whole group.) for sidecar in sidecars: diff --git a/plugins/nemo-deployments/tests/unit/backends/docker/test_containers.py b/plugins/nemo-deployments/tests/unit/backends/docker/test_containers.py index 2e2f2e9c7e..ffad0564c3 100644 --- a/plugins/nemo-deployments/tests/unit/backends/docker/test_containers.py +++ b/plugins/nemo-deployments/tests/unit/backends/docker/test_containers.py @@ -5,6 +5,8 @@ from __future__ import annotations +from unittest.mock import patch + import pytest from backends.docker.docker_helpers import lora_config, sample_config from nemo_deployments_plugin.backends.docker.containers import ( @@ -14,6 +16,8 @@ ) from nemo_deployments_plugin.entities import Container, ContainerPort, DeploymentConfig +_AUTH_PROXY_MOD = "nemo_deployments_plugin.auth_proxy" + def test_single_container_plan_has_no_init_or_sidecars() -> None: plan = build_docker_plan(sample_config()) @@ -23,6 +27,47 @@ def test_single_container_plan_has_no_init_or_sidecars() -> None: assert plan.is_multi_container is False +def test_auth_proxy_injected_as_docker_sidecar_when_auth_on() -> None: + config = sample_config() + config = config.model_copy(update={"auth_proxy_sidecar": True, "auth_proxy_sidecar_identity": "agents"}) + with ( + patch(f"{_AUTH_PROXY_MOD}.platform_auth_enabled", return_value=True), + patch(f"{_AUTH_PROXY_MOD}.get_qualified_image", return_value="my-registry/nmp-api:local"), + patch(f"{_AUTH_PROXY_MOD}._upstream_base_url", return_value="http://host.docker.internal:8080"), + ): + plan = build_docker_plan(config) + # Injected as a sidecar (shares primary netns), not the primary. + assert plan.primary.name == "main" + proxy = next(c for c in plan.sidecars if c.name == "auth-proxy") + env = {e.name: e.value for e in proxy.env} + assert env["NMP_AUTH_PROXY_PRINCIPAL"] == "agents" + assert env["NMP_BASE_URL"] == "http://host.docker.internal:8080" + # Sidecar must not declare ports (shares netns). + assert proxy.ports == [] + + +def test_auth_proxy_not_injected_when_auth_off() -> None: + config = sample_config().model_copy(update={"auth_proxy_sidecar": True, "auth_proxy_sidecar_identity": "agents"}) + with patch(f"{_AUTH_PROXY_MOD}.platform_auth_enabled", return_value=False): + plan = build_docker_plan(config) + assert plan.sidecars == [] + + +def test_auth_proxy_docker_upstream_rewrites_loopback() -> None: + # docker=True + loopback base_url -> host.docker.internal substitution. + # _upstream_base_url imports these lazily from nemo_platform_plugin.config. + from nemo_deployments_plugin.auth_proxy import _upstream_base_url + + with ( + patch("nemo_platform_plugin.config.determine_loopback_override", return_value="host.docker.internal"), + patch("nemo_platform_plugin.config.get_platform_config") as get_cfg, + ): + get_cfg.return_value.base_url = "http://localhost:8080" + assert _upstream_base_url(docker=True) == "http://host.docker.internal:8080" + # k8s path leaves it verbatim. + assert _upstream_base_url(docker=False) == "http://localhost:8080" + + def test_lora_plan_splits_init_primary_and_sidecar() -> None: plan = build_docker_plan(lora_config()) assert plan.primary.name == "server" From ee535eb413f8759a4c5ad3c501eea8ccda6681e4 Mon Sep 17 00:00:00 2001 From: Ben McCown Date: Mon, 27 Jul 2026 13:12:31 -0600 Subject: [PATCH 4/5] refactor(deployments): require auth-proxy identity; trim sidecar header stripping Addresses PR feedback. - Remove the default service-principal identity. A DeploymentConfig with auth_proxy_sidecar=True now requires auth_proxy_sidecar_identity: a model_validator rejects the invalid combination, surfaced as a 4xx at the create endpoint. The sidecar's run() likewise requires NMP_AUTH_PROXY_PRINCIPAL rather than defaulting. No silent "agents" fallback anywhere. - Trim the auth-proxy's header sanitization to only what would be actively wrong. Request: host + content-length (httpx recomputes), authorization + x-nmp-principal-id (we stamp the identity; must not be spoofed/conflict). Response: content-length + transfer-encoding (responses are streamed). Dropped the hop-by-hop policing that a trusted loopback sidecar doesn't need. - Tests: replace the identity-default test with a validation-rejection test. Signed-off-by: Ben McCown --- .../nmp/common/auth/workload_proxy/main.py | 29 +++++++------------ .../src/nemo_deployments_plugin/auth_proxy.py | 4 +-- .../src/nemo_deployments_plugin/entities.py | 6 ++++ .../tests/unit/test_auth_proxy.py | 17 +++++------ 4 files changed, 26 insertions(+), 30 deletions(-) diff --git a/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py b/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py index 3e072745c2..9ab5d82c01 100644 --- a/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py +++ b/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py @@ -38,37 +38,28 @@ AUTH_PROXY_PRINCIPAL_ENVVAR = "NMP_AUTH_PROXY_PRINCIPAL" DEFAULT_AUTH_PROXY_HOST = "127.0.0.1" DEFAULT_AUTH_PROXY_PORT = 8090 -DEFAULT_AUTH_PROXY_PRINCIPAL = "agents" _READ_TIMEOUT_ENVVAR = "NMP_AUTH_PROXY_READ_TIMEOUT" _PRINCIPAL_ID_HEADER = "x-nmp-principal-id" -# Hop-by-hop headers must not be forwarded (RFC 7230). We also drop the -# workload's own authorization and principal headers (we set the principal), and -# host/content-length (rewritten by httpx / chunked transfer). +# Minimal request-header sanitization. We only drop what would be actively wrong: +# - the workload's own credential / principal header (we set the identity), so it +# can't be spoofed or conflict with what we stamp; +# - host and content-length, which httpx recomputes for the upstream request +# (a stale value corrupts routing / the body). _STRIP_REQUEST_HEADERS = frozenset( { "host", - "connection", - "keep-alive", - "proxy-authenticate", - "proxy-authorization", - "te", - "trailers", - "transfer-encoding", - "upgrade", - "authorization", "content-length", + "authorization", _PRINCIPAL_ID_HEADER, - "x-nmp-principal-on-behalf-of", } ) +# We stream the response, so the upstream's framing headers no longer apply. _STRIP_RESPONSE_HEADERS = frozenset( { - "connection", - "keep-alive", - "transfer-encoding", "content-length", + "transfer-encoding", } ) @@ -127,7 +118,9 @@ async def _body() -> AsyncIterator[bytes]: def run(parent_stop_signal: threading.Event | None = None) -> None: """Sidecar entrypoint. Serves the loopback auth-proxy until stopped.""" base_url = _upstream_base_url() - principal = os.environ.get(AUTH_PROXY_PRINCIPAL_ENVVAR, DEFAULT_AUTH_PROXY_PRINCIPAL) + principal = os.environ.get(AUTH_PROXY_PRINCIPAL_ENVVAR) + if not principal: + raise RuntimeError(f"{AUTH_PROXY_PRINCIPAL_ENVVAR} is required for the auth-proxy sidecar") host = os.environ.get(AUTH_PROXY_HOST_ENVVAR, DEFAULT_AUTH_PROXY_HOST) port = int(os.environ.get(AUTH_PROXY_PORT_ENVVAR, str(DEFAULT_AUTH_PROXY_PORT))) app = build_app(base_url=base_url, principal=principal) diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py index d15febdad7..b1027b98c1 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/auth_proxy.py @@ -37,7 +37,6 @@ _AUTH_PROXY_PRINCIPAL_ENVVAR = "NMP_AUTH_PROXY_PRINCIPAL" _AUTH_PROXY_HOST_ENVVAR = "NMP_AUTH_PROXY_HOST" _AUTH_PROXY_PORT_ENVVAR = "NMP_AUTH_PROXY_PORT" -_DEFAULT_IDENTITY = "agents" def auth_proxy_port() -> int: @@ -83,7 +82,8 @@ def build_auth_proxy_container(config: DeploymentConfig, *, docker: bool = False return None deployments_config = get_nemo_config(DeploymentsConfig) - identity = config.auth_proxy_sidecar_identity or _DEFAULT_IDENTITY + # Guaranteed present: DeploymentConfig validates identity when the sidecar is enabled. + identity = config.auth_proxy_sidecar_identity port = deployments_config.auth_proxy_port image = deployments_config.auth_proxy_image or get_qualified_image(deployments_config.auth_proxy_image_name) diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py index 4f810b4bfb..5d94482d0a 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py @@ -295,6 +295,12 @@ class DeploymentConfig(NemoEntity, entity_type=ENTITY_TYPE_DEPLOYMENT_CONFIG): model_config = {"populate_by_name": True} + @model_validator(mode="after") + def _validate_auth_proxy_identity(self) -> DeploymentConfig: + if self.auth_proxy_sidecar and not self.auth_proxy_sidecar_identity: + raise ValueError("auth_proxy_sidecar_identity is required when auth_proxy_sidecar is True") + return self + class Deployment(NemoEntity, entity_type=ENTITY_TYPE_DEPLOYMENT): """Desired and observed deployment state.""" diff --git a/plugins/nemo-deployments/tests/unit/test_auth_proxy.py b/plugins/nemo-deployments/tests/unit/test_auth_proxy.py index d94e5ccb8c..07f1bf2510 100644 --- a/plugins/nemo-deployments/tests/unit/test_auth_proxy.py +++ b/plugins/nemo-deployments/tests/unit/test_auth_proxy.py @@ -7,8 +7,10 @@ from unittest.mock import patch +import pytest from nemo_deployments_plugin.auth_proxy import AUTH_PROXY_CONTAINER_NAME, build_auth_proxy_container from nemo_deployments_plugin.entities import DeploymentConfig +from pydantic import ValidationError _MOD = "nemo_deployments_plugin.auth_proxy" @@ -51,13 +53,8 @@ def test_builds_sidecar_when_requested_and_auth_on() -> None: assert "127.0.0.1" in " ".join(container.readiness_probe.exec_action.command) -def test_identity_defaults_to_agents_when_unset() -> None: - with ( - patch(f"{_MOD}.platform_auth_enabled", return_value=True), - patch(f"{_MOD}.get_qualified_image", return_value="img"), - patch(f"{_MOD}._upstream_base_url", return_value="http://x:8080"), - ): - container = build_auth_proxy_container(_config(auth_proxy_sidecar=True)) - assert container is not None - env = {e.name: e.value for e in container.env} - assert env["NMP_AUTH_PROXY_PRINCIPAL"] == "agents" +def test_sidecar_without_identity_is_rejected() -> None: + # No default identity: a config requesting the sidecar without an identity + # is invalid and fails validation (surfaced as a 4xx at the create endpoint). + with pytest.raises(ValidationError, match="auth_proxy_sidecar_identity is required"): + _config(auth_proxy_sidecar=True) From dba859b62011e50c1b75c5140126bfd7c3e6442c Mon Sep 17 00:00:00 2001 From: Ben McCown Date: Mon, 27 Jul 2026 13:35:39 -0600 Subject: [PATCH 5/5] refactor(auth-proxy): close upstream client on shutdown; drop redundant cleanup Addresses review feedback on the auth-proxy sidecar. - Add a FastAPI lifespan that closes the shared httpx AsyncClient on shutdown so its connection pool is released gracefully instead of relying on process exit. - Remove the redundant BackgroundTask(response.aclose): the finally in the streaming generator already runs on normal completion, exception, and client disconnect, so the response is always closed there. httpx's aclose() is idempotent, so the duplicate was harmless but implied cleanup was missing. - Test that the lifespan actually closes the upstream client. Not changed: the request body is still buffered before forwarding. Requests through this proxy are a workload's outbound platform API calls (bounded in size), and streaming them would force chunked transfer-encoding on every request since we strip content-length. Revisit if large payloads are routed through the proxy. Signed-off-by: Ben McCown --- .../nmp/common/auth/workload_proxy/main.py | 16 ++++++++++--- .../tests/auth/test_workload_proxy.py | 24 +++++++++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py b/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py index 9ab5d82c01..6d05a0f067 100644 --- a/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py +++ b/packages/nmp_common/src/nmp/common/auth/workload_proxy/main.py @@ -22,12 +22,12 @@ import os import threading from collections.abc import AsyncIterator +from contextlib import asynccontextmanager import httpx import uvicorn from fastapi import FastAPI, Request from fastapi.responses import StreamingResponse -from starlette.background import BackgroundTask logger = logging.getLogger(__name__) @@ -75,12 +75,20 @@ def _upstream_base_url() -> str: def build_app(*, base_url: str, principal: str) -> FastAPI: """Build the forwarding FastAPI app for the given upstream and service principal.""" - app = FastAPI(title="nmp-auth-proxy") principal_id = principal if principal.startswith("service:") else f"service:{principal}" read_timeout = float(os.environ.get(_READ_TIMEOUT_ENVVAR, "300")) timeout = httpx.Timeout(connect=10.0, read=read_timeout, write=60.0, pool=10.0) client = httpx.AsyncClient(base_url=base_url, timeout=timeout, follow_redirects=False) + @asynccontextmanager + async def lifespan(_: FastAPI) -> AsyncIterator[None]: + try: + yield + finally: + await client.aclose() + + app = FastAPI(title="nmp-auth-proxy", lifespan=lifespan) + @app.get("/healthz") async def healthz() -> dict[str, str]: return {"status": "ok"} @@ -98,6 +106,9 @@ async def forward(request: Request, path: str) -> StreamingResponse: response = await client.send(upstream, stream=True) async def _body() -> AsyncIterator[bytes]: + # The finally runs on normal completion, exception, and client + # disconnect (Starlette closes the generator), so this is the only + # cleanup the response needs. try: async for chunk in response.aiter_raw(): yield chunk @@ -109,7 +120,6 @@ async def _body() -> AsyncIterator[bytes]: _body(), status_code=response.status_code, headers=resp_headers, - background=BackgroundTask(response.aclose), ) return app diff --git a/packages/nmp_common/tests/auth/test_workload_proxy.py b/packages/nmp_common/tests/auth/test_workload_proxy.py index 346bac8083..14546abc31 100644 --- a/packages/nmp_common/tests/auth/test_workload_proxy.py +++ b/packages/nmp_common/tests/auth/test_workload_proxy.py @@ -5,6 +5,8 @@ from __future__ import annotations +from unittest.mock import patch + import httpx import respx from fastapi.testclient import TestClient @@ -63,3 +65,25 @@ def test_healthz_does_not_require_upstream() -> None: resp = client.get("/healthz") assert resp.status_code == 200 assert resp.json() == {"status": "ok"} + + +def test_lifespan_closes_upstream_client() -> None: + # Entering TestClient as a context manager runs the lifespan; the shared + # httpx client's connection pool must be closed on shutdown. + created: list[httpx.AsyncClient] = [] + real_async_client = httpx.AsyncClient + + def _tracking_client(*args, **kwargs) -> httpx.AsyncClient: + client = real_async_client(*args, **kwargs) + created.append(client) + return client + + with patch("nmp.common.auth.workload_proxy.main.httpx.AsyncClient", side_effect=_tracking_client): + app = build_app(base_url="http://nemo-platform-api:8080", principal="agents") + + assert len(created) == 1 + upstream_client = created[0] + with TestClient(app) as test_client: + assert test_client.get("/healthz").status_code == 200 + assert upstream_client.is_closed is False + assert upstream_client.is_closed is True