Fill out boilerplate

Author: dandavison

temporalio/bridge/worker.py

@@ -165,6 +165,7 @@ class TunerHolder:
     workflow_slot_supplier: SlotSupplier
     activity_slot_supplier: SlotSupplier
     local_activity_slot_supplier: SlotSupplier
+    nexus_slot_supplier: SlotSupplier
 
 
 class Worker:

temporalio/worker/_replayer.py

@@ -264,6 +264,9 @@ def on_eviction_hook(
                         local_activity_slot_supplier=temporalio.bridge.worker.FixedSizeSlotSupplier(
                             1
                         ),
+                        nexus_slot_supplier=temporalio.bridge.worker.FixedSizeSlotSupplier(
+                            1
+                        ),
                     ),
                     nonsticky_to_sticky_poll_ratio=1,
                     no_remote_activities=True,

temporalio/worker/_tuning.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import logging
 from abc import ABC, abstractmethod
@@ -303,7 +305,8 @@ def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
 
 
 def _to_bridge_slot_supplier(
-    slot_supplier: SlotSupplier, kind: Literal["workflow", "activity", "local_activity"]
+    slot_supplier: SlotSupplier,
+    kind: Literal["workflow", "activity", "local_activity", "nexus"],
 ) -> temporalio.bridge.worker.SlotSupplier:
     if isinstance(slot_supplier, FixedSizeSlotSupplier):
         return temporalio.bridge.worker.FixedSizeSlotSupplier(slot_supplier.num_slots)
@@ -347,7 +350,8 @@ def create_resource_based(
         workflow_config: Optional[ResourceBasedSlotConfig] = None,
         activity_config: Optional[ResourceBasedSlotConfig] = None,
         local_activity_config: Optional[ResourceBasedSlotConfig] = None,
-    ) -> "WorkerTuner":
+        nexus_config: Optional[ResourceBasedSlotConfig] = None,
+    ) -> WorkerTuner:
         """Create a resource-based tuner with the provided options."""
         resource_cfg = ResourceBasedTunerConfig(target_memory_usage, target_cpu_usage)
         wf = ResourceBasedSlotSupplier(
@@ -359,10 +363,14 @@ def create_resource_based(
         local_act = ResourceBasedSlotSupplier(
             local_activity_config or ResourceBasedSlotConfig(), resource_cfg
         )
+        nexus = ResourceBasedSlotSupplier(
+            nexus_config or ResourceBasedSlotConfig(), resource_cfg
+        )
         return _CompositeTuner(
             wf,
             act,
             local_act,
+            nexus,
         )
 
     @staticmethod
@@ -371,14 +379,16 @@ def create_fixed(
         workflow_slots: Optional[int],
         activity_slots: Optional[int],
         local_activity_slots: Optional[int],
-    ) -> "WorkerTuner":
+        nexus_slots: Optional[int],
+    ) -> WorkerTuner:
         """Create a fixed-size tuner with the provided number of slots. Any unspecified slots will default to 100."""
         return _CompositeTuner(
             FixedSizeSlotSupplier(workflow_slots if workflow_slots else 100),
             FixedSizeSlotSupplier(activity_slots if activity_slots else 100),
             FixedSizeSlotSupplier(
                 local_activity_slots if local_activity_slots else 100
             ),
+            FixedSizeSlotSupplier(nexus_slots if nexus_slots else 100),
         )
 
     @staticmethod
@@ -387,12 +397,14 @@ def create_composite(
         workflow_supplier: SlotSupplier,
         activity_supplier: SlotSupplier,
         local_activity_supplier: SlotSupplier,
-    ) -> "WorkerTuner":
+        nexus_supplier: SlotSupplier,
+    ) -> WorkerTuner:
         """Create a tuner composed of the provided slot suppliers."""
         return _CompositeTuner(
             workflow_supplier,
             activity_supplier,
             local_activity_supplier,
+            nexus_supplier,
         )
 
     @abstractmethod
@@ -407,6 +419,10 @@ def _get_activity_task_slot_supplier(self) -> SlotSupplier:
     def _get_local_activity_task_slot_supplier(self) -> SlotSupplier:
         raise NotImplementedError
 
+    @abstractmethod
+    def _get_nexus_slot_supplier(self) -> SlotSupplier:
+        raise NotImplementedError
+
     def _to_bridge_tuner(self) -> temporalio.bridge.worker.TunerHolder:
         return temporalio.bridge.worker.TunerHolder(
             _to_bridge_slot_supplier(
@@ -418,6 +434,7 @@ def _to_bridge_tuner(self) -> temporalio.bridge.worker.TunerHolder:
             _to_bridge_slot_supplier(
                 self._get_local_activity_task_slot_supplier(), "local_activity"
             ),
+            _to_bridge_slot_supplier(self._get_nexus_slot_supplier(), "nexus"),
         )
 
     def _get_activities_max(self) -> Optional[int]:
@@ -436,6 +453,7 @@ class _CompositeTuner(WorkerTuner):
     workflow_slot_supplier: SlotSupplier
     activity_slot_supplier: SlotSupplier
     local_activity_slot_supplier: SlotSupplier
+    nexus_slot_supplier: SlotSupplier
 
     def _get_workflow_task_slot_supplier(self) -> SlotSupplier:
         return self.workflow_slot_supplier
@@ -445,3 +463,6 @@ def _get_activity_task_slot_supplier(self) -> SlotSupplier:
 
     def _get_local_activity_task_slot_supplier(self) -> SlotSupplier:
         return self.local_activity_slot_supplier
+
+    def _get_nexus_slot_supplier(self) -> SlotSupplier:
+        return self.nexus_slot_supplier

temporalio/worker/_worker.py

@@ -119,6 +119,7 @@ def __init__(
         max_concurrent_workflow_tasks: Optional[int] = None,
         max_concurrent_activities: Optional[int] = None,
         max_concurrent_local_activities: Optional[int] = None,
+        max_concurrent_nexus_tasks: Optional[int] = None,
         tuner: Optional[WorkerTuner] = None,
         max_concurrent_workflow_task_polls: Optional[int] = None,
         nonsticky_to_sticky_poll_ratio: float = 0.2,
@@ -214,13 +215,16 @@ def __init__(
             max_concurrent_workflow_tasks: Maximum allowed number of workflow
                 tasks that will ever be given to this worker at one time. Mutually exclusive with
                 ``tuner``. Must be set to at least two if ``max_cached_workflows`` is nonzero.
-            max_concurrent_activities: Maximum number of activity tasks that
-                will ever be given to the activity worker concurrently. Mutually exclusive with ``tuner``.
+            max_concurrent_activities: Maximum number of activity tasks that will ever be given to
+                the activity worker concurrently. Mutually exclusive with ``tuner``.
             max_concurrent_local_activities: Maximum number of local activity
-                tasks that will ever be given to the activity worker concurrently. Mutually exclusive with ``tuner``.
+                tasks that will ever be given to the activity worker concurrently. Mutually
+                exclusive with ``tuner``.
+            max_concurrent_nexus_tasks: Maximum number of Nexus tasks that will ever be given to
+                the Nexus worker concurrently. Mutually exclusive with ``tuner``.
             tuner:  Provide a custom :py:class:`WorkerTuner`. Mutually exclusive with the
-                ``max_concurrent_workflow_tasks``, ``max_concurrent_activities``, and
-                ``max_concurrent_local_activities`` arguments.
+                ``max_concurrent_workflow_tasks``, ``max_concurrent_activities``,
+                ``max_concurrent_local_activities``, and ``max_concurrent_nexus_tasks`` arguments.
 
                 Defaults to fixed-size 100 slots for each slot kind if unset and none of the
                 max_* arguments are provided.
@@ -337,6 +341,7 @@ def __init__(
             max_concurrent_workflow_tasks=max_concurrent_workflow_tasks,
             max_concurrent_activities=max_concurrent_activities,
             max_concurrent_local_activities=max_concurrent_local_activities,
+            max_concurrent_nexus_tasks=max_concurrent_nexus_tasks,
             tuner=tuner,
             max_concurrent_workflow_task_polls=max_concurrent_workflow_task_polls,
             nonsticky_to_sticky_poll_ratio=nonsticky_to_sticky_poll_ratio,
@@ -387,7 +392,6 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf
         """
         self._config = config
 
-        # TODO(nexus-preview): max_concurrent_nexus_tasks / tuner support
         if not (
             config["activities"]
             or config["nexus_service_handlers"]
@@ -462,6 +466,7 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf
             )
         self._nexus_worker: Optional[_NexusWorker] = None
         if config["nexus_service_handlers"]:
+            # TODO(nexus-preview) analogous calculations to those done for activity?
             self._nexus_worker = _NexusWorker(
                 bridge_worker=lambda: self._bridge_worker,
                 client=config["client"],
@@ -522,16 +527,19 @@ def check_activity(activity):
                 config["max_concurrent_workflow_tasks"]
                 or config["max_concurrent_activities"]
                 or config["max_concurrent_local_activities"]
+                or config["max_concurrent_nexus_tasks"]
             ):
                 raise ValueError(
                     "Cannot specify max_concurrent_workflow_tasks, max_concurrent_activities, "
-                    "or max_concurrent_local_activities when also specifying tuner"
+                    "max_concurrent_local_activities, or max_concurrent_nexus_tasks when also "
+                    "specifying tuner"
                 )
         else:
             tuner = WorkerTuner.create_fixed(
                 workflow_slots=config["max_concurrent_workflow_tasks"],
                 activity_slots=config["max_concurrent_activities"],
                 local_activity_slots=config["max_concurrent_local_activities"],
+                nexus_slots=config["max_concurrent_nexus_tasks"],
             )
 
         bridge_tuner = tuner._to_bridge_tuner()
@@ -883,6 +891,7 @@ class WorkerConfig(TypedDict, total=False):
     max_concurrent_workflow_tasks: Optional[int]
     max_concurrent_activities: Optional[int]
     max_concurrent_local_activities: Optional[int]
+    max_concurrent_nexus_tasks: Optional[int]
     tuner: Optional[WorkerTuner]
     max_concurrent_workflow_task_polls: Optional[int]
     nonsticky_to_sticky_poll_ratio: float