nvidia-nat 1.4.0a20251120__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

nat/front_ends/fastapi/async_job.py

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, 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.
+"""
+The functions in this module are intentionally written to be submitted as Dask tasks, as such they are self-contained.
+"""
+
+import asyncio
+import logging
+import typing
+
+
+def _configure_logging(configure_logging: bool, log_level: int) -> logging.Logger:
+    from nat.utils.log_utils import setup_logging
+    if configure_logging:
+        setup_logging(log_level)
+
+    return logging.getLogger(__name__)
+
+
+async def run_generation(configure_logging: bool,
+                         log_level: int,
+                         scheduler_address: str,
+                         db_url: str,
+                         config_file_path: str,
+                         job_id: str,
+                         payload: typing.Any):
+    """
+    Background async task to run the workflow.
+
+    Parameters
+    ----------
+    configure_logging : bool
+        Whether to configure logging.
+    log_level : int
+        The log level to use when `configure_logging` is `True`, ignored otherwise.
+    scheduler_address : str
+        The Dask scheduler address.
+    db_url : str
+        The database URL for the job store.
+    config_file_path : str
+        The path to the workflow configuration file.
+    job_id : str
+        The job ID.
+    payload : typing.Any
+        The input payload for the workflow.
+    """
+    from nat.front_ends.fastapi.job_store import JobStatus
+    from nat.front_ends.fastapi.job_store import JobStore
+    from nat.front_ends.fastapi.response_helpers import generate_single_response
+    from nat.runtime.loader import load_workflow
+
+    logger = _configure_logging(configure_logging, log_level)
+
+    job_store = None
+    try:
+        job_store = JobStore(scheduler_address=scheduler_address, db_url=db_url)
+        await job_store.update_status(job_id, JobStatus.RUNNING)
+        async with load_workflow(config_file_path) as local_session_manager:
+            async with local_session_manager.session() as session:
+                result = await generate_single_response(payload,
+                                                        session,
+                                                        result_type=session.workflow.single_output_schema)
+
+        await job_store.update_status(job_id, JobStatus.SUCCESS, output=result)
+    except asyncio.CancelledError:
+        logger.info("Async job %s cancelled", job_id)
+        if job_store is not None:
+            await job_store.update_status(job_id, JobStatus.INTERRUPTED, error="cancelled")
+    except Exception as e:
+        logger.exception("Error in async job %s", job_id)
+        if job_store is not None:
+            await job_store.update_status(job_id, JobStatus.FAILURE, error=str(e))
+
+
+async def periodic_cleanup(*,
+                           scheduler_address: str,
+                           db_url: str,
+                           sleep_time_sec: int = 300,
+                           configure_logging: bool = True,
+                           log_level: int = logging.INFO):
+    """
+    Dask task to periodically clean up expired jobs from the job store. This task is intended to be submitted only
+    once to the Dask cluster and run indefinitely.
+
+    Parameters
+    ----------
+    scheduler_address : str
+        The Dask scheduler address.
+    db_url : str
+        The database URL for the job store.
+    sleep_time_sec : int
+        The sleep time between cleanup operations in seconds.
+    configure_logging : bool
+        Whether to configure logging.
+    log_level : int
+        The log level to use when `configure_logging` is `True`, ignored otherwise.
+    """
+    from nat.front_ends.fastapi.job_store import JobStore
+
+    logger = _configure_logging(configure_logging, log_level)
+
+    job_store = None
+
+    logger.info("Starting periodic cleanup of expired jobs every %d seconds", sleep_time_sec)
+    while True:
+        await asyncio.sleep(sleep_time_sec)
+
+        try:
+            if job_store is None:
+                job_store = JobStore(scheduler_address=scheduler_address, db_url=db_url)
+
+            num_expired = await job_store.cleanup_expired_jobs()
+            logger.info("Expired jobs cleaned up: %d", num_expired)
+        except:  # noqa: E722
+            logger.exception("Error during job cleanup")
+            job_store = None  # Reset job store to attempt re-creation on next iteration

nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -49,14 +49,18 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
     def __init__(self,
                  add_flow_cb: Callable[[str, FlowState], Awaitable[None]],
                  remove_flow_cb: Callable[[str], Awaitable[None]],
-                 web_socket_message_handler: WebSocketMessageHandler):
+                 web_socket_message_handler: WebSocketMessageHandler,
+                 auth_timeout_seconds: float = 300.0):
 
         self._add_flow_cb: Callable[[str, FlowState], Awaitable[None]] = add_flow_cb
         self._remove_flow_cb: Callable[[str], Awaitable[None]] = remove_flow_cb
         self._web_socket_message_handler: WebSocketMessageHandler = web_socket_message_handler
+        self._auth_timeout_seconds: float = auth_timeout_seconds
 
-    async def authenticate(self, config: OAuth2AuthCodeFlowProviderConfig,
-                           method: AuthFlowType) -> AuthenticatedContext:
+    async def authenticate(
+            self,
+            config: OAuth2AuthCodeFlowProviderConfig,  # type: ignore[override]
+            method: AuthFlowType) -> AuthenticatedContext:
         if method == AuthFlowType.OAUTH2_AUTHORIZATION_CODE:
             return await self._handle_oauth2_auth_code_flow(config)
 
