diff --git a/plugins/nemo-agents/README.md b/plugins/nemo-agents/README.md index 598095587f..10f62e966a 100644 --- a/plugins/nemo-agents/README.md +++ b/plugins/nemo-agents/README.md @@ -352,6 +352,23 @@ The injected URL format: --- +## Performance tips + +### First-deploy cold start + +The first `nemo agents deploy` after installing packages is noticeably slower +than subsequent deploys because Python compiles `.pyc` bytecache files on first +import. Pre-compiling NAT's dependencies eliminates this overhead: + +```bash +python -m compileall -q $(python -c "import nat; print(nat.__path__[0])") 2>/dev/null +python -m compileall -q .venv/lib/ 2>/dev/null +``` + +This can cut 20--40 seconds off the first deploy. + +--- + ## Notes and known limitations - **`tool_calling_agent`** is broken with `langchain-openai==1.1.x` due to a diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/config.py b/plugins/nemo-agents/src/nemo_agents_plugin/config.py index f00abe0870..1199cbabad 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/config.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/config.py @@ -16,7 +16,7 @@ class ControllerConfig(BaseModel): """Configuration for the AgentDeploymentController reconcile loop.""" - interval_seconds: int = Field(default=5, description="Reconciliation loop interval in seconds.") + interval_seconds: int = Field(default=2, description="Reconciliation loop interval in seconds.") health_check_timeout_seconds: int = Field( default=120, description="Maximum time to wait for agent health check to succeed." ) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/runner/backend.py b/plugins/nemo-agents/src/nemo_agents_plugin/runner/backend.py index a3f6d901fa..97baaca487 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/runner/backend.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/runner/backend.py @@ -97,9 +97,9 @@ async def health_check(self, endpoint: str) -> bool: ... @abstractmethod - def shutdown(self) -> None: + async def shutdown(self) -> None: """Terminate all managed processes and release resources. - Called synchronously during service shutdown. Must be idempotent. + Called during service shutdown. Must be idempotent. """ ... diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/runner/controller.py b/plugins/nemo-agents/src/nemo_agents_plugin/runner/controller.py index 5b33463e63..8a5ff5e718 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/runner/controller.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/runner/controller.py @@ -56,7 +56,7 @@ def __init__(self) -> None: self._entities: NemoEntitiesClient | None = None self._controller_config: ControllerConfig | None = None self._starting_since: dict[str, float] = {} - self._interval_seconds: float = 5.0 # default; overwritten in on_startup + self._interval_seconds: float = 2.0 # overwritten in on_startup # ------------------------------------------------------------------ # Narrowing properties — raise clearly if accessed before on_startup() @@ -90,6 +90,10 @@ def interval_seconds(self) -> float: async def on_startup(self) -> None: """Initialise the entity client and runner backend from config.""" + # Imports deferred intentionally: these modules pull in the SDK, + # entity-store client, and HTTP machinery. Importing at module level + # would add ~1s to every `nemo` CLI invocation during plugin discovery, + # even when the agents controller is never started. Do not hoist. from nemo_agents_plugin.config import AgentsConfig from nemo_agents_plugin.runner.registry import RunnerBackendRegistry from nemo_platform.resources.entities import AsyncEntitiesResource @@ -120,7 +124,7 @@ async def on_startup(self) -> None: async def on_shutdown(self) -> None: """Shut down the runner backend.""" if self._backend is not None: - self._backend.shutdown() + await self._backend.shutdown() logger.info("AgentDeploymentController shut down.") async def list_objects(self) -> list: @@ -146,7 +150,7 @@ async def reconcile_one(self, obj: object) -> None: logger.debug("Optimistic lock conflict on '%s' — will retry next cycle.", dep.name) # ------------------------------------------------------------------ - # Internal state-machine helpers (unchanged from original) + # Internal state-machine helpers # ------------------------------------------------------------------ async def _reconcile_one(self, dep: AgentDeployment) -> None: @@ -160,7 +164,8 @@ async def _reconcile_one(self, dep: AgentDeployment) -> None: await self._delete_deployment(dep) async def _start_deployment(self, dep: AgentDeployment) -> None: - """pending → starting: allocate port and spawn the agent process.""" + """pending -> starting: allocate port and spawn the agent process.""" + t0 = time.perf_counter() port = self.backend.allocate_port() try: info = await self.backend.create_deployment( @@ -175,6 +180,7 @@ async def _start_deployment(self, dep: AgentDeployment) -> None: await self._save(dep) return + spawn_ms = (time.perf_counter() - t0) * 1000 dep.status = "starting" dep.port = info.port dep.pid = info.pid @@ -183,23 +189,35 @@ async def _start_deployment(self, dep: AgentDeployment) -> None: self._starting_since[dep.name] = time.monotonic() await self._save(dep) logger.info( - "Deployment '%s' started (pid=%d, port=%d, log=%s).", + "Deployment '%s' spawned (pid=%d, port=%d, spawn=%.0fms, log=%s).", dep.name, dep.pid, dep.port, + spawn_ms, info.log_path or "", ) async def _check_health(self, dep: AgentDeployment) -> None: - """starting → running | failed: poll the health endpoint. + """starting -> running | failed: single-shot health check per reconcile cycle. - Process death takes precedence over a successful health check: if - the subprocess has already exited, surface the failure immediately - with the exit code in the error message, instead of letting a stale - ``/health`` reply mark the deployment as ``running``. + Checks once and returns so the reconcile loop can service other + deployments promptly. The ``_starting_since`` timestamp persists + across cycles so the overall ``health_check_timeout_seconds`` budget + is enforced across many cycles. """ since = self._starting_since.get(dep.name, time.monotonic()) + timeout = self.controller_config.health_check_timeout_seconds elapsed = time.monotonic() - since + remaining = timeout - elapsed + + if remaining <= 0: + dep.status = "failed" + dep.error = f"Health check timed out after {timeout}s." + await self.backend.delete_deployment(dep.name) + self._starting_since.pop(dep.name, None) + await self._save(dep) + logger.warning("Deployment '%s' health check timed out.", dep.name) + return info = await self.backend.get_deployment_status(dep.name) if info is not None and info.status == "failed": @@ -216,23 +234,19 @@ async def _check_health(self, dep: AgentDeployment) -> None: return healthy = bool(dep.endpoint) and await self.backend.health_check(dep.endpoint) + if healthy: dep.status = "running" self._starting_since.pop(dep.name, None) await self._save(dep) - logger.info("Deployment '%s' is running at %s.", dep.name, dep.endpoint) - elif elapsed > self.controller_config.health_check_timeout_seconds: - dep.status = "failed" - dep.error = f"Health check timed out after {self.controller_config.health_check_timeout_seconds}s." - self._starting_since.pop(dep.name, None) - log_path = info.log_path if info is not None else "" - await self.backend.delete_deployment(dep.name) - await self._save(dep) - logger.warning( - "Deployment '%s' health check timed out (log: %s).", + logger.info( + "Deployment '%s' is running at %s (took %.1fs).", dep.name, - log_path or "", + dep.endpoint, + time.monotonic() - since, ) + else: + logger.debug("Deployment '%s' not healthy yet (%.1fs elapsed).", dep.name, elapsed) async def _verify_running(self, dep: AgentDeployment) -> None: """mark failed if the process has exited or pending if process is not found to attempt to restart.""" 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 6efca48309..a3d0726c15 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 @@ -101,6 +101,7 @@ def __init__(self, config: ControllerConfig) -> None: self._deployments: dict[str, DeploymentInfo] = {} self._next_port: int = config.port_range_start self._temp_files: dict[str, Path] = {} + self._http_client: httpx.AsyncClient | None = None @property def output_base_dir(self) -> Path: @@ -227,21 +228,38 @@ async def list_deployments(self) -> list[DeploymentInfo]: async def health_check(self, endpoint: str) -> bool: url = endpoint.rstrip("/") + "/health" try: - async with httpx.AsyncClient(timeout=5.0) as client: - resp = await client.get(url) - return resp.status_code < 400 + client = self._get_http_client() + resp = await client.get(url) + return resp.status_code < 400 except Exception: return False - def shutdown(self) -> None: - """Terminate all managed processes synchronously.""" - for name, proc in list(self._processes.items()): - self._terminate(name, proc) + def _get_http_client(self) -> httpx.AsyncClient: + if self._http_client is None or self._http_client.is_closed: + self._http_client = httpx.AsyncClient(timeout=5.0) + return self._http_client + + async def shutdown(self) -> None: + """Terminate all managed processes (best-effort).""" + names = list(self._processes.keys()) + results = await asyncio.gather( + *(asyncio.to_thread(self._terminate, name, proc) for name, proc in list(self._processes.items())), + return_exceptions=True, + ) + for name, result in zip(names, results, strict=False): + if isinstance(result, Exception): + logger.warning("Error terminating '%s' during shutdown", name, exc_info=result) self._processes.clear() self._deployments.clear() for path in self._temp_files.values(): path.unlink(missing_ok=True) self._temp_files.clear() + if self._http_client is not None and not self._http_client.is_closed: + try: + await self._http_client.aclose() + except Exception: + logger.warning("Error closing HTTP client during shutdown", exc_info=True) + self._http_client = None logger.info("InMemoryRunnerBackend shut down — all processes terminated.") def _write_config(self, name: str, config: dict[str, Any]) -> Path: diff --git a/plugins/nemo-agents/tests/unit/test_controller.py b/plugins/nemo-agents/tests/unit/test_controller.py new file mode 100644 index 0000000000..8ab48c18cf --- /dev/null +++ b/plugins/nemo-agents/tests/unit/test_controller.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for AgentDeploymentController health checks and state transitions.""" + +from __future__ import annotations + +import time +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from nemo_agents_plugin.config import ControllerConfig +from nemo_agents_plugin.entities import AgentDeployment, DeploymentStatus +from nemo_agents_plugin.runner.backend import DeploymentInfo +from nemo_agents_plugin.runner.controller import AgentDeploymentController + + +def _make_deployment( + name: str = "test-dep", + workspace: str = "default", + agent: str = "test-agent", + status: DeploymentStatus = "pending", + port: int = 0, + pid: int = 0, + endpoint: str = "", +) -> AgentDeployment: + dep = AgentDeployment(name=name, workspace=workspace, agent=agent, status=status) + dep.port = port + dep.pid = pid + dep.endpoint = endpoint + return dep + + +def _make_controller( + interval_seconds: int = 2, + health_check_timeout_seconds: int = 10, + health_check_interval_seconds: int = 1, +) -> Any: + """Return a controller with mocked backend and entities. + + Returns ``Any`` so ty does not flag ``AsyncMock`` attribute access + (e.g. ``.return_value``, ``.side_effect``) on the mock backend/entities. + """ + ctrl = AgentDeploymentController() + ctrl._backend = AsyncMock() + ctrl._entities = AsyncMock() + ctrl._controller_config = ControllerConfig( + interval_seconds=interval_seconds, + health_check_timeout_seconds=health_check_timeout_seconds, + health_check_interval_seconds=health_check_interval_seconds, + ) + ctrl._interval_seconds = float(interval_seconds) + return ctrl + + +class TestDefaultInterval: + def test_default_interval_is_2s(self) -> None: + cfg = ControllerConfig() + assert cfg.interval_seconds == 2 + + def test_default_health_check_interval_is_2s(self) -> None: + cfg = ControllerConfig() + assert cfg.health_check_interval_seconds == 2 + + +class TestStartDeployment: + @pytest.mark.asyncio + async def test_start_transitions_to_starting(self) -> None: + ctrl = _make_controller() + dep = _make_deployment(status="pending") + ctrl.backend.allocate_port.return_value = 50000 + ctrl.backend.create_deployment.return_value = DeploymentInfo( + name="test-dep", + status="starting", + port=50000, + pid=12345, + endpoint="http://127.0.0.1:50000", + ) + + await ctrl._start_deployment(dep) + + assert dep.status == "starting" + assert dep.port == 50000 + assert dep.pid == 12345 + assert dep.endpoint == "http://127.0.0.1:50000" + assert dep.name in ctrl._starting_since + + @pytest.mark.asyncio + async def test_start_failure_transitions_to_failed(self) -> None: + ctrl = _make_controller() + dep = _make_deployment(status="pending") + ctrl.backend.allocate_port.return_value = 50000 + ctrl.backend.create_deployment.side_effect = RuntimeError("spawn failed") + + await ctrl._start_deployment(dep) + + assert dep.status == "failed" + assert "spawn failed" in dep.error + + +class TestCheckHealth: + @pytest.mark.asyncio + async def test_healthy_on_first_check(self) -> None: + ctrl = _make_controller(health_check_interval_seconds=1) + dep = _make_deployment( + status="starting", + port=50000, + pid=123, + endpoint="http://127.0.0.1:50000", + ) + ctrl._starting_since["test-dep"] = time.monotonic() + + ctrl.backend.get_deployment_status.return_value = DeploymentInfo( + name="test-dep", + status="starting", + ) + ctrl.backend.health_check.return_value = True + + await ctrl._check_health(dep) + + assert dep.status == "running" + assert "test-dep" not in ctrl._starting_since + + @pytest.mark.asyncio + async def test_not_healthy_stays_starting(self) -> None: + """Single-shot: one failed probe leaves status as 'starting' for the next cycle.""" + ctrl = _make_controller() + dep = _make_deployment( + status="starting", + port=50000, + pid=123, + endpoint="http://127.0.0.1:50000", + ) + ctrl._starting_since["test-dep"] = time.monotonic() + + ctrl.backend.get_deployment_status.return_value = DeploymentInfo( + name="test-dep", + status="starting", + ) + ctrl.backend.health_check.return_value = False + + await ctrl._check_health(dep) + + assert dep.status == "starting" + assert "test-dep" in ctrl._starting_since + ctrl.backend.health_check.assert_called_once() + + @pytest.mark.asyncio + async def test_process_exit_during_health_check(self) -> None: + ctrl = _make_controller(health_check_interval_seconds=0) + dep = _make_deployment( + status="starting", + port=50000, + pid=123, + endpoint="http://127.0.0.1:50000", + ) + ctrl._starting_since["test-dep"] = time.monotonic() + + ctrl.backend.get_deployment_status.return_value = DeploymentInfo( + name="test-dep", + status="failed", + error="Process exited with code 1", + ) + + await ctrl._check_health(dep) + + assert dep.status == "failed" + assert "exited" in dep.error.lower() + + @pytest.mark.asyncio + async def test_timeout_transitions_to_failed(self) -> None: + ctrl = _make_controller( + health_check_timeout_seconds=0, + health_check_interval_seconds=0, + ) + dep = _make_deployment( + status="starting", + port=50000, + pid=123, + endpoint="http://127.0.0.1:50000", + ) + ctrl._starting_since["test-dep"] = 0 # started long ago + + ctrl.backend.get_deployment_status.return_value = DeploymentInfo( + name="test-dep", + status="starting", + ) + ctrl.backend.health_check.return_value = False + + await ctrl._check_health(dep) + + assert dep.status == "failed" + assert "timed out" in dep.error.lower() + ctrl.backend.delete_deployment.assert_called_once_with("test-dep") + + @pytest.mark.asyncio + async def test_no_endpoint_skips_health_check(self) -> None: + """If endpoint is empty, health_check should not be called.""" + ctrl = _make_controller( + health_check_timeout_seconds=0, + health_check_interval_seconds=0, + ) + dep = _make_deployment(status="starting", endpoint="") + ctrl._starting_since["test-dep"] = 0 + + ctrl.backend.get_deployment_status.return_value = DeploymentInfo( + name="test-dep", + status="starting", + ) + ctrl.backend.health_check.return_value = False + + await ctrl._check_health(dep) + + assert dep.status == "failed" + ctrl.backend.health_check.assert_not_called() + + +class TestReconcileOne: + @pytest.mark.asyncio + async def test_pending_triggers_start(self) -> None: + ctrl = _make_controller() + dep = _make_deployment(status="pending") + ctrl.backend.allocate_port.return_value = 50000 + ctrl.backend.create_deployment.return_value = DeploymentInfo( + name="test-dep", + status="starting", + port=50000, + pid=99, + endpoint="http://127.0.0.1:50000", + ) + + await ctrl._reconcile_one(dep) + + assert dep.status == "starting" + + @pytest.mark.asyncio + async def test_running_verify_process_gone_resets_to_pending(self) -> None: + ctrl = _make_controller() + dep = _make_deployment(status="running") + ctrl.backend.get_deployment_status.return_value = None + + await ctrl._reconcile_one(dep) + + assert dep.status == "pending" + + @pytest.mark.asyncio + async def test_deleting_removes_deployment(self) -> None: + ctrl = _make_controller() + dep = _make_deployment(status="deleting") + + await ctrl._reconcile_one(dep) + + ctrl.backend.delete_deployment.assert_called_once_with("test-dep") diff --git a/plugins/nemo-agents/tests/unit/test_runner_controller.py b/plugins/nemo-agents/tests/unit/test_runner_controller.py index 36bce17eab..2a16361274 100644 --- a/plugins/nemo-agents/tests/unit/test_runner_controller.py +++ b/plugins/nemo-agents/tests/unit/test_runner_controller.py @@ -16,6 +16,7 @@ from __future__ import annotations +import time from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -35,6 +36,7 @@ def _make_controller() -> tuple[AgentDeploymentController, Any]: """ ctrl = AgentDeploymentController() backend = MagicMock() + backend.delete_deployment = AsyncMock() # Bypass on_startup() — wire stubs directly. ctrl._backend = backend ctrl._entities = MagicMock() @@ -112,7 +114,7 @@ async def test_check_health_marks_failed_when_subprocess_exited() -> None: status="starting", endpoint="http://127.0.0.1:49200", ) - ctrl._starting_since["dep-1"] = 0.0 # arbitrary + ctrl._starting_since["dep-1"] = time.monotonic() await ctrl._check_health(dep) @@ -136,7 +138,7 @@ async def test_check_health_marks_running_when_healthy() -> None: status="starting", endpoint="http://127.0.0.1:49200", ) - ctrl._starting_since["dep-1"] = 0.0 + ctrl._starting_since["dep-1"] = time.monotonic() await ctrl._check_health(dep)