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

nat/cli/type_registry.py

@@ -64,6 +64,8 @@ from nat.data_models.logging import LoggingBaseConfig
 from nat.data_models.logging import LoggingMethodConfigT
 from nat.data_models.memory import MemoryBaseConfig
 from nat.data_models.memory import MemoryBaseConfigT
+from nat.data_models.middleware import MiddlewareBaseConfig
+from nat.data_models.middleware import MiddlewareBaseConfigT
 from nat.data_models.object_store import ObjectStoreBaseConfig
 from nat.data_models.object_store import ObjectStoreBaseConfigT
 from nat.data_models.registry_handler import RegistryHandlerBaseConfig
@@ -76,6 +78,7 @@ from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfigT
 from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
 from nat.memory.interfaces import MemoryEditor
+from nat.middleware.middleware import Middleware
 from nat.object_store.interfaces import ObjectStore
 from nat.observability.exporter.base_exporter import BaseExporter
 from nat.registry_handlers.registry_handler_base import AbstractRegistryHandler
@@ -89,6 +92,7 @@ EvaluatorBuildCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AsyncIte
 FrontEndBuildCallableT = Callable[[FrontEndConfigT, Config], AsyncIterator[FrontEndBase]]
 FunctionBuildCallableT = Callable[[FunctionConfigT, Builder], AsyncIterator[FunctionInfo | Callable | FunctionBase]]
 FunctionGroupBuildCallableT = Callable[[FunctionGroupConfigT, Builder], AsyncIterator[FunctionGroup]]
+MiddlewareBuildCallableT = Callable[[MiddlewareBaseConfigT, Builder], AsyncIterator[Middleware]]
 TTCStrategyBuildCallableT = Callable[[TTCStrategyBaseConfigT, Builder], AsyncIterator[StrategyBase]]
 LLMClientBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[typing.Any]]
 LLMProviderBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[LLMProviderInfo]]
@@ -111,6 +115,7 @@ FrontEndRegisteredCallableT = Callable[[FrontEndConfigT, Config], AbstractAsyncC
 FunctionRegisteredCallableT = Callable[[FunctionConfigT, Builder],
                                        AbstractAsyncContextManager[FunctionInfo | Callable | FunctionBase]]
 FunctionGroupRegisteredCallableT = Callable[[FunctionGroupConfigT, Builder], AbstractAsyncContextManager[FunctionGroup]]
+MiddlewareRegisteredCallableT = Callable[[MiddlewareBaseConfigT, Builder], AbstractAsyncContextManager[Middleware]]
 TTCStrategyRegisterCallableT = Callable[[TTCStrategyBaseConfigT, Builder], AbstractAsyncContextManager[StrategyBase]]
 LLMClientRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
 LLMProviderRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[LLMProviderInfo]]
@@ -179,6 +184,8 @@ class RegisteredFunctionInfo(RegisteredInfo[FunctionBaseConfig]):
     and a description.
     """
 
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
+
     build_fn: FunctionRegisteredCallableT = Field(repr=False)
     framework_wrappers: list[str] = Field(default_factory=list)
 
@@ -193,6 +200,15 @@ class RegisteredFunctionGroupInfo(RegisteredInfo[FunctionGroupBaseConfig]):
     framework_wrappers: list[str] = Field(default_factory=list)
 
 
+class RegisteredMiddlewareInfo(RegisteredInfo[MiddlewareBaseConfig]):
+    """
+    Represents registered middleware. Middleware provides middleware-style wrapping of
+    calls with preprocessing and postprocessing logic.
+    """
+
+    build_fn: MiddlewareRegisteredCallableT = Field(repr=False)
+
+
 class RegisteredLLMProviderInfo(RegisteredInfo[LLMBaseConfig]):
     """
     Represents a registered LLM provider. LLM Providers are the operators of the LLMs. i.e. NIMs, OpenAI, Anthropic,
@@ -331,6 +347,9 @@ class TypeRegistry:
         # Function Groups
         self._registered_function_groups: dict[type[FunctionGroupBaseConfig], RegisteredFunctionGroupInfo] = {}
 
+        # Middleware
+        self._registered_middleware: dict[type[MiddlewareBaseConfig], RegisteredMiddlewareInfo] = {}
+
         # LLMs
         self._registered_llm_provider_infos: dict[type[LLMBaseConfig], RegisteredLLMProviderInfo] = {}
         self._llm_client_provider_to_framework: dict[type[LLMBaseConfig], dict[str, RegisteredLLMClientInfo]] = {}
@@ -540,6 +559,49 @@ class TypeRegistry:
         """
         return list(self._registered_function_groups.values())
 
+    def register_middleware(self, registration: RegisteredMiddlewareInfo):
+        """Register middleware with the type registry.
+
+        Args:
+            registration: The middleware registration information
+
+        Raises:
+            ValueError: If middleware with the same config type is already registered
+        """
+        if (registration.config_type in self._registered_middleware):
+            raise ValueError(f"Middleware with the same config type `{registration.config_type}` has already been "
+                             "registered.")
+
+        self._registered_middleware[registration.config_type] = registration
+
+        self._registration_changed()
+
+    def get_middleware(self, config_type: type[MiddlewareBaseConfig]) -> RegisteredMiddlewareInfo:
+        """Get registered middleware by its config type.
+
+        Args:
+            config_type: The middleware configuration type
+
+        Returns:
+            RegisteredMiddlewareInfo: The registered middleware information
+
+        Raises:
+            KeyError: If no middleware is registered for the given config type
+        """
+        try:
+            return self._registered_middleware[config_type]
+        except KeyError as err:
+            raise KeyError(f"Could not find registered middleware for config `{config_type}`. "
+                           f"Registered configs: {set(self._registered_middleware.keys())}") from err
+
+    def get_registered_middleware(self) -> list[RegisteredInfo[MiddlewareBaseConfig]]:
+        """Get all registered middleware.
+
+        Returns:
+            list[RegisteredInfo[MiddlewareBaseConfig]]: List of all registered middleware
+        """
+        return list(self._registered_middleware.values())
+
     def register_llm_provider(self, info: RegisteredLLMProviderInfo):
 
         if (info.config_type in self._registered_llm_provider_infos):
@@ -912,6 +974,9 @@ class TypeRegistry:
         if component_type == ComponentEnum.TTC_STRATEGY:
             return self._registered_ttc_strategies
 
+        if component_type == ComponentEnum.MIDDLEWARE:
+            return self._registered_middleware
+
         raise ValueError(f"Supplied an unsupported component type {component_type}")
 
     def get_registered_types_by_component_type(self, component_type: ComponentEnum) -> list[str]:
@@ -1038,6 +1103,9 @@ class TypeRegistry:
         if issubclass(cls, TTCStrategyBaseConfig):
             return self._do_compute_annotation(cls, self.get_registered_ttc_strategies())
 
+        if issubclass(cls, MiddlewareBaseConfig):
+            return self._do_compute_annotation(cls, self.get_registered_middleware())
+
         raise ValueError(f"Supplied an unsupported component type {cls}")
 
 

nat/data_models/component.py

@@ -28,6 +28,7 @@ class ComponentEnum(StrEnum):
     FRONT_END = "front_end"
     FUNCTION = "function"
     FUNCTION_GROUP = "function_group"
+    MIDDLEWARE = "middleware"
     TTC_STRATEGY = "ttc_strategy"
     LLM_CLIENT = "llm_client"
     LLM_PROVIDER = "llm_provider"
@@ -49,6 +50,7 @@ class ComponentGroup(StrEnum):
     EMBEDDERS = "embedders"
     FUNCTIONS = "functions"
     FUNCTION_GROUPS = "function_groups"
+    MIDDLEWARE = "middleware"
     TTC_STRATEGIES = "ttc_strategies"
     LLMS = "llms"
     MEMORY = "memory"

nat/data_models/component_ref.py

@@ -177,3 +177,14 @@ class TTCStrategyRef(ComponentRef):
     @override
     def component_group(self):
         return ComponentGroup.TTC_STRATEGIES
+
+
+class MiddlewareRef(ComponentRef):
+    """
+    A reference to middleware in a NAT configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.MIDDLEWARE

nat/data_models/config.py

@@ -43,6 +43,7 @@ from .common import TypedBaseModel
 from .embedder import EmbedderBaseConfig
 from .llm import LLMBaseConfig
 from .memory import MemoryBaseConfig
+from .middleware import FunctionMiddlewareBaseConfig
 from .object_store import ObjectStoreBaseConfig
 from .retriever import RetrieverBaseConfig
 
@@ -86,6 +87,8 @@ def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWr
                 registered_keys = GlobalTypeRegistry.get().get_registered_front_ends()
             elif (info.field_name == "ttc_strategies"):
                 registered_keys = GlobalTypeRegistry.get().get_registered_ttc_strategies()
+            elif (info.field_name == "middleware"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_middleware()
 
             else:
                 assert False, f"Unknown field name {info.field_name} in validator"
@@ -253,6 +256,9 @@ class Config(HashableBaseModel):
     # Function Groups Configuration
     function_groups: dict[str, FunctionGroupBaseConfig] = Field(default_factory=dict)
 
+    # Middleware Configuration
+    middleware: dict[str, FunctionMiddlewareBaseConfig] = Field(default_factory=dict)
+
     # LLMs Configuration
     llms: dict[str, LLMBaseConfig] = Field(default_factory=dict)
 
@@ -303,6 +309,7 @@ class Config(HashableBaseModel):
 
     @field_validator("functions",
                      "function_groups",
+                     "middleware",
                      "llms",
                      "embedders",
                      "memory",
@@ -348,6 +355,10 @@ class Config(HashableBaseModel):
                                         typing.Annotated[type_registry.compute_annotation(FunctionGroupBaseConfig),
                                                          Discriminator(TypedBaseModel.discriminator)]]
 
+        MiddlewareAnnotation = dict[str,
+                                    typing.Annotated[type_registry.compute_annotation(FunctionMiddlewareBaseConfig),
+                                                     Discriminator(TypedBaseModel.discriminator)]]
+
         MemoryAnnotation = dict[str,
                                 typing.Annotated[type_registry.compute_annotation(MemoryBaseConfig),
                                                  Discriminator(TypedBaseModel.discriminator)]]
@@ -393,6 +404,11 @@ class Config(HashableBaseModel):
             function_groups_field.annotation = FunctionGroupsAnnotation
             should_rebuild = True
 
+        middleware_field = cls.model_fields.get("middleware")
+        if (middleware_field is not None and middleware_field.annotation != MiddlewareAnnotation):
+            middleware_field.annotation = MiddlewareAnnotation
+            should_rebuild = True
+
         memory_field = cls.model_fields.get("memory")
         if memory_field is not None and memory_field.annotation != MemoryAnnotation:
             memory_field.annotation = MemoryAnnotation

nat/data_models/function.py

@@ -24,7 +24,16 @@ from .common import TypedBaseModel
 
 
 class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
-    pass
+    """Base configuration for functions.
+
+    Attributes:
+        middleware: List of function middleware names to apply to this function.
+            These must match names defined in the `middleware` section of the YAML configuration.
+    """
+    middleware: list[str] = Field(
+        default_factory=list,
+        description="List of function middleware names to apply to this function in order",
+    )
 
 
 class FunctionGroupBaseConfig(TypedBaseModel, BaseModelRegistryTag):
@@ -40,6 +49,10 @@ class FunctionGroupBaseConfig(TypedBaseModel, BaseModelRegistryTag):
         default_factory=list,
         description="The list of function names which should be excluded from default access to the group",
     )
+    middleware: list[str] = Field(
+        default_factory=list,
+        description="List of function middleware names to apply to all functions in this group",
+    )
 
     @field_validator("include", "exclude")
     @classmethod

nat/data_models/middleware.py

@@ -0,0 +1,35 @@
+# 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 .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class MiddlewareBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """The base level config object for middleware.
+
+    Middleware provides middleware-style wrapping of calls with
+    preprocessing and postprocessing logic.
+    """
+    pass
+
+
+MiddlewareBaseConfigT = typing.TypeVar("MiddlewareBaseConfigT", bound=MiddlewareBaseConfig)
+
+# Specialized type for function-specific middleware
+FunctionMiddlewareBaseConfig = MiddlewareBaseConfig
+FunctionMiddlewareBaseConfigT = MiddlewareBaseConfigT

nat/data_models/runtime_enum.py

@@ -0,0 +1,26 @@
+# 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 enum
+
+
+class RuntimeTypeEnum(str, enum.Enum):
+    """
+    Enum representing different runtime types.
+    """
+
+    RUN_OR_SERVE = "run_or_serve"
+    EVALUATE = "evaluate"
+    OTHER = "other"

nat/eval/dataset_handler/dataset_filter.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import fnmatch
+
 import pandas as pd
 
 from nat.data_models.dataset_handler import EvalFilterConfig
@@ -24,6 +26,7 @@ class DatasetFilter:
         - If a allowlist is provided, only keep rows matching the filter values.
         - If a denylist is provided, remove rows matching the filter values.
         - If the filter column does not exist in the DataFrame, the filtering is skipped for that column.
+        - Supports Unix shell-style wildcards (``*``, ``?``, ``[seq]``, ``[!seq]``) for string matching.
 
     This is a utility class that is dataset agnostic and can be used to filter any DataFrame based on the provided
     filter configuration.
@@ -33,6 +36,33 @@ class DatasetFilter:
 
         self.filter_config = filter_config
 
+    @staticmethod
+    def _match_wildcard_patterns(series: pd.Series, patterns: list[str | int | float]) -> pd.Series:
+        """
+        Match series values against wildcard patterns and exact values.
+
+        Args:
+            series (pd.Series): pandas Series to match against
+            patterns (list[str | int | float]): List of patterns/values
+
+        Returns:
+            pd.Series: Boolean Series indicating matches
+        """
+        # Convert series to string for pattern matching
+        str_series = series.astype(str)
+
+        # Initialize boolean mask
+        matches = pd.Series([False] * len(series), index=series.index)
+
+        # Check each pattern using fnmatch with list comprehension to avoid lambda capture
+        for pattern in patterns:
+            pattern_str = str(pattern)
+            pattern_matches = pd.Series([fnmatch.fnmatch(val, pattern_str) for val in str_series],
+                                        index=str_series.index)
+            matches |= pattern_matches
+
+        return matches
+
     def apply_filters(self, df) -> pd.DataFrame:
 
         filtered_df = df.copy()
@@ -41,12 +71,14 @@ class DatasetFilter:
         if self.filter_config.allowlist:
             for column, values in self.filter_config.allowlist.field.items():
                 if column in filtered_df.columns:
-                    filtered_df = filtered_df[filtered_df[column].isin(values)]
+                    matches = self._match_wildcard_patterns(filtered_df[column], values)
+                    filtered_df = filtered_df[matches]
 
         # Apply denylist (remove specified rows)
         if self.filter_config.denylist:
             for column, values in self.filter_config.denylist.field.items():
                 if column in filtered_df.columns:
-                    filtered_df = filtered_df[~filtered_df[column].isin(values)]
+                    matches = self._match_wildcard_patterns(filtered_df[column], values)
+                    filtered_df = filtered_df[~matches]
 
         return filtered_df

nat/eval/evaluate.py

@@ -16,6 +16,7 @@
 import asyncio
 import logging
 import shutil
+import warnings
 from pathlib import Path
 from typing import Any
 from uuid import uuid4
@@ -25,6 +26,7 @@ from tqdm import tqdm
 
 from nat.data_models.evaluate import EvalConfig
 from nat.data_models.evaluate import JobEvictionPolicy
+from nat.data_models.runtime_enum import RuntimeTypeEnum
 from nat.eval.config import EvaluationRunConfig
 from nat.eval.config import EvaluationRunOutput
 from nat.eval.dataset_handler.dataset_handler import DatasetHandler
@@ -67,7 +69,13 @@ class EvaluationRun:
         # Create evaluation trace context
         try:
             from nat.eval.utils.eval_trace_ctx import WeaveEvalTraceContext
-            self.eval_trace_context = WeaveEvalTraceContext()
+            with warnings.catch_warnings():
+                # Ignore deprecation warnings being triggered by weave. https://github.com/wandb/weave/issues/3666
+                warnings.filterwarnings("ignore",
+                                        category=DeprecationWarning,
+                                        message=r"`sentry_sdk\.Hub` is deprecated")
+
+                self.eval_trace_context = WeaveEvalTraceContext()
         except Exception:
             from nat.eval.utils.eval_trace_ctx import EvalTraceContext
             self.eval_trace_context = EvalTraceContext()
@@ -161,7 +169,7 @@ class EvaluationRun:
             if stop_event.is_set():
                 return "", []
 
-            async with session_manager.run(item.input_obj) as runner:
+            async with session_manager.run(item.input_obj, runtime_type=RuntimeTypeEnum.EVALUATE) as runner:
                 if not session_manager.workflow.has_single_output:
                     # raise an error if the workflow has multiple outputs
                     raise NotImplementedError("Multiple outputs are not supported")
@@ -514,7 +522,7 @@ class EvaluationRun:
         # Run workflow and evaluate
         async with WorkflowEvalBuilder.from_config(config=config) as eval_workflow:
             # Initialize Weave integration
-            self.weave_eval.initialize_logger(workflow_alias, self.eval_input, config)
+            self.weave_eval.initialize_logger(workflow_alias, self.eval_input, config, job_id=job_id)
 
             with self.eval_trace_context.evaluation_context():
                 # Run workflow

nat/eval/utils/weave_eval.py

@@ -82,7 +82,7 @@ class WeaveEvaluationIntegration:
         """Get the full dataset for Weave."""
         return [item.full_dataset_entry for item in eval_input.eval_input_items]
 
-    def initialize_logger(self, workflow_alias: str, eval_input: EvalInput, config: Any):
+    def initialize_logger(self, workflow_alias: str, eval_input: EvalInput, config: Any, job_id: str | None = None):
         """Initialize the Weave evaluation logger."""
         if not self.client and not self.initialize_client():
             # lazy init the client
@@ -92,10 +92,16 @@ class WeaveEvaluationIntegration:
             weave_dataset = self._get_weave_dataset(eval_input)
             config_dict = config.model_dump(mode="json")
             config_dict["name"] = workflow_alias
+
+            # Include job_id in eval_attributes if provided
+            eval_attributes = {}
+            if job_id:
+                eval_attributes["job_id"] = job_id
+
             self.eval_logger = self.evaluation_logger_cls(model=config_dict,
                                                           dataset=weave_dataset,
                                                           name=workflow_alias,
-                                                          eval_attributes={})
+                                                          eval_attributes=eval_attributes)
             self.pred_loggers = {}
 
             # Capture the current evaluation call for context propagation
@@ -136,9 +142,17 @@ class WeaveEvaluationIntegration:
         coros = []
         for eval_output_item in eval_output.eval_output_items:
             if eval_output_item.id in self.pred_loggers:
+                # Structure the score as a dict and include reasoning if available
+                score_value = {
+                    "score": eval_output_item.score,
+                }
+
+                if eval_output_item.reasoning is not None:
+                    score_value["reasoning"] = eval_output_item.reasoning
+
                 coros.append(self.pred_loggers[eval_output_item.id].alog_score(
                     scorer=evaluator_name,
-                    score=eval_output_item.score,
+                    score=score_value,
                 ))
 
         # Execute all coroutines concurrently

nat/front_ends/fastapi/fastapi_front_end_config.py

@@ -27,6 +27,8 @@ from pydantic import field_validator
 from nat.data_models.component_ref import ObjectStoreRef
 from nat.data_models.front_end import FrontEndBaseConfig
 from nat.data_models.step_adaptor import StepAdaptorConfig
+from nat.eval.evaluator.evaluator_model import EvalInputItem
+from nat.eval.evaluator.evaluator_model import EvalOutputItem
 
 logger = logging.getLogger(__name__)
 
@@ -133,6 +135,19 @@ class AsyncGenerationStatusResponse(BaseAsyncStatusResponse):
         description="Output of the generate request, this is only available if the job completed successfully.")
 
 
+class EvaluateItemRequest(BaseModel):
+    """Request model for single-item evaluation endpoint."""
+    item: EvalInputItem = Field(description="Single evaluation input item to evaluate")
+    evaluator_name: str = Field(description="Name of the evaluator to use (must match config)")
+
+
+class EvaluateItemResponse(BaseModel):
+    """Response model for single-item evaluation endpoint."""
+    success: bool = Field(description="Whether the evaluation completed successfully")
+    result: EvalOutputItem | None = Field(default=None, description="Evaluation result if successful")
+    error: str | None = Field(default=None, description="Error message if evaluation failed")
+
+
 class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
     """
     A FastAPI based front end that allows a NAT workflow to be served as a microservice.
@@ -211,6 +226,13 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
             "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.",
@@ -232,6 +254,13 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
         description="Evaluates the performance and accuracy of the workflow on a dataset",
     )
 
+    evaluate_item: typing.Annotated[EndpointBase,
+                                    Field(description="Endpoint for evaluating a single item.")] = EndpointBase(
+                                        method="POST",
+                                        path="/evaluate/item",
+                                        description="Evaluate a single item with a specified evaluator",
+                                    )
+
     oauth2_callback_path: str | None = Field(
         default="/auth/redirect",
         description="OAuth2.0 authentication callback endpoint. If None, no OAuth2 callback endpoint is created.")

nat/front_ends/fastapi/fastapi_front_end_plugin.py

@@ -120,18 +120,24 @@ class FastApiFrontEndPlugin(DaskClientMixin, FrontEndBase[FastApiFrontEndConfig]
 
                     from dask.distributed import LocalCluster
 
-                    self._cluster = LocalCluster(processes=True,
+                    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,
-                                                 n_workers=self.front_end_config.max_running_async_jobs,
-                                                 threads_per_worker=1)
+                                                 protocol="tcp",
+                                                 n_workers=self.front_end_config.max_running_async_jobs + 1)
 
                     self._scheduler_address = self._cluster.scheduler.address
 
-                    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)
+                    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", self._scheduler_address)
+                    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.")