@@ -65,7 +69,7 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
     def create_oauth_client(self, config: OAuth2AuthCodeFlowProviderConfig) -> AsyncOAuth2Client:
         try:
             return AsyncOAuth2Client(client_id=config.client_id,
-                                     client_secret=config.client_secret,
+                                     client_secret=config.client_secret.get_secret_value(),
                                      redirect_uri=config.redirect_uri,
                                      scope=" ".join(config.scopes) if config.scopes else None,
                                      token_endpoint=config.token_url,
@@ -80,8 +84,8 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
                                   client: AsyncOAuth2Client,
                                   config: OAuth2AuthCodeFlowProviderConfig,
                                   state: str,
-                                  verifier: str = None,
-                                  challenge: str = None) -> str:
+                                  verifier: str | None = None,
+                                  challenge: str | None = None) -> str:
         """
         Create OAuth authorization URL with proper error handling.
 
@@ -129,9 +133,9 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
         await self._web_socket_message_handler.create_websocket_message(_HumanPromptOAuthConsent(text=authorization_url)
                                                                         )
         try:
-            token = await asyncio.wait_for(flow_state.future, timeout=300)
+            token = await asyncio.wait_for(flow_state.future, timeout=self._auth_timeout_seconds)
         except TimeoutError as exc:
-            raise RuntimeError("Authentication flow timed out after 5 minutes.") from exc
+            raise RuntimeError(f"Authentication flow timed out after {self._auth_timeout_seconds} seconds.") from exc
         finally:
 
             await self._remove_flow_cb(state)

nat/front_ends/fastapi/dask_client_mixin.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/fastapi_front_end_config.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/fastapi_front_end_controller.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/fastapi_front_end_plugin.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
+import copy
 import logging
 import os
 import sys
@@ -21,6 +21,7 @@ import tempfile
 import typing
 
 from nat.builder.front_end import FrontEndBase
+from nat.front_ends.fastapi.async_job import periodic_cleanup
 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
@@ -28,6 +29,7 @@ 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
+from nat.utils.log_utils import LOG_DATE_FORMAT
 
 if (typing.TYPE_CHECKING):
     from nat.data_models.config import Config
@@ -44,6 +46,7 @@ class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]
         self._cluster = None
         self._periodic_cleanup_future = None
         self._scheduler_address = None
+        self._use_dask_threads = False
 
     def get_worker_class(self) -> type[FastApiFrontEndPluginWorkerBase]:
         from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorker
@@ -60,34 +63,15 @@ class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]
 
         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)
+        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,
+            self._periodic_cleanup_future = client.submit(periodic_cleanup,
                                                           scheduler_address=self._scheduler_address,
                                                           db_url=db_url,
-                                                          log_level=log_level)
+                                                          log_level=log_level,
+                                                          configure_logging=not self._use_dask_threads)
 
     @staticmethod
     def _setup_worker():
@@ -99,6 +83,7 @@ class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]
         os.setsid()
 
     async def run(self):
+        log_level = logger.getEffectiveLevel()
 
         # Write the entire config to a temporary file
         with tempfile.NamedTemporaryFile(mode="w", prefix="nat_config", suffix=".yml", delete=False) as config_file:
@@ -120,17 +105,17 @@ class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]
 
                     from dask.distributed import LocalCluster
 
-                    use_threads = self.front_end_config.dask_workers == 'threads'
+                    self._use_dask_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,
+                    self._cluster = LocalCluster(processes=not self._use_dask_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":
+                    if not self._use_dask_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)
@@ -156,13 +141,15 @@ class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]
                 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)
+                                                log_level=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,
+                    "NAT_USE_DASK_THREADS": str(int(self._use_dask_threads)),
+                    "NAT_FASTAPI_LOG_LEVEL": str(log_level),
                 })
 
             # Write to YAML file
@@ -194,6 +181,12 @@ class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]
                     # For non-macOS platforms
                     event_loop_policy = "auto"
 
+                # Start with the default uvicorn logging config, but override with our desired format
+                log_config = copy.deepcopy(uvicorn.config.LOGGING_CONFIG)
+                for formatter in log_config.get("formatters", {}).values():
+                    formatter["fmt"] = f"%(asctime)s - {formatter['fmt']}"
+                    formatter["datefmt"] = LOG_DATE_FORMAT
+
                 uvicorn.run("nat.front_ends.fastapi.main:get_app",
                             host=self.front_end_config.host,
                             port=self.front_end_config.port,
@@ -201,7 +194,9 @@ class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]
                             reload=self.front_end_config.reload,
                             factory=True,
                             reload_excludes=reload_excludes,
-                            loop=event_loop_policy)
+                            loop=event_loop_policy,
+                            log_level=log_level,
+                            log_config=log_config)
 
             else:
                 app = get_app()