From 800de5a51aaabb177f827977a7237f8cb4b3a789 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Sun, 26 Jul 2026 18:00:16 -0500 Subject: [PATCH 01/16] add fastapi server Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/server.py | 76 ++++++++++++++ .../tests/unit/test_fabric_server.py | 98 +++++++++++++++++++ 2 files changed, 174 insertions(+) create mode 100644 plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py create mode 100644 plugins/nemo-agents/tests/unit/test_fabric_server.py diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py new file mode 100644 index 0000000000..1e7dd1d837 --- /dev/null +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Local HTTP server for Platform-managed Fabric agent runtimes.""" + +from __future__ import annotations + +import argparse +import logging +import sys +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +from fastapi import FastAPI, HTTPException +from nemo_agents_plugin.agent_config import AgentConfig, load_agent_config + +logger = logging.getLogger(__name__) + + +async def _validate_agent_config(config: AgentConfig, *, base_dir: Path) -> Any: + from nemo_agents_plugin.fabric.validation import validate_platform_agent_config + + return await validate_platform_agent_config(config, base_dir=base_dir) + + +def create_fabric_serving_app(agent_config_path: str | Path) -> FastAPI: + """Create a serving app that validates its agent definition at startup.""" + config_path = Path(agent_config_path).resolve() + + @asynccontextmanager + async def lifespan(app: FastAPI) -> AsyncIterator[None]: + agent_config = load_agent_config(config_path) + validation_result = await _validate_agent_config(agent_config, base_dir=config_path.parent) + app.state.agent_config = agent_config + app.state.base_dir = config_path.parent + app.state.validation_result = validation_result + logger.info("Validated Fabric-backed agent config at %s", config_path) + yield + + app = FastAPI(title="NeMo Agents Fabric Server", lifespan=lifespan) + + @app.get("/health") + async def health() -> dict[str, str]: + return {"status": "ok"} + + @app.post("/v1/chat/completions") + async def chat_completions() -> None: + raise HTTPException(status_code=503, detail="Fabric runtime session manager is not initialized.") + + return app + + +def main(argv: list[str] | None = None) -> int: + """Run the local Fabric agent server.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--agent-config", required=True, type=Path, help="Path to an agent YAML config file.") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, required=True) + args = parser.parse_args(argv) + + import uvicorn + + logging.basicConfig(level=logging.INFO) + uvicorn.run( + create_fabric_serving_app(args.agent_config), + host=args.host, + port=args.port, + log_config=None, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py new file mode 100644 index 0000000000..d8c1143407 --- /dev/null +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest +import yaml +from fastapi.testclient import TestClient +from nemo_agents_plugin.agent_config import AgentConfig, AgentConfigLoadError +from nemo_agents_plugin.fabric import server +from nemo_agents_plugin.fabric.server import create_fabric_serving_app + + +@pytest.fixture() +def mock_validate_agent_config(monkeypatch: pytest.MonkeyPatch) -> list[tuple[AgentConfig, Path]]: + validation_calls: list[tuple[AgentConfig, Path]] = [] + + async def validate(config: AgentConfig, *, base_dir: Path) -> object: + validation_calls.append((config, base_dir)) + return object() + + monkeypatch.setattr(server, "_validate_agent_config", validate) + return validation_calls + + +def _example_config() -> dict[str, Any]: + return { + "config_format": "nemo-agents-spec-v1", + "name": "test-agent", + "default_harness": "hermes", + "harnesses": { + "hermes": { + "kind": "hermes", + "model": { + "provider": "nvidia", + "model": "nvidia/test-model", + }, + } + }, + } + + +def _write_agent_config(tmp_path: Path, config: dict[str, Any] | None = None) -> Path: + config_path = tmp_path / "agent.yaml" + config_path.write_text(yaml.safe_dump(config or _example_config()), encoding="utf-8") + return config_path + + +def test_startup_loads_and_validates_agent_config( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], +) -> None: + config_path = _write_agent_config(tmp_path) + app = create_fabric_serving_app(config_path) + + with TestClient(app) as client: + response = client.get("/health") + + assert response.status_code == 200 + assert response.json() == {"status": "ok"} + assert app.state.agent_config.name == "test-agent" + assert app.state.base_dir == tmp_path + assert app.state.validation_result is not None + + assert mock_validate_agent_config == [(app.state.agent_config, tmp_path)] + + +def test_startup_fails_for_invalid_agent_config( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], +) -> None: + config_path = _write_agent_config(tmp_path, {"name": "invalid"}) + app = create_fabric_serving_app(config_path) + + with pytest.raises(AgentConfigLoadError), TestClient(app): + pass + + assert mock_validate_agent_config == [] + + +def test_chat_completions_is_unavailable_until_session_manager_is_added( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], +) -> None: + config_path = _write_agent_config(tmp_path) + app = create_fabric_serving_app(config_path) + + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status_code == 503 + assert response.json() == {"detail": "Fabric runtime session manager is not initialized."} From 77f8a2da90fc560dd5b1adfc22361707e939df94 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Sun, 26 Jul 2026 18:14:03 -0500 Subject: [PATCH 02/16] add session registry Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/server.py | 2 + .../fabric/session_registry.py | 83 +++++++++++++++++++ .../tests/unit/test_fabric_server.py | 2 + .../unit/test_fabric_session_registry.py | 70 ++++++++++++++++ 4 files changed, 157 insertions(+) create mode 100644 plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py create mode 100644 plugins/nemo-agents/tests/unit/test_fabric_session_registry.py diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 1e7dd1d837..7a8c9bbc0a 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -15,6 +15,7 @@ from fastapi import FastAPI, HTTPException from nemo_agents_plugin.agent_config import AgentConfig, load_agent_config +from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry logger = logging.getLogger(__name__) @@ -36,6 +37,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: app.state.agent_config = agent_config app.state.base_dir = config_path.parent app.state.validation_result = validation_result + app.state.session_registry = FabricSessionRegistry() logger.info("Validated Fabric-backed agent config at %s", config_path) yield diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py new file mode 100644 index 0000000000..d59b3b380a --- /dev/null +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Process-local logical session registry for Fabric runtimes.""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from nemo_fabric import Runtime + + +@dataclass(slots=True) +class FabricRuntimeSession: + """Platform session bound to one stateful Fabric runtime.""" + + session_id: str + runtime: Runtime + created_at: float + last_accessed_at: float + + +class FabricSessionNotFoundError(LookupError): + """Raised when a logical Fabric session is not registered.""" + + +class FabricSessionAlreadyExistsError(ValueError): + """Raised when a logical Fabric session ID is registered twice.""" + + +class FabricSessionRegistry: + """Maintain the process-local mapping from Platform sessions to runtimes.""" + + def __init__(self) -> None: + self._sessions: dict[str, FabricRuntimeSession] = {} + self._lock = asyncio.Lock() + + async def register( + self, + runtime: Runtime, + *, + session_id: str | None = None, + ) -> FabricRuntimeSession: + """Bind a new opaque Platform session ID to a Fabric runtime.""" + resolved_session_id = session_id or str(uuid.uuid4()) + now = time.monotonic() + session = FabricRuntimeSession( + session_id=resolved_session_id, + runtime=runtime, + created_at=now, + last_accessed_at=now, + ) + + async with self._lock: + if resolved_session_id in self._sessions: + raise FabricSessionAlreadyExistsError(f"Fabric session '{resolved_session_id}' is already registered.") + self._sessions[resolved_session_id] = session + return session + + async def get(self, session_id: str) -> FabricRuntimeSession: + """Return an active session and update its last-accessed time.""" + async with self._lock: + try: + session = self._sessions[session_id] + except KeyError as error: + raise FabricSessionNotFoundError(f"Fabric session '{session_id}' was not found.") from error + session.last_accessed_at = time.monotonic() + return session + + async def remove(self, session_id: str) -> FabricRuntimeSession | None: + """Remove and return a session without stopping its runtime.""" + async with self._lock: + return self._sessions.pop(session_id, None) + + async def count(self) -> int: + """Return the number of registered logical sessions.""" + async with self._lock: + return len(self._sessions) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index d8c1143407..57aed523d0 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -12,6 +12,7 @@ from nemo_agents_plugin.agent_config import AgentConfig, AgentConfigLoadError from nemo_agents_plugin.fabric import server from nemo_agents_plugin.fabric.server import create_fabric_serving_app +from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry @pytest.fixture() @@ -64,6 +65,7 @@ def test_startup_loads_and_validates_agent_config( assert app.state.agent_config.name == "test-agent" assert app.state.base_dir == tmp_path assert app.state.validation_result is not None + assert isinstance(app.state.session_registry, FabricSessionRegistry) assert mock_validate_agent_config == [(app.state.agent_config, tmp_path)] diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py b/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py new file mode 100644 index 0000000000..4e6c83e84d --- /dev/null +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import uuid +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from nemo_agents_plugin.fabric import session_registry +from nemo_agents_plugin.fabric.session_registry import ( + FabricSessionAlreadyExistsError, + FabricSessionNotFoundError, + FabricSessionRegistry, +) + + +@pytest.mark.asyncio +async def test_register_generates_opaque_session_id() -> None: + registry = FabricSessionRegistry() + runtime = object() + + session = await registry.register(cast(Any, runtime)) + + assert uuid.UUID(session.session_id) + assert session.runtime is runtime + assert session.created_at == session.last_accessed_at + assert await registry.count() == 1 + + +@pytest.mark.asyncio +async def test_get_returns_session_and_updates_last_accessed_at(monkeypatch: pytest.MonkeyPatch) -> None: + clock = iter([10.0, 20.0]) + monkeypatch.setattr(session_registry, "time", SimpleNamespace(monotonic=lambda: next(clock))) + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, object()), session_id="session-1") + + resolved = await registry.get("session-1") + + assert resolved is session + assert resolved.created_at == 10.0 + assert resolved.last_accessed_at == 20.0 + + +@pytest.mark.asyncio +async def test_register_rejects_duplicate_session_id() -> None: + registry = FabricSessionRegistry() + await registry.register(cast(Any, object()), session_id="session-1") + + with pytest.raises(FabricSessionAlreadyExistsError, match="already registered"): + await registry.register(cast(Any, object()), session_id="session-1") + + +@pytest.mark.asyncio +async def test_get_rejects_unknown_session_id() -> None: + registry = FabricSessionRegistry() + + with pytest.raises(FabricSessionNotFoundError, match="was not found"): + await registry.get("missing") + + +@pytest.mark.asyncio +async def test_remove_returns_session_and_is_idempotent() -> None: + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, object()), session_id="session-1") + + assert await registry.remove("session-1") is session + assert await registry.remove("session-1") is None + assert await registry.count() == 0 From fdc28ce741798c0db75a509c429741b618668e69 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Sun, 26 Jul 2026 18:37:42 -0500 Subject: [PATCH 03/16] lazy fabric session initializiation per session Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/server.py | 11 +- .../fabric/session_manager.py | 63 ++++++++++ .../tests/unit/test_fabric_server.py | 6 +- .../tests/unit/test_fabric_session_manager.py | 113 ++++++++++++++++++ 4 files changed, 189 insertions(+), 4 deletions(-) create mode 100644 plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py create mode 100644 plugins/nemo-agents/tests/unit/test_fabric_session_manager.py diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 7a8c9bbc0a..111fb53ab3 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -15,6 +15,7 @@ from fastapi import FastAPI, HTTPException from nemo_agents_plugin.agent_config import AgentConfig, load_agent_config +from nemo_agents_plugin.fabric.session_manager import FabricSessionManager from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry logger = logging.getLogger(__name__) @@ -37,7 +38,13 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: app.state.agent_config = agent_config app.state.base_dir = config_path.parent app.state.validation_result = validation_result - app.state.session_registry = FabricSessionRegistry() + session_registry = FabricSessionRegistry() + app.state.session_registry = session_registry + app.state.session_manager = FabricSessionManager( + agent_config, + base_dir=config_path.parent, + session_registry=session_registry, + ) logger.info("Validated Fabric-backed agent config at %s", config_path) yield @@ -49,7 +56,7 @@ async def health() -> dict[str, str]: @app.post("/v1/chat/completions") async def chat_completions() -> None: - raise HTTPException(status_code=503, detail="Fabric runtime session manager is not initialized.") + raise HTTPException(status_code=503, detail="Fabric runtime invocation is not initialized.") return app diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py new file mode 100644 index 0000000000..baedab3e89 --- /dev/null +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lifecycle coordination for Platform-managed Fabric runtime sessions.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any + +from nemo_agents_plugin.agent_config import AgentConfig +from nemo_agents_plugin.fabric.session_registry import FabricRuntimeSession, FabricSessionRegistry +from nemo_agents_plugin.fabric.translator import FabricTranslationError, translate_agent_config + +# CI type-checks this plugin via ty extra-paths without installing nemo-agents deps. +from nemo_fabric import Fabric, FabricError # ty: ignore[unresolved-import] + +logger = logging.getLogger(__name__) + + +class FabricSessionStartError(RuntimeError): + """Raised when a Fabric runtime cannot be started for a Platform session.""" + + +class FabricSessionManager: + """Create Fabric runtimes lazily from one reusable Platform agent definition.""" + + def __init__( + self, + agent_config: AgentConfig, + *, + base_dir: Path, + session_registry: FabricSessionRegistry, + fabric: Any | None = None, + ) -> None: + self._agent_config = agent_config + self._base_dir = base_dir + self._session_registry = session_registry + self._fabric = fabric + + async def open_session(self) -> FabricRuntimeSession: + """Materialize a Fabric config, start its runtime, and register the session.""" + try: + fabric_config = translate_agent_config(self._agent_config) + except FabricTranslationError as error: + raise FabricSessionStartError(f"Fabric config translation failed: {error}") from error + + fabric = self._fabric or Fabric() + try: + runtime = await fabric.start_runtime(fabric_config, base_dir=self._base_dir) + except FabricError as error: + raise FabricSessionStartError(f"Fabric runtime startup failed: {error}") from error + + try: + return await self._session_registry.register(runtime) + except BaseException: + # A started runtime must not leak if registration fails or is cancelled. + try: + await runtime.stop() + except FabricError: + logger.exception("Failed to stop Fabric runtime after session registration failed.") + raise diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index 57aed523d0..a8716255de 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -12,6 +12,7 @@ from nemo_agents_plugin.agent_config import AgentConfig, AgentConfigLoadError from nemo_agents_plugin.fabric import server from nemo_agents_plugin.fabric.server import create_fabric_serving_app +from nemo_agents_plugin.fabric.session_manager import FabricSessionManager from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry @@ -66,6 +67,7 @@ def test_startup_loads_and_validates_agent_config( assert app.state.base_dir == tmp_path assert app.state.validation_result is not None assert isinstance(app.state.session_registry, FabricSessionRegistry) + assert isinstance(app.state.session_manager, FabricSessionManager) assert mock_validate_agent_config == [(app.state.agent_config, tmp_path)] @@ -83,7 +85,7 @@ def test_startup_fails_for_invalid_agent_config( assert mock_validate_agent_config == [] -def test_chat_completions_is_unavailable_until_session_manager_is_added( +def test_chat_completions_is_unavailable_until_invocation_is_added( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], ) -> None: @@ -97,4 +99,4 @@ def test_chat_completions_is_unavailable_until_session_manager_is_added( ) assert response.status_code == 503 - assert response.json() == {"detail": "Fabric runtime session manager is not initialized."} + assert response.json() == {"detail": "Fabric runtime invocation is not initialized."} diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py new file mode 100644 index 0000000000..3be2b58d5f --- /dev/null +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest +from nemo_agents_plugin.agent_config import AgentConfig +from nemo_agents_plugin.fabric import session_manager +from nemo_agents_plugin.fabric.session_manager import FabricSessionManager +from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry + + +class _FakeRuntime: + def __init__(self) -> None: + self.stop_calls = 0 + + async def stop(self) -> None: + self.stop_calls += 1 + + +class _FakeFabric: + def __init__(self, runtime: _FakeRuntime) -> None: + self.runtime = runtime + self.start_calls: list[tuple[Any, Path]] = [] + + async def start_runtime(self, config: Any, *, base_dir: Path) -> _FakeRuntime: + self.start_calls.append((config, base_dir)) + return self.runtime + + +def _agent_config() -> AgentConfig: + return AgentConfig.model_validate( + { + "config_format": "nemo-agents-spec-v1", + "name": "test-agent", + "default_harness": "hermes", + "harnesses": { + "hermes": { + "kind": "hermes", + "model": { + "provider": "nvidia", + "model": "nvidia/test-model", + }, + } + }, + } + ) + + +@pytest.mark.asyncio +async def test_open_session_materializes_config_and_starts_runtime( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fabric_config = object() + translation_calls: list[AgentConfig] = [] + + def translate(config: AgentConfig) -> Any: + translation_calls.append(config) + return fabric_config + + monkeypatch.setattr(session_manager, "translate_agent_config", translate) + runtime = _FakeRuntime() + fabric = _FakeFabric(runtime) + registry = FabricSessionRegistry() + agent_config = _agent_config() + manager = FabricSessionManager( + agent_config, + base_dir=tmp_path, + session_registry=registry, + fabric=fabric, + ) + + assert translation_calls == [] + assert fabric.start_calls == [] + + session = await manager.open_session() + + assert translation_calls == [agent_config] + assert fabric.start_calls == [(fabric_config, tmp_path)] + assert session.runtime is runtime + assert await registry.get(session.session_id) is session + + +@pytest.mark.asyncio +async def test_open_session_stops_runtime_when_registration_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(session_manager, "translate_agent_config", lambda config: object()) + runtime = _FakeRuntime() + fabric = _FakeFabric(runtime) + registry = FabricSessionRegistry() + + async def fail_registration(runtime: Any) -> None: + raise RuntimeError("registration failed") + + monkeypatch.setattr(registry, "register", fail_registration) + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + fabric=fabric, + ) + + with pytest.raises(RuntimeError, match="registration failed"): + await manager.open_session() + + assert runtime.stop_calls == 1 + assert await registry.count() == 0 From c1b22740b3e92c764630667e3fbda8deae04a4c4 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Sun, 26 Jul 2026 19:01:59 -0500 Subject: [PATCH 04/16] route requests with sessions Signed-off-by: Manjesh Mogallapalli --- .../nemo_agents_plugin/fabric/invocation.py | 4 +- .../src/nemo_agents_plugin/fabric/runtime.py | 51 ++++- .../src/nemo_agents_plugin/fabric/server.py | 120 +++++++++- .../fabric/serving_models.py | 61 ++++++ .../fabric/session_manager.py | 6 + .../tests/unit/test_fabric_runtime.py | 47 +++- .../tests/unit/test_fabric_server.py | 205 +++++++++++++++++- .../tests/unit/test_fabric_serving_models.py | 33 +++ .../tests/unit/test_fabric_session_manager.py | 46 +++- 9 files changed, 544 insertions(+), 29 deletions(-) create mode 100644 plugins/nemo-agents/src/nemo_agents_plugin/fabric/serving_models.py create mode 100644 plugins/nemo-agents/tests/unit/test_fabric_serving_models.py diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/invocation.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/invocation.py index fe2538afff..0b0a916978 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/invocation.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/invocation.py @@ -11,7 +11,7 @@ from typing import Any from nemo_agents_plugin.agent_config import AgentConfig -from nemo_agents_plugin.fabric.runtime import FabricRuntimeRequest, FabricRuntimeResult, run_fabric_agent_once +from nemo_agents_plugin.fabric.runtime import FabricOneShotRequest, FabricRuntimeResult, run_fabric_agent_once from nemo_agents_plugin.fabric.translator import translate_agent_config @@ -29,7 +29,7 @@ async def invoke_agent_config_once( for item in inputs: results.append( await run_fabric_agent_once( - FabricRuntimeRequest( + FabricOneShotRequest( fabric_config=fabric_config, base_dir=base_dir, input=item, diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py index eb8eaeeb38..188aa8fa7a 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py @@ -20,12 +20,29 @@ from typing import Any # CI type-checks this plugin via ty extra-paths without installing nemo-agents deps. -from nemo_fabric import Fabric, FabricConfig, FabricError, RunRequest, RunResult # ty: ignore[unresolved-import] +from nemo_fabric import ( # ty: ignore[unresolved-import] + Fabric, + FabricConfig, + FabricError, + RunRequest, + RunResult, + Runtime, +) @dataclass(frozen=True, slots=True) -class FabricRuntimeRequest: - """Platform-owned request for one Fabric runtime invocation. +class FabricInvocationRequest: + """Platform-owned request for one invocation on an active Fabric runtime.""" + + input: Any = "" + request_id: str | None = None + caller_context: dict[str, Any] = field(default_factory=dict) + timeout_seconds: float | None = None + + +@dataclass(frozen=True, slots=True) +class FabricOneShotRequest: + """Platform-owned request for one ephemeral Fabric runtime invocation. This is an internal bridge type. The fields are intentionally close to Fabric's ``RunRequest`` while preserving Platform-owned lifecycle inputs @@ -70,8 +87,30 @@ class FabricRuntimeTimeoutError(FabricRuntimeExecutionError): """Raised when a Fabric runtime invocation exceeds the Platform timeout.""" +async def invoke_fabric_runtime( + runtime: Runtime, + request: FabricInvocationRequest, +) -> FabricRuntimeResult: + """Invoke an active Fabric runtime without changing its lifecycle.""" + try: + result = await asyncio.wait_for( + runtime.invoke(request=_with_platform_invocation_context(request)), + timeout=request.timeout_seconds, + ) + except TimeoutError as error: + raise FabricRuntimeTimeoutError( + f"Fabric runtime invocation timed out after {request.timeout_seconds:g}s.", + ) from error + except FabricError as error: + raise FabricRuntimeExecutionError( + f"Fabric runtime invocation failed: {error}", + ) from error + + return _normalize_fabric_run_result(result) + + async def run_fabric_agent_once( - request: FabricRuntimeRequest, + request: FabricOneShotRequest, *, fabric: Any | None = None, ) -> FabricRuntimeResult: @@ -96,7 +135,7 @@ async def run_fabric_agent_once( async def _invoke_fabric_agent_once( - request: FabricRuntimeRequest, + request: FabricOneShotRequest, *, fabric: Any, ) -> RunResult: @@ -108,7 +147,7 @@ async def _invoke_fabric_agent_once( return await runtime.invoke(request=_with_platform_invocation_context(request)) -def _with_platform_invocation_context(request: FabricRuntimeRequest) -> RunRequest: +def _with_platform_invocation_context(request: FabricInvocationRequest | FabricOneShotRequest) -> RunRequest: """Preserve Platform invocation metadata when calling Fabric.""" request_kwargs: dict[str, Any] = { "context": request.caller_context, diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 111fb53ab3..792b59c258 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -8,18 +8,78 @@ import argparse import logging import sys -from collections.abc import AsyncIterator +import uuid +from collections.abc import AsyncIterator, Mapping from contextlib import asynccontextmanager from pathlib import Path -from typing import Any +from typing import Annotated, Any -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI, Header, HTTPException, Response from nemo_agents_plugin.agent_config import AgentConfig, load_agent_config -from nemo_agents_plugin.fabric.session_manager import FabricSessionManager -from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry +from nemo_agents_plugin.fabric.runtime import ( + FabricInvocationRequest, + FabricRuntimeExecutionError, + FabricRuntimeResult, + FabricRuntimeTimeoutError, + invoke_fabric_runtime, +) +from nemo_agents_plugin.fabric.serving_models import ( + ChatCompletionChoice, + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseMessage, +) +from nemo_agents_plugin.fabric.session_manager import FabricSessionManager, FabricSessionStartError +from nemo_agents_plugin.fabric.session_registry import FabricSessionNotFoundError, FabricSessionRegistry logger = logging.getLogger(__name__) +SESSION_ID_HEADER = "X-Nemo-Session-Id" + + +def _to_fabric_invocation_request( + request: ChatCompletionRequest, + *, + session_id: str, +) -> FabricInvocationRequest: + """Translate the current chat turn into a Platform-owned Fabric request.""" + return FabricInvocationRequest( + input=request.messages[-1].content, + caller_context={"session_id": session_id}, + ) + + +def _to_chat_completion_response(result: FabricRuntimeResult) -> ChatCompletionResponse: + """Convert a successful Fabric result into an OpenAI-compatible response.""" + if not isinstance(result.response, str): + raise ValueError("Fabric invocation did not return a text response.") + + usage = None + if isinstance(result.output, Mapping) and isinstance(result.output.get("usage"), Mapping): + usage = dict(result.output["usage"]) + + return ChatCompletionResponse( + id=result.invocation_id or result.request_id or result.runtime_id or f"chatcmpl-{uuid.uuid4().hex}", + choices=[ + ChatCompletionChoice( + message=ChatCompletionResponseMessage(content=result.response), + ) + ], + usage=usage, + ) + + +def _session_headers(session_id: str) -> dict[str, str]: + return {SESSION_ID_HEADER: session_id} + + +def _failed_result_detail(result: FabricRuntimeResult) -> str: + if isinstance(result.error, Mapping): + message = result.error.get("message") + if isinstance(message, str): + return message + return f"Fabric invocation returned status {result.status!r}." + async def _validate_agent_config(config: AgentConfig, *, base_dir: Path) -> Any: from nemo_agents_plugin.fabric.validation import validate_platform_agent_config @@ -54,9 +114,53 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: async def health() -> dict[str, str]: return {"status": "ok"} - @app.post("/v1/chat/completions") - async def chat_completions() -> None: - raise HTTPException(status_code=503, detail="Fabric runtime invocation is not initialized.") + @app.post("/v1/chat/completions", response_model_exclude_none=True) + async def chat_completions( + request: ChatCompletionRequest, + response: Response, + session_id: Annotated[str | None, Header(alias=SESSION_ID_HEADER)] = None, + ) -> ChatCompletionResponse: + try: + session = await app.state.session_manager.resolve_session(session_id) + except FabricSessionNotFoundError as error: + raise HTTPException(status_code=404, detail=str(error)) from error + except FabricSessionStartError as error: + raise HTTPException(status_code=503, detail=str(error)) from error + + invocation_request = _to_fabric_invocation_request(request, session_id=session.session_id) + try: + result = await invoke_fabric_runtime(session.runtime, invocation_request) + except FabricRuntimeTimeoutError as error: + raise HTTPException( + status_code=504, + detail=str(error), + headers=_session_headers(session.session_id), + ) from error + except FabricRuntimeExecutionError as error: + raise HTTPException( + status_code=502, + detail=str(error), + headers=_session_headers(session.session_id), + ) from error + + if result.status != "succeeded": + raise HTTPException( + status_code=502, + detail=_failed_result_detail(result), + headers=_session_headers(session.session_id), + ) + + try: + completion = _to_chat_completion_response(result) + except ValueError as error: + raise HTTPException( + status_code=502, + detail=str(error), + headers=_session_headers(session.session_id), + ) from error + + response.headers[SESSION_ID_HEADER] = session.session_id + return completion return app diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/serving_models.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/serving_models.py new file mode 100644 index 0000000000..e9a802e42d --- /dev/null +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/serving_models.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""HTTP request and response models for Fabric-backed agent serving.""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class ChatCompletionMessage(BaseModel): + """Supported OpenAI-compatible chat message.""" + + model_config = ConfigDict(extra="allow") + + role: Literal["assistant", "developer", "system", "tool", "user"] + content: str + + +class ChatCompletionRequest(BaseModel): + """Supported subset of an OpenAI chat-completions request.""" + + model_config = ConfigDict(extra="allow") + + messages: list[ChatCompletionMessage] = Field(min_length=1) + stream: bool = False + + @model_validator(mode="after") + def validate_current_turn(self) -> ChatCompletionRequest: + if self.stream: + raise ValueError("Streaming chat completions are not supported.") + if self.messages[-1].role != "user": + raise ValueError("The final chat message must have role 'user'.") + return self + + +class ChatCompletionResponseMessage(BaseModel): + """OpenAI-compatible assistant response message.""" + + role: Literal["assistant"] = "assistant" + content: str + + +class ChatCompletionChoice(BaseModel): + """OpenAI-compatible chat-completion choice.""" + + index: int = 0 + message: ChatCompletionResponseMessage + finish_reason: Literal["stop"] = "stop" + + +class ChatCompletionResponse(BaseModel): + """OpenAI-compatible response for one Fabric runtime invocation.""" + + id: str + object: Literal["chat.completion"] = "chat.completion" + model: str = "unknown-model" + choices: list[ChatCompletionChoice] + usage: dict[str, Any] | None = None diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py index baedab3e89..5d08e19a66 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -61,3 +61,9 @@ async def open_session(self) -> FabricRuntimeSession: except FabricError: logger.exception("Failed to stop Fabric runtime after session registration failed.") raise + + async def resolve_session(self, session_id: str | None) -> FabricRuntimeSession: + """Open a new session or resolve an existing session by its opaque ID.""" + if session_id is None: + return await self.open_session() + return await self._session_registry.get(session_id) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py index 218a02289e..3b905be01e 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py @@ -12,9 +12,11 @@ import pytest from nemo_agents_plugin.fabric import runtime as fabric_runtime from nemo_agents_plugin.fabric.runtime import ( + FabricInvocationRequest, + FabricOneShotRequest, FabricRuntimeExecutionError, - FabricRuntimeRequest, FabricRuntimeTimeoutError, + invoke_fabric_runtime, run_fabric_agent_once, ) from nemo_fabric import FabricConfig # ty: ignore[unresolved-import] @@ -121,7 +123,7 @@ async def test_starts_invokes_and_cleans_up_ephemeral_runtime(self) -> None: fabric_config = cast(FabricConfig, object()) fake_runtime = _FakeRuntime() fake_fabric = _FakeFabric(runtime=fake_runtime) - request = FabricRuntimeRequest( + request = FabricOneShotRequest( fabric_config=fabric_config, base_dir=Path("/tmp/agent"), input={"prompt": "hi"}, @@ -154,7 +156,7 @@ async def test_starts_invokes_and_cleans_up_ephemeral_runtime(self) -> None: async def test_wraps_timeout(self) -> None: fake_runtime = _FakeRuntime(invoke_delay=1.0) fake_fabric = _FakeFabric(runtime=fake_runtime) - request = FabricRuntimeRequest( + request = FabricOneShotRequest( fabric_config=cast(FabricConfig, object()), base_dir=Path("/tmp/agent"), timeout_seconds=0.01, @@ -167,7 +169,7 @@ async def test_wraps_timeout(self) -> None: async def test_wraps_fabric_lifecycle_errors(self) -> None: fake_fabric = _FakeFabric(start_error=fabric_runtime.FabricError("native unavailable")) - request = FabricRuntimeRequest( + request = FabricOneShotRequest( fabric_config=cast(FabricConfig, object()), base_dir=Path("/tmp/agent"), ) @@ -182,7 +184,7 @@ async def test_failed_run_result_is_returned_as_normalized_result(self) -> None: error=_FabricMapping({"stage": "invoke", "message": "adapter failed"}), ) fake_fabric = _FakeFabric(runtime=_FakeRuntime(result=failed_result)) - request = FabricRuntimeRequest( + request = FabricOneShotRequest( fabric_config=cast(FabricConfig, object()), base_dir=Path("/tmp/agent"), ) @@ -200,7 +202,7 @@ async def test_normalizes_fabric_mapping_fields_to_plain_values(self) -> None: }, ) fake_fabric = _FakeFabric(runtime=_FakeRuntime(result=fake_result)) - request = FabricRuntimeRequest( + request = FabricOneShotRequest( fabric_config=cast(FabricConfig, object()), base_dir=Path("/tmp/agent"), ) @@ -216,3 +218,36 @@ async def test_normalizes_fabric_mapping_fields_to_plain_values(self) -> None: assert result.telemetry == [{"provider": "relay", "kind": "trace"}] assert result.events == [{"kind": "runtime_start", "message": "started"}] assert result.metadata == {"adapter_runner": "python"} + + +@pytest.mark.asyncio +class TestInvokeFabricRuntime: + async def test_invokes_active_runtime_without_changing_its_lifecycle(self) -> None: + fake_runtime = _FakeRuntime() + request = FabricInvocationRequest( + input={"prompt": "hi"}, + request_id="platform-request-1", + caller_context={"session_id": "session-1"}, + ) + + result = await invoke_fabric_runtime(cast(Any, fake_runtime), request) + + assert fake_runtime.entered is False + assert fake_runtime.exited is False + fabric_request = fake_runtime.invoke_requests[0] + assert fabric_request.input == {"prompt": "hi"} + assert fabric_request.request_id == "platform-request-1" + assert fabric_request.context == {"session_id": "session-1"} + assert result.status == "succeeded" + assert result.response == "hello" + assert result.runtime_id == "runtime-1" + + async def test_wraps_timeout_without_stopping_runtime(self) -> None: + fake_runtime = _FakeRuntime(invoke_delay=1.0) + request = FabricInvocationRequest(timeout_seconds=0.01) + + with pytest.raises(FabricRuntimeTimeoutError, match="timed out after 0.01s"): + await invoke_fabric_runtime(cast(Any, fake_runtime), request) + + assert fake_runtime.entered is False + assert fake_runtime.exited is False diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index a8716255de..ed4ee72d02 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -4,6 +4,7 @@ from __future__ import annotations from pathlib import Path +from types import SimpleNamespace from typing import Any import pytest @@ -11,9 +12,15 @@ from fastapi.testclient import TestClient from nemo_agents_plugin.agent_config import AgentConfig, AgentConfigLoadError from nemo_agents_plugin.fabric import server -from nemo_agents_plugin.fabric.server import create_fabric_serving_app -from nemo_agents_plugin.fabric.session_manager import FabricSessionManager -from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry +from nemo_agents_plugin.fabric.runtime import ( + FabricRuntimeExecutionError, + FabricRuntimeResult, + FabricRuntimeTimeoutError, +) +from nemo_agents_plugin.fabric.server import SESSION_ID_HEADER, create_fabric_serving_app +from nemo_agents_plugin.fabric.serving_models import ChatCompletionRequest +from nemo_agents_plugin.fabric.session_manager import FabricSessionManager, FabricSessionStartError +from nemo_agents_plugin.fabric.session_registry import FabricSessionNotFoundError, FabricSessionRegistry @pytest.fixture() @@ -85,18 +92,204 @@ def test_startup_fails_for_invalid_agent_config( assert mock_validate_agent_config == [] -def test_chat_completions_is_unavailable_until_invocation_is_added( +def test_chat_completion_without_session_id_opens_and_returns_session( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, ) -> None: config_path = _write_agent_config(tmp_path) app = create_fabric_serving_app(config_path) + resolve_calls: list[str | None] = [] + invocation_calls: list[tuple[Any, Any]] = [] + runtime = object() + + async def resolve_session(session_id: str | None) -> Any: + resolve_calls.append(session_id) + return SimpleNamespace(session_id="session-1", runtime=runtime) + + async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRuntimeResult: + invocation_calls.append((active_runtime, request)) + return FabricRuntimeResult( + status="succeeded", + output={"response": "hello", "usage": {"total_tokens": 3}}, + response="hello", + invocation_id="invocation-1", + ) + + monkeypatch.setattr(server, "invoke_fabric_runtime", invoke_fabric_runtime) + + with TestClient(app) as client: + monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + response = client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status_code == 200 + assert response.headers[SESSION_ID_HEADER] == "session-1" + assert response.json() == { + "id": "invocation-1", + "object": "chat.completion", + "model": "unknown-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + } + ], + "usage": {"total_tokens": 3}, + } + assert resolve_calls == [None] + active_runtime, invocation_request = invocation_calls[0] + assert active_runtime is runtime + assert invocation_request.input == "hello" + assert invocation_request.caller_context == {"session_id": "session-1"} + + +def test_chat_completion_with_session_id_reuses_session( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = _write_agent_config(tmp_path) + app = create_fabric_serving_app(config_path) + resolve_calls: list[str | None] = [] + + async def resolve_session(session_id: str | None) -> Any: + resolve_calls.append(session_id) + return SimpleNamespace(session_id="session-1", runtime=object()) + + async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRuntimeResult: + return FabricRuntimeResult(status="succeeded", response="hello again") + + monkeypatch.setattr(server, "invoke_fabric_runtime", invoke_fabric_runtime) + + with TestClient(app) as client: + monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + response = client.post( + "/v1/chat/completions", + headers={SESSION_ID_HEADER: "session-1"}, + json={"messages": [{"role": "user", "content": "hello again"}]}, + ) + + assert response.status_code == 200 + assert response.headers[SESSION_ID_HEADER] == "session-1" + assert resolve_calls == ["session-1"] + + +@pytest.mark.parametrize( + ("error", "status_code"), + [ + (FabricRuntimeTimeoutError("timed out"), 504), + (FabricRuntimeExecutionError("invoke failed"), 502), + ], +) +def test_chat_completion_maps_runtime_errors( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, + error: Exception, + status_code: int, +) -> None: + app = create_fabric_serving_app(_write_agent_config(tmp_path)) + + async def resolve_session(session_id: str | None) -> Any: + return SimpleNamespace(session_id="session-1", runtime=object()) + + async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRuntimeResult: + raise error + + monkeypatch.setattr(server, "invoke_fabric_runtime", invoke_fabric_runtime) + + with TestClient(app) as client: + monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + response = client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status_code == status_code + assert response.headers[SESSION_ID_HEADER] == "session-1" + assert response.json() == {"detail": str(error)} + + +@pytest.mark.parametrize( + ("error", "status_code"), + [ + (FabricSessionNotFoundError("missing session"), 404), + (FabricSessionStartError("startup failed"), 503), + ], +) +def test_chat_completion_maps_session_resolution_errors( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, + error: Exception, + status_code: int, +) -> None: + app = create_fabric_serving_app(_write_agent_config(tmp_path)) + + async def resolve_session(session_id: str | None) -> Any: + raise error with TestClient(app) as client: + monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) response = client.post( "/v1/chat/completions", + headers={SESSION_ID_HEADER: "missing"}, json={"messages": [{"role": "user", "content": "hello"}]}, ) - assert response.status_code == 503 - assert response.json() == {"detail": "Fabric runtime invocation is not initialized."} + assert response.status_code == status_code + assert SESSION_ID_HEADER not in response.headers + assert response.json() == {"detail": str(error)} + + +def test_chat_completion_maps_failed_run_result( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + app = create_fabric_serving_app(_write_agent_config(tmp_path)) + + async def resolve_session(session_id: str | None) -> Any: + return SimpleNamespace(session_id="session-1", runtime=object()) + + async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRuntimeResult: + return FabricRuntimeResult( + status="failed", + error={"stage": "invoke", "message": "adapter failed"}, + ) + + monkeypatch.setattr(server, "invoke_fabric_runtime", invoke_fabric_runtime) + + with TestClient(app) as client: + monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + response = client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status_code == 502 + assert response.headers[SESSION_ID_HEADER] == "session-1" + assert response.json() == {"detail": "adapter failed"} + + +def test_chat_completion_request_translates_final_user_turn() -> None: + request = ChatCompletionRequest.model_validate( + { + "messages": [ + {"role": "system", "content": "Be concise."}, + {"role": "assistant", "content": "How can I help?"}, + {"role": "user", "content": "Say hello."}, + ], + "model": "test-model", + "stream": False, + } + ) + + invocation_request = server._to_fabric_invocation_request(request, session_id="session-1") + + assert invocation_request.input == "Say hello." + assert invocation_request.caller_context == {"session_id": "session-1"} diff --git a/plugins/nemo-agents/tests/unit/test_fabric_serving_models.py b/plugins/nemo-agents/tests/unit/test_fabric_serving_models.py new file mode 100644 index 0000000000..34a6935e87 --- /dev/null +++ b/plugins/nemo-agents/tests/unit/test_fabric_serving_models.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +import pytest +from nemo_agents_plugin.fabric.serving_models import ChatCompletionRequest +from pydantic import ValidationError + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ({"messages": []}, "List should have at least 1 item"), + ( + {"messages": [{"role": "assistant", "content": "hello"}]}, + "The final chat message must have role 'user'.", + ), + ( + {"messages": [{"role": "user", "content": "hello"}], "stream": True}, + "Streaming chat completions are not supported.", + ), + ( + {"messages": [{"role": "user", "content": [{"type": "text", "text": "hello"}]}]}, + "Input should be a valid string", + ), + ], +) +def test_chat_completion_request_rejects_unsupported_inputs(payload: dict[str, Any], message: str) -> None: + with pytest.raises(ValidationError, match=message): + ChatCompletionRequest.model_validate(payload) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py index 3be2b58d5f..b9ef2fc5ce 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py @@ -4,7 +4,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any +from typing import Any, cast import pytest from nemo_agents_plugin.agent_config import AgentConfig @@ -111,3 +111,47 @@ async def fail_registration(runtime: Any) -> None: assert runtime.stop_calls == 1 assert await registry.count() == 0 + + +@pytest.mark.asyncio +async def test_resolve_session_opens_session_when_id_is_absent( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(session_manager, "translate_agent_config", lambda config: object()) + runtime = _FakeRuntime() + fabric = _FakeFabric(runtime) + registry = FabricSessionRegistry() + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + fabric=fabric, + ) + + session = await manager.resolve_session(None) + + assert session.runtime is runtime + assert len(fabric.start_calls) == 1 + + +@pytest.mark.asyncio +async def test_resolve_session_reuses_registered_runtime( + tmp_path: Path, +) -> None: + runtime = _FakeRuntime() + fabric = _FakeFabric(_FakeRuntime()) + registry = FabricSessionRegistry() + registered = await registry.register(cast(Any, runtime), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + fabric=fabric, + ) + + session = await manager.resolve_session("session-1") + + assert session is registered + assert session.runtime is runtime + assert fabric.start_calls == [] From f42fd0ef613567167914e6b07745b504211cf5ee Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Sun, 26 Jul 2026 19:13:26 -0500 Subject: [PATCH 05/16] add per session invocation lock Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/server.py | 3 +- .../fabric/session_manager.py | 10 ++ .../fabric/session_registry.py | 3 +- .../tests/unit/test_fabric_server.py | 26 ++-- .../tests/unit/test_fabric_session_manager.py | 121 ++++++++++++++++++ .../unit/test_fabric_session_registry.py | 10 ++ 6 files changed, 155 insertions(+), 18 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 792b59c258..27afae4d9b 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -21,7 +21,6 @@ FabricRuntimeExecutionError, FabricRuntimeResult, FabricRuntimeTimeoutError, - invoke_fabric_runtime, ) from nemo_agents_plugin.fabric.serving_models import ( ChatCompletionChoice, @@ -129,7 +128,7 @@ async def chat_completions( invocation_request = _to_fabric_invocation_request(request, session_id=session.session_id) try: - result = await invoke_fabric_runtime(session.runtime, invocation_request) + result = await app.state.session_manager.invoke_session(session, invocation_request) except FabricRuntimeTimeoutError as error: raise HTTPException( status_code=504, diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py index 5d08e19a66..6c387345b3 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -10,6 +10,7 @@ from typing import Any from nemo_agents_plugin.agent_config import AgentConfig +from nemo_agents_plugin.fabric.runtime import FabricInvocationRequest, FabricRuntimeResult, invoke_fabric_runtime from nemo_agents_plugin.fabric.session_registry import FabricRuntimeSession, FabricSessionRegistry from nemo_agents_plugin.fabric.translator import FabricTranslationError, translate_agent_config @@ -67,3 +68,12 @@ async def resolve_session(self, session_id: str | None) -> FabricRuntimeSession: if session_id is None: return await self.open_session() return await self._session_registry.get(session_id) + + async def invoke_session( + self, + session: FabricRuntimeSession, + request: FabricInvocationRequest, + ) -> FabricRuntimeResult: + """Serialize and invoke one turn on a session's active runtime.""" + async with session.invocation_lock: + return await invoke_fabric_runtime(session.runtime, request) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py index d59b3b380a..bdd3a67f40 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py @@ -8,7 +8,7 @@ import asyncio import time import uuid -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -23,6 +23,7 @@ class FabricRuntimeSession: runtime: Runtime created_at: float last_accessed_at: float + invocation_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) class FabricSessionNotFoundError(LookupError): diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index ed4ee72d02..69357a6c25 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -107,8 +107,8 @@ async def resolve_session(session_id: str | None) -> Any: resolve_calls.append(session_id) return SimpleNamespace(session_id="session-1", runtime=runtime) - async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRuntimeResult: - invocation_calls.append((active_runtime, request)) + async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: + invocation_calls.append((session, request)) return FabricRuntimeResult( status="succeeded", output={"response": "hello", "usage": {"total_tokens": 3}}, @@ -116,10 +116,9 @@ async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRunt invocation_id="invocation-1", ) - monkeypatch.setattr(server, "invoke_fabric_runtime", invoke_fabric_runtime) - with TestClient(app) as client: monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + monkeypatch.setattr(app.state.session_manager, "invoke_session", invoke_session) response = client.post( "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hello"}]}, @@ -141,8 +140,8 @@ async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRunt "usage": {"total_tokens": 3}, } assert resolve_calls == [None] - active_runtime, invocation_request = invocation_calls[0] - assert active_runtime is runtime + resolved_session, invocation_request = invocation_calls[0] + assert resolved_session.runtime is runtime assert invocation_request.input == "hello" assert invocation_request.caller_context == {"session_id": "session-1"} @@ -160,13 +159,12 @@ async def resolve_session(session_id: str | None) -> Any: resolve_calls.append(session_id) return SimpleNamespace(session_id="session-1", runtime=object()) - async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRuntimeResult: + async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: return FabricRuntimeResult(status="succeeded", response="hello again") - monkeypatch.setattr(server, "invoke_fabric_runtime", invoke_fabric_runtime) - with TestClient(app) as client: monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + monkeypatch.setattr(app.state.session_manager, "invoke_session", invoke_session) response = client.post( "/v1/chat/completions", headers={SESSION_ID_HEADER: "session-1"}, @@ -197,13 +195,12 @@ def test_chat_completion_maps_runtime_errors( async def resolve_session(session_id: str | None) -> Any: return SimpleNamespace(session_id="session-1", runtime=object()) - async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRuntimeResult: + async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: raise error - monkeypatch.setattr(server, "invoke_fabric_runtime", invoke_fabric_runtime) - with TestClient(app) as client: monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + monkeypatch.setattr(app.state.session_manager, "invoke_session", invoke_session) response = client.post( "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hello"}]}, @@ -256,16 +253,15 @@ def test_chat_completion_maps_failed_run_result( async def resolve_session(session_id: str | None) -> Any: return SimpleNamespace(session_id="session-1", runtime=object()) - async def invoke_fabric_runtime(active_runtime: Any, request: Any) -> FabricRuntimeResult: + async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: return FabricRuntimeResult( status="failed", error={"stage": "invoke", "message": "adapter failed"}, ) - monkeypatch.setattr(server, "invoke_fabric_runtime", invoke_fabric_runtime) - with TestClient(app) as client: monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + monkeypatch.setattr(app.state.session_manager, "invoke_session", invoke_session) response = client.post( "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hello"}]}, diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py index b9ef2fc5ce..056c25115d 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py @@ -3,12 +3,14 @@ from __future__ import annotations +import asyncio from pathlib import Path from typing import Any, cast import pytest from nemo_agents_plugin.agent_config import AgentConfig from nemo_agents_plugin.fabric import session_manager +from nemo_agents_plugin.fabric.runtime import FabricInvocationRequest, FabricRuntimeResult from nemo_agents_plugin.fabric.session_manager import FabricSessionManager from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry @@ -155,3 +157,122 @@ async def test_resolve_session_reuses_registered_runtime( assert session is registered assert session.runtime is runtime assert fabric.start_calls == [] + + +@pytest.mark.asyncio +async def test_invoke_session_serializes_turns_for_same_runtime( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + fabric=_FakeFabric(_FakeRuntime()), + ) + first_started = asyncio.Event() + release_first = asyncio.Event() + invocation_order: list[str] = [] + active_invocations = 0 + max_active_invocations = 0 + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + nonlocal active_invocations, max_active_invocations + active_invocations += 1 + max_active_invocations = max(max_active_invocations, active_invocations) + invocation_order.append(request.input) + if request.input == "first": + first_started.set() + await release_first.wait() + active_invocations -= 1 + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + + first = asyncio.create_task( + manager.invoke_session(session, FabricInvocationRequest(input="first")), + ) + await first_started.wait() + second = asyncio.create_task( + manager.invoke_session(session, FabricInvocationRequest(input="second")), + ) + await asyncio.sleep(0) + + assert invocation_order == ["first"] + + release_first.set() + results = await asyncio.gather(first, second) + + assert invocation_order == ["first", "second"] + assert max_active_invocations == 1 + assert [result.response for result in results] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_invoke_session_releases_lock_after_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + fabric=_FakeFabric(_FakeRuntime()), + ) + invocation_count = 0 + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + nonlocal invocation_count + invocation_count += 1 + if invocation_count == 1: + raise RuntimeError("invoke failed") + return FabricRuntimeResult(status="succeeded", response="recovered") + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + + with pytest.raises(RuntimeError, match="invoke failed"): + await manager.invoke_session(session, FabricInvocationRequest(input="first")) + + result = await manager.invoke_session(session, FabricInvocationRequest(input="second")) + + assert result.response == "recovered" + + +@pytest.mark.asyncio +async def test_invoke_session_releases_lock_after_cancellation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + fabric=_FakeFabric(_FakeRuntime()), + ) + invocation_started = asyncio.Event() + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + if request.input == "cancel": + invocation_started.set() + await asyncio.Event().wait() + return FabricRuntimeResult(status="succeeded", response="recovered") + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + cancelled = asyncio.create_task( + manager.invoke_session(session, FabricInvocationRequest(input="cancel")), + ) + await invocation_started.wait() + + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + + result = await manager.invoke_session(session, FabricInvocationRequest(input="next")) + + assert result.response == "recovered" diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py b/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py index 4e6c83e84d..c2b1570366 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py @@ -29,6 +29,16 @@ async def test_register_generates_opaque_session_id() -> None: assert await registry.count() == 1 +@pytest.mark.asyncio +async def test_each_session_has_an_independent_invocation_lock() -> None: + registry = FabricSessionRegistry() + + first = await registry.register(cast(Any, object()), session_id="session-1") + second = await registry.register(cast(Any, object()), session_id="session-2") + + assert first.invocation_lock is not second.invocation_lock + + @pytest.mark.asyncio async def test_get_returns_session_and_updates_last_accessed_at(monkeypatch: pytest.MonkeyPatch) -> None: clock = iter([10.0, 20.0]) From 8aa67505ab69af641378916b73549df3b24552c8 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 11:14:07 -0500 Subject: [PATCH 06/16] concurrency & session lifecycle management Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/server.py | 102 +++++- .../fabric/session_manager.py | 71 +++- .../fabric/session_registry.py | 44 ++- .../tests/unit/test_fabric_server.py | 117 ++++++- .../tests/unit/test_fabric_session_manager.py | 317 +++++++++++++++++- .../unit/test_fabric_session_registry.py | 53 +++ 6 files changed, 692 insertions(+), 12 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 27afae4d9b..16994d13ed 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -6,6 +6,7 @@ from __future__ import annotations import argparse +import asyncio import logging import sys import uuid @@ -28,7 +29,14 @@ ChatCompletionResponse, ChatCompletionResponseMessage, ) -from nemo_agents_plugin.fabric.session_manager import FabricSessionManager, FabricSessionStartError +from nemo_agents_plugin.fabric.session_manager import ( + DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS, + DEFAULT_MAX_CONCURRENT_INVOCATIONS, + DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS, + FabricSessionManager, + FabricSessionStartError, + FabricSessionStopError, +) from nemo_agents_plugin.fabric.session_registry import FabricSessionNotFoundError, FabricSessionRegistry logger = logging.getLogger(__name__) @@ -86,8 +94,37 @@ async def _validate_agent_config(config: AgentConfig, *, base_dir: Path) -> Any: return await validate_platform_agent_config(config, base_dir=base_dir) -def create_fabric_serving_app(agent_config_path: str | Path) -> FastAPI: +async def _run_idle_session_cleanup( + manager: FabricSessionManager, + *, + idle_timeout_seconds: float, + cleanup_interval_seconds: float, + shutdown_event: asyncio.Event, +) -> None: + """Periodically expire inactive logical sessions until shutdown.""" + while not shutdown_event.is_set(): + try: + await asyncio.wait_for(shutdown_event.wait(), timeout=cleanup_interval_seconds) + except TimeoutError: + try: + await manager.expire_idle_sessions(idle_timeout_seconds=idle_timeout_seconds) + except Exception: + logger.exception("Failed to expire idle Fabric sessions.") + + +def create_fabric_serving_app( + agent_config_path: str | Path, + *, + max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS, + idle_session_timeout_seconds: float = DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS, + session_cleanup_interval_seconds: float = DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS, +) -> FastAPI: """Create a serving app that validates its agent definition at startup.""" + if idle_session_timeout_seconds <= 0: + raise ValueError("idle_session_timeout_seconds must be greater than zero.") + if session_cleanup_interval_seconds <= 0: + raise ValueError("session_cleanup_interval_seconds must be greater than zero.") + config_path = Path(agent_config_path).resolve() @asynccontextmanager @@ -99,13 +136,29 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: app.state.validation_result = validation_result session_registry = FabricSessionRegistry() app.state.session_registry = session_registry - app.state.session_manager = FabricSessionManager( + session_manager = FabricSessionManager( agent_config, base_dir=config_path.parent, session_registry=session_registry, + max_concurrent_invocations=max_concurrent_invocations, + ) + app.state.session_manager = session_manager + cleanup_shutdown = asyncio.Event() + cleanup_task = asyncio.create_task( + _run_idle_session_cleanup( + session_manager, + idle_timeout_seconds=idle_session_timeout_seconds, + cleanup_interval_seconds=session_cleanup_interval_seconds, + shutdown_event=cleanup_shutdown, + ) ) logger.info("Validated Fabric-backed agent config at %s", config_path) - yield + try: + yield + finally: + cleanup_shutdown.set() + await cleanup_task + await session_manager.close_all_sessions() app = FastAPI(title="NeMo Agents Fabric Server", lifespan=lifespan) @@ -129,6 +182,12 @@ async def chat_completions( invocation_request = _to_fabric_invocation_request(request, session_id=session.session_id) try: result = await app.state.session_manager.invoke_session(session, invocation_request) + except FabricSessionNotFoundError as error: + raise HTTPException( + status_code=404, + detail=str(error), + headers=_session_headers(session.session_id), + ) from error except FabricRuntimeTimeoutError as error: raise HTTPException( status_code=504, @@ -161,6 +220,16 @@ async def chat_completions( response.headers[SESSION_ID_HEADER] = session.session_id return completion + @app.delete("/v1/sessions/{session_id}", status_code=204) + async def close_session(session_id: str) -> Response: + try: + await app.state.session_manager.close_session(session_id) + except FabricSessionNotFoundError as error: + raise HTTPException(status_code=404, detail=str(error)) from error + except FabricSessionStopError as error: + raise HTTPException(status_code=502, detail=str(error)) from error + return Response(status_code=204) + return app @@ -170,13 +239,36 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--agent-config", required=True, type=Path, help="Path to an agent YAML config file.") parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, required=True) + parser.add_argument( + "--max-concurrent-invocations", + type=int, + default=DEFAULT_MAX_CONCURRENT_INVOCATIONS, + help="Maximum concurrent Fabric invocations; use 0 for unlimited.", + ) + parser.add_argument( + "--idle-session-timeout-seconds", + type=float, + default=DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS, + help="Seconds of inactivity before a logical session expires.", + ) + parser.add_argument( + "--session-cleanup-interval-seconds", + type=float, + default=DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS, + help="Seconds between idle-session cleanup checks.", + ) args = parser.parse_args(argv) import uvicorn logging.basicConfig(level=logging.INFO) uvicorn.run( - create_fabric_serving_app(args.agent_config), + create_fabric_serving_app( + args.agent_config, + max_concurrent_invocations=args.max_concurrent_invocations, + idle_session_timeout_seconds=args.idle_session_timeout_seconds, + session_cleanup_interval_seconds=args.session_cleanup_interval_seconds, + ), host=args.host, port=args.port, log_config=None, diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py index 6c387345b3..42524fd15a 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -5,13 +5,18 @@ from __future__ import annotations +import asyncio import logging from pathlib import Path from typing import Any from nemo_agents_plugin.agent_config import AgentConfig from nemo_agents_plugin.fabric.runtime import FabricInvocationRequest, FabricRuntimeResult, invoke_fabric_runtime -from nemo_agents_plugin.fabric.session_registry import FabricRuntimeSession, FabricSessionRegistry +from nemo_agents_plugin.fabric.session_registry import ( + FabricRuntimeSession, + FabricSessionNotFoundError, + FabricSessionRegistry, +) from nemo_agents_plugin.fabric.translator import FabricTranslationError, translate_agent_config # CI type-checks this plugin via ty extra-paths without installing nemo-agents deps. @@ -19,11 +24,19 @@ logger = logging.getLogger(__name__) +DEFAULT_MAX_CONCURRENT_INVOCATIONS = 8 +DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS = 30 * 60 +DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS = 5 * 60 + class FabricSessionStartError(RuntimeError): """Raised when a Fabric runtime cannot be started for a Platform session.""" +class FabricSessionStopError(RuntimeError): + """Raised when a Fabric runtime cannot be stopped for a Platform session.""" + + class FabricSessionManager: """Create Fabric runtimes lazily from one reusable Platform agent definition.""" @@ -34,11 +47,18 @@ def __init__( base_dir: Path, session_registry: FabricSessionRegistry, fabric: Any | None = None, + max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS, ) -> None: + if max_concurrent_invocations < 0: + raise ValueError("max_concurrent_invocations must be greater than or equal to zero.") + self._agent_config = agent_config self._base_dir = base_dir self._session_registry = session_registry self._fabric = fabric + self._invocation_semaphore = ( + asyncio.Semaphore(max_concurrent_invocations) if max_concurrent_invocations > 0 else None + ) async def open_session(self) -> FabricRuntimeSession: """Materialize a Fabric config, start its runtime, and register the session.""" @@ -76,4 +96,51 @@ async def invoke_session( ) -> FabricRuntimeResult: """Serialize and invoke one turn on a session's active runtime.""" async with session.invocation_lock: - return await invoke_fabric_runtime(session.runtime, request) + if session.closing: + raise FabricSessionNotFoundError(f"Fabric session '{session.session_id}' was not found.") + try: + if self._invocation_semaphore is None: + return await invoke_fabric_runtime(session.runtime, request) + async with self._invocation_semaphore: + return await invoke_fabric_runtime(session.runtime, request) + finally: + await self._session_registry.refresh_activity(session) + + async def close_session(self, session_id: str) -> None: + """Remove a session and stop its runtime after any active turn finishes.""" + session = await self._session_registry.remove(session_id) + if session is None: + raise FabricSessionNotFoundError(f"Fabric session '{session_id}' was not found.") + + await self._stop_session(session) + + async def expire_idle_sessions(self, *, idle_timeout_seconds: float) -> int: + """Stop and remove sessions that have exceeded the idle timeout.""" + expired = await self._session_registry.remove_expired(idle_timeout_seconds=idle_timeout_seconds) + for session in expired: + try: + await self._stop_session(session) + except FabricSessionStopError: + logger.exception("Failed to stop expired Fabric session %s.", session.session_id) + return len(expired) + + async def close_all_sessions(self) -> int: + """Drain the registry and stop every remaining runtime.""" + sessions = await self._session_registry.drain() + + async def stop_session(session: FabricRuntimeSession) -> None: + try: + await self._stop_session(session) + except FabricSessionStopError: + logger.exception("Failed to stop Fabric session %s during shutdown.", session.session_id) + + await asyncio.gather(*(stop_session(session) for session in sessions)) + return len(sessions) + + async def _stop_session(self, session: FabricRuntimeSession) -> None: + """Stop one session after any active invocation releases its lock.""" + async with session.invocation_lock: + try: + await session.runtime.stop() + except FabricError as error: + raise FabricSessionStopError(f"Fabric runtime shutdown failed: {error}") from error diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py index bdd3a67f40..7287d99be9 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py @@ -23,6 +23,7 @@ class FabricRuntimeSession: runtime: Runtime created_at: float last_accessed_at: float + closing: bool = False invocation_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) @@ -34,12 +35,17 @@ class FabricSessionAlreadyExistsError(ValueError): """Raised when a logical Fabric session ID is registered twice.""" +class FabricSessionRegistryClosedError(RuntimeError): + """Raised when registering a session after shutdown has started.""" + + class FabricSessionRegistry: """Maintain the process-local mapping from Platform sessions to runtimes.""" def __init__(self) -> None: self._sessions: dict[str, FabricRuntimeSession] = {} self._lock = asyncio.Lock() + self._closed = False async def register( self, @@ -58,6 +64,8 @@ async def register( ) async with self._lock: + if self._closed: + raise FabricSessionRegistryClosedError("Fabric session registry is closed.") if resolved_session_id in self._sessions: raise FabricSessionAlreadyExistsError(f"Fabric session '{resolved_session_id}' is already registered.") self._sessions[resolved_session_id] = session @@ -73,10 +81,42 @@ async def get(self, session_id: str) -> FabricRuntimeSession: session.last_accessed_at = time.monotonic() return session + async def refresh_activity(self, session: FabricRuntimeSession) -> None: + """Refresh activity for a session that is still registered.""" + async with self._lock: + if self._sessions.get(session.session_id) is session: + session.last_accessed_at = time.monotonic() + async def remove(self, session_id: str) -> FabricRuntimeSession | None: - """Remove and return a session without stopping its runtime.""" + """Mark a session as closing, then remove and return it.""" + async with self._lock: + session = self._sessions.pop(session_id, None) + if session is not None: + session.closing = True + return session + + async def remove_expired(self, *, idle_timeout_seconds: float) -> list[FabricRuntimeSession]: + """Remove inactive sessions that are not currently invoking.""" + cutoff = time.monotonic() - idle_timeout_seconds + expired: list[FabricRuntimeSession] = [] + async with self._lock: + for session_id, session in list(self._sessions.items()): + if session.last_accessed_at > cutoff or session.invocation_lock.locked(): + continue + session.closing = True + expired.append(session) + del self._sessions[session_id] + return expired + + async def drain(self) -> list[FabricRuntimeSession]: + """Close the registry and remove all remaining sessions.""" async with self._lock: - return self._sessions.pop(session_id, None) + self._closed = True + sessions = list(self._sessions.values()) + self._sessions.clear() + for session in sessions: + session.closing = True + return sessions async def count(self) -> int: """Return the number of registered logical sessions.""" diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index 69357a6c25..600bfa38c8 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -3,9 +3,10 @@ from __future__ import annotations +import asyncio from pathlib import Path from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest import yaml @@ -19,7 +20,11 @@ ) from nemo_agents_plugin.fabric.server import SESSION_ID_HEADER, create_fabric_serving_app from nemo_agents_plugin.fabric.serving_models import ChatCompletionRequest -from nemo_agents_plugin.fabric.session_manager import FabricSessionManager, FabricSessionStartError +from nemo_agents_plugin.fabric.session_manager import ( + FabricSessionManager, + FabricSessionStartError, + FabricSessionStopError, +) from nemo_agents_plugin.fabric.session_registry import FabricSessionNotFoundError, FabricSessionRegistry @@ -79,6 +84,31 @@ def test_startup_loads_and_validates_agent_config( assert mock_validate_agent_config == [(app.state.agent_config, tmp_path)] +def test_shutdown_stops_all_registered_runtimes( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], +) -> None: + class _Runtime: + def __init__(self) -> None: + self.stop_calls = 0 + + async def stop(self) -> None: + self.stop_calls += 1 + + runtime = _Runtime() + app = create_fabric_serving_app(_write_agent_config(tmp_path)) + + with TestClient(app) as client: + registry = app.state.session_registry + + async def register_runtime() -> None: + await registry.register(cast(Any, runtime), session_id="session-1") + + client.portal.call(register_runtime) + + assert runtime.stop_calls == 1 + + def test_startup_fails_for_invalid_agent_config( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], @@ -92,6 +122,15 @@ def test_startup_fails_for_invalid_agent_config( assert mock_validate_agent_config == [] +def test_rejects_non_positive_session_cleanup_settings(tmp_path: Path) -> None: + config_path = _write_agent_config(tmp_path) + + with pytest.raises(ValueError, match="idle_session_timeout_seconds"): + create_fabric_serving_app(config_path, idle_session_timeout_seconds=0) + with pytest.raises(ValueError, match="session_cleanup_interval_seconds"): + create_fabric_serving_app(config_path, session_cleanup_interval_seconds=0) + + def test_chat_completion_without_session_id_opens_and_returns_session( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], @@ -289,3 +328,77 @@ def test_chat_completion_request_translates_final_user_turn() -> None: assert invocation_request.input == "Say hello." assert invocation_request.caller_context == {"session_id": "session-1"} + + +def test_close_session_stops_registered_runtime( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + app = create_fabric_serving_app(_write_agent_config(tmp_path)) + close_calls: list[str] = [] + + async def close_session(session_id: str) -> None: + close_calls.append(session_id) + + with TestClient(app) as client: + monkeypatch.setattr(app.state.session_manager, "close_session", close_session) + response = client.delete("/v1/sessions/session-1") + + assert response.status_code == 204 + assert response.content == b"" + assert close_calls == ["session-1"] + + +@pytest.mark.parametrize( + ("error", "status_code"), + [ + (FabricSessionNotFoundError("missing session"), 404), + (FabricSessionStopError("shutdown failed"), 502), + ], +) +def test_close_session_maps_errors( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, + error: Exception, + status_code: int, +) -> None: + app = create_fabric_serving_app(_write_agent_config(tmp_path)) + + async def close_session(session_id: str) -> None: + raise error + + with TestClient(app) as client: + monkeypatch.setattr(app.state.session_manager, "close_session", close_session) + response = client.delete("/v1/sessions/session-1") + + assert response.status_code == status_code + assert response.json() == {"detail": str(error)} + + +@pytest.mark.asyncio +async def test_idle_cleanup_runs_periodically_until_shutdown() -> None: + cleanup_calls: list[float] = [] + cleanup_ran = asyncio.Event() + shutdown = asyncio.Event() + + class _Manager: + async def expire_idle_sessions(self, *, idle_timeout_seconds: float) -> int: + cleanup_calls.append(idle_timeout_seconds) + cleanup_ran.set() + return 0 + + cleanup = asyncio.create_task( + server._run_idle_session_cleanup( + cast(Any, _Manager()), + idle_timeout_seconds=30.0, + cleanup_interval_seconds=0.01, + shutdown_event=shutdown, + ) + ) + await asyncio.wait_for(cleanup_ran.wait(), timeout=1) + shutdown.set() + await cleanup + + assert cleanup_calls == [30.0] diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py index 056c25115d..7ffc023f86 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py @@ -12,7 +12,10 @@ from nemo_agents_plugin.fabric import session_manager from nemo_agents_plugin.fabric.runtime import FabricInvocationRequest, FabricRuntimeResult from nemo_agents_plugin.fabric.session_manager import FabricSessionManager -from nemo_agents_plugin.fabric.session_registry import FabricSessionRegistry +from nemo_agents_plugin.fabric.session_registry import ( + FabricSessionNotFoundError, + FabricSessionRegistry, +) class _FakeRuntime: @@ -159,6 +162,178 @@ async def test_resolve_session_reuses_registered_runtime( assert fabric.start_calls == [] +@pytest.mark.asyncio +async def test_close_session_removes_session_and_stops_runtime(tmp_path: Path) -> None: + runtime = _FakeRuntime() + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, runtime), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + ) + + await manager.close_session(session.session_id) + + assert runtime.stop_calls == 1 + assert await registry.count() == 0 + with pytest.raises(FabricSessionNotFoundError, match="session-1"): + await manager.resolve_session(session.session_id) + + +@pytest.mark.asyncio +async def test_close_session_waits_for_active_invocation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + runtime = _FakeRuntime() + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, runtime), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + ) + invocation_started = asyncio.Event() + release_invocation = asyncio.Event() + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + invocation_started.set() + await release_invocation.wait() + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + invocation = asyncio.create_task( + manager.invoke_session(session, FabricInvocationRequest(input="hello")), + ) + await invocation_started.wait() + + close = asyncio.create_task(manager.close_session(session.session_id)) + await asyncio.sleep(0) + + assert runtime.stop_calls == 0 + with pytest.raises(FabricSessionNotFoundError, match="session-1"): + await manager.resolve_session(session.session_id) + + release_invocation.set() + await invocation + await close + + assert runtime.stop_calls == 1 + + +@pytest.mark.asyncio +async def test_resolved_session_cannot_invoke_after_close_starts(tmp_path: Path) -> None: + runtime = _FakeRuntime() + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, runtime), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + ) + + await manager.close_session(session.session_id) + + with pytest.raises(FabricSessionNotFoundError, match="session-1"): + await manager.invoke_session(session, FabricInvocationRequest(input="hello")) + + +@pytest.mark.asyncio +async def test_expire_idle_sessions_stops_expired_runtime( + tmp_path: Path, +) -> None: + runtime = _FakeRuntime() + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, runtime), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + ) + session.last_accessed_at = float("-inf") + + expired_count = await manager.expire_idle_sessions(idle_timeout_seconds=30.0) + + assert expired_count == 1 + assert runtime.stop_calls == 1 + assert await registry.count() == 0 + + +@pytest.mark.asyncio +async def test_close_all_sessions_stops_every_runtime(tmp_path: Path) -> None: + first_runtime = _FakeRuntime() + second_runtime = _FakeRuntime() + registry = FabricSessionRegistry() + await registry.register(cast(Any, first_runtime), session_id="session-1") + await registry.register(cast(Any, second_runtime), session_id="session-2") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + ) + + closed_count = await manager.close_all_sessions() + + assert closed_count == 2 + assert first_runtime.stop_calls == 1 + assert second_runtime.stop_calls == 1 + assert await registry.count() == 0 + + +@pytest.mark.asyncio +async def test_close_all_sessions_continues_after_stop_failure(tmp_path: Path) -> None: + class _FailingRuntime(_FakeRuntime): + async def stop(self) -> None: + self.stop_calls += 1 + raise session_manager.FabricError("stop failed") + + failing_runtime = _FailingRuntime() + healthy_runtime = _FakeRuntime() + registry = FabricSessionRegistry() + await registry.register(cast(Any, failing_runtime), session_id="session-1") + await registry.register(cast(Any, healthy_runtime), session_id="session-2") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + ) + + closed_count = await manager.close_all_sessions() + + assert closed_count == 2 + assert failing_runtime.stop_calls == 1 + assert healthy_runtime.stop_calls == 1 + + +@pytest.mark.asyncio +async def test_invoke_session_refreshes_activity( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + ) + refresh_calls: list[Any] = [] + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + return FabricRuntimeResult(status="succeeded", response=request.input) + + async def refresh_activity(resolved_session: Any) -> None: + refresh_calls.append(resolved_session) + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + monkeypatch.setattr(registry, "refresh_activity", refresh_activity) + + await manager.invoke_session(session, FabricInvocationRequest(input="hello")) + + assert refresh_calls == [session] + + @pytest.mark.asyncio async def test_invoke_session_serializes_turns_for_same_runtime( tmp_path: Path, @@ -276,3 +451,143 @@ async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) result = await manager.invoke_session(session, FabricInvocationRequest(input="next")) assert result.response == "recovered" + + +def test_rejects_negative_concurrency_limit(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="must be greater than or equal to zero"): + FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=FabricSessionRegistry(), + max_concurrent_invocations=-1, + ) + + +@pytest.mark.asyncio +async def test_invoke_session_limits_concurrency_across_sessions( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = FabricSessionRegistry() + first_session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-1") + second_session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-2") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + max_concurrent_invocations=1, + ) + first_started = asyncio.Event() + release_first = asyncio.Event() + invocation_order: list[str] = [] + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + invocation_order.append(request.input) + if request.input == "first": + first_started.set() + await release_first.wait() + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + first = asyncio.create_task( + manager.invoke_session(first_session, FabricInvocationRequest(input="first")), + ) + await first_started.wait() + second = asyncio.create_task( + manager.invoke_session(second_session, FabricInvocationRequest(input="second")), + ) + await asyncio.sleep(0) + + assert invocation_order == ["first"] + + release_first.set() + await asyncio.gather(first, second) + + assert invocation_order == ["first", "second"] + + +@pytest.mark.asyncio +async def test_cancelled_capacity_waiter_does_not_leak_capacity( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = FabricSessionRegistry() + first_session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-1") + second_session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-2") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + max_concurrent_invocations=1, + ) + first_started = asyncio.Event() + release_first = asyncio.Event() + invocation_order: list[str] = [] + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + invocation_order.append(request.input) + if request.input == "first": + first_started.set() + await release_first.wait() + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + first = asyncio.create_task( + manager.invoke_session(first_session, FabricInvocationRequest(input="first")), + ) + await first_started.wait() + cancelled = asyncio.create_task( + manager.invoke_session(second_session, FabricInvocationRequest(input="cancelled")), + ) + await asyncio.sleep(0) + + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + + release_first.set() + await first + result = await manager.invoke_session(second_session, FabricInvocationRequest(input="next")) + + assert invocation_order == ["first", "next"] + assert result.response == "next" + + +@pytest.mark.asyncio +async def test_zero_concurrency_limit_allows_parallel_sessions( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = FabricSessionRegistry() + first_session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-1") + second_session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-2") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + max_concurrent_invocations=0, + ) + both_started = asyncio.Event() + release = asyncio.Event() + active_invocations = 0 + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + nonlocal active_invocations + active_invocations += 1 + if active_invocations == 2: + both_started.set() + await release.wait() + active_invocations -= 1 + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + first = asyncio.create_task( + manager.invoke_session(first_session, FabricInvocationRequest(input="first")), + ) + second = asyncio.create_task( + manager.invoke_session(second_session, FabricInvocationRequest(input="second")), + ) + + await asyncio.wait_for(both_started.wait(), timeout=1) + release.set() + await asyncio.gather(first, second) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py b/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py index c2b1570366..ddf5a791d8 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py @@ -13,6 +13,7 @@ FabricSessionAlreadyExistsError, FabricSessionNotFoundError, FabricSessionRegistry, + FabricSessionRegistryClosedError, ) @@ -78,3 +79,55 @@ async def test_remove_returns_session_and_is_idempotent() -> None: assert await registry.remove("session-1") is session assert await registry.remove("session-1") is None assert await registry.count() == 0 + + +@pytest.mark.asyncio +async def test_refresh_activity_updates_registered_session(monkeypatch: pytest.MonkeyPatch) -> None: + clock = iter([10.0, 20.0]) + monkeypatch.setattr(session_registry, "time", SimpleNamespace(monotonic=lambda: next(clock))) + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, object()), session_id="session-1") + + await registry.refresh_activity(session) + + assert session.last_accessed_at == 20.0 + + +@pytest.mark.asyncio +async def test_remove_expired_removes_only_idle_sessions(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(session_registry, "time", SimpleNamespace(monotonic=lambda: 100.0)) + registry = FabricSessionRegistry() + expired = await registry.register(cast(Any, object()), session_id="expired") + active = await registry.register(cast(Any, object()), session_id="active") + recent = await registry.register(cast(Any, object()), session_id="recent") + expired.last_accessed_at = 10.0 + active.last_accessed_at = 10.0 + recent.last_accessed_at = 90.0 + + await active.invocation_lock.acquire() + try: + removed = await registry.remove_expired(idle_timeout_seconds=30.0) + finally: + active.invocation_lock.release() + + assert removed == [expired] + assert expired.closing is True + assert active.closing is False + assert recent.closing is False + assert await registry.count() == 2 + + +@pytest.mark.asyncio +async def test_drain_removes_sessions_and_rejects_new_registrations() -> None: + registry = FabricSessionRegistry() + first = await registry.register(cast(Any, object()), session_id="session-1") + second = await registry.register(cast(Any, object()), session_id="session-2") + + drained = await registry.drain() + + assert drained == [first, second] + assert first.closing is True + assert second.closing is True + assert await registry.count() == 0 + with pytest.raises(FabricSessionRegistryClosedError, match="registry is closed"): + await registry.register(cast(Any, object()), session_id="session-3") From 57fe33b54096767fb619414c34eb76ab05fd96c5 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 11:19:40 -0500 Subject: [PATCH 07/16] shared config validation protocol Signed-off-by: Manjesh Mogallapalli --- .../agent_config_formats.py | 112 ++++++++++++++++++ .../src/nemo_agents_plugin/api/v2/agents.py | 24 +--- .../nemo_agents_plugin/api/v2/deployments.py | 28 ++--- .../tests/unit/test_agent_config_formats.py | 100 ++++++++++++++++ 4 files changed, 228 insertions(+), 36 deletions(-) create mode 100644 plugins/nemo-agents/src/nemo_agents_plugin/agent_config_formats.py create mode 100644 plugins/nemo-agents/tests/unit/test_agent_config_formats.py diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/agent_config_formats.py b/plugins/nemo-agents/src/nemo_agents_plugin/agent_config_formats.py new file mode 100644 index 0000000000..f1be7353c8 --- /dev/null +++ b/plugins/nemo-agents/src/nemo_agents_plugin/agent_config_formats.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared validation and deployment resolution for agent config formats.""" + +from __future__ import annotations + +from typing import Any, Protocol + +from nemo_agents_plugin.agent_config import AgentConfig +from nemo_agents_plugin.entities import NAT_WORKFLOW_CONFIG_FORMAT, NEMO_AGENTS_SPEC_CONFIG_FORMAT +from nemo_agents_plugin.utils import inject_default_model, inject_gateway_url, inject_nemo_trace_fields +from pydantic import ValidationError + + +class AgentConfigFormatError(ValueError): + """Base error for unsupported or invalid agent config formats.""" + + +class UnsupportedAgentConfigFormatError(AgentConfigFormatError): + """Raised when no handler exists for an agent config format.""" + + +class InvalidAgentConfigError(AgentConfigFormatError): + """Raised when an agent config does not satisfy its format contract.""" + + +class AgentConfigFormatHandler(Protocol): + """Validate and resolve one persisted agent config format.""" + + def validate(self, config: dict[str, Any]) -> dict[str, Any]: ... + + def resolve_for_deployment( + self, + config: dict[str, Any], + *, + workspace: str, + agent_name: str, + ) -> dict[str, Any]: ... + + +class _NatWorkflowConfigHandler: + def validate(self, config: dict[str, Any]) -> dict[str, Any]: + return config + + def resolve_for_deployment( + self, + config: dict[str, Any], + *, + workspace: str, + agent_name: str, + ) -> dict[str, Any]: + resolved = inject_gateway_url(config, workspace) + resolved = inject_default_model(resolved) + inject_nemo_trace_fields(resolved, workspace=workspace, agent_name=agent_name) + return resolved + + +class _NemoAgentsSpecConfigHandler: + def validate(self, config: dict[str, Any]) -> dict[str, Any]: + return self._normalize(config) + + def resolve_for_deployment( + self, + config: dict[str, Any], + *, + workspace: str, + agent_name: str, + ) -> dict[str, Any]: + del workspace, agent_name + return self._normalize(config) + + @staticmethod + def _normalize(config: dict[str, Any]) -> dict[str, Any]: + try: + return AgentConfig.model_validate(config).model_dump(exclude_none=True) + except ValidationError as error: + raise InvalidAgentConfigError(f"Invalid agent config: {error}") from error + + +_AGENT_CONFIG_FORMAT_HANDLERS: dict[str, AgentConfigFormatHandler] = { + NAT_WORKFLOW_CONFIG_FORMAT: _NatWorkflowConfigHandler(), + NEMO_AGENTS_SPEC_CONFIG_FORMAT: _NemoAgentsSpecConfigHandler(), +} + + +def get_agent_config_format_handler(config_format: str) -> AgentConfigFormatHandler: + """Return the handler registered for an agent config format.""" + try: + return _AGENT_CONFIG_FORMAT_HANDLERS[config_format] + except KeyError as error: + raise UnsupportedAgentConfigFormatError(f"Unsupported config_format {config_format!r}.") from error + + +def validate_agent_config(config_format: str, config: dict[str, Any]) -> dict[str, Any]: + """Validate and normalize an agent config before persistence.""" + return get_agent_config_format_handler(config_format).validate(config) + + +def resolve_agent_config_for_deployment( + config_format: str, + config: dict[str, Any], + *, + workspace: str, + agent_name: str, +) -> dict[str, Any]: + """Resolve a persisted agent config for deployment.""" + return get_agent_config_format_handler(config_format).resolve_for_deployment( + config, + workspace=workspace, + agent_name=agent_name, + ) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/api/v2/agents.py b/plugins/nemo-agents/src/nemo_agents_plugin/api/v2/agents.py index 23b24f2801..e294576939 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/api/v2/agents.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/api/v2/agents.py @@ -14,16 +14,11 @@ from typing import Any from fastapi import APIRouter, Depends, HTTPException, Query -from nemo_agents_plugin.agent_config import AgentConfig +from nemo_agents_plugin.agent_config_formats import AgentConfigFormatError, validate_agent_config from nemo_agents_plugin.api.v2._perms import AgentPerms from nemo_agents_plugin.api.v2.dependencies import get_entity_client from nemo_agents_plugin.authz import scope -from nemo_agents_plugin.entities import ( - NAT_WORKFLOW_CONFIG_FORMAT, - NEMO_AGENTS_SPEC_CONFIG_FORMAT, - Agent, - AgentDeployment, -) +from nemo_agents_plugin.entities import Agent, AgentDeployment from nemo_agents_plugin.schema import ( AgentFilter, AgentPage, @@ -33,7 +28,6 @@ from nemo_platform_plugin.authz import CallerKind, path_rule from nemo_platform_plugin.entity_client import NemoEntitiesClient, NemoEntityConflictError, NemoEntityNotFoundError from nemo_platform_plugin.schema import PaginationData -from pydantic import ValidationError # Deployment statuses that block agent deletion. # "failed" and "deleting" are excluded — they are terminal/in-cleanup and @@ -192,13 +186,7 @@ async def delete_agent( def _validate_agent_config_for_create(body: CreateAgentRequest) -> dict[str, Any]: - if body.config_format == NAT_WORKFLOW_CONFIG_FORMAT: - return body.config - - if body.config_format == NEMO_AGENTS_SPEC_CONFIG_FORMAT: - try: - return AgentConfig.model_validate(body.config).model_dump(exclude_none=True) - except ValidationError as exc: - raise HTTPException(status_code=400, detail=f"Invalid agent config: {exc}") from exc - - raise HTTPException(status_code=400, detail=f"Unsupported config_format {body.config_format!r}.") + try: + return validate_agent_config(body.config_format, body.config) + except AgentConfigFormatError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/api/v2/deployments.py b/plugins/nemo-agents/src/nemo_agents_plugin/api/v2/deployments.py index a0cbba077f..e9bacf33e5 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/api/v2/deployments.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/api/v2/deployments.py @@ -21,13 +21,11 @@ from typing import Any from fastapi import APIRouter, Depends, HTTPException, Query -from nemo_agents_plugin.agent_config import AgentConfig +from nemo_agents_plugin.agent_config_formats import AgentConfigFormatError, resolve_agent_config_for_deployment from nemo_agents_plugin.api.v2._perms import DeploymentPerms from nemo_agents_plugin.api.v2.dependencies import get_entity_client from nemo_agents_plugin.authz import scope from nemo_agents_plugin.entities import ( - NAT_WORKFLOW_CONFIG_FORMAT, - NEMO_AGENTS_SPEC_CONFIG_FORMAT, Agent, AgentDeployment, is_container_deployment_mode, @@ -37,12 +35,10 @@ DeploymentFilter, DeploymentPage, ) -from nemo_agents_plugin.utils import inject_default_model, inject_gateway_url, inject_nemo_trace_fields from nemo_platform_plugin.api.filters import make_filter_obj_dep from nemo_platform_plugin.authz import CallerKind, path_rule from nemo_platform_plugin.entity_client import NemoEntitiesClient, NemoEntityConflictError, NemoEntityNotFoundError from nemo_platform_plugin.schema import PaginationData -from pydantic import ValidationError logger = logging.getLogger(__name__) @@ -112,19 +108,15 @@ async def create_deployment( def _resolve_deployment_config(agent: Agent, *, workspace: str) -> dict[str, Any]: - if agent.config_format == NAT_WORKFLOW_CONFIG_FORMAT: - resolved_config = inject_gateway_url(agent.config, workspace) - resolved_config = inject_default_model(resolved_config) - inject_nemo_trace_fields(resolved_config, workspace=workspace, agent_name=agent.name) - return resolved_config - - if agent.config_format == NEMO_AGENTS_SPEC_CONFIG_FORMAT: - try: - return AgentConfig.model_validate(agent.config).model_dump(exclude_none=True) - except ValidationError as exc: - raise HTTPException(status_code=400, detail=f"Invalid agent config: {exc}") from exc - - raise HTTPException(status_code=400, detail=f"Unsupported config_format {agent.config_format!r}.") + try: + return resolve_agent_config_for_deployment( + agent.config_format, + agent.config, + workspace=workspace, + agent_name=agent.name, + ) + except AgentConfigFormatError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc @router.get("/deployments", response_model=DeploymentPage, tags=["Agent Deployments"]) diff --git a/plugins/nemo-agents/tests/unit/test_agent_config_formats.py b/plugins/nemo-agents/tests/unit/test_agent_config_formats.py new file mode 100644 index 0000000000..e5e25d2816 --- /dev/null +++ b/plugins/nemo-agents/tests/unit/test_agent_config_formats.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +import pytest +from nemo_agents_plugin import agent_config_formats +from nemo_agents_plugin.agent_config_formats import ( + InvalidAgentConfigError, + UnsupportedAgentConfigFormatError, + resolve_agent_config_for_deployment, + validate_agent_config, +) +from nemo_agents_plugin.entities import NAT_WORKFLOW_CONFIG_FORMAT, NEMO_AGENTS_SPEC_CONFIG_FORMAT + + +def _nemo_agents_config() -> dict[str, Any]: + return { + "config_format": NEMO_AGENTS_SPEC_CONFIG_FORMAT, + "name": "test-agent", + "default_harness": "hermes", + "harnesses": {"hermes": {"kind": "hermes"}}, + } + + +def test_nat_config_validation_preserves_legacy_payload() -> None: + config = {"workflow": {"_type": "chat_completion"}} + + validated = validate_agent_config(NAT_WORKFLOW_CONFIG_FORMAT, config) + + assert validated is config + + +def test_nemo_agents_config_validation_normalizes_payload() -> None: + validated = validate_agent_config(NEMO_AGENTS_SPEC_CONFIG_FORMAT, _nemo_agents_config()) + + assert validated["config_format"] == NEMO_AGENTS_SPEC_CONFIG_FORMAT + assert validated["environment"]["provider"] == "local" + + +def test_nemo_agents_config_validation_rejects_invalid_payload() -> None: + config = _nemo_agents_config() + config["default_harness"] = "missing" + + with pytest.raises(InvalidAgentConfigError, match="Invalid agent config"): + validate_agent_config(NEMO_AGENTS_SPEC_CONFIG_FORMAT, config) + + +def test_unknown_config_format_is_rejected() -> None: + with pytest.raises(UnsupportedAgentConfigFormatError, match="Unsupported config_format 'custom-v2'"): + validate_agent_config("custom-v2", {}) + + +def test_nat_deployment_resolution_applies_legacy_injections(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[tuple[str, Any]] = [] + + def inject_gateway(config: dict[str, Any], workspace: str) -> dict[str, Any]: + calls.append(("gateway", workspace)) + return {**config, "gateway": True} + + def inject_model(config: dict[str, Any]) -> dict[str, Any]: + calls.append(("model", None)) + return {**config, "model": True} + + def inject_trace(config: dict[str, Any], *, workspace: str, agent_name: str) -> None: + calls.append(("trace", (workspace, agent_name))) + config["trace"] = True + + monkeypatch.setattr(agent_config_formats, "inject_gateway_url", inject_gateway) + monkeypatch.setattr(agent_config_formats, "inject_default_model", inject_model) + monkeypatch.setattr(agent_config_formats, "inject_nemo_trace_fields", inject_trace) + + resolved = resolve_agent_config_for_deployment( + NAT_WORKFLOW_CONFIG_FORMAT, + {"workflow": {}}, + workspace="test-workspace", + agent_name="test-agent", + ) + + assert resolved == {"workflow": {}, "gateway": True, "model": True, "trace": True} + assert calls == [ + ("gateway", "test-workspace"), + ("model", None), + ("trace", ("test-workspace", "test-agent")), + ] + + +def test_nemo_agents_deployment_resolution_only_normalizes_payload() -> None: + resolved = resolve_agent_config_for_deployment( + NEMO_AGENTS_SPEC_CONFIG_FORMAT, + _nemo_agents_config(), + workspace="test-workspace", + agent_name="test-agent", + ) + + assert resolved["config_format"] == NEMO_AGENTS_SPEC_CONFIG_FORMAT + assert resolved["environment"]["provider"] == "local" + assert "workflow" not in resolved From e6fb7c048c5e8cb851c2248e2dd54152700ef66d Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 11:33:48 -0500 Subject: [PATCH 08/16] wiring fabric deployment Signed-off-by: Manjesh Mogallapalli --- .../nemo_agents_plugin/runner/in_memory.py | 123 +++++++++++------- .../nemo-agents/tests/unit/test_controller.py | 1 - .../tests/unit/test_runner_in_memory.py | 96 ++++++++++++-- 3 files changed, 161 insertions(+), 59 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/runner/in_memory.py b/plugins/nemo-agents/src/nemo_agents_plugin/runner/in_memory.py index 0c2b62d6eb..3b1010e2ca 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/runner/in_memory.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/runner/in_memory.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""InMemoryRunnerBackend — spawns ``nat serve`` subprocesses for agent deployments. +"""InMemoryRunnerBackend — spawns local agent-server subprocesses. Agent processes run as local subprocesses on the same machine as the platform server (``deployment_mode=subprocess``). Process state is tracked in memory; @@ -12,7 +12,7 @@ :class:`~nemo_agents_plugin.runner.deployments_backend.DeploymentsRunnerBackend` instead. -Module-level helpers ``system_dir_for_workspace`` and ``log_path_for_deployment`` +Module-level helpers ``system_dir`` and ``log_path_for_deployment`` encode the on-disk layout convention so out-of-process callers (e.g. the ``nemo agents logs`` CLI) can locate a deployment's log file without instantiating the backend. The convention is intentionally narrow: it is @@ -36,7 +36,7 @@ import httpx import yaml from nemo_agents_plugin.config import AgentsConfig, ControllerConfig -from nemo_agents_plugin.entities import NEMO_AGENTS_SPEC_CONFIG_FORMAT, DeploymentMode +from nemo_agents_plugin.entities import AGENT_CONFIG_FILENAME, NEMO_AGENTS_SPEC_CONFIG_FORMAT, DeploymentMode from nemo_agents_plugin.runner.backend import DeploymentInfo, LocalLog, LogLocation, NotYetAvailable, RunnerBackend # Match characters not safe for filesystem paths. Deployment names are @@ -116,6 +116,16 @@ def config_path_for_deployment(workspace: str, name: str, workspace_dir: Path | return base / f"{_sanitize_filename(name)}.yaml" +def _write_yaml_config(path: Path, config: dict[str, Any]) -> Path: + """Atomically write a YAML config to an absolute path.""" + path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = path.with_suffix(path.suffix + ".tmp") + with tmp_path.open("w", encoding="utf-8") as fh: + yaml.safe_dump(config, fh) + tmp_path.replace(path) + return path + + logger = logging.getLogger(__name__) @@ -213,7 +223,7 @@ async def create_deployment( """Start a local deployment for NAT workflows or Platform-owned agent specs.""" del image, deployment_mode if config.get("config_format") == NEMO_AGENTS_SPEC_CONFIG_FORMAT: - return await self._create_fabric_deployment(workspace, name, config) + return await self._create_fabric_deployment(workspace, name, config, port) key = (workspace, name) config_path = await asyncio.to_thread(self._write_config, workspace, name, config) @@ -245,35 +255,45 @@ async def create_deployment( ) return info - async def _create_fabric_deployment(self, workspace: str, name: str, config: dict[str, Any]) -> DeploymentInfo: - """Validate and prepare a Platform-owned Fabric-backed deployment.""" + async def _create_fabric_deployment( + self, + workspace: str, + name: str, + config: dict[str, Any], + port: int, + ) -> DeploymentInfo: + """Validate and start a Platform-owned Fabric-backed agent server.""" + key = (workspace, name) base_dir = self._fabric_base_dir_for(workspace, name) await asyncio.to_thread(base_dir.mkdir, parents=True, exist_ok=True) try: - validation_result = await validate_platform_agent_config(config, base_dir=base_dir) + config_path = await asyncio.to_thread(self._write_fabric_config, base_dir, config) + await validate_platform_agent_config(config, base_dir=base_dir) + log_path = self.log_path_for(workspace, name) + proc = await asyncio.to_thread(self._spawn_fabric, name, config_path, log_path, port) except Exception: await asyncio.to_thread(shutil.rmtree, base_dir, ignore_errors=True) raise - log_path = self.log_path_for(workspace, name) - await asyncio.to_thread(self._write_fabric_validation_log, workspace, name, log_path, validation_result) - info = DeploymentInfo( name=name, - status="running", + status="starting", + port=port, + pid=proc.pid, + endpoint=f"http://127.0.0.1:{port}", log_path=str(log_path), - extra={ - "runtime": "fabric", - "base_dir": str(base_dir), - "prepared": True, - }, + extra={"base_dir": str(base_dir)}, ) - self._deployments[(workspace, name)] = info + self._processes[key] = proc + self._deployments[key] = info logger.info( - "Prepared Fabric-backed deployment for '%s/%s' (base_dir=%s)", + "Spawned Fabric-backed deployment for '%s/%s' (pid=%d, port=%d, base_dir=%s, log=%s)", workspace, name, + proc.pid, + port, base_dir, + log_path, ) return info @@ -282,8 +302,6 @@ async def get_deployment_status(self, workspace: str, name: str) -> DeploymentIn info = self._deployments.get(key) if info is None: return None - if info.extra.get("runtime") == "fabric": - return info proc = self._processes.get(key) if proc is not None and proc.poll() is not None: @@ -335,6 +353,11 @@ def _get_http_client(self) -> httpx.AsyncClient: async def shutdown(self) -> None: """Terminate all managed processes (best-effort).""" items = list(self._processes.items()) + fabric_base_dirs: list[Path] = [] + for info in self._deployments.values(): + base_dir = info.extra.get("base_dir") + if base_dir is not None: + fabric_base_dirs.append(Path(base_dir)) labels = [f"{ws}/{nm}" for (ws, nm), _ in items] results = await asyncio.gather( *(asyncio.to_thread(self._terminate, f"{ws}/{nm}", proc) for (ws, nm), proc in items), @@ -348,6 +371,7 @@ async def shutdown(self) -> None: for path in self._temp_files.values(): path.unlink(missing_ok=True) self._temp_files.clear() + await asyncio.gather(*(asyncio.to_thread(shutil.rmtree, path, ignore_errors=True) for path in fabric_base_dirs)) if self._http_client is not None and not self._http_client.is_closed: try: await self._http_client.aclose() @@ -357,37 +381,15 @@ async def shutdown(self) -> None: logger.info("InMemoryRunnerBackend shut down — all processes terminated.") def _write_config(self, workspace: str, name: str, config: dict[str, Any]) -> Path: - config_path = self.config_path_for(workspace, name) - config_path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = config_path.with_suffix(config_path.suffix + ".tmp") - with tmp_path.open("w", encoding="utf-8") as fh: - yaml.safe_dump(config, fh) - tmp_path.replace(config_path) - return config_path + return _write_yaml_config(self.config_path_for(workspace, name), config) def _fabric_base_dir_for(self, workspace: str, name: str) -> Path: - """Return the local base directory used for Fabric validation/preparation.""" + """Return the local base directory used by a Fabric-backed deployment.""" return self.system_dir / _sanitize_filename(workspace) / f"{_sanitize_filename(name)}-fabric" - def _write_fabric_validation_log( - self, - workspace: str, - name: str, - log_path: Path, - validation_result: Any, - ) -> None: - log_path.parent.mkdir(parents=True, exist_ok=True) - log_path.write_text( - "\n".join( - [ - f"Validated Fabric-backed deployment for {workspace}/{name}.", - f"agent={validation_result.agent_config.name}", - f"base_dir={self._fabric_base_dir_for(workspace, name)}", - "", - ] - ), - encoding="utf-8", - ) + def _write_fabric_config(self, base_dir: Path, config: dict[str, Any]) -> Path: + """Write a Platform-owned agent config into its deployment directory.""" + return _write_yaml_config(base_dir / AGENT_CONFIG_FILENAME, config) def _spawn( self, @@ -426,6 +428,33 @@ def _spawn( # copy for the lifetime of the subprocess. log_file.close() + def _spawn_fabric( + self, + name: str, + config_path: Path, + log_path: Path, + port: int, + ) -> subprocess.Popen[bytes]: + """Spawn the Platform-owned Fabric server on a loopback port.""" + cmd = [ + sys.executable, + "-m", + "nemo_agents_plugin.fabric.server", + "--agent-config", + str(config_path), + "--host", + "127.0.0.1", + "--port", + str(port), + ] + log_path.parent.mkdir(parents=True, exist_ok=True) + logger.info("Spawning: %s (log: %s)", " ".join(cmd), log_path) + log_file = log_path.open("a") + try: + return subprocess.Popen(cmd, stdout=log_file, stderr=subprocess.STDOUT) + finally: + log_file.close() + def _terminate(self, name: str, proc: subprocess.Popen[bytes]) -> None: if proc.poll() is None: proc.terminate() diff --git a/plugins/nemo-agents/tests/unit/test_controller.py b/plugins/nemo-agents/tests/unit/test_controller.py index 0564ea1ed6..75345fd9d3 100644 --- a/plugins/nemo-agents/tests/unit/test_controller.py +++ b/plugins/nemo-agents/tests/unit/test_controller.py @@ -130,7 +130,6 @@ async def test_start_accepts_backend_running_status(self) -> None: ctrl.backend.create_deployment.return_value = DeploymentInfo( name="test-dep", status="running", - extra={"runtime": "fabric", "prepared": True}, ) await ctrl._start_deployment(dep) diff --git a/plugins/nemo-agents/tests/unit/test_runner_in_memory.py b/plugins/nemo-agents/tests/unit/test_runner_in_memory.py index ad162b2ffa..7a29ab785c 100644 --- a/plugins/nemo-agents/tests/unit/test_runner_in_memory.py +++ b/plugins/nemo-agents/tests/unit/test_runner_in_memory.py @@ -25,11 +25,12 @@ from pathlib import Path from types import SimpleNamespace from typing import Any -from unittest.mock import patch +from unittest.mock import ANY, patch import pytest import yaml from nemo_agents_plugin.config import AgentsConfig, ControllerConfig +from nemo_agents_plugin.runner.backend import DeploymentInfo from nemo_agents_plugin.runner.in_memory import InMemoryRunnerBackend, _resolve_nat_bin from nemo_platform_plugin.config import Configuration, nmp_user_data_dir @@ -205,21 +206,33 @@ async def _validate_platform_agent_config(config_: dict[str, Any], *, base_dir: validation_calls.append({"config": config_, "base_dir": base_dir}) return SimpleNamespace(agent_config=SimpleNamespace(name="fabric-agent")) - with patch("nemo_agents_plugin.runner.in_memory.validate_platform_agent_config", _validate_platform_agent_config): - info = await backend.create_deployment("ws", "fabric-dep", config, port=0) + fake_process = SimpleNamespace(pid=4242, returncode=None, poll=lambda: None) + + def _spawn_fabric(self_, name, config_path, log_path, port): # noqa: ANN001 + del self_, name, config_path, port + log_path.parent.mkdir(parents=True, exist_ok=True) + log_path.write_text("") + return fake_process + + with ( + patch("nemo_agents_plugin.runner.in_memory.validate_platform_agent_config", _validate_platform_agent_config), + patch.object(InMemoryRunnerBackend, "_spawn_fabric", _spawn_fabric), + ): + info = await backend.create_deployment("ws", "fabric-dep", config, port=49210) - assert info.status == "running" - assert info.extra["runtime"] == "fabric" - assert info.extra["prepared"] is True + assert info.status == "starting" + assert info.endpoint == "http://127.0.0.1:49210" + assert info.port == 49210 + assert info.pid == 4242 assert Path(info.log_path).exists() assert validation_calls == [{"config": config, "base_dir": tmp_path / "system" / "ws" / "fabric-dep-fabric"}] - assert "Validated Fabric-backed deployment" in Path(info.log_path).read_text() + assert yaml.safe_load((Path(info.extra["base_dir"]) / "agent.yaml").read_text()) == config status = await backend.get_deployment_status("ws", "fabric-dep") assert status is info @pytest.mark.asyncio -async def test_delete_deployment_removes_prepared_fabric_deployment(tmp_path: Path) -> None: +async def test_delete_deployment_removes_fabric_deployment(tmp_path: Path) -> None: backend = _backend(tmp_path) config = { "config_format": "nemo-agents-spec-v1", @@ -233,13 +246,31 @@ async def _validate_platform_agent_config(config_: dict[str, Any], *, base_dir: del config_, base_dir return SimpleNamespace(agent_config=SimpleNamespace(name="fabric-agent")) - with patch("nemo_agents_plugin.runner.in_memory.validate_platform_agent_config", _validate_platform_agent_config): - info = await backend.create_deployment("ws", "fabric-dep", config, port=0) + fake_process = SimpleNamespace(pid=4242, returncode=None, poll=lambda: None) + terminate_calls: list[tuple[str, Any]] = [] + + def _spawn_fabric(self_, name, config_path, log_path, port): # noqa: ANN001 + del self_, name, config_path, port + log_path.parent.mkdir(parents=True, exist_ok=True) + log_path.write_text("") + return fake_process + + def _terminate(self_, name, proc): # noqa: ANN001 + del self_ + terminate_calls.append((name, proc)) + + with ( + patch("nemo_agents_plugin.runner.in_memory.validate_platform_agent_config", _validate_platform_agent_config), + patch.object(InMemoryRunnerBackend, "_spawn_fabric", _spawn_fabric), + patch.object(InMemoryRunnerBackend, "_terminate", _terminate), + ): + info = await backend.create_deployment("ws", "fabric-dep", config, port=49211) base_dir = Path(info.extra["base_dir"]) assert base_dir.exists() cleaned = await backend.delete_deployment("ws", "fabric-dep") assert cleaned is True + assert terminate_calls == [("fabric-dep", fake_process)] assert not base_dir.exists() assert await backend.get_deployment_status("ws", "fabric-dep") is None @@ -258,7 +289,7 @@ async def test_create_deployment_cleans_fabric_base_dir_on_validation_failure(tm async def _validate_platform_agent_config(config_: dict[str, Any], *, base_dir: Path) -> Any: del config_ - (base_dir / "prepared.txt").write_text("created during validation") + (base_dir / "validation.txt").write_text("created during validation") raise ValueError("bad fabric config") with patch("nemo_agents_plugin.runner.in_memory.validate_platform_agent_config", _validate_platform_agent_config): @@ -269,6 +300,49 @@ async def _validate_platform_agent_config(config_: dict[str, Any], *, base_dir: assert await backend.get_deployment_status("ws", "fabric-dep") is None +def test_spawn_fabric_uses_current_python_and_platform_server(tmp_path: Path) -> None: + backend = _backend(tmp_path) + config_path = tmp_path / "agent.yaml" + config_path.write_text("name: test-agent\n") + log_path = tmp_path / "agent.log" + process = SimpleNamespace() + + with patch("nemo_agents_plugin.runner.in_memory.subprocess.Popen", return_value=process) as popen: + spawned = backend._spawn_fabric("fabric-dep", config_path, log_path, 49212) + + assert spawned is process + popen.assert_called_once_with( + [ + sys.executable, + "-m", + "nemo_agents_plugin.fabric.server", + "--agent-config", + str(config_path), + "--host", + "127.0.0.1", + "--port", + "49212", + ], + stdout=ANY, + stderr=subprocess.STDOUT, + ) + + +@pytest.mark.asyncio +async def test_shutdown_removes_fabric_deployment_directory(tmp_path: Path) -> None: + backend = _backend(tmp_path) + base_dir = backend._fabric_base_dir_for("ws", "fabric-dep") + base_dir.mkdir(parents=True) + backend._deployments[("ws", "fabric-dep")] = DeploymentInfo( + name="fabric-dep", + extra={"base_dir": str(base_dir)}, + ) + + await backend.shutdown() + + assert not base_dir.exists() + + # --------------------------------------------------------------------------- # get_deployment_status surfaces subprocess exit code # --------------------------------------------------------------------------- From 0098e9afa99e04aaf716122e3d18f632e21a4749 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 11:49:15 -0500 Subject: [PATCH 09/16] add common env setup helper Signed-off-by: Manjesh Mogallapalli --- .../nemo_agents_plugin/fabric/environment.py | 19 +++++++++++++++++++ .../nemo_agents_plugin/fabric/invocation.py | 13 ++----------- .../fabric/session_manager.py | 2 ++ .../tests/unit/test_fabric_session_manager.py | 1 + 4 files changed, 24 insertions(+), 11 deletions(-) create mode 100644 plugins/nemo-agents/src/nemo_agents_plugin/fabric/environment.py diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/environment.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/environment.py new file mode 100644 index 0000000000..d84b507441 --- /dev/null +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/environment.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prepare Platform-owned environment paths before Fabric runtime startup.""" + +from pathlib import Path + +from nemo_agents_plugin.agent_config import AgentConfig + + +def ensure_local_workspace_dir(agent_config: AgentConfig, base_dir: Path) -> None: + """Create the configured local workspace relative to the agent base directory.""" + if agent_config.environment.provider != "local": + return + + workspace = Path(agent_config.environment.workspace) + if not workspace.is_absolute(): + workspace = base_dir / workspace + workspace.mkdir(parents=True, exist_ok=True) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/invocation.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/invocation.py index 0b0a916978..040b913cf6 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/invocation.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/invocation.py @@ -11,6 +11,7 @@ from typing import Any from nemo_agents_plugin.agent_config import AgentConfig +from nemo_agents_plugin.fabric.environment import ensure_local_workspace_dir from nemo_agents_plugin.fabric.runtime import FabricOneShotRequest, FabricRuntimeResult, run_fabric_agent_once from nemo_agents_plugin.fabric.translator import translate_agent_config @@ -23,7 +24,7 @@ async def invoke_agent_config_once( ) -> list[FabricRuntimeResult]: """Translate a Platform agent config and run each input through Fabric once.""" fabric_config = translate_agent_config(agent_config) - await asyncio.to_thread(_ensure_local_workspace_dir, agent_config, base_dir) + await asyncio.to_thread(ensure_local_workspace_dir, agent_config, base_dir) results: list[FabricRuntimeResult] = [] for item in inputs: @@ -37,13 +38,3 @@ async def invoke_agent_config_once( ) ) return results - - -def _ensure_local_workspace_dir(agent_config: AgentConfig, base_dir: Path) -> None: - if agent_config.environment.provider != "local": - return - - workspace = Path(agent_config.environment.workspace) - if not workspace.is_absolute(): - workspace = base_dir / workspace - workspace.mkdir(parents=True, exist_ok=True) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py index 42524fd15a..1f13b142da 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -11,6 +11,7 @@ from typing import Any from nemo_agents_plugin.agent_config import AgentConfig +from nemo_agents_plugin.fabric.environment import ensure_local_workspace_dir from nemo_agents_plugin.fabric.runtime import FabricInvocationRequest, FabricRuntimeResult, invoke_fabric_runtime from nemo_agents_plugin.fabric.session_registry import ( FabricRuntimeSession, @@ -67,6 +68,7 @@ async def open_session(self) -> FabricRuntimeSession: except FabricTranslationError as error: raise FabricSessionStartError(f"Fabric config translation failed: {error}") from error + await asyncio.to_thread(ensure_local_workspace_dir, self._agent_config, self._base_dir) fabric = self._fabric or Fabric() try: runtime = await fabric.start_runtime(fabric_config, base_dir=self._base_dir) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py index 7ffc023f86..507ca63909 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py @@ -85,6 +85,7 @@ def translate(config: AgentConfig) -> Any: session = await manager.open_session() assert translation_calls == [agent_config] + assert (tmp_path / "workspace").is_dir() assert fabric.start_calls == [(fabric_config, tmp_path)] assert session.runtime is runtime assert await registry.get(session.session_id) is session From 7ba35ca721f62ae2cb8427a7538ce573186d0ab6 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 12:36:34 -0500 Subject: [PATCH 10/16] followups Signed-off-by: Manjesh Mogallapalli --- .../nemo_agents_plugin/fabric/environment.py | 12 +- .../src/nemo_agents_plugin/fabric/runtime.py | 12 +- .../src/nemo_agents_plugin/fabric/server.py | 6 +- .../tests/unit/test_fabric_invocation.py | 20 +++ .../tests/unit/test_fabric_runtime.py | 23 +++ .../tests/unit/test_fabric_server.py | 29 ++++ pr_description.md | 147 ++++++++++++++++++ 7 files changed, 241 insertions(+), 8 deletions(-) create mode 100644 pr_description.md diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/environment.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/environment.py index d84b507441..7a2aa3483f 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/environment.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/environment.py @@ -13,7 +13,13 @@ def ensure_local_workspace_dir(agent_config: AgentConfig, base_dir: Path) -> Non if agent_config.environment.provider != "local": return - workspace = Path(agent_config.environment.workspace) - if not workspace.is_absolute(): - workspace = base_dir / workspace + configured_workspace = Path(agent_config.environment.workspace) + if configured_workspace.is_absolute(): + raise ValueError("Local workspace path must be relative to the agent base directory.") + + resolved_base_dir = base_dir.resolve() + workspace = (resolved_base_dir / configured_workspace).resolve() + if not workspace.is_relative_to(resolved_base_dir): + raise ValueError("Local workspace path must remain within the agent base directory.") + workspace.mkdir(parents=True, exist_ok=True) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py index 188aa8fa7a..ff5815f1e4 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py @@ -84,7 +84,13 @@ class FabricRuntimeExecutionError(RuntimeError): class FabricRuntimeTimeoutError(FabricRuntimeExecutionError): - """Raised when a Fabric runtime invocation exceeds the Platform timeout.""" + """Raised when a Fabric runtime invocation times out.""" + + +def _timeout_error_message(timeout_seconds: float | None) -> str: + if timeout_seconds is None: + return "Fabric runtime invocation timed out." + return f"Fabric runtime invocation timed out after {timeout_seconds:g}s." async def invoke_fabric_runtime( @@ -99,7 +105,7 @@ async def invoke_fabric_runtime( ) except TimeoutError as error: raise FabricRuntimeTimeoutError( - f"Fabric runtime invocation timed out after {request.timeout_seconds:g}s.", + _timeout_error_message(request.timeout_seconds), ) from error except FabricError as error: raise FabricRuntimeExecutionError( @@ -124,7 +130,7 @@ async def run_fabric_agent_once( ) except TimeoutError as error: raise FabricRuntimeTimeoutError( - f"Fabric runtime invocation timed out after {request.timeout_seconds:g}s.", + _timeout_error_message(request.timeout_seconds), ) from error except FabricError as error: raise FabricRuntimeExecutionError( diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 16994d13ed..3bf292bfcb 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -157,8 +157,10 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: yield finally: cleanup_shutdown.set() - await cleanup_task - await session_manager.close_all_sessions() + try: + await cleanup_task + finally: + await session_manager.close_all_sessions() app = FastAPI(title="NeMo Agents Fabric Server", lifespan=lifespan) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_invocation.py b/plugins/nemo-agents/tests/unit/test_fabric_invocation.py index 521a34fff6..e72e4344fd 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_invocation.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_invocation.py @@ -64,3 +64,23 @@ async def _run_fabric_agent_once(request: Any) -> FabricRuntimeResult: agent_config = AgentConfig.model_validate(config) with patch("nemo_agents_plugin.fabric.invocation.run_fabric_agent_once", _run_fabric_agent_once): await invoke_agent_config_once(agent_config, ["one"], base_dir=tmp_path) + + +@pytest.mark.asyncio +async def test_invoke_agent_config_once_rejects_absolute_local_workspace(tmp_path: Path) -> None: + config = _agent_config() + config["environment"] = {"workspace": str(tmp_path.parent / "outside")} + agent_config = AgentConfig.model_validate(config) + + with pytest.raises(ValueError, match="Local workspace path must be relative"): + await invoke_agent_config_once(agent_config, ["one"], base_dir=tmp_path) + + +@pytest.mark.asyncio +async def test_invoke_agent_config_once_rejects_workspace_traversal(tmp_path: Path) -> None: + config = _agent_config() + config["environment"] = {"workspace": "../../outside"} + agent_config = AgentConfig.model_validate(config) + + with pytest.raises(ValueError, match="Local workspace path must remain within"): + await invoke_agent_config_once(agent_config, ["one"], base_dir=tmp_path) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py index 3b905be01e..37a8c53624 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py @@ -167,6 +167,20 @@ async def test_wraps_timeout(self) -> None: assert fake_runtime.exited is True + async def test_wraps_runtime_timeout_without_configured_deadline(self) -> None: + timeout_error = TimeoutError("adapter timed out") + fake_runtime = _FakeRuntime(invoke_error=timeout_error) + fake_fabric = _FakeFabric(runtime=fake_runtime) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + with pytest.raises(FabricRuntimeTimeoutError, match=r"Fabric runtime invocation timed out\.$") as exc_info: + await run_fabric_agent_once(request, fabric=fake_fabric) + + assert exc_info.value.__cause__ is timeout_error + async def test_wraps_fabric_lifecycle_errors(self) -> None: fake_fabric = _FakeFabric(start_error=fabric_runtime.FabricError("native unavailable")) request = FabricOneShotRequest( @@ -249,5 +263,14 @@ async def test_wraps_timeout_without_stopping_runtime(self) -> None: with pytest.raises(FabricRuntimeTimeoutError, match="timed out after 0.01s"): await invoke_fabric_runtime(cast(Any, fake_runtime), request) + async def test_wraps_active_runtime_timeout_without_configured_deadline(self) -> None: + timeout_error = TimeoutError("adapter timed out") + fake_runtime = _FakeRuntime(invoke_error=timeout_error) + + with pytest.raises(FabricRuntimeTimeoutError, match=r"Fabric runtime invocation timed out\.$") as exc_info: + await invoke_fabric_runtime(cast(Any, fake_runtime), FabricInvocationRequest()) + + assert exc_info.value.__cause__ is timeout_error + assert fake_runtime.entered is False assert fake_runtime.exited is False diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index 600bfa38c8..eab3c2c7e1 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -109,6 +109,35 @@ async def register_runtime() -> None: assert runtime.stop_calls == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize("cleanup_error", [RuntimeError("cleanup failed"), asyncio.CancelledError()]) +async def test_shutdown_stops_sessions_when_cleanup_task_fails( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, + cleanup_error: BaseException, +) -> None: + close_calls = 0 + + async def fail_cleanup(*args: Any, **kwargs: Any) -> None: + raise cleanup_error + + async def close_all_sessions(self: FabricSessionManager) -> int: + nonlocal close_calls + close_calls += 1 + return 0 + + monkeypatch.setattr(server, "_run_idle_session_cleanup", fail_cleanup) + monkeypatch.setattr(FabricSessionManager, "close_all_sessions", close_all_sessions) + app = create_fabric_serving_app(_write_agent_config(tmp_path)) + + with pytest.raises(type(cleanup_error)): + async with app.router.lifespan_context(app): + pass + + assert close_calls == 1 + + def test_startup_fails_for_invalid_agent_config( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], diff --git a/pr_description.md b/pr_description.md new file mode 100644 index 0000000000..6306ef35a6 --- /dev/null +++ b/pr_description.md @@ -0,0 +1,147 @@ +## Summary + +This PR implements the first Platform-managed serving lifecycle for Fabric-backed NeMo Agents as part of AIRCORE-932. + +It adds a local FastAPI serving process that creates one Fabric runtime per logical user session and exposes it through the existing agent deployment and gateway flow: + +```text +Platform Agent entity + -> persisted agent.yaml + -> local Fabric serving process + -> logical session + -> translated FabricConfig + -> FabricRuntime + -> ordered invoke calls + -> runtime stop +``` + +Fabric continues to own harness execution and the runtime `start` / `invoke` / `stop` lifecycle. NeMo Platform owns the multi-user server, logical session identity, runtime registry, request routing, concurrency policy, expiration, and cleanup. + +## Changes + +- Added an OpenAI-compatible Fabric serving application with: + - `GET /health` + - `POST /v1/chat/completions` + - `DELETE /v1/sessions/{session_id}` +- Added typed request and response models for the chat-completions boundary. +- Added a runtime session registry that maps opaque Platform session IDs to active Fabric runtimes. +- Added a session manager responsible for: + - Lazy `FabricConfig` translation and runtime startup. + - Reusing the same runtime for later turns in a logical session. + - Serializing invocations within one session. + - Limiting concurrent invocations across independent sessions. + - Explicit session closure. + - Idle-session expiration. + - Draining and stopping all runtimes during server shutdown. +- Added invocation support for an already-active Fabric runtime while retaining the existing one-shot invocation path. +- Added shared local-environment preparation so configured workspaces exist before either one-shot or managed runtime startup. +- Added a shared agent-config format registry/protocol used by agent creation and deployment config resolution. +- Updated the in-memory runner to deploy `nemo-agents-spec-v1` agents by: + - Persisting a canonical `agent.yaml` under the Platform agent workspace. + - Running Fabric plan/doctor validation before spawning the server. + - Starting the Fabric server as a managed local subprocess. + - Reusing existing port allocation, readiness polling, log handling, process termination, and deployment cleanup behavior. +- Preserved the existing NAT deployment path for `nat-workflow-v1`. +- Added focused tests for format handling, HTTP routing, session lifecycle, concurrency, expiration, runtime cleanup, local deployment, and failure paths. + +## Design Choices + +### Platform owns sessions; Fabric owns runtimes + +A Platform session is a logical conversation rather than an HTTP connection. Platform maintains the `session_id -> FabricRuntime` mapping, while Fabric remains unaware of users, HTTP routing, and other runtimes. + +This keeps the integration aligned with Fabric's public lifecycle contract instead of adding lifecycle behavior inside adapters. + +### Runtimes are created lazily + +Server startup loads and validates the reusable Platform agent definition but does not create a Fabric runtime. A complete `FabricConfig` is translated when the first request opens a logical session, and that config is bound to the resulting runtime. + +This avoids allocating harness resources for sessions that never invoke the agent and leaves room for future per-session policy, environment, and profile resolution. + +### Session identity uses a response header + +The first request may omit `X-Nemo-Session-Id`. Platform generates an opaque session ID and returns it in that response header. Later requests provide the same header to reuse the runtime. + +Using a header keeps the request body compatible with the OpenAI chat-completions shape. Supplying an unknown or closed session ID returns `404`; it does not silently create a replacement runtime. + +### One runtime processes one turn at a time + +Invocations for the same session are serialized with a per-session lock because ordered turns share harness state. Different sessions may run concurrently, subject to a server-wide semaphore. The initial default permits eight concurrent invocations. + +### The runtime owns conversation state + +Each HTTP request passes the current user message to the existing runtime. Prior turns are not replayed from the HTTP payload because the Fabric runtime and selected harness adapter own the session's conversation state. + +### Cleanup is explicit and bounded + +Clients can close sessions explicitly. The server also expires idle sessions after 30 minutes, checks every five minutes, and drains all remaining runtimes during shutdown. Runtime registration failures also stop any runtime that was already started. + +### Local deployment builds on the existing runner + +The first implementation uses the existing in-memory subprocess backend rather than introducing a second process-management system. Fabric and NAT deployments therefore share port allocation, health polling, logs, status transitions, termination, and filesystem cleanup. + +Docker/Kubernetes runtime placement and distributed session ownership remain separate follow-up work. + +### Agent formats share a small internal protocol + +Agent creation and deployment now resolve behavior through config-format handlers. This keeps `nat-workflow-v1` as the default and adds `nemo-agents-spec-v1` without spreading format-specific branches through the API layer. + +The registry is intentionally narrow and internal while the RFC 122 entity shapes are still being finalized. + +## Error Behavior + +- Unknown or closed session: `404` +- Fabric runtime startup failure: `503` +- Fabric invocation failure or failed result: `502` +- Fabric invocation timeout: `504` +- Runtime shutdown failure: `502` + +Errors for an existing session preserve the session ID header where appropriate. + +## Out of Scope + +- Durable session recovery after a Platform/server restart. +- Distributed session registries or routing across replicas. +- Docker/Kubernetes Fabric server deployment and remote runtime placement. +- Authentication and authorization inside the Fabric serving process; those remain Platform gateway concerns. +- Streaming chat completions. +- User-facing cancellation APIs. +- Per-user concurrency quotas or configurable limits through public API fields. +- Final RFC 122 `AgentRun`, input, output, environment, sandbox, and harness entity shapes. + +## Validation + +Focused branch coverage: + +```text +222 passed +``` + +This includes agent config handling, format dispatch, Fabric translation/validation, one-shot and active-runtime invocation, serving routes, session registry/manager behavior, deployment APIs, controller behavior, and the in-memory runner. + +Additional nested NeMo Agents suites: + +```text +104 passed +``` + +Repository Python style and formatting: + +```text +All checks passed +2793 files already formatted +``` + +Manual end-to-end validation: + +1. Registered a `nemo-agents-spec-v1` Agent from `agent.yaml`. +2. Deployed it through the Platform API/CLI. +3. Observed `pending -> starting -> running`. +4. Invoked it through the Platform agent gateway. +5. Opened a logical session and received `X-Nemo-Session-Id`. +6. Reused that session for a second turn and confirmed conversation state was preserved. +7. Closed the session and received `204`. +8. Confirmed reuse of the closed session returned `404`. +9. Deleted the deployment and verified its process and prepared base directory were removed. + +A broader top-level NeMo Agents unit sweep produced `710 passed` and 16 failures in existing CLI delete/list/optimize tests outside the files changed by this branch. From 0f588be8e414df0061ba2a591f677a8a3f2a1dc4 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 12:38:51 -0500 Subject: [PATCH 11/16] lint Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/session_registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py index 7287d99be9..1c7ed20230 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from nemo_fabric import Runtime + from nemo_fabric import Runtime # ty: ignore[unresolved-import] @dataclass(slots=True) From 4f66513ad9ad0eeaa26a120b0794e8efb6af803c Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 12:56:58 -0500 Subject: [PATCH 12/16] remove file Signed-off-by: Manjesh Mogallapalli --- pr_description.md | 147 ---------------------------------------------- 1 file changed, 147 deletions(-) delete mode 100644 pr_description.md diff --git a/pr_description.md b/pr_description.md deleted file mode 100644 index 6306ef35a6..0000000000 --- a/pr_description.md +++ /dev/null @@ -1,147 +0,0 @@ -## Summary - -This PR implements the first Platform-managed serving lifecycle for Fabric-backed NeMo Agents as part of AIRCORE-932. - -It adds a local FastAPI serving process that creates one Fabric runtime per logical user session and exposes it through the existing agent deployment and gateway flow: - -```text -Platform Agent entity - -> persisted agent.yaml - -> local Fabric serving process - -> logical session - -> translated FabricConfig - -> FabricRuntime - -> ordered invoke calls - -> runtime stop -``` - -Fabric continues to own harness execution and the runtime `start` / `invoke` / `stop` lifecycle. NeMo Platform owns the multi-user server, logical session identity, runtime registry, request routing, concurrency policy, expiration, and cleanup. - -## Changes - -- Added an OpenAI-compatible Fabric serving application with: - - `GET /health` - - `POST /v1/chat/completions` - - `DELETE /v1/sessions/{session_id}` -- Added typed request and response models for the chat-completions boundary. -- Added a runtime session registry that maps opaque Platform session IDs to active Fabric runtimes. -- Added a session manager responsible for: - - Lazy `FabricConfig` translation and runtime startup. - - Reusing the same runtime for later turns in a logical session. - - Serializing invocations within one session. - - Limiting concurrent invocations across independent sessions. - - Explicit session closure. - - Idle-session expiration. - - Draining and stopping all runtimes during server shutdown. -- Added invocation support for an already-active Fabric runtime while retaining the existing one-shot invocation path. -- Added shared local-environment preparation so configured workspaces exist before either one-shot or managed runtime startup. -- Added a shared agent-config format registry/protocol used by agent creation and deployment config resolution. -- Updated the in-memory runner to deploy `nemo-agents-spec-v1` agents by: - - Persisting a canonical `agent.yaml` under the Platform agent workspace. - - Running Fabric plan/doctor validation before spawning the server. - - Starting the Fabric server as a managed local subprocess. - - Reusing existing port allocation, readiness polling, log handling, process termination, and deployment cleanup behavior. -- Preserved the existing NAT deployment path for `nat-workflow-v1`. -- Added focused tests for format handling, HTTP routing, session lifecycle, concurrency, expiration, runtime cleanup, local deployment, and failure paths. - -## Design Choices - -### Platform owns sessions; Fabric owns runtimes - -A Platform session is a logical conversation rather than an HTTP connection. Platform maintains the `session_id -> FabricRuntime` mapping, while Fabric remains unaware of users, HTTP routing, and other runtimes. - -This keeps the integration aligned with Fabric's public lifecycle contract instead of adding lifecycle behavior inside adapters. - -### Runtimes are created lazily - -Server startup loads and validates the reusable Platform agent definition but does not create a Fabric runtime. A complete `FabricConfig` is translated when the first request opens a logical session, and that config is bound to the resulting runtime. - -This avoids allocating harness resources for sessions that never invoke the agent and leaves room for future per-session policy, environment, and profile resolution. - -### Session identity uses a response header - -The first request may omit `X-Nemo-Session-Id`. Platform generates an opaque session ID and returns it in that response header. Later requests provide the same header to reuse the runtime. - -Using a header keeps the request body compatible with the OpenAI chat-completions shape. Supplying an unknown or closed session ID returns `404`; it does not silently create a replacement runtime. - -### One runtime processes one turn at a time - -Invocations for the same session are serialized with a per-session lock because ordered turns share harness state. Different sessions may run concurrently, subject to a server-wide semaphore. The initial default permits eight concurrent invocations. - -### The runtime owns conversation state - -Each HTTP request passes the current user message to the existing runtime. Prior turns are not replayed from the HTTP payload because the Fabric runtime and selected harness adapter own the session's conversation state. - -### Cleanup is explicit and bounded - -Clients can close sessions explicitly. The server also expires idle sessions after 30 minutes, checks every five minutes, and drains all remaining runtimes during shutdown. Runtime registration failures also stop any runtime that was already started. - -### Local deployment builds on the existing runner - -The first implementation uses the existing in-memory subprocess backend rather than introducing a second process-management system. Fabric and NAT deployments therefore share port allocation, health polling, logs, status transitions, termination, and filesystem cleanup. - -Docker/Kubernetes runtime placement and distributed session ownership remain separate follow-up work. - -### Agent formats share a small internal protocol - -Agent creation and deployment now resolve behavior through config-format handlers. This keeps `nat-workflow-v1` as the default and adds `nemo-agents-spec-v1` without spreading format-specific branches through the API layer. - -The registry is intentionally narrow and internal while the RFC 122 entity shapes are still being finalized. - -## Error Behavior - -- Unknown or closed session: `404` -- Fabric runtime startup failure: `503` -- Fabric invocation failure or failed result: `502` -- Fabric invocation timeout: `504` -- Runtime shutdown failure: `502` - -Errors for an existing session preserve the session ID header where appropriate. - -## Out of Scope - -- Durable session recovery after a Platform/server restart. -- Distributed session registries or routing across replicas. -- Docker/Kubernetes Fabric server deployment and remote runtime placement. -- Authentication and authorization inside the Fabric serving process; those remain Platform gateway concerns. -- Streaming chat completions. -- User-facing cancellation APIs. -- Per-user concurrency quotas or configurable limits through public API fields. -- Final RFC 122 `AgentRun`, input, output, environment, sandbox, and harness entity shapes. - -## Validation - -Focused branch coverage: - -```text -222 passed -``` - -This includes agent config handling, format dispatch, Fabric translation/validation, one-shot and active-runtime invocation, serving routes, session registry/manager behavior, deployment APIs, controller behavior, and the in-memory runner. - -Additional nested NeMo Agents suites: - -```text -104 passed -``` - -Repository Python style and formatting: - -```text -All checks passed -2793 files already formatted -``` - -Manual end-to-end validation: - -1. Registered a `nemo-agents-spec-v1` Agent from `agent.yaml`. -2. Deployed it through the Platform API/CLI. -3. Observed `pending -> starting -> running`. -4. Invoked it through the Platform agent gateway. -5. Opened a logical session and received `X-Nemo-Session-Id`. -6. Reused that session for a second turn and confirmed conversation state was preserved. -7. Closed the session and received `204`. -8. Confirmed reuse of the closed session returned `404`. -9. Deleted the deployment and verified its process and prepared base directory were removed. - -A broader top-level NeMo Agents unit sweep produced `710 passed` and 16 failures in existing CLI delete/list/optimize tests outside the files changed by this branch. From e76e524f113700e171ab40e2022055bae046224d Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 13:31:48 -0500 Subject: [PATCH 13/16] nemo agents run Signed-off-by: Manjesh Mogallapalli --- .../nemo-agents/src/nemo_agents_plugin/cli.py | 32 ++++++++-- plugins/nemo-agents/tests/unit/test_cli.py | 59 +++++++++++++++++++ 2 files changed, 86 insertions(+), 5 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/cli.py b/plugins/nemo-agents/src/nemo_agents_plugin/cli.py index d60f9ce15a..1e429ba716 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/cli.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/cli.py @@ -12,7 +12,7 @@ instance. - ``invoke`` — single invocation -- ``run`` — start a persistent local FastAPI server for NAT configs +- ``run`` — start a persistent local FastAPI server The ``evaluate`` and ``optimize`` commands are auto-generated from the ``EvaluateAgentJob`` and ``OptimizeAgentJob`` registered under the @@ -38,6 +38,7 @@ import logging import os import re +import sys import time from dataclasses import asdict from datetime import datetime @@ -205,7 +206,7 @@ def run( ..., "--agent-config", "-c", - help="Path to a NAT workflow YAML config file.", + help="Path to an agent YAML config file.", exists=True, file_okay=True, dir_okay=False, @@ -213,10 +214,28 @@ def run( host: str = typer.Option("0.0.0.0", "--host"), port: int = typer.Option(8080, "--port", "-p"), ) -> None: - """Run an agent locally as a persistent FastAPI server (wraps ``nat start fastapi``).""" + """Run an agent locally as a persistent FastAPI server.""" import subprocess - cmd = ["nat", "start", "fastapi", "--config_file", agent_config.name, "--host", host, "--port", str(port)] + config_format = _load_yaml(agent_config).get("config_format", NAT_WORKFLOW_CONFIG_FORMAT) + if config_format == NEMO_AGENTS_SPEC_CONFIG_FORMAT: + cmd = [ + sys.executable, + "-m", + "nemo_agents_plugin.fabric.server", + "--agent-config", + agent_config.name, + "--host", + host, + "--port", + str(port), + ] + elif config_format == NAT_WORKFLOW_CONFIG_FORMAT: + cmd = ["nat", "start", "fastapi", "--config_file", agent_config.name, "--host", host, "--port", str(port)] + else: + typer.echo(f"Error: unsupported config_format {config_format!r}", err=True) + raise typer.Exit(code=1) + typer.echo(f"Starting agent server: {' '.join(cmd)}") try: subprocess.run(cmd, check=True, cwd=agent_config.parent) @@ -224,7 +243,10 @@ def run( typer.echo(f"Agent server exited with code {exc.returncode}.", err=True) raise typer.Exit(code=exc.returncode) except FileNotFoundError: - typer.echo("Error: 'nat' command not found. Install nvidia-nat-core.", err=True) + if config_format == NAT_WORKFLOW_CONFIG_FORMAT: + typer.echo("Error: 'nat' command not found. Install nvidia-nat-core.", err=True) + else: + typer.echo(f"Error: server command {cmd[0]!r} was not found.", err=True) raise typer.Exit(code=1) diff --git a/plugins/nemo-agents/tests/unit/test_cli.py b/plugins/nemo-agents/tests/unit/test_cli.py index 4445e74d1d..42c9920f51 100644 --- a/plugins/nemo-agents/tests/unit/test_cli.py +++ b/plugins/nemo-agents/tests/unit/test_cli.py @@ -3,6 +3,7 @@ from __future__ import annotations +import sys from collections.abc import Callable from contextlib import AbstractContextManager from pathlib import Path @@ -47,6 +48,64 @@ def test_no_args_prints_help_successfully() -> None: assert "Agent lifecycle management" in result.stdout +def test_run_starts_nat_server_for_nat_config(tmp_path: Path) -> None: + config = tmp_path / "workflow.yaml" + config.write_text("workflow:\n _type: chat_completion\n") + + app = AgentsCLI().get_cli() + with patch("subprocess.run") as run: + result = CliRunner().invoke( + app, + ["run", "--agent-config", str(config), "--host", "127.0.0.1", "--port", "8081"], + ) + + assert result.exit_code == 0, result.stderr + run.assert_called_once_with( + [ + "nat", + "start", + "fastapi", + "--config_file", + config.name, + "--host", + "127.0.0.1", + "--port", + "8081", + ], + check=True, + cwd=tmp_path, + ) + + +def test_run_starts_fabric_server_for_platform_config(tmp_path: Path) -> None: + config = tmp_path / "agent.yaml" + config.write_text("config_format: nemo-agents-spec-v1\nname: fabric-agent\n") + + app = AgentsCLI().get_cli() + with patch("subprocess.run") as run: + result = CliRunner().invoke( + app, + ["run", "--agent-config", str(config), "--host", "127.0.0.1", "--port", "8081"], + ) + + assert result.exit_code == 0, result.stderr + run.assert_called_once_with( + [ + sys.executable, + "-m", + "nemo_agents_plugin.fabric.server", + "--agent-config", + config.name, + "--host", + "127.0.0.1", + "--port", + "8081", + ], + check=True, + cwd=tmp_path, + ) + + def test_list_404_prints_request_context_and_hint() -> None: def handler(req: httpx.Request) -> httpx.Response: return httpx.Response(404, json={"detail": "Not Found"}) From 3c20a6b1c9683659fa6b7c4e3bdee317cb7aa23a Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Mon, 27 Jul 2026 13:45:03 -0500 Subject: [PATCH 14/16] cr followup Signed-off-by: Manjesh Mogallapalli --- plugins/nemo-agents/src/nemo_agents_plugin/cli.py | 7 ++++++- plugins/nemo-agents/tests/unit/test_cli.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/cli.py b/plugins/nemo-agents/src/nemo_agents_plugin/cli.py index 1e429ba716..8f5ffcc6c5 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/cli.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/cli.py @@ -217,7 +217,12 @@ def run( """Run an agent locally as a persistent FastAPI server.""" import subprocess - config_format = _load_yaml(agent_config).get("config_format", NAT_WORKFLOW_CONFIG_FORMAT) + config = _load_yaml(agent_config) + if not isinstance(config, dict): + typer.echo(f"Error: agent config {agent_config} root must be a YAML mapping.", err=True) + raise typer.Exit(code=1) + + config_format = config.get("config_format", NAT_WORKFLOW_CONFIG_FORMAT) if config_format == NEMO_AGENTS_SPEC_CONFIG_FORMAT: cmd = [ sys.executable, diff --git a/plugins/nemo-agents/tests/unit/test_cli.py b/plugins/nemo-agents/tests/unit/test_cli.py index 42c9920f51..76eb40c622 100644 --- a/plugins/nemo-agents/tests/unit/test_cli.py +++ b/plugins/nemo-agents/tests/unit/test_cli.py @@ -106,6 +106,19 @@ def test_run_starts_fabric_server_for_platform_config(tmp_path: Path) -> None: ) +def test_run_rejects_empty_yaml_config(tmp_path: Path) -> None: + config = tmp_path / "agent.yaml" + config.write_text("") + + app = AgentsCLI().get_cli() + with patch("subprocess.run") as run: + result = CliRunner().invoke(app, ["run", "--agent-config", str(config)]) + + assert result.exit_code == 1 + assert f"agent config {config} root must be a YAML mapping" in result.stderr + run.assert_not_called() + + def test_list_404_prints_request_context_and_hint() -> None: def handler(req: httpx.Request) -> httpx.Response: return httpx.Response(404, json={"detail": "Not Found"}) From f69cf41a401825ba4fdf38fd78bfbb5706aed40d Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Tue, 28 Jul 2026 14:13:21 -0500 Subject: [PATCH 15/16] reviewer feedback Signed-off-by: Manjesh Mogallapalli --- .../agent_config_formats.py | 11 +++-- .../fabric/session_manager.py | 15 +++++-- .../nemo_agents_plugin/runner/in_memory.py | 2 +- .../src/nemo_agents_plugin/utils.py | 33 +++++++++++++++ .../tests/unit/test_agent_config_formats.py | 13 +++++- .../tests/unit/test_fabric_session_manager.py | 39 ++++++++++++++---- plugins/nemo-agents/tests/unit/test_utils.py | 40 +++++++++++++++++++ 7 files changed, 138 insertions(+), 15 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/agent_config_formats.py b/plugins/nemo-agents/src/nemo_agents_plugin/agent_config_formats.py index f1be7353c8..d5620f2bf9 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/agent_config_formats.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/agent_config_formats.py @@ -9,7 +9,12 @@ from nemo_agents_plugin.agent_config import AgentConfig from nemo_agents_plugin.entities import NAT_WORKFLOW_CONFIG_FORMAT, NEMO_AGENTS_SPEC_CONFIG_FORMAT -from nemo_agents_plugin.utils import inject_default_model, inject_gateway_url, inject_nemo_trace_fields +from nemo_agents_plugin.utils import ( + inject_default_model, + inject_fabric_gateway_url, + inject_gateway_url, + inject_nemo_trace_fields, +) from pydantic import ValidationError @@ -67,8 +72,8 @@ def resolve_for_deployment( workspace: str, agent_name: str, ) -> dict[str, Any]: - del workspace, agent_name - return self._normalize(config) + del agent_name + return self._normalize(inject_fabric_gateway_url(config, workspace)) @staticmethod def _normalize(config: dict[str, Any]) -> dict[str, Any]: diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py index 1f13b142da..491a534397 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -8,7 +8,6 @@ import asyncio import logging from pathlib import Path -from typing import Any from nemo_agents_plugin.agent_config import AgentConfig from nemo_agents_plugin.fabric.environment import ensure_local_workspace_dir @@ -47,7 +46,7 @@ def __init__( *, base_dir: Path, session_registry: FabricSessionRegistry, - fabric: Any | None = None, + fabric: Fabric | None = None, max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS, ) -> None: if max_concurrent_invocations < 0: @@ -136,7 +135,17 @@ async def stop_session(session: FabricRuntimeSession) -> None: except FabricSessionStopError: logger.exception("Failed to stop Fabric session %s during shutdown.", session.session_id) - await asyncio.gather(*(stop_session(session) for session in sessions)) + results = await asyncio.gather( + *(stop_session(session) for session in sessions), + return_exceptions=True, + ) + for session, result in zip(sessions, results, strict=True): + if isinstance(result, BaseException): + logger.error( + "Unexpected error stopping Fabric session %s during shutdown.", + session.session_id, + exc_info=(type(result), result, result.__traceback__), + ) return len(sessions) async def _stop_session(self, session: FabricRuntimeSession) -> None: diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/runner/in_memory.py b/plugins/nemo-agents/src/nemo_agents_plugin/runner/in_memory.py index 3b1010e2ca..f8a0b54270 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/runner/in_memory.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/runner/in_memory.py @@ -449,7 +449,7 @@ def _spawn_fabric( ] log_path.parent.mkdir(parents=True, exist_ok=True) logger.info("Spawning: %s (log: %s)", " ".join(cmd), log_path) - log_file = log_path.open("a") + log_file = log_path.open("w") try: return subprocess.Popen(cmd, stdout=log_file, stderr=subprocess.STDOUT) finally: diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/utils.py b/plugins/nemo-agents/src/nemo_agents_plugin/utils.py index 7df1b7daf0..5d7ec4613f 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/utils.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/utils.py @@ -234,6 +234,39 @@ def inject_gateway_url( return config +def inject_fabric_gateway_url( + config: dict[str, Any], + workspace: str, + base_url: str | None = None, +) -> dict[str, Any]: + """Deep-copy a Platform agent config and bind its models to the Inference Gateway.""" + if base_url is None: + base_url = get_base_url() + gateway_url = f"{base_url.rstrip('/')}/apis/inference-gateway/v2/workspaces/{workspace}/openai/-/v1" + + resolved = copy.deepcopy(config) + model_configs: list[Any] = [] + + models = resolved.get("models") + if isinstance(models, dict): + model_configs.extend(models.values()) + + harnesses = resolved.get("harnesses") + if isinstance(harnesses, dict): + for harness in harnesses.values(): + if isinstance(harness, dict) and "model" in harness: + model_configs.append(harness["model"]) + + for model_config in model_configs: + if not isinstance(model_config, dict): + continue + settings = model_config.setdefault("settings", {}) + if isinstance(settings, dict): + settings.setdefault("base_url", gateway_url) + + return resolved + + def inject_nemo_trace_fields( config: dict[str, Any], workspace: str, diff --git a/plugins/nemo-agents/tests/unit/test_agent_config_formats.py b/plugins/nemo-agents/tests/unit/test_agent_config_formats.py index e5e25d2816..16b0c94daa 100644 --- a/plugins/nemo-agents/tests/unit/test_agent_config_formats.py +++ b/plugins/nemo-agents/tests/unit/test_agent_config_formats.py @@ -87,7 +87,17 @@ def inject_trace(config: dict[str, Any], *, workspace: str, agent_name: str) -> ] -def test_nemo_agents_deployment_resolution_only_normalizes_payload() -> None: +def test_nemo_agents_deployment_resolution_injects_gateway_and_normalizes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + def inject_gateway(config: dict[str, Any], workspace: str) -> dict[str, Any]: + calls.append(workspace) + return config + + monkeypatch.setattr(agent_config_formats, "inject_fabric_gateway_url", inject_gateway) + resolved = resolve_agent_config_for_deployment( NEMO_AGENTS_SPEC_CONFIG_FORMAT, _nemo_agents_config(), @@ -98,3 +108,4 @@ def test_nemo_agents_deployment_resolution_only_normalizes_payload() -> None: assert resolved["config_format"] == NEMO_AGENTS_SPEC_CONFIG_FORMAT assert resolved["environment"]["provider"] == "local" assert "workflow" not in resolved + assert calls == ["test-workspace"] diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py index 507ca63909..8ff6ed8d92 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py @@ -76,7 +76,7 @@ def translate(config: AgentConfig) -> Any: agent_config, base_dir=tmp_path, session_registry=registry, - fabric=fabric, + fabric=cast(Any, fabric), ) assert translation_calls == [] @@ -109,7 +109,7 @@ async def fail_registration(runtime: Any) -> None: _agent_config(), base_dir=tmp_path, session_registry=registry, - fabric=fabric, + fabric=cast(Any, fabric), ) with pytest.raises(RuntimeError, match="registration failed"): @@ -132,7 +132,7 @@ async def test_resolve_session_opens_session_when_id_is_absent( _agent_config(), base_dir=tmp_path, session_registry=registry, - fabric=fabric, + fabric=cast(Any, fabric), ) session = await manager.resolve_session(None) @@ -153,7 +153,7 @@ async def test_resolve_session_reuses_registered_runtime( _agent_config(), base_dir=tmp_path, session_registry=registry, - fabric=fabric, + fabric=cast(Any, fabric), ) session = await manager.resolve_session("session-1") @@ -307,6 +307,31 @@ async def stop(self) -> None: assert healthy_runtime.stop_calls == 1 +@pytest.mark.asyncio +async def test_close_all_sessions_continues_after_unexpected_stop_failure(tmp_path: Path) -> None: + class _FailingRuntime(_FakeRuntime): + async def stop(self) -> None: + self.stop_calls += 1 + raise RuntimeError("unexpected stop failure") + + failing_runtime = _FailingRuntime() + healthy_runtime = _FakeRuntime() + registry = FabricSessionRegistry() + await registry.register(cast(Any, failing_runtime), session_id="session-1") + await registry.register(cast(Any, healthy_runtime), session_id="session-2") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + ) + + closed_count = await manager.close_all_sessions() + + assert closed_count == 2 + assert failing_runtime.stop_calls == 1 + assert healthy_runtime.stop_calls == 1 + + @pytest.mark.asyncio async def test_invoke_session_refreshes_activity( tmp_path: Path, @@ -346,7 +371,7 @@ async def test_invoke_session_serializes_turns_for_same_runtime( _agent_config(), base_dir=tmp_path, session_registry=registry, - fabric=_FakeFabric(_FakeRuntime()), + fabric=cast(Any, _FakeFabric(_FakeRuntime())), ) first_started = asyncio.Event() release_first = asyncio.Event() @@ -397,7 +422,7 @@ async def test_invoke_session_releases_lock_after_failure( _agent_config(), base_dir=tmp_path, session_registry=registry, - fabric=_FakeFabric(_FakeRuntime()), + fabric=cast(Any, _FakeFabric(_FakeRuntime())), ) invocation_count = 0 @@ -429,7 +454,7 @@ async def test_invoke_session_releases_lock_after_cancellation( _agent_config(), base_dir=tmp_path, session_registry=registry, - fabric=_FakeFabric(_FakeRuntime()), + fabric=cast(Any, _FakeFabric(_FakeRuntime())), ) invocation_started = asyncio.Event() diff --git a/plugins/nemo-agents/tests/unit/test_utils.py b/plugins/nemo-agents/tests/unit/test_utils.py index 166577a34c..92a07cdbe2 100644 --- a/plugins/nemo-agents/tests/unit/test_utils.py +++ b/plugins/nemo-agents/tests/unit/test_utils.py @@ -7,6 +7,8 @@ - inject_gateway_url: IGW URL construction, setdefault semantics, explicit base_url override, NEMO_BASE_URL env var, non-openai LLMs left unchanged, original config dict not mutated. +- inject_fabric_gateway_url: shared and harness-specific Fabric models are + bound to the IGW without overriding explicit endpoints. - merge_agent_config: per-section merge semantics for component dicts vs. scalar/dict sections, workflow ownership, and input immutability. - temp_injected_config: temp file written to same directory as source, @@ -30,6 +32,7 @@ from nemo_agents_plugin.utils import ( get_internal_base_url, inject_default_model, + inject_fabric_gateway_url, inject_gateway_url, merge_agent_config, preflight_validate_llm_models, @@ -140,6 +143,43 @@ def test_trailing_slash_stripped_from_base_url(self) -> None: assert "//apis" not in result["llms"]["llm"]["base_url"] +class TestInjectFabricGatewayUrl: + def test_injects_shared_and_harness_model_settings(self) -> None: + config = { + "models": {"default": {"provider": "openai", "model": "test-model"}}, + "harnesses": { + "hermes": { + "kind": "hermes", + "model": {"provider": "nvidia", "model": "test-harness-model"}, + } + }, + } + + result = inject_fabric_gateway_url(config, "test-workspace", base_url="http://platform:8080") + expected_url = "http://platform:8080/apis/inference-gateway/v2/workspaces/test-workspace/openai/-/v1" + + assert result["models"]["default"]["settings"]["base_url"] == expected_url + assert result["harnesses"]["hermes"]["model"]["settings"]["base_url"] == expected_url + assert "settings" not in config["models"]["default"] + assert "settings" not in config["harnesses"]["hermes"]["model"] + + def test_preserves_explicit_endpoint_and_input(self) -> None: + config = { + "models": { + "default": { + "provider": "openai", + "model": "test-model", + "settings": {"base_url": "http://explicit:8080/v1"}, + } + } + } + + result = inject_fabric_gateway_url(config, "test-workspace", base_url="http://platform:8080") + + assert result["models"]["default"]["settings"]["base_url"] == "http://explicit:8080/v1" + assert config["models"]["default"]["settings"]["base_url"] == "http://explicit:8080/v1" + + class TestGetInternalBaseUrl: def test_returns_none_when_unset(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("NEMO_INTERNAL_BASE_URL", raising=False) From 3d1d94b379afddee01883cca3c95332c09f918c7 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Tue, 28 Jul 2026 14:29:19 -0500 Subject: [PATCH 16/16] feedback pt2 Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/server.py | 42 ++++++++++++------- .../tests/unit/test_fabric_server.py | 12 +++--- 2 files changed, 34 insertions(+), 20 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 3bf292bfcb..943f7ad7ba 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -12,6 +12,7 @@ import uuid from collections.abc import AsyncIterator, Mapping from contextlib import asynccontextmanager +from dataclasses import dataclass from pathlib import Path from typing import Annotated, Any @@ -44,6 +45,23 @@ SESSION_ID_HEADER = "X-Nemo-Session-Id" +@dataclass(frozen=True, slots=True) +class FabricServingSettings: + """Operational settings for the Platform-owned Fabric server.""" + + max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS + idle_session_timeout_seconds: float = DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS + session_cleanup_interval_seconds: float = DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS + + def __post_init__(self) -> None: + if self.max_concurrent_invocations < 0: + raise ValueError("max_concurrent_invocations must be greater than or equal to zero.") + if self.idle_session_timeout_seconds <= 0: + raise ValueError("idle_session_timeout_seconds must be greater than zero.") + if self.session_cleanup_interval_seconds <= 0: + raise ValueError("session_cleanup_interval_seconds must be greater than zero.") + + def _to_fabric_invocation_request( request: ChatCompletionRequest, *, @@ -115,16 +133,10 @@ async def _run_idle_session_cleanup( def create_fabric_serving_app( agent_config_path: str | Path, *, - max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS, - idle_session_timeout_seconds: float = DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS, - session_cleanup_interval_seconds: float = DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS, + settings: FabricServingSettings | None = None, ) -> FastAPI: """Create a serving app that validates its agent definition at startup.""" - if idle_session_timeout_seconds <= 0: - raise ValueError("idle_session_timeout_seconds must be greater than zero.") - if session_cleanup_interval_seconds <= 0: - raise ValueError("session_cleanup_interval_seconds must be greater than zero.") - + settings = settings or FabricServingSettings() config_path = Path(agent_config_path).resolve() @asynccontextmanager @@ -140,15 +152,15 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: agent_config, base_dir=config_path.parent, session_registry=session_registry, - max_concurrent_invocations=max_concurrent_invocations, + max_concurrent_invocations=settings.max_concurrent_invocations, ) app.state.session_manager = session_manager cleanup_shutdown = asyncio.Event() cleanup_task = asyncio.create_task( _run_idle_session_cleanup( session_manager, - idle_timeout_seconds=idle_session_timeout_seconds, - cleanup_interval_seconds=session_cleanup_interval_seconds, + idle_timeout_seconds=settings.idle_session_timeout_seconds, + cleanup_interval_seconds=settings.session_cleanup_interval_seconds, shutdown_event=cleanup_shutdown, ) ) @@ -267,9 +279,11 @@ def main(argv: list[str] | None = None) -> int: uvicorn.run( create_fabric_serving_app( args.agent_config, - max_concurrent_invocations=args.max_concurrent_invocations, - idle_session_timeout_seconds=args.idle_session_timeout_seconds, - session_cleanup_interval_seconds=args.session_cleanup_interval_seconds, + settings=FabricServingSettings( + max_concurrent_invocations=args.max_concurrent_invocations, + idle_session_timeout_seconds=args.idle_session_timeout_seconds, + session_cleanup_interval_seconds=args.session_cleanup_interval_seconds, + ), ), host=args.host, port=args.port, diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index eab3c2c7e1..c5b509cf8d 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -18,7 +18,7 @@ FabricRuntimeResult, FabricRuntimeTimeoutError, ) -from nemo_agents_plugin.fabric.server import SESSION_ID_HEADER, create_fabric_serving_app +from nemo_agents_plugin.fabric.server import SESSION_ID_HEADER, FabricServingSettings, create_fabric_serving_app from nemo_agents_plugin.fabric.serving_models import ChatCompletionRequest from nemo_agents_plugin.fabric.session_manager import ( FabricSessionManager, @@ -151,13 +151,13 @@ def test_startup_fails_for_invalid_agent_config( assert mock_validate_agent_config == [] -def test_rejects_non_positive_session_cleanup_settings(tmp_path: Path) -> None: - config_path = _write_agent_config(tmp_path) - +def test_rejects_invalid_serving_settings() -> None: + with pytest.raises(ValueError, match="max_concurrent_invocations"): + FabricServingSettings(max_concurrent_invocations=-1) with pytest.raises(ValueError, match="idle_session_timeout_seconds"): - create_fabric_serving_app(config_path, idle_session_timeout_seconds=0) + FabricServingSettings(idle_session_timeout_seconds=0) with pytest.raises(ValueError, match="session_cleanup_interval_seconds"): - create_fabric_serving_app(config_path, session_cleanup_interval_seconds=0) + FabricServingSettings(session_cleanup_interval_seconds=0) def test_chat_completion_without_session_id_opens_and_returns_session(