nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.3.0a20250917__py3-none-any.whl

nat/data_models/temperature_mixin.py

@@ -16,9 +16,10 @@
 import re
 
 from pydantic import BaseModel
-from pydantic import Field
 
 from nat.data_models.gated_field_mixin import GatedFieldMixin
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import SearchSpace
 
 
 class TemperatureMixin(
@@ -35,9 +36,9 @@ class TemperatureMixin(
     Attributes:
         temperature: Sampling temperature in [0, 1]. Defaults to 0.0 when supported on the model.
     """
-    temperature: float | None = Field(
+    temperature: float | None = OptimizableField(
         default=None,
         ge=0.0,
         le=1.0,
         description="Sampling temperature in [0, 1]. Defaults to 0.0 when supported on the model.",
-    )
+        space=SearchSpace(high=0.9, low=0.1, step=0.2))

nat/data_models/top_p_mixin.py

@@ -16,9 +16,10 @@
 import re
 
 from pydantic import BaseModel
-from pydantic import Field
 
 from nat.data_models.gated_field_mixin import GatedFieldMixin
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import SearchSpace
 
 
 class TopPMixin(
@@ -35,9 +36,9 @@ class TopPMixin(
     Attributes:
         top_p: Top-p for distribution sampling. Defaults to 1.0 when supported on the model.
     """
-    top_p: float | None = Field(
+    top_p: float | None = OptimizableField(
         default=None,
         ge=0.0,
         le=1.0,
         description="Top-p for distribution sampling. Defaults to 1.0 when supported on the model.",
-    )
+        space=SearchSpace(high=1.0, low=0.5, step=0.1))

nat/embedder/nim_embedder.py

@@ -50,7 +50,7 @@ class NIMEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="nim"):
                                        description=("The truncation strategy if the input on the "
                                                     "server side if it's too large."))
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
 
 @register_embedder_provider(config_type=NIMEmbedderModelConfig)

nat/embedder/openai_embedder.py

@@ -27,7 +27,7 @@ from nat.data_models.retry_mixin import RetryMixin
 class OpenAIEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
     api_key: str | None = Field(default=None, description="OpenAI API key to interact with hosted model.")
     base_url: str | None = Field(default=None, description="Base url to the hosted model.")

nat/eval/config.py

@@ -27,7 +27,7 @@ class EvaluationRunConfig(BaseModel):
     """
     Parameters used for a single evaluation run.
     """
-    config_file: Path
+    config_file: Path | BaseModel
     dataset: str | None = None  # dataset file path can be specified in the config file
     result_json_path: str = "$"
     skip_workflow: bool = False

nat/eval/evaluate.py

@@ -449,10 +449,14 @@ class EvaluationRun:
         from nat.runtime.loader import load_config
 
         # Load and override the config
-        if self.config.override:
+        config = None
+        if isinstance(self.config.config_file, BaseModel):
+            config = self.config.config_file
+        elif self.config.override:
             config = self.apply_overrides()
         else:
             config = load_config(self.config.config_file)
+
         self.eval_config = config.eval
         workflow_alias = self._get_workflow_alias(config.workflow.type)
         logger.debug("Loaded %s evaluation configuration: %s", workflow_alias, self.eval_config)

nat/eval/register.py

@@ -17,6 +17,10 @@
 
 # Import evaluators which need to be automatically registered here
 from .rag_evaluator.register import register_ragas_evaluator
+from .runtime_evaluator.register import register_avg_llm_latency_evaluator
+from .runtime_evaluator.register import register_avg_num_llm_calls_evaluator
+from .runtime_evaluator.register import register_avg_tokens_per_llm_end_evaluator
+from .runtime_evaluator.register import register_avg_workflow_runtime_evaluator
 from .swe_bench_evaluator.register import register_swe_bench_evaluator
 from .trajectory_evaluator.register import register_trajectory_evaluator
 from .tunable_rag_evaluator.register import register_tunable_rag_evaluator

nat/eval/runtime_evaluator/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/eval/runtime_evaluator/evaluate.py

@@ -0,0 +1,123 @@
+# 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 __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass
+
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.eval.evaluator.base_evaluator import BaseEvaluator
+from nat.eval.evaluator.evaluator_model import EvalInputItem
+from nat.eval.evaluator.evaluator_model import EvalOutputItem
+from nat.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
+
+
+@dataclass
+class _CallTiming:
+    start_ts: float | None = None
+    end_ts: float | None = None
+
+    @property
+    def latency(self) -> float | None:
+        if self.start_ts is None or self.end_ts is None:
+            return None
+        return max(0.0, self.end_ts - self.start_ts)
+
+
+class AverageLLMLatencyEvaluator(BaseEvaluator):
+    """
+    Mean difference between connected LLM_START and LLM_END events (same UUID).
+    The score is the average latency in seconds for the item. Reasoning contains per-call latencies.
+    """
+
+    def __init__(self, max_concurrency: int = 8):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Avg LLM Latency")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:  # noqa: D401
+        calls: dict[str, _CallTiming] = defaultdict(_CallTiming)
+
+        for step in (IntermediatePropertyAdaptor.from_intermediate_step(s) for s in item.trajectory):
+            if step.event_type == IntermediateStepType.LLM_START:
+                calls[step.UUID].start_ts = step.event_timestamp
+            elif step.event_type == IntermediateStepType.LLM_END:
+                calls[step.UUID].end_ts = step.event_timestamp
+
+        latencies = [ct.latency for ct in calls.values() if ct.latency is not None]
+        avg_latency = sum(latencies) / len(latencies) if latencies else 0.0
+
+        reasoning = {
+            "num_llm_calls": len(latencies),
+            "latencies": latencies,
+        }
+        return EvalOutputItem(id=item.id, score=round(avg_latency, 4), reasoning=reasoning)
+
+
+class AverageWorkflowRuntimeEvaluator(BaseEvaluator):
+    """
+    Average workflow runtime per item: max(event_timestamp) - min(event_timestamp) across the trajectory.
+    The score is the runtime in seconds for the item.
+    """
+
+    def __init__(self, max_concurrency: int = 8):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Avg Workflow Runtime")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:  # noqa: D401
+        if not item.trajectory:
+            return EvalOutputItem(id=item.id, score=0.0, reasoning={"note": "no steps"})
+
+        timestamps = [s.event_timestamp for s in item.trajectory]
+        runtime = max(timestamps) - min(timestamps)
+        return EvalOutputItem(id=item.id, score=round(max(0.0, runtime), 4), reasoning={"steps": len(timestamps)})
+
+
+class AverageNumberOfLLMCallsEvaluator(BaseEvaluator):
+    """
+    Average number of LLM calls per item. The score is the count for the item.
+    """
+
+    def __init__(self, max_concurrency: int = 8):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Avg # LLM Calls")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:  # noqa: D401
+        num_calls = sum(1 for s in item.trajectory if s.event_type == IntermediateStepType.LLM_END)
+        return EvalOutputItem(id=item.id, score=float(num_calls), reasoning={"num_llm_end": num_calls})
+
+
+class AverageTokensPerLLMEndEvaluator(BaseEvaluator):
+    """
+    Average total tokens per LLM_END event: sum of prompt and completion tokens if available.
+    The score is the average tokens per LLM_END for the item (0 if none).
+    """
+
+    def __init__(self, max_concurrency: int = 8):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Avg Tokens/LLM_END")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:  # noqa: D401
+        totals: list[int] = []
+        for step in (IntermediatePropertyAdaptor.from_intermediate_step(s) for s in item.trajectory):
+            if step.event_type == IntermediateStepType.LLM_END:
+                total_tokens = step.token_usage.total_tokens
+                # If framework doesn't set total, compute from prompt+completion
+                if total_tokens == 0:
+                    total_tokens = step.token_usage.prompt_tokens + step.token_usage.completion_tokens
+                totals.append(total_tokens)
+
+        avg_tokens = (sum(totals) / len(totals)) if totals else 0.0
+        reasoning = {
+            "num_llm_end": len(totals),
+            "totals": totals,
+        }
+        return EvalOutputItem(id=item.id, score=round(avg_tokens, 2), reasoning=reasoning)

nat/eval/runtime_evaluator/register.py

@@ -0,0 +1,100 @@
+# 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 pydantic import Field
+
+from nat.builder.builder import EvalBuilder
+from nat.builder.evaluator import EvaluatorInfo
+from nat.cli.register_workflow import register_evaluator
+from nat.data_models.evaluator import EvaluatorBaseConfig
+from nat.eval.evaluator.evaluator_model import EvalInput
+from nat.eval.evaluator.evaluator_model import EvalOutput
+
+
+class AverageLLMLatencyConfig(EvaluatorBaseConfig, name="avg_llm_latency"):
+    """Mean difference between connected LLM_START and LLM_END events (same UUID)."""
+
+    max_concurrency: int = Field(default=8, description="Max concurrency for evaluation.")
+
+
+class AverageWorkflowRuntimeConfig(EvaluatorBaseConfig, name="avg_workflow_runtime"):
+    """Average workflow runtime per item (max timestamp - min timestamp)."""
+
+    max_concurrency: int = Field(default=8, description="Max concurrency for evaluation.")
+
+
+class AverageNumberOfLLMCallsConfig(EvaluatorBaseConfig, name="avg_num_llm_calls"):
+    """Average number of LLM calls per item (count of LLM_END)."""
+
+    max_concurrency: int = Field(default=8, description="Max concurrency for evaluation.")
+
+
+class AverageTokensPerLLMEndConfig(EvaluatorBaseConfig, name="avg_tokens_per_llm_end"):
+    """Average total tokens per LLM_END event (prompt + completion if available)."""
+
+    max_concurrency: int = Field(default=8, description="Max concurrency for evaluation.")
+
+
+@register_evaluator(config_type=AverageLLMLatencyConfig)
+async def register_avg_llm_latency_evaluator(config: AverageLLMLatencyConfig, builder: EvalBuilder):
+    from .evaluate import AverageLLMLatencyEvaluator
+
+    evaluator = AverageLLMLatencyEvaluator(max_concurrency=config.max_concurrency or builder.get_max_concurrency())
+
+    async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:
+        return await evaluator.evaluate(eval_input)
+
+    yield EvaluatorInfo(config=config,
+                        evaluate_fn=evaluate_fn,
+                        description="Average LLM latency (s) from LLM_START to LLM_END")
+
+
+@register_evaluator(config_type=AverageWorkflowRuntimeConfig)
+async def register_avg_workflow_runtime_evaluator(config: AverageWorkflowRuntimeConfig, builder: EvalBuilder):
+    from .evaluate import AverageWorkflowRuntimeEvaluator
+
+    evaluator = AverageWorkflowRuntimeEvaluator(max_concurrency=config.max_concurrency or builder.get_max_concurrency())
+
+    async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:
+        return await evaluator.evaluate(eval_input)
+
+    yield EvaluatorInfo(config=config, evaluate_fn=evaluate_fn, description="Average workflow runtime (s)")
+
+
+@register_evaluator(config_type=AverageNumberOfLLMCallsConfig)
+async def register_avg_num_llm_calls_evaluator(config: AverageNumberOfLLMCallsConfig, builder: EvalBuilder):
+    from .evaluate import AverageNumberOfLLMCallsEvaluator
+
+    evaluator = AverageNumberOfLLMCallsEvaluator(
+        max_concurrency=config.max_concurrency or builder.get_max_concurrency())
+
+    async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:
+        return await evaluator.evaluate(eval_input)
+
+    yield EvaluatorInfo(config=config, evaluate_fn=evaluate_fn, description="Average number of LLM calls")
+
+
+@register_evaluator(config_type=AverageTokensPerLLMEndConfig)
+async def register_avg_tokens_per_llm_end_evaluator(config: AverageTokensPerLLMEndConfig, builder: EvalBuilder):
+    from .evaluate import AverageTokensPerLLMEndEvaluator
+
+    evaluator = AverageTokensPerLLMEndEvaluator(max_concurrency=config.max_concurrency or builder.get_max_concurrency())
+
+    async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:
+        return await evaluator.evaluate(eval_input)
+
+    yield EvaluatorInfo(config=config,
+                        evaluate_fn=evaluate_fn,
+                        description="Average total tokens per LLM_END (prompt + completion)")

nat/experimental/test_time_compute/functions/plan_select_execute_function.py

@@ -97,7 +97,11 @@ async def plan_select_execute_function(config: PlanSelectExecuteFunctionConfig,
                          f"function without a description.")
 
     # Get the function dependencies of the augmented function
-    function_used_tools = builder.get_function_dependencies(config.augmented_fn).functions
+    function_dependencies = builder.get_function_dependencies(config.augmented_fn)
+    function_used_tools = set(function_dependencies.functions)
+    for function_group in function_dependencies.function_groups:
+        function_used_tools.update(builder.get_function_group_dependencies(function_group).functions)
+
     tool_list = "Tool: Description\n"
 
     for tool in function_used_tools:

nat/front_ends/fastapi/dask_client_mixin.py

@@ -0,0 +1,43 @@
+# 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 typing
+from abc import ABC
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+if typing.TYPE_CHECKING:
+    from dask.distributed import Client
+
+
+class DaskClientMixin(ABC):
+
+    @asynccontextmanager
+    async def client(self, address: str) -> AsyncGenerator["Client"]:
+        """
+        Async context manager for obtaining a Dask client connection.
+
+        Yields
+        ------
+        Client
+            An active Dask client connected to the scheduler. The client is automatically closed when exiting the
+            context manager.
+        """
+        from dask.distributed import Client
+        client = await Client(address=address, asynchronous=True)
+
+        yield client
+
+        await client.close()

nat/front_ends/fastapi/fastapi_front_end_config.py

@@ -197,9 +197,20 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
     port: int = Field(default=8000, description="Port to bind the server to", ge=0, le=65535)
     reload: bool = Field(default=False, description="Enable auto-reload for development")
     workers: int = Field(default=1, description="Number of workers to run", ge=1)
-    max_running_async_jobs: int = Field(default=10,
-                                        description="Maximum number of async jobs to run concurrently",
-                                        ge=1)
+    scheduler_address: str | None = Field(
+        default=None,
+        description=("Address of the Dask scheduler to use for async jobs. If None, a Dask local cluster is created. "
+                     "Note: This requires the optional dask dependency to be installed."))
+    db_url: str | None = Field(
+        default=None,
+        description=
+        "SQLAlchemy database URL for storing async job metadata, if unset a temporary SQLite database is used.")
+    max_running_async_jobs: int = Field(
+        default=10,
+        description=(
+            "Maximum number of async jobs to run concurrently, this controls the number of dask workers created. "
+            "This parameter is only used when scheduler_address is `None` and a Dask local cluster is created."),
+        ge=1)
     step_adaptor: StepAdaptorConfig = StepAdaptorConfig()
 
     workflow: typing.Annotated[EndpointBase, Field(description="Endpoint for the default workflow.")] = EndpointBase(

nat/front_ends/fastapi/fastapi_front_end_plugin.py

@@ -13,21 +13,36 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import logging
 import os
+import sys
 import tempfile
 import typing
 
 from nat.builder.front_end import FrontEndBase
+from nat.front_ends.fastapi.dask_client_mixin import DaskClientMixin
 from nat.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
 from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorkerBase
 from nat.front_ends.fastapi.main import get_app
+from nat.front_ends.fastapi.utils import get_class_name
 from nat.utils.io.yaml_tools import yaml_dump
 
+if (typing.TYPE_CHECKING):
+    from nat.data_models.config import Config
+
 logger = logging.getLogger(__name__)
 
 
-class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
+class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]):
+
+    def __init__(self, full_config: "Config"):
+        super().__init__(full_config)
+
+        # This attribute is set if dask is installed, and an external cluster is not used (scheduler_address is None)
+        self._cluster = None
+        self._periodic_cleanup_future = None
+        self._scheduler_address = None
 
     def get_worker_class(self) -> type[FastApiFrontEndPluginWorkerBase]:
         from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorker
@@ -42,7 +57,38 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
 
         worker_class = self.get_worker_class()
 
-        return f"{worker_class.__module__}.{worker_class.__qualname__}"
+        return get_class_name(worker_class)
+
+    @staticmethod
+    async def _periodic_cleanup(scheduler_address: str,
+                                db_url: str,
+                                sleep_time_sec: int = 300,
+                                log_level: int = logging.INFO):
+        from nat.front_ends.fastapi.job_store import JobStore
+
+        job_store = JobStore(scheduler_address=scheduler_address, db_url=db_url)
+
+        logging.basicConfig(level=log_level)
+        logger.info("Starting periodic cleanup of expired jobs every %d seconds", sleep_time_sec)
+        while True:
+            await asyncio.sleep(sleep_time_sec)
+
+            try:
+                await job_store.cleanup_expired_jobs()
+                logger.debug("Expired jobs cleaned up")
+            except:  # noqa: E722
+                logger.exception("Error during job cleanup")
+
+    async def _submit_cleanup_task(self, scheduler_address: str, db_url: str):
+        """Submit a cleanup task to the cluster to remove the job after expiry."""
+        logger.info("Submitting periodic cleanup task to Dask cluster at %s", scheduler_address)
+        async with self.client(self._scheduler_address) as client:
+            self._periodic_cleanup_future = client.submit(self._periodic_cleanup,
+                                                          scheduler_address=self._scheduler_address,
+                                                          db_url=db_url,
+                                                          log_level=logger.getEffectiveLevel())
+
+        logger.info("Submitted periodic cleanup task to Dask cluster at %s", scheduler_address)
 
     async def run(self):
 
@@ -52,6 +98,45 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
             # Get as dict
             config_dict = self.full_config.model_dump(mode="json", by_alias=True, round_trip=True)
 
+            # Three possible cases:
+            # 1. Dask is installed and scheduler_address is None, we create a LocalCluster
+            # 2. Dask is installed and scheduler_address is set, we use the existing cluster
+            # 3. Dask is not installed, we skip the cluster setup
+            self._scheduler_address = self.front_end_config.scheduler_address
+            if self._scheduler_address is None:
+                try:
+                    from dask.distributed import LocalCluster
+
+                    self._cluster = LocalCluster(n_workers=self.front_end_config.max_running_async_jobs,
+                                                 threads_per_worker=1)
+
+                    self._scheduler_address = self._cluster.scheduler.address
+                    logger.info("Created local Dask cluster with scheduler at %s", self._scheduler_address)
+
+                except ImportError:
+                    logger.warning("Dask is not installed, async execution and evaluation will not be available.")
+
+            if self._scheduler_address is not None:
+                # If we are here then either the user provided a scheduler address, or we created a LocalCluster
+
+                from nat.front_ends.fastapi.job_store import Base
+                from nat.front_ends.fastapi.job_store import get_db_engine
+
+                db_engine = get_db_engine(self.front_end_config.db_url, use_async=True)
+                async with db_engine.begin() as conn:
+                    await conn.run_sync(Base.metadata.create_all, checkfirst=True)  # create tables if they do not exist
+
+                # If self.front_end_config.db_url is None, then we need to get the actual url from the engine
+                db_url = str(db_engine.url)
+                await self._submit_cleanup_task(scheduler_address=self._scheduler_address, db_url=db_url)
+
+                # Set environment variabls such that the worker subprocesses will know how to connect to dask and to
+                # the database
+                os.environ.update({
+                    "NAT_DASK_SCHEDULER_ADDRESS": self._scheduler_address,
+                    "NAT_JOB_STORE_DB_URL": db_url,
+                })
+
             # Write to YAML file
             yaml_dump(config_dict, config_file)
 
@@ -70,13 +155,25 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
 
                 reload_excludes = ["./.*"]
 
+                # By default, Uvicorn uses "auto" event loop policy, which prefers `uvloop` if installed. However,
+                # uvloop’s event loop policy for macOS doesn’t provide a child watcher (which is needed for MCP server),
+                # so setting loop="asyncio" forces Uvicorn to use the standard event loop, which includes child-watcher
+                # support.
+                if sys.platform == "darwin" or sys.platform.startswith("linux"):
+                    # For macOS
+                    event_loop_policy = "asyncio"
+                else:
+                    # For non-macOS platforms
+                    event_loop_policy = "auto"
+
                 uvicorn.run("nat.front_ends.fastapi.main:get_app",
                             host=self.front_end_config.host,
                             port=self.front_end_config.port,
                             workers=self.front_end_config.workers,
                             reload=self.front_end_config.reload,
                             factory=True,
-                            reload_excludes=reload_excludes)
+                            reload_excludes=reload_excludes,
+                            loop=event_loop_policy)
 
             else:
                 app = get_app()
@@ -110,6 +207,17 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
                 StandaloneApplication(app, options=options).run()
 
         finally:
+            logger.debug("Shutting down")
+            if self._periodic_cleanup_future is not None:
+                logger.info("Cancelling periodic cleanup task.")
+                # Use the scheduler address, because self._cluster is None if an external cluster is used
+                async with self.client(self._scheduler_address) as client:
+                    await client.cancel([self._periodic_cleanup_future], asynchronous=True, force=True)
+
+            if self._cluster is not None:
+                # Only shut down the cluster if we created it
+                logger.info("Closing Local Dask cluster.")
+                self._cluster.close()
             try:
                 os.remove(config_file_name)
             except OSError as e: