nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__py3-none-any.whl

nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py

@@ -22,6 +22,7 @@ from dataclasses import dataclass
 from dataclasses import field
 
 import pkce
+from authlib.common.errors import AuthlibBaseError as OAuthError
 from authlib.integrations.httpx_client import AsyncOAuth2Client
 
 from nat.authentication.interfaces import FlowHandlerBase
@@ -61,14 +62,50 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
 
         raise NotImplementedError(f"Authentication method '{method}' is not supported by the websocket frontend.")
 
-    def create_oauth_client(self, config: OAuth2AuthCodeFlowProviderConfig):
-        return AsyncOAuth2Client(client_id=config.client_id,
-                                 client_secret=config.client_secret,
-                                 redirect_uri=config.redirect_uri,
-                                 scope=" ".join(config.scopes) if config.scopes else None,
-                                 token_endpoint=config.token_url,
-                                 code_challenge_method='S256' if config.use_pkce else None,
-                                 token_endpoint_auth_method=config.token_endpoint_auth_method)
+    def create_oauth_client(self, config: OAuth2AuthCodeFlowProviderConfig) -> AsyncOAuth2Client:
+        try:
+            return AsyncOAuth2Client(client_id=config.client_id,
+                                     client_secret=config.client_secret,
+                                     redirect_uri=config.redirect_uri,
+                                     scope=" ".join(config.scopes) if config.scopes else None,
+                                     token_endpoint=config.token_url,
+                                     code_challenge_method='S256' if config.use_pkce else None,
+                                     token_endpoint_auth_method=config.token_endpoint_auth_method)
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Invalid OAuth2 configuration: {e}") from e
+        except Exception as e:
+            raise RuntimeError(f"Failed to create OAuth2 client: {e}") from e
+
+    def _create_authorization_url(self,
+                                  client: AsyncOAuth2Client,
+                                  config: OAuth2AuthCodeFlowProviderConfig,
+                                  state: str,
+                                  verifier: str = None,
+                                  challenge: str = None) -> str:
+        """
+        Create OAuth authorization URL with proper error handling.
+
+        Args:
+            client: The OAuth2 client instance
+            config: OAuth2 configuration
+            state: OAuth state parameter
+            verifier: PKCE verifier (if using PKCE)
+            challenge: PKCE challenge (if using PKCE)
+
+        Returns:
+            The authorization URL
+        """
+        try:
+            authorization_url, _ = client.create_authorization_url(
+                config.authorization_url,
+                state=state,
+                code_verifier=verifier if config.use_pkce else None,
+                code_challenge=challenge if config.use_pkce else None,
+                **(config.authorization_kwargs or {})
+            )
+            return authorization_url
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Error creating OAuth authorization URL: {e}") from e
 
     async def _handle_oauth2_auth_code_flow(self, config: OAuth2AuthCodeFlowProviderConfig) -> AuthenticatedContext:
 
@@ -82,21 +119,19 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
             flow_state.verifier = verifier
             flow_state.challenge = challenge
 
-        authorization_url, _ = flow_state.client.create_authorization_url(
-            config.authorization_url,
-            state=state,
-            code_verifier=flow_state.verifier if config.use_pkce else None,
-            code_challenge=flow_state.challenge if config.use_pkce else None,
-            **(config.authorization_kwargs or {})
-        )
+        authorization_url = self._create_authorization_url(client=flow_state.client,
+                                                           config=config,
+                                                           state=state,
+                                                           verifier=flow_state.verifier,
+                                                           challenge=flow_state.challenge)
 
         await self._add_flow_cb(state, flow_state)
         await self._web_socket_message_handler.create_websocket_message(_HumanPromptOAuthConsent(text=authorization_url)
                                                                         )
         try:
             token = await asyncio.wait_for(flow_state.future, timeout=300)
-        except asyncio.TimeoutError:
-            raise RuntimeError("Authentication flow timed out after 5 minutes.")
+        except TimeoutError as exc:
+            raise RuntimeError("Authentication flow timed out after 5 minutes.") from exc
         finally:
 
             await self._remove_flow_cb(state)

nat/front_ends/fastapi/dask_client_mixin.py

@@ -0,0 +1,65 @@
+# 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 collections.abc import Generator
+from contextlib import asynccontextmanager
+from contextlib import contextmanager
+
+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.
+
+        Yields
+        ------
+        Client
+            An async 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)
+
+        try:
+            yield client
+        finally:
+            await client.close()
+
+    @contextmanager
+    def blocking_client(self, address: str) -> Generator["Client"]:
+        """
+        context manager for obtaining a blocking Dask client.
+
+        Yields
+        ------
+        Client
+            A blocking Dask client connected to the scheduler. The client is automatically closed when exiting the
+            context manager.
+        """
+        from dask.distributed import Client
+        client = Client(address=address)
+
+        try:
+            yield client
+        finally:
+            client.close()

nat/front_ends/fastapi/fastapi_front_end_config.py

@@ -197,9 +197,31 @@ 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)
+    dask_workers: typing.Literal["threads", "processes"] = Field(
+        default="processes",
+        description=(
+            "Type of Dask workers to use. Options are 'threads' for Threaded Dask workers or 'processes' for "
+            "Process based Dask workers. This parameter is only used when scheduler_address is `None` and a local Dask "
+            "cluster is created."),
+    )
+    dask_log_level: str = Field(
+        default="WARNING",
+        description="Logging level for Dask.",
+    )
     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,37 @@
 # 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
+from nat.utils.log_levels import LOG_LEVELS
+
+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 +58,45 @@ 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, log_level: int = logging.INFO):
+        """Submit a cleanup task to the cluster to remove the job after expiry."""
+        logger.debug("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=log_level)
+
+    @staticmethod
+    def _setup_worker():
+        """
+        Setup function to be run in each worker process. This moves each worker into it's own process group.
+        This fixes an issue where a Ctrl-C in the terminal sends a SIGINT to all workers, which then causes the
+        workers to exit before the main process can shutdown the cluster gracefully.
+        """
+        os.setsid()
 
     async def run(self):
 
@@ -52,6 +106,65 @@ 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
+            dask_log_level = LOG_LEVELS.get(self.front_end_config.dask_log_level.upper(), logging.WARNING)
+            dask_logger = logging.getLogger("distributed")
+            dask_logger.setLevel(dask_log_level)
+
+            self._scheduler_address = self.front_end_config.scheduler_address
+            if self._scheduler_address is None:
+                try:
+
+                    from dask.distributed import LocalCluster
+
+                    use_threads = self.front_end_config.dask_workers == 'threads'
+
+                    # set n_workers to max_running_async_jobs + 1 to allow for one worker to handle the cleanup task
+                    self._cluster = LocalCluster(processes=not use_threads,
+                                                 silence_logs=dask_log_level,
+                                                 protocol="tcp",
+                                                 n_workers=self.front_end_config.max_running_async_jobs + 1)
+
+                    self._scheduler_address = self._cluster.scheduler.address
+
+                    if not use_threads and sys.platform != "win32":
+                        with self.blocking_client(self._scheduler_address) as client:
+                            # Client.run submits a function to be run on each worker
+                            client.run(self._setup_worker)
+
+                    logger.info("Created local Dask cluster with scheduler at %s using %s workers",
+                                self._scheduler_address,
+                                self.front_end_config.dask_workers)
+
+                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,
+                                                log_level=dask_log_level)
+
+                # 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 +183,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 +235,18 @@ 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.debug("Closing Local Dask cluster.")
+                self._cluster.close()
+
             try:
                 os.remove(config_file_name)
             except OSError as e: