nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/front_ends/mcp/mcp_front_end_plugin_worker.py

@@ -0,0 +1,268 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Mapping
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+from starlette.exceptions import HTTPException
+from starlette.requests import Request
+
+from nat.builder.function import Function
+from nat.builder.function_base import FunctionBase
+from nat.builder.workflow import Workflow
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.data_models.config import Config
+from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
+from nat.front_ends.mcp.memory_profiler import MemoryProfiler
+
+logger = logging.getLogger(__name__)
+
+
+class MCPFrontEndPluginWorkerBase(ABC):
+    """Base class for MCP front end plugin workers."""
+
+    def __init__(self, config: Config):
+        """Initialize the MCP worker with configuration.
+
+        Args:
+            config: The full NAT configuration
+        """
+        self.full_config = config
+        self.front_end_config: MCPFrontEndConfig = config.general.front_end
+
+        # Initialize memory profiler if enabled
+        self.memory_profiler = MemoryProfiler(enabled=self.front_end_config.enable_memory_profiling,
+                                              log_interval=self.front_end_config.memory_profile_interval,
+                                              top_n=self.front_end_config.memory_profile_top_n,
+                                              log_level=self.front_end_config.memory_profile_log_level)
+
+    def _setup_health_endpoint(self, mcp: FastMCP):
+        """Set up the HTTP health endpoint that exercises MCP ping handler."""
+
+        @mcp.custom_route("/health", methods=["GET"])
+        async def health_check(_request: Request):
+            """HTTP health check using server's internal ping handler"""
+            from starlette.responses import JSONResponse
+
+            try:
+                from mcp.types import PingRequest
+
+                # Create a ping request
+                ping_request = PingRequest(method="ping")
+
+                # Call the ping handler directly (same one that responds to MCP pings)
+                await mcp._mcp_server.request_handlers[PingRequest](ping_request)
+
+                return JSONResponse({
+                    "status": "healthy",
+                    "error": None,
+                    "server_name": mcp.name,
+                })
+
+            except Exception as e:
+                return JSONResponse({
+                    "status": "unhealthy",
+                    "error": str(e),
+                    "server_name": mcp.name,
+                },
+                                    status_code=503)
+
+    @abstractmethod
+    async def add_routes(self, mcp: FastMCP, builder: WorkflowBuilder):
+        """Add routes to the MCP server.
+
+        Args:
+            mcp: The FastMCP server instance
+            builder (WorkflowBuilder): The workflow builder instance
+        """
+        pass
+
+    async def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]:
+        """Get all functions from the workflow.
+
+        Args:
+            workflow: The NAT workflow.
+
+        Returns:
+            Dict mapping function names to Function objects.
+        """
+        functions: dict[str, Function] = {}
+
+        # Extract all functions from the workflow
+        functions.update(workflow.functions)
+        for function_group in workflow.function_groups.values():
+            functions.update(await function_group.get_accessible_functions())
+
+        if workflow.config.workflow.workflow_alias:
+            functions[workflow.config.workflow.workflow_alias] = workflow
+        else:
+            functions[workflow.config.workflow.type] = workflow
+
+        return functions
+
+    def _setup_debug_endpoints(self, mcp: FastMCP, functions: Mapping[str, FunctionBase]) -> None:
+        """Set up HTTP debug endpoints for introspecting tools and schemas.
+
+        Exposes:
+          - GET /debug/tools/list: List tools. Optional query param `name` (one or more, repeatable or comma separated)
+            selects a subset and returns details for those tools.
+          - GET /debug/memory/stats: Get current memory profiling statistics (read-only)
+        """
+
+        @mcp.custom_route("/debug/tools/list", methods=["GET"])
+        async def list_tools(request: Request):
+            """HTTP list tools endpoint."""
+
+            from starlette.responses import JSONResponse
+
+            from nat.front_ends.mcp.tool_converter import get_function_description
+
+            # Query params
+            # Support repeated names and comma-separated lists
+            names_param_list = set(request.query_params.getlist("name"))
+            names: list[str] = []
+            for raw in names_param_list:
+                # if p.strip() is empty, it won't be included in the list!
+                parts = [p.strip() for p in raw.split(",") if p.strip()]
+                names.extend(parts)
+            detail_raw = request.query_params.get("detail")
+
+            def _parse_detail_param(detail_param: str | None, has_names: bool) -> bool:
+                if detail_param is None:
+                    if has_names:
+                        return True
+                    return False
+                v = detail_param.strip().lower()
+                if v in ("0", "false", "no", "off"):
+                    return False
+                if v in ("1", "true", "yes", "on"):
+                    return True
+                # For invalid values, default based on whether names are present
+                return has_names
+
+            # Helper function to build the input schema info
+            def _build_schema_info(fn: FunctionBase) -> dict[str, Any] | None:
+                schema = getattr(fn, "input_schema", None)
+                if schema is None:
+                    return None
+
+                # check if schema is a ChatRequest
+                schema_name = getattr(schema, "__name__", "")
+                schema_qualname = getattr(schema, "__qualname__", "")
+                if "ChatRequest" in schema_name or "ChatRequest" in schema_qualname:
+                    # Simplified interface used by MCP wrapper for ChatRequest
+                    return {
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string", "description": "User query string"
+                            }
+                        },
+                        "required": ["query"],
+                        "title": "ChatRequestQuery",
+                    }
+
+                # Pydantic models provide model_json_schema
+                if schema is not None and hasattr(schema, "model_json_schema"):
+                    return schema.model_json_schema()
+
+                return None
+
+            def _build_final_json(functions_to_include: Mapping[str, FunctionBase],
+                                  include_schemas: bool = False) -> dict[str, Any]:
+                tools = []
+                for name, fn in functions_to_include.items():
+                    list_entry: dict[str, Any] = {
+                        "name": name, "description": get_function_description(fn), "is_workflow": hasattr(fn, "run")
+                    }
+                    if include_schemas:
+                        list_entry["schema"] = _build_schema_info(fn)
+                    tools.append(list_entry)
+
+                return {
+                    "count": len(tools),
+                    "tools": tools,
+                    "server_name": mcp.name,
+                }
+
+            if names:
+                # Return selected tools
+                try:
+                    functions_to_include = {n: functions[n] for n in names}
+                except KeyError as e:
+                    raise HTTPException(status_code=404, detail=f"Tool \"{e.args[0]}\" not found.") from e
+            else:
+                functions_to_include = functions
+
+            # Default for listing all: detail defaults to False unless explicitly set true
+            return JSONResponse(
+                _build_final_json(functions_to_include, _parse_detail_param(detail_raw, has_names=bool(names))))
+
+        # Memory profiling endpoint (read-only)
+        @mcp.custom_route("/debug/memory/stats", methods=["GET"])
+        async def get_memory_stats(_request: Request):
+            """Get current memory profiling statistics."""
+            from starlette.responses import JSONResponse
+
+            stats = self.memory_profiler.get_stats()
+            return JSONResponse(stats)
+
+
+class MCPFrontEndPluginWorker(MCPFrontEndPluginWorkerBase):
+    """Default MCP front end plugin worker implementation."""
+
+    async def add_routes(self, mcp: FastMCP, builder: WorkflowBuilder):
+        """Add default routes to the MCP server.
+
+        Args:
+            mcp: The FastMCP server instance
+            builder (WorkflowBuilder): The workflow builder instance
+        """
+        from nat.front_ends.mcp.tool_converter import register_function_with_mcp
+
+        # Set up the health endpoint
+        self._setup_health_endpoint(mcp)
+
+        # Build the workflow and register all functions with MCP
+        workflow = await builder.build()
+
+        # Get all functions from the workflow
+        functions = await self._get_all_functions(workflow)
+
+        # Filter functions based on tool_names if provided
+        if self.front_end_config.tool_names:
+            logger.info("Filtering functions based on tool_names: %s", self.front_end_config.tool_names)
+            filtered_functions: dict[str, Function] = {}
+            for function_name, function in functions.items():
+                if function_name in self.front_end_config.tool_names:
+                    filtered_functions[function_name] = function
+                else:
+                    logger.debug("Skipping function %s as it's not in tool_names", function_name)
+            functions = filtered_functions
+
+        # Register each function with MCP, passing workflow context for observability
+        for function_name, function in functions.items():
+            register_function_with_mcp(mcp, function_name, function, workflow, self.memory_profiler)
+
+        # Add a simple fallback function if no functions were found
+        if not functions:
+            raise RuntimeError("No functions found in workflow. Please check your configuration.")
+
+        # After registration, expose debug endpoints for tool/schema inspection
+        self._setup_debug_endpoints(mcp, functions)

nat/front_ends/mcp/memory_profiler.py

@@ -0,0 +1,320 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Memory profiling utilities for MCP frontend."""
+
+import gc
+import logging
+import tracemalloc
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryProfiler:
+    """Memory profiler for tracking memory usage and potential leaks."""
+
+    def __init__(self, enabled: bool = False, log_interval: int = 50, top_n: int = 10, log_level: str = "DEBUG"):
+        """Initialize the memory profiler.
+
+        Args:
+            enabled: Whether memory profiling is enabled
+            log_interval: Log stats every N requests
+            top_n: Number of top allocations to log
+            log_level: Log level for memory profiling output (e.g., "DEBUG", "INFO")
+        """
+        self.enabled = enabled
+        # normalize interval to avoid modulo-by-zero
+        self.log_interval = max(1, int(log_interval))
+        self.top_n = top_n
+        self.log_level = getattr(logging, log_level.upper(), logging.DEBUG)
+        self.request_count = 0
+        self.baseline_snapshot = None
+
+        # Track whether this instance started tracemalloc (to avoid resetting external tracing)
+        self._we_started_tracemalloc = False
+
+        if self.enabled:
+            logger.info("Memory profiling ENABLED (interval=%d, top_n=%d, log_level=%s)",
+                        self.log_interval,
+                        top_n,
+                        log_level)
+            try:
+                if not tracemalloc.is_tracing():
+                    tracemalloc.start()
+                    self._we_started_tracemalloc = True
+                # Take baseline snapshot
+                gc.collect()
+                self.baseline_snapshot = tracemalloc.take_snapshot()
+            except RuntimeError as e:
+                logger.warning("tracemalloc unavailable or not tracing: %s", e)
+        else:
+            logger.info("Memory profiling DISABLED")
+
+    def _log(self, message: str, *args: Any) -> None:
+        """Log a message at the configured log level.
+
+        Args:
+            message: Log message format string
+            args: Arguments for the format string
+        """
+        logger.log(self.log_level, message, *args)
+
+    def on_request_complete(self) -> None:
+        """Called after each request completes."""
+        if not self.enabled:
+            return
+        self.request_count += 1
+        if self.request_count % self.log_interval == 0:
+            self.log_memory_stats()
+
+    def _ensure_tracing(self) -> bool:
+        """Ensure tracemalloc is running if we started it originally.
+
+        Returns:
+            True if tracemalloc is active, False otherwise
+        """
+        if tracemalloc.is_tracing():
+            return True
+
+        # Only restart if we started it originally (respect external control)
+        if not self._we_started_tracemalloc:
+            return False
+
+        # Attempt to restart
+        try:
+            logger.warning("tracemalloc was stopped externally; restarting (we started it originally)")
+            tracemalloc.start()
+
+            # Reset baseline since old tracking data is lost
+            gc.collect()
+            self.baseline_snapshot = tracemalloc.take_snapshot()
+            logger.info("Baseline snapshot reset after tracemalloc restart")
+
+            return True
+        except RuntimeError as e:
+            logger.error("Failed to restart tracemalloc: %s", e)
+            return False
+
+    def _safe_traced_memory(self) -> tuple[float, float] | None:
+        """Return (current, peak usage in MB) if tracemalloc is active, else None."""
+        if not self._ensure_tracing():
+            return None
+
+        try:
+            current, peak = tracemalloc.get_traced_memory()
+            megabyte = (1 << 20)
+            return (current / megabyte, peak / megabyte)
+        except RuntimeError:
+            return None
+
+    def _safe_snapshot(self) -> tracemalloc.Snapshot | None:
+        """Return a tracemalloc Snapshot if available, else None."""
+        if not self._ensure_tracing():
+            return None
+
+        try:
+            return tracemalloc.take_snapshot()
+        except RuntimeError:
+            return None
+
+    def log_memory_stats(self) -> dict[str, Any]:
+        """Log current memory statistics and return them."""
+        if not self.enabled:
+            return {}
+
+        # Force garbage collection first
+        gc.collect()
+
+        # Get current memory usage
+        mem = self._safe_traced_memory()
+        if mem is None:
+            logger.info("tracemalloc is not active; cannot collect memory stats.")
+            # still return structural fields
+            stats = {
+                "request_count": self.request_count,
+                "current_memory_mb": None,
+                "peak_memory_mb": None,
+                "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
+                "outstanding_steps": self._safe_outstanding_step_count(),
+                "active_exporters": self._safe_exporter_count(),
+                "isolated_exporters": self._safe_isolated_exporter_count(),
+                "subject_instances": self._count_instances_of_type("Subject"),
+            }
+            return stats
+
+        current_mb, peak_mb = mem
+
+        # Take snapshot and compare to baseline
+        snapshot = self._safe_snapshot()
+
+        # Track BaseExporter instances (observability layer)
+        exporter_count = self._safe_exporter_count()
+        isolated_exporter_count = self._safe_isolated_exporter_count()
+
+        # Track Subject instances (event streams)
+        subject_count = self._count_instances_of_type("Subject")
+
+        stats = {
+            "request_count": self.request_count,
+            "current_memory_mb": round(current_mb, 2),
+            "peak_memory_mb": round(peak_mb, 2),
+            "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
+            "outstanding_steps": self._safe_outstanding_step_count(),
+            "active_exporters": exporter_count,
+            "isolated_exporters": isolated_exporter_count,
+            "subject_instances": subject_count,
+        }
+
+        self._log("=" * 80)
+        self._log("MEMORY PROFILE AFTER %d REQUESTS:", self.request_count)
+        self._log("  Current Memory: %.2f MB", current_mb)
+        self._log("  Peak Memory: %.2f MB", peak_mb)
+        self._log("")
+        self._log("NAT COMPONENT INSTANCES:")
+        self._log("  IntermediateStepManagers: %d active (%d outstanding steps)",
+                  stats["active_intermediate_managers"],
+                  stats["outstanding_steps"])
+        self._log("  BaseExporters: %d active (%d isolated)", stats["active_exporters"], stats["isolated_exporters"])
+        self._log("  Subject (event streams): %d instances", stats["subject_instances"])
+
+        # Show top allocations
+        if snapshot is None:
+            self._log("tracemalloc snapshot unavailable.")
+        else:
+            if self.baseline_snapshot:
+                self._log("TOP %d MEMORY GROWTH SINCE BASELINE:", self.top_n)
+                top_stats = snapshot.compare_to(self.baseline_snapshot, 'lineno')
+            else:
+                self._log("TOP %d MEMORY ALLOCATIONS:", self.top_n)
+                top_stats = snapshot.statistics('lineno')
+
+            for i, stat in enumerate(top_stats[:self.top_n], 1):
+                self._log("  #%d: %s", i, stat)
+
+        self._log("=" * 80)
+
+        return stats
+
+    def _count_instances_of_type(self, type_name: str) -> int:
+        """Count instances of a specific type in memory."""
+        count = 0
+        for obj in gc.get_objects():
+            try:
+                if type(obj).__name__ == type_name:
+                    count += 1
+            except Exception:
+                pass
+        return count
+
+    def _safe_exporter_count(self) -> int:
+        try:
+            from nat.observability.exporter.base_exporter import BaseExporter
+            return BaseExporter.get_active_instance_count()
+        except Exception as e:
+            logger.debug("Could not get BaseExporter stats: %s", e)
+            return 0
+
+    def _safe_isolated_exporter_count(self) -> int:
+        try:
+            from nat.observability.exporter.base_exporter import BaseExporter
+            return BaseExporter.get_isolated_instance_count()
+        except Exception:
+            return 0
+
+    def _safe_intermediate_step_manager_count(self) -> int:
+        try:
+            from nat.builder.intermediate_step_manager import IntermediateStepManager
+            # len() is atomic in CPython, but catch RuntimeError just in case
+            try:
+                return IntermediateStepManager.get_active_instance_count()
+            except RuntimeError:
+                # Set was modified during len() - very rare
+                logger.debug("Set changed during count, returning 0")
+                return 0
+        except Exception as e:
+            logger.debug("Could not get IntermediateStepManager stats: %s", e)
+            return 0
+
+    def _safe_outstanding_step_count(self) -> int:
+        """Get total outstanding steps across all active IntermediateStepManager instances."""
+        try:
+            from nat.builder.intermediate_step_manager import IntermediateStepManager
+
+            # Make a snapshot to avoid "Set changed size during iteration" if GC runs
+            try:
+                instances_snapshot = list(IntermediateStepManager._active_instances)
+            except RuntimeError:
+                # Set changed during list() call - rare but possible
+                logger.debug("Set changed during snapshot, returning 0 for outstanding steps")
+                return 0
+
+            total_outstanding = 0
+            # Iterate through snapshot safely
+            for ref in instances_snapshot:
+                try:
+                    manager = ref()
+                    if manager is not None:
+                        total_outstanding += manager.get_outstanding_step_count()
+                except (ReferenceError, AttributeError):
+                    # Manager was GC'd or in invalid state - skip it
+                    continue
+            return total_outstanding
+        except Exception as e:
+            logger.debug("Could not get outstanding step count: %s", e)
+            return 0
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get current memory statistics without logging."""
+        if not self.enabled:
+            return {"enabled": False}
+
+        mem = self._safe_traced_memory()
+        if mem is None:
+            return {
+                "enabled": True,
+                "request_count": self.request_count,
+                "current_memory_mb": None,
+                "peak_memory_mb": None,
+                "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
+                "outstanding_steps": self._safe_outstanding_step_count(),
+                "active_exporters": self._safe_exporter_count(),
+                "isolated_exporters": self._safe_isolated_exporter_count(),
+                "subject_instances": self._count_instances_of_type("Subject"),
+            }
+
+        current_mb, peak_mb = mem
+        return {
+            "enabled": True,
+            "request_count": self.request_count,
+            "current_memory_mb": round(current_mb, 2),
+            "peak_memory_mb": round(peak_mb, 2),
+            "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
+            "outstanding_steps": self._safe_outstanding_step_count(),
+            "active_exporters": self._safe_exporter_count(),
+            "isolated_exporters": self._safe_isolated_exporter_count(),
+            "subject_instances": self._count_instances_of_type("Subject"),
+        }
+
+    def reset_baseline(self) -> None:
+        """Reset the baseline snapshot to current state."""
+        if not self.enabled:
+            return
+        gc.collect()
+        snap = self._safe_snapshot()
+        if snap is None:
+            logger.info("Cannot reset baseline: tracemalloc is not active.")
+            return
+        self.baseline_snapshot = snap
+        logger.info("Memory profiling baseline reset at request %d", self.request_count)

nat/front_ends/mcp/register.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import AsyncIterator
+
+from nat.cli.register_workflow import register_front_end
+from nat.data_models.config import Config
+from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
+
+
+@register_front_end(config_type=MCPFrontEndConfig)
+async def register_mcp_front_end(config: MCPFrontEndConfig, full_config: Config) -> AsyncIterator:
+    from nat.front_ends.mcp.mcp_front_end_plugin import MCPFrontEndPlugin
+
+    yield MCPFrontEndPlugin(full_config=full_config)