Moving dts logic into its own module

Signed-off-by: Ryan Lettieri <ryanLettieri@microsoft.com>

Author: RyanLettieri

durabletask/worker.py

@@ -16,7 +16,7 @@
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
-from durabletask.accessTokenManager import AccessTokenManager
+
 from durabletask import task
 
 TInput = TypeVar('TInput')
@@ -89,36 +89,15 @@ def __init__(self, *,
                  metadata: Optional[list[tuple[str, str]]] = None,
                  log_handler=None,
                  log_formatter: Optional[logging.Formatter] = None,
-                 secure_channel: bool = False,
-                 access_token_manager: AccessTokenManager = None):
+                 secure_channel: bool = False):
         self._registry = _Registry()
         self._host_address = host_address if host_address else shared.get_default_host_address()
         self._metadata = metadata
         self._logger = shared.get_logger("worker", log_handler, log_formatter)
         self._shutdown = Event()
         self._is_running = False
         self._secure_channel = secure_channel
-        self._access_token_manager = access_token_manager
-        self.__update_metadata_with_token()
-
-    def __update_metadata_with_token(self):
-        """
-        Add or update the `authorization` key in the metadata with the current access token.
-        """
-        if self._access_token_manager is not None:
-            token = self._access_token_manager.get_access_token()
-            
-            # Check if "authorization" already exists in the metadata
-            updated = False
-            for i, (key, _) in enumerate(self._metadata):
-                if key == "authorization":
-                    self._metadata[i] = ("authorization", token)
-                    updated = True
-                    break
-            
-            # If not updated, add a new entry
-            if not updated:
-                self._metadata.append(("authorization", token))
+
 
     def __enter__(self):
         return self
@@ -153,7 +132,6 @@ def run_loop():
             with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
                 while not self._shutdown.is_set():
                     try:
-                        self.__update_metadata_with_token()
                         # send a "Hello" message to the sidecar to ensure that it's listening
                         stub.Hello(empty_pb2.Empty())
 

examples/dts/dts_activity_sequence.py

@@ -3,8 +3,9 @@
 
 """End-to-end sample that demonstrates how to configure an orchestrator
 that calls an activity function in a sequence and prints the outputs."""
-from durabletask import client, task, worker
-from durabletask.accessTokenManager import AccessTokenManager
+from durabletask import client, task
+from externalpackages.durabletaskscheduler.durabletask_scheduler_worker import DurableTaskSchedulerWorker
+from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
 
 def hello(ctx: task.ActivityContext, name: str) -> str:
     """Activity function that returns a greeting"""
@@ -56,7 +57,7 @@ def sequence(ctx: task.OrchestrationContext, _):
 ]
 
 # configure and start the worker
-with worker.TaskHubGrpcWorker(host_address=endpoint, metadata=metaData, secure_channel=True, access_token_manager=token_manager) as w:
+with DurableTaskSchedulerWorker(host_address=endpoint, metadata=metaData, secure_channel=True, access_token_manager=token_manager) as w:
     w.add_orchestrator(sequence)
     w.add_activity(hello)
     w.start()

externalpackages/durabletaskscheduler/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Durable Task SDK for Python"""
+
+
+PACKAGE_NAME = "durabletaskscheduler"

durabletask/accessTokenManager.py renamed to externalpackages/durabletaskscheduler/access_token_manager.py

@@ -0,0 +1,6 @@
+from durabletask import TaskHubGrpcClient
+
+class DurableTaskSchedulerClient(TaskHubGrpcClient):
+    def __init__(self, *args, **kwargs):
+        # Initialize the base class
+        super().__init__(*args, **kwargs)

externalpackages/durabletaskscheduler/durabletask_scheduler_client.py

@@ -0,0 +1,94 @@
+import concurrent.futures
+from threading import Thread
+from google.protobuf import empty_pb2
+import grpc
+import durabletask.internal.orchestrator_service_pb2 as pb
+import durabletask.internal.orchestrator_service_pb2_grpc as stubs
+import durabletask.internal.shared as shared
+
+from durabletask.worker import TaskHubGrpcWorker
+from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
+
+class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
+    def __init__(self, *args, access_token_manager: AccessTokenManager = None, **kwargs):
+        # Initialize the base class
+        super().__init__(*args, **kwargs)
+        self._access_token_manager = access_token_manager
+        self.__update_metadata_with_token()
+
+    def __update_metadata_with_token(self):
+        """
+        Add or update the `authorization` key in the metadata with the current access token.
+        """
+        if self._access_token_manager is not None:
+            token = self._access_token_manager.get_access_token()
+            
+            # Check if "authorization" already exists in the metadata
+            updated = False
+            for i, (key, _) in enumerate(self._metadata):
+                if key == "authorization":
+                    self._metadata[i] = ("authorization", token)
+                    updated = True
+                    break
+            
+            # If not updated, add a new entry
+            if not updated:
+                self._metadata.append(("authorization", token))
+
+    def start(self):
+        """Starts the worker on a background thread and begins listening for work items."""
+        channel = shared.get_grpc_channel(self._host_address, self._metadata, self._secure_channel)
+        stub = stubs.TaskHubSidecarServiceStub(channel)
+
+        if self._is_running:
+            raise RuntimeError('The worker is already running.')
+
+        def run_loop():
+            # TODO: Investigate whether asyncio could be used to enable greater concurrency for async activity
+            #       functions. We'd need to know ahead of time whether a function is async or not.
+            # TODO: Max concurrency configuration settings
+            with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
+                while not self._shutdown.is_set():
+                    try:
+                        self.__update_metadata_with_token()
+                        # send a "Hello" message to the sidecar to ensure that it's listening
+                        stub.Hello(empty_pb2.Empty())
+
+                        # stream work items
+                        self._response_stream = stub.GetWorkItems(pb.GetWorkItemsRequest())
+                        self._logger.info(f'Successfully connected to {self._host_address}. Waiting for work items...')
+
+                        # The stream blocks until either a work item is received or the stream is canceled
+                        # by another thread (see the stop() method).
+                        for work_item in self._response_stream:  # type: ignore
+                            request_type = work_item.WhichOneof('request')
+                            self._logger.debug(f'Received "{request_type}" work item')
+                            if work_item.HasField('orchestratorRequest'):
+                                executor.submit(self._execute_orchestrator, work_item.orchestratorRequest, stub, work_item.completionToken)
+                            elif work_item.HasField('activityRequest'):
+                                executor.submit(self._execute_activity, work_item.activityRequest, stub, work_item.completionToken)
+                            elif work_item.HasField('healthPing'):
+                                pass # no-op
+                            else:
+                                self._logger.warning(f'Unexpected work item type: {request_type}')
+
+                    except grpc.RpcError as rpc_error:
+                        if rpc_error.code() == grpc.StatusCode.CANCELLED:  # type: ignore
+                            self._logger.info(f'Disconnected from {self._host_address}')
+                        elif rpc_error.code() == grpc.StatusCode.UNAVAILABLE:  # type: ignore
+                            self._logger.warning(
+                                f'The sidecar at address {self._host_address} is unavailable - will continue retrying')
+                        else:
+                            self._logger.warning(f'Unexpected error: {rpc_error}')
+                    except Exception as ex:
+                        self._logger.warning(f'Unexpected error: {ex}')
+
+                    # CONSIDER: exponential backoff
+                    self._shutdown.wait(5)
+                self._logger.info("No longer listening for work items")
+                return
+
+        self._logger.info(f"Starting gRPC worker that connects to {self._host_address}")
+        self._runLoop = Thread(target=run_loop)
+        self._runLoop.start()
+        self._is_running = True