add new deterministic functions, non-retryable errors, and shutdown helpers

Signed-off-by: Filinto Duran <[email protected]>

Author: filintod

dev-requirements.txt

@@ -1 +0,0 @@
-grpcio-tools==1.62.3 # 1.62.X is the latest version before protobuf 1.26.X is used which has breaking changes for Python # supports protobuf 6.x and aligns with generated code

durabletask/client.py

@@ -127,9 +127,28 @@ def __init__(
             interceptors=interceptors,
             options=channel_options,
         )
+        self._channel = channel
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger("client", log_handler, log_formatter)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        try:
+            self.close()
+        finally:
+            return False
+
+    def close(self) -> None:
+        """Close the underlying gRPC channel."""
+        try:
+            # grpc.Channel.close() is idempotent
+            self._channel.close()
+        except Exception:
+            # Best-effort cleanup
+            pass
+
     def schedule_new_orchestration(
         self,
         orchestrator: Union[task.Orchestrator[TInput, TOutput], str],
@@ -188,10 +207,59 @@ def wait_for_orchestration_completion(
     ) -> Optional[OrchestrationState]:
         req = pb.GetInstanceRequest(instanceId=instance_id, getInputsAndOutputs=fetch_payloads)
         try:
-            grpc_timeout = None if timeout == 0 else timeout
-            self._logger.info(
-                f"Waiting {'indefinitely' if timeout == 0 else f'up to {timeout}s'} for instance '{instance_id}' to complete."
-            )
+            # gRPC timeout mapping (pytest unit tests may pass None explicitly)
+            grpc_timeout = None if (timeout is None or timeout == 0) else timeout
+
+            # If timeout is None or 0, skip pre-checks/polling and call server-side wait directly
+            if timeout is None or timeout == 0:
+                self._logger.info(
+                    f"Waiting {'indefinitely' if not timeout else f'up to {timeout}s'} for instance '{instance_id}' to complete."
+                )
+                res: pb.GetInstanceResponse = self._stub.WaitForInstanceCompletion(
+                    req, timeout=grpc_timeout
+                )
+                state = new_orchestration_state(req.instanceId, res)
+                return state
+
+            # For positive timeout, best-effort pre-check and short polling to avoid long server waits
+            try:
+                # First check if the orchestration is already completed
+                current_state = self.get_orchestration_state(
+                    instance_id, fetch_payloads=fetch_payloads
+                )
+                if current_state and current_state.runtime_status in [
+                    OrchestrationStatus.COMPLETED,
+                    OrchestrationStatus.FAILED,
+                    OrchestrationStatus.TERMINATED,
+                ]:
+                    return current_state
+
+                # Poll for completion with exponential backoff to handle eventual consistency
+                import time
+
+                poll_timeout = min(timeout, 10)
+                poll_start = time.time()
+                poll_interval = 0.1
+
+                while time.time() - poll_start < poll_timeout:
+                    current_state = self.get_orchestration_state(
+                        instance_id, fetch_payloads=fetch_payloads
+                    )
+
+                    if current_state and current_state.runtime_status in [
+                        OrchestrationStatus.COMPLETED,
+                        OrchestrationStatus.FAILED,
+                        OrchestrationStatus.TERMINATED,
+                    ]:
+                        return current_state
+
+                    time.sleep(poll_interval)
+                    poll_interval = min(poll_interval * 1.5, 1.0)  # Exponential backoff, max 1s
+            except Exception:
+                # Ignore pre-check/poll issues (e.g., mocked stubs in unit tests) and fall back
+                pass
+
+            self._logger.info(f"Waiting up to {timeout}s for instance '{instance_id}' to complete.")
             res: pb.GetInstanceResponse = self._stub.WaitForInstanceCompletion(
                 req, timeout=grpc_timeout
             )

durabletask/deterministic.py

@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""
+Deterministic utilities for Durable Task workflows (async and generator).
+
+This module provides deterministic alternatives to non-deterministic Python
+functions, ensuring workflow replay consistency across different executions.
+It is shared by both the asyncio authoring model and the generator-based model.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import random
+import string as _string
+import uuid
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional, Protocol, TypeVar, runtime_checkable
+
+
+@dataclass
+class DeterminismSeed:
+    """Seed data for deterministic operations."""
+
+    instance_id: str
+    orchestration_unix_ts: int
+
+    def to_int(self) -> int:
+        """Convert seed to integer for PRNG initialization."""
+        combined = f"{self.instance_id}:{self.orchestration_unix_ts}"
+        hash_bytes = hashlib.sha256(combined.encode("utf-8")).digest()
+        return int.from_bytes(hash_bytes[:8], byteorder="big")
+
+
+def derive_seed(instance_id: str, orchestration_time: datetime) -> int:
+    """
+    Derive a deterministic seed from instance ID and orchestration time.
+    """
+    ts = int(orchestration_time.timestamp())
+    return DeterminismSeed(instance_id=instance_id, orchestration_unix_ts=ts).to_int()
+
+
+def deterministic_random(instance_id: str, orchestration_time: datetime) -> random.Random:
+    """
+    Create a deterministic random number generator.
+    """
+    seed = derive_seed(instance_id, orchestration_time)
+    return random.Random(seed)
+
+
+def deterministic_uuid4(rnd: random.Random) -> uuid.UUID:
+    """Generate a deterministic UUID4 using the provided random generator."""
+    bytes_ = bytes(rnd.randrange(0, 256) for _ in range(16))
+    bytes_list = list(bytes_)
+    bytes_list[6] = (bytes_list[6] & 0x0F) | 0x40  # Version 4
+    bytes_list[8] = (bytes_list[8] & 0x3F) | 0x80  # Variant bits
+    return uuid.UUID(bytes=bytes(bytes_list))
+
+
+@runtime_checkable
+class DeterministicContextProtocol(Protocol):
+    """Protocol for contexts that provide deterministic operations."""
+
+    @property
+    def instance_id(self) -> str: ...
+
+    @property
+    def current_utc_datetime(self) -> datetime: ...
+
+
+class DeterministicContextMixin:
+    """
+    Mixin providing deterministic helpers for workflow contexts.
+
+    Assumes the inheriting class exposes `instance_id` and `current_utc_datetime` attributes.
+    """
+
+    def now(self) -> datetime:
+        """Return orchestration time (deterministic UTC)."""
+        value = self.current_utc_datetime  # type: ignore[attr-defined]
+        assert isinstance(value, datetime)
+        return value
+
+    def random(self) -> random.Random:
+        """Return a PRNG seeded deterministically from instance id and orchestration time."""
+        rnd = deterministic_random(
+            self.instance_id,  # type: ignore[attr-defined]
+            self.current_utc_datetime,  # type: ignore[attr-defined]
+        )
+        # Mark as deterministic for sandbox detector whitelisting of bound methods
+        try:
+            rnd._dt_deterministic = True
+        except Exception:
+            pass
+        return rnd
+
+    def uuid4(self) -> uuid.UUID:
+        """Return a deterministically generated UUID using the deterministic PRNG."""
+        rnd = self.random()
+        return deterministic_uuid4(rnd)
+
+    def new_guid(self) -> uuid.UUID:
+        """Alias for uuid4 for API parity with other SDKs."""
+        return self.uuid4()
+
+    def random_string(self, length: int, *, alphabet: Optional[str] = None) -> str:
+        """Return a deterministically generated random string of the given length."""
+        if length < 0:
+            raise ValueError("length must be non-negative")
+        chars = alphabet if alphabet is not None else (_string.ascii_letters + _string.digits)
+        if not chars:
+            raise ValueError("alphabet must not be empty")
+        rnd = self.random()
+        size = len(chars)
+        return "".join(chars[rnd.randrange(0, size)] for _ in range(length))
+
+    def random_int(self, min_value: int = 0, max_value: int = 2**31 - 1) -> int:
+        """Return a deterministic random integer in the specified range."""
+        if min_value > max_value:
+            raise ValueError("min_value must be <= max_value")
+        rnd = self.random()
+        return rnd.randint(min_value, max_value)
+
+    T = TypeVar("T")
+
+    def random_choice(self, sequence: Sequence[T]) -> T:
+        """Return a deterministic random element from a non-empty sequence."""
+        if not sequence:
+            raise IndexError("Cannot choose from empty sequence")
+        rnd = self.random()
+        return rnd.choice(sequence)

durabletask/internal/shared.py

@@ -102,7 +102,7 @@ def get_logger(
     # Add a default log handler if none is provided
     if log_handler is None:
         log_handler = logging.StreamHandler()
-        log_handler.setLevel(logging.INFO)
+        log_handler.setLevel(logging.DEBUG)
     logger.handlers.append(log_handler)
 
     # Set a default log formatter to our handler if none is provided

durabletask/task.py

@@ -7,7 +7,7 @@
 import math
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from typing import Any, Callable, Generator, Generic, Optional, TypeVar, Union
+from typing import Any, Callable, Generic, Generator, Optional, TypeVar, Union, cast
 
 import durabletask.internal.helpers as pbh
 import durabletask.internal.orchestrator_service_pb2 as pb
@@ -233,6 +233,16 @@ class OrchestrationStateError(Exception):
     pass
 
 
+class NonRetryableError(Exception):
+    """Exception indicating the operation should not be retried.
+
+    If an activity or sub-orchestration raises this exception, retry logic will be
+    bypassed and the failure will be returned immediately to the orchestrator.
+    """
+
+    pass
+
+
 class Task(ABC, Generic[T]):
     """Abstract base class for asynchronous tasks in a durable orchestration."""
 
@@ -395,7 +405,7 @@ def compute_next_delay(self) -> Optional[timedelta]:
                 next_delay_f = min(
                     next_delay_f, self._retry_policy.max_retry_interval.total_seconds()
                 )
-                return timedelta(seconds=next_delay_f)
+            return timedelta(seconds=next_delay_f)
 
         return None
 
@@ -486,6 +496,7 @@ def __init__(
         backoff_coefficient: Optional[float] = 1.0,
         max_retry_interval: Optional[timedelta] = None,
         retry_timeout: Optional[timedelta] = None,
+        non_retryable_error_types: Optional[list[Union[str, type]]] = None,
     ):
         """Creates a new RetryPolicy instance.
 
@@ -501,6 +512,11 @@ def __init__(
             The maximum retry interval to use for any retry attempt.
         retry_timeout : Optional[timedelta]
             The maximum amount of time to spend retrying the operation.
+        non_retryable_error_types : Optional[list[Union[str, type]]]
+            A list of exception type names or classes that should not be retried.
+            If a failure's error type matches any of these, the task fails immediately.
+            The built-in NonRetryableError is always treated as non-retryable regardless
+            of this setting.
         """
         # validate inputs
         if first_retry_interval < timedelta(seconds=0):
@@ -519,6 +535,17 @@ def __init__(
         self._backoff_coefficient = backoff_coefficient
         self._max_retry_interval = max_retry_interval
         self._retry_timeout = retry_timeout
+        # Normalize non-retryable error type names to a set of strings
+        names: Optional[set[str]] = None
+        if non_retryable_error_types:
+            names = set()
+            for t in non_retryable_error_types:
+                if isinstance(t, str):
+                    if t:
+                        names.add(t)
+                elif isinstance(t, type):
+                    names.add(t.__name__)
+        self._non_retryable_error_types = names
 
     @property
     def first_retry_interval(self) -> timedelta:
@@ -545,6 +572,14 @@ def retry_timeout(self) -> Optional[timedelta]:
         """The maximum amount of time to spend retrying the operation."""
         return self._retry_timeout
 
+    @property
+    def non_retryable_error_types(self) -> Optional[set[str]]:
+        """Set of error type names that should not be retried.
+
+        Comparison is performed against the errorType string provided by the
+        backend (typically the exception class name).
+        """
+        return self._non_retryable_error_types
 
 def get_name(fn: Callable) -> str:
     """Returns the name of the provided function"""