Creation of DTS example and passing of completionToken

durabletask/client.py

@@ -97,6 +97,7 @@ def __init__(self, *,
                  log_handler: Optional[logging.Handler] = None,
                  log_formatter: Optional[logging.Formatter] = None,
                  secure_channel: bool = False):
+        self._metadata = metadata
         channel = shared.get_grpc_channel(host_address, metadata, secure_channel=secure_channel)
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger("client", log_handler, log_formatter)

durabletask/internal/shared.py

@@ -10,6 +10,7 @@
 import grpc
 
 from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
+from externalpackages.durabletaskscheduler.durabletask_grpc_interceptor import DTSDefaultClientInterceptorImpl
 
 # Field name used to indicate that an object was automatically serialized
 # and should be deserialized as a SimpleNamespace
@@ -50,7 +51,13 @@ def get_grpc_channel(
         channel = grpc.insecure_channel(host_address)
 
     if metadata is not None and len(metadata) > 0:
-        interceptors = [DefaultClientInterceptorImpl(metadata)]
+        for key, _ in metadata:
+            # Check if we are using DTS as the backend and if so, construct the DTS specific interceptors
+            if key == "dts":
+                interceptors = [DTSDefaultClientInterceptorImpl(metadata)]
+                break
+            else:
+                interceptors = [DefaultClientInterceptorImpl(metadata)]
         channel = grpc.intercept_channel(channel, *interceptors)
     return channel
 

durabletask/worker.py

@@ -143,9 +143,11 @@ def run_loop():
                             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)
+                                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)
+                                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}')
 
@@ -184,39 +186,42 @@ def stop(self):
         self._logger.info("Worker shutdown completed")
         self._is_running = False
 
-    def _execute_orchestrator(self, req: pb.OrchestratorRequest, stub: stubs.TaskHubSidecarServiceStub):
+    def _execute_orchestrator(self, req: pb.OrchestratorRequest, stub: stubs.TaskHubSidecarServiceStub, completionToken):
         try:
             executor = _OrchestrationExecutor(self._registry, self._logger)
             result = executor.execute(req.instanceId, req.pastEvents, req.newEvents)
             res = pb.OrchestratorResponse(
                 instanceId=req.instanceId,
                 actions=result.actions,
-                customStatus=pbh.get_string_value(result.encoded_custom_status))
+                customStatus=pbh.get_string_value(result.encoded_custom_status),
+                completionToken=completionToken)
         except Exception as ex:
             self._logger.exception(f"An error occurred while trying to execute instance '{req.instanceId}': {ex}")
             failure_details = pbh.new_failure_details(ex)
             actions = [pbh.new_complete_orchestration_action(-1, pb.ORCHESTRATION_STATUS_FAILED, "", failure_details)]
-            res = pb.OrchestratorResponse(instanceId=req.instanceId, actions=actions)
+            res = pb.OrchestratorResponse(instanceId=req.instanceId, actions=actions, completionToken=completionToken)
 
         try:
             stub.CompleteOrchestratorTask(res)
         except Exception as ex:
             self._logger.exception(f"Failed to deliver orchestrator response for '{req.instanceId}' to sidecar: {ex}")
 
-    def _execute_activity(self, req: pb.ActivityRequest, stub: stubs.TaskHubSidecarServiceStub):
+    def _execute_activity(self, req: pb.ActivityRequest, stub: stubs.TaskHubSidecarServiceStub, completionToken):
         instance_id = req.orchestrationInstance.instanceId
         try:
             executor = _ActivityExecutor(self._registry, self._logger)
             result = executor.execute(instance_id, req.name, req.taskId, req.input.value)
             res = pb.ActivityResponse(
                 instanceId=instance_id,
                 taskId=req.taskId,
-                result=pbh.get_string_value(result))
+                result=pbh.get_string_value(result),
+                completionToken=completionToken)
         except Exception as ex:
             res = pb.ActivityResponse(
                 instanceId=instance_id,
                 taskId=req.taskId,
-                failureDetails=pbh.new_failure_details(ex))
+                failureDetails=pbh.new_failure_details(ex),
+                completionToken=completionToken)
 
         try:
             stub.CompleteActivityTask(res)

examples/README.md

@@ -8,7 +8,7 @@ All the examples assume that you have a Durable Task-compatible sidecar running
 
 1. Install the latest version of the [Dapr CLI](https://docs.dapr.io/getting-started/install-dapr-cli/), which contains and exposes an embedded version of the Durable Task engine. The setup process (which requires Docker) will configure the workflow engine to store state in a local Redis container.
 
-1. Clone and run the [Durable Task Sidecar](https://github.com/microsoft/durabletask-go) project locally (requires Go 1.18 or higher). Orchestration state will be stored in a local sqlite database.
+2. Clone and run the [Durable Task Sidecar](https://github.com/microsoft/durabletask-go) project locally (requires Go 1.18 or higher). Orchestration state will be stored in a local sqlite database.
 
 ## Running the examples
 

examples/dts/README.md

@@ -0,0 +1,41 @@
+# Examples
+
+This directory contains examples of how to author durable orchestrations using the Durable Task Python SDK in conjunction with the Durable Task Scheduler (DTS).
+
+## Prerequisites
+
+All the examples assume that you have a Durable Task Scheduler taskhub created.
+
+The simplest way to create a taskhub is by using the az cli commands:
+
+1. Create a scheduler:
+    az durabletask scheduler create --resource-group <testrg> --name <testscheduler> --location <eastus> --ip-allowlist "[0.0.0.0/0]" --sku-capacity 1 --sku-name "Dedicated" --tags "{}"
+
+2. Create your taskhub
+    az durabletask taskhub create --resource-group <testrg> --scheduler-name <testscheduler> --name <testtaskhub>
+
+3. Retrieve the endpoint for the scheduler. This can be done by locating the taskhub in the portal.
+
+4. Set the appropriate environment variables for the TASKHUB and ENDPOINT
+
+```sh
+export TASKHUB=<taskhubname>
+```
+
+```sh
+export ENDPOINT=<taskhubEndpoint>
+```
+
+5. Since the samples rely on azure identity, ensure the package is installed and up-to-date
+
+```sh
+python3 -m pip install azure-identity
+```
+
+## Running the examples
+
+With one of the sidecars running, you can simply execute any of the examples in this directory using `python3`:
+
+```sh
+python3 dts_activity_sequence.py
+```

examples/dts/dts_activity_sequence.py

@@ -0,0 +1,62 @@
+"""End-to-end sample that demonstrates how to configure an orchestrator
+that calls an activity function in a sequence and prints the outputs."""
+import os
+from durabletask import client, task
+from externalpackages.durabletaskscheduler.durabletask_scheduler_worker import DurableTaskSchedulerWorker
+from externalpackages.durabletaskscheduler.durabletask_scheduler_client import DurableTaskSchedulerClient
+
+def hello(ctx: task.ActivityContext, name: str) -> str:
+    """Activity function that returns a greeting"""
+    return f'Hello {name}!'
+
+
+def sequence(ctx: task.OrchestrationContext, _):
+    """Orchestrator function that calls the 'hello' activity function in a sequence"""
+    # call "hello" activity function in a sequence
+    result1 = yield ctx.call_activity(hello, input='Tokyo')
+    result2 = yield ctx.call_activity(hello, input='Seattle')
+    result3 = yield ctx.call_activity(hello, input='London')
+
+    # return an array of results
+    return [result1, result2, result3]
+
+
+# Read the environment variable
+taskhub_name = os.getenv("TASKHUB")
+
+# Check if the variable exists
+if taskhub_name:
+    print(f"The value of TASKHUB is: {taskhub_name}")
+else:
+    print("TASKHUB is not set. Please set the TASKHUB environment variable to the name of the taskhub you wish to use")
+    print("If you are using windows powershell, run the following: $env:TASKHUB=\"<taskhubname>\"")
+    print("If you are using bash, run the following: export TASKHUB=\"<taskhubname>\"")
+    exit()
+
+# Read the environment variable
+endpoint = os.getenv("ENDPOINT")
+
+# Check if the variable exists
+if endpoint:
+    print(f"The value of ENDPOINT is: {endpoint}")
+else:
+    print("ENDPOINT is not set. Please set the ENDPOINT environment variable to the endpoint of the scheduler")
+    print("If you are using windows powershell, run the following: $env:ENDPOINT=\"<schedulerEndpoint>\"")
+    print("If you are using bash, run the following: export ENDPOINT=\"<schedulerEndpoint>\"")
+    exit()
+
+
+# configure and start the worker
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, use_managed_identity=False, client_id="", taskhub=taskhub_name) as w:
+    w.add_orchestrator(sequence)
+    w.add_activity(hello)
+    w.start()
+
+    # Construct the client and run the orchestrations
+    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True, taskhub=taskhub_name)
+    instance_id = c.schedule_new_orchestration(sequence)
+    state = c.wait_for_orchestration_completion(instance_id, timeout=60)
+    if state and state.runtime_status == client.OrchestrationStatus.COMPLETED:
+        print(f'Orchestration completed! Result: {state.serialized_output}')
+    elif state:
+        print(f'Orchestration failed: {state.failure_details}')

examples/dts/dts_fanout_fanin.py

@@ -0,0 +1,90 @@
+"""End-to-end sample that demonstrates how to configure an orchestrator
+that a dynamic number activity functions in parallel, waits for them all
+to complete, and prints an aggregate summary of the outputs."""
+import random
+import time
+import os
+from durabletask import client, task
+from durabletask import client, task
+from externalpackages.durabletaskscheduler.durabletask_scheduler_worker import DurableTaskSchedulerWorker
+from externalpackages.durabletaskscheduler.durabletask_scheduler_client import DurableTaskSchedulerClient
+
+
+def get_work_items(ctx: task.ActivityContext, _) -> list[str]:
+    """Activity function that returns a list of work items"""
+    # return a random number of work items
+    count = random.randint(2, 10)
+    print(f'generating {count} work items...')
+    return [f'work item {i}' for i in range(count)]
+
+
+def process_work_item(ctx: task.ActivityContext, item: str) -> int:
+    """Activity function that returns a result for a given work item"""
+    print(f'processing work item: {item}')
+
+    # simulate some work that takes a variable amount of time
+    time.sleep(random.random() * 5)
+
+    # return a result for the given work item, which is also a random number in this case
+    return random.randint(0, 10)
+
+
+def orchestrator(ctx: task.OrchestrationContext, _):
+    """Orchestrator function that calls the 'get_work_items' and 'process_work_item'
+    activity functions in parallel, waits for them all to complete, and prints
+    an aggregate summary of the outputs"""
+
+    work_items: list[str] = yield ctx.call_activity(get_work_items)
+
+    # execute the work-items in parallel and wait for them all to return
+    tasks = [ctx.call_activity(process_work_item, input=item) for item in work_items]
+    results: list[int] = yield task.when_all(tasks)
+
+    # return an aggregate summary of the results
+    return {
+        'work_items': work_items,
+        'results': results,
+        'total': sum(results),
+    }
+
+
+# Read the environment variable
+taskhub_name = os.getenv("TASKHUB")
+
+# Check if the variable exists
+if taskhub_name:
+    print(f"The value of TASKHUB is: {taskhub_name}")
+else:
+    print("TASKHUB is not set. Please set the TASKHUB environment variable to the name of the taskhub you wish to use")
+    print("If you are using windows powershell, run the following: $env:TASKHUB=\"<taskhubname>\"")
+    print("If you are using bash, run the following: export TASKHUB=\"<taskhubname>\"")
+    exit()
+
+# Read the environment variable
+endpoint = os.getenv("ENDPOINT")
+
+# Check if the variable exists
+if endpoint:
+    print(f"The value of ENDPOINT is: {endpoint}")
+else:
+    print("ENDPOINT is not set. Please set the ENDPOINT environment variable to the endpoint of the scheduler")
+    print("If you are using windows powershell, run the following: $env:ENDPOINT=\"<schedulerEndpoint>\"")
+    print("If you are using bash, run the following: export ENDPOINT=\"<schedulerEndpoint>\"")
+    exit()
+
+# configure and start the worker
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, taskhub=taskhub_name) as w:
+    w.add_orchestrator(orchestrator)
+    w.add_activity(process_work_item)
+    w.add_activity(get_work_items)
+    w.start()
+
+    # create a client, start an orchestration, and wait for it to finish
+    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True, taskhub=taskhub_name)
+    instance_id = c.schedule_new_orchestration(orchestrator)
+    state = c.wait_for_orchestration_completion(instance_id, timeout=30)
+    if state and state.runtime_status == client.OrchestrationStatus.COMPLETED:
+        print(f'Orchestration completed! Result: {state.serialized_output}')
+    elif state:
+        print(f'Orchestration failed: {state.failure_details}')
+    exit()

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"

externalpackages/durabletaskscheduler/access_token_manager.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+import durabletask.internal.shared as shared
+
+# By default, when there's 10minutes left before the token expires, refresh the token
+class AccessTokenManager:
+    def __init__(self, refresh_buffer: int = 600, metadata: Optional[list[tuple[str, str]]] = None):
+        self.scope = "https://durabletask.io/.default"
+        self.refresh_buffer = refresh_buffer
+        self._use_managed_identity = False
+        self._metadata = metadata
+        self._client_id = None
+        self._logger = shared.get_logger("token_manager")
+
+        if metadata:  # Ensure metadata is not None
+            for key, value in metadata:
+                if key == "use_managed_identity":
+                    self._use_managed_identity = value.lower() == "true"  # Properly convert string to bool
+                elif key == "client_id":
+                    self._client_id = value  # Directly assign string
+
+        # Choose the appropriate credential based on use_managed_identity
+        if self._use_managed_identity:
+            if not self._client_id:
+                self._logger.debug("Using System Assigned Managed Identity for authentication.")
+                self.credential = ManagedIdentityCredential()
+            else:
+                self._logger.debug("Using User Assigned Managed Identity for authentication.")
+                self.credential = ManagedIdentityCredential(client_id=self._client_id)
+        else:
+            self.credential = DefaultAzureCredential()
+            self._logger.debug("Using Default Azure Credentials for authentication.")
+
+        self.token = None
+        self.expiry_time = None
+
+    def get_access_token(self) -> str:
+        if self.token is None or self.is_token_expired():
+            self.refresh_token()
+        return self.token
+
+    # Checks if the token is expired, or if it will expire in the next "refresh_buffer" seconds.
+    # For example, if the token is created to have a lifespan of 2 hours, and the refresh buffer is set to 30 minutes,
+    # We will grab a new token when there're 30minutes left on the lifespan of the token
+    def is_token_expired(self) -> bool:
+        if self.expiry_time is None:
+            return True
+        return datetime.now(timezone.utc) >= (self.expiry_time - timedelta(seconds=self.refresh_buffer))
+
+    def refresh_token(self):
+        new_token = self.credential.get_token(self.scope)
+        self.token = f"Bearer {new_token.token}"
+
+        # Convert UNIX timestamp to timezone-aware datetime
+        self.expiry_time = datetime.fromtimestamp(new_token.expires_on, tz=timezone.utc)
+        self._logger.debug(f"Token refreshed. Expires at: {self.expiry_time}")

externalpackages/durabletaskscheduler/durabletask_grpc_interceptor.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from durabletask.internal.grpc_interceptor import _ClientCallDetails, DefaultClientInterceptorImpl
+from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
+
+import grpc
+
+class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
+    """The class implements a UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor,
+    StreamUnaryClientInterceptor and StreamStreamClientInterceptor from grpc to add an 
+    interceptor to add additional headers to all calls as needed."""
+
+    def __init__(self, metadata: list[tuple[str, str]]):
+        super().__init__(metadata)
+        self._token_manager = AccessTokenManager(metadata=self._metadata)
+
+    def _intercept_call(
+                self, client_call_details: _ClientCallDetails) -> grpc.ClientCallDetails:
+        """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
+            call details."""
+        # Refresh the auth token if it is present and needed
+        if self._metadata is not None:
+            for i, (key, _) in enumerate(self._metadata):
+                if key.lower() == "authorization":  # Ensure case-insensitive comparison
+                    new_token = self._token_manager.get_access_token()  # Get the new token
+                    self._metadata[i] = ("authorization", new_token)  # Update the token
+
+        return super()._intercept_call(client_call_details)