conductor-python 1.3.9__py3-none-any.whl → 1.3.11__py3-none-any.whl

conductor/client/__init__.py

@@ -0,0 +1,29 @@
+# Convenience re-exports for common symbols
+# Allows: from conductor.client import Configuration, TaskHandler, ...
+from conductor.client.configuration.configuration import Configuration
+from conductor.client.automator.task_handler import TaskHandler
+from conductor.client.automator.task_runner import TaskRunner
+from conductor.client.orkes_clients import OrkesClients
+from conductor.client.workflow.conductor_workflow import ConductorWorkflow
+from conductor.client.workflow.executor.workflow_executor import WorkflowExecutor
+from conductor.client.worker.worker_task import worker_task
+from conductor.client.worker.worker_interface import WorkerInterface
+from conductor.client.http.models.task import Task
+from conductor.client.http.models.task_result import TaskResult
+from conductor.client.http.models.task_result_status import TaskResultStatus
+from conductor.client.http.models.start_workflow_request import StartWorkflowRequest
+
+__all__ = [
+    "Configuration",
+    "TaskHandler",
+    "TaskRunner",
+    "OrkesClients",
+    "ConductorWorkflow",
+    "WorkflowExecutor",
+    "worker_task",
+    "WorkerInterface",
+    "Task",
+    "TaskResult",
+    "TaskResultStatus",
+    "StartWorkflowRequest",
+]

conductor/client/automator/async_task_runner.py

@@ -32,6 +32,7 @@ from conductor.client.worker.worker_interface import WorkerInterface
 from conductor.client.worker.worker_config import resolve_worker_config, get_worker_config_oneline
 from conductor.client.worker.exception import NonRetryableException
 from conductor.client.automator.json_schema_generator import generate_json_schema_from_function
+from conductor.client.automator.lease_tracker import LeaseManager
 
 logger = logging.getLogger(
     Configuration.get_logging_formatted_name(
@@ -112,6 +113,9 @@ class AsyncTaskRunner:
         self._semaphore = None
         self._shutdown = False  # Flag to indicate graceful shutdown
         self._use_update_v2 = True  # Will be set to False if server doesn't support v2 endpoint
+        self._lease_manager = LeaseManager.get_instance()
+        self._tracked_task_ids = set()  # Local set for cleanup on shutdown
+        self._sync_task_client = None  # Created after fork for LeaseManager heartbeats
 
     async def run(self) -> None:
         """Main async loop - runs continuously in single event loop."""
@@ -131,6 +135,17 @@ class AsyncTaskRunner:
             api_client=self.async_api_client
         )
 
+        # Create a sync TaskResourceApi for LeaseManager heartbeats
+        # (LeaseManager sends heartbeats from its own ThreadPoolExecutor)
+        from conductor.client.http.api.task_resource_api import TaskResourceApi
+        from conductor.client.http.api_client import ApiClient
+        self._sync_task_client = TaskResourceApi(
+            ApiClient(
+                configuration=self.configuration,
+                metrics_collector=self.metrics_collector
+            )
+        )
+
         # Create semaphore in the event loop (must be created within the loop)
         self._semaphore = asyncio.Semaphore(self._max_workers)
 
@@ -166,6 +181,11 @@ class AsyncTaskRunner:
         """Clean up async resources."""
         logger.debug("Cleaning up AsyncTaskRunner resources...")
 
+        # Untrack all tasks this runner was tracking from the shared LeaseManager
+        for task_id in list(self._tracked_task_ids):
+            self._lease_manager.untrack(task_id)
+        self._tracked_task_ids.clear()
+
         # Cancel any running tasks (EAFP style)
         try:
             for task in list(self._running_tasks):
@@ -182,6 +202,13 @@ class AsyncTaskRunner:
             except (IOError, OSError) as e:
                 logger.warning(f"Error closing async client: {e}")
 
+        # Close sync HTTP client used for lease heartbeats
+        if self._sync_task_client:
+            try:
+                self._sync_task_client.api_client.rest_client.connection.close()
+            except Exception:
+                pass
+
         # Clear event listeners
         self.event_dispatcher = None
 
@@ -227,7 +254,18 @@ class AsyncTaskRunner:
             output_schema_name = None
             schema_registry_available = True
 
-            if hasattr(self.worker, 'execute_function'):
+            # Check if schema registration is enabled for this worker
+            register_schema = getattr(self.worker, 'register_schema', False)
+            # Also check global Configuration default
+            if hasattr(self.configuration, 'register_schema') and self.configuration.register_schema is not None:
+                # Worker-level setting takes precedence if explicitly set (not default)
+                if not hasattr(self.worker, 'register_schema'):
+                    register_schema = self.configuration.register_schema
+
+            if not register_schema:
+                logger.debug(f"Schema registration disabled for {task_name} (register_schema=False)")
+
+            if register_schema and hasattr(self.worker, 'execute_function'):
                 logger.debug(f"Generating JSON schemas from function signature...")
                 # Pass strict_schema flag to control additionalProperties
                 strict_mode = getattr(self.worker, 'strict_schema', False)
@@ -309,6 +347,8 @@ class AsyncTaskRunner:
                             logger.debug(f"Could not register schemas for {task_name}: {e}")
                 else:
                     logger.debug(f"  ⚠ No schemas generated (unable to analyze function signature)")
+            elif not register_schema:
+                pass  # Already logged above
             else:
                 logger.debug(f"  ⚠ Class-based worker (no execute_function) - registering task without schemas")
 
@@ -573,15 +613,13 @@ class AsyncTaskRunner:
         # Acquire semaphore for entire task lifecycle (execution + update)
         # This ensures we never exceed thread_count tasks in any stage of processing
         async with self._semaphore:
+            self._track_lease(task)
             try:
                 while task is not None and not self._shutdown:
                     task_result = await self.__async_execute_task(task)
-                    # If task returned TaskInProgress, don't update yet
-                    if isinstance(task_result, TaskInProgress):
-                        logger.debug("Task %s is in progress, will update when complete", task.task_id)
-                        return
                     if task_result is None:
                         return
+                    self._untrack_lease(task.task_id)
                     # Update task and get next task from v2 response
                     task = await self.__async_update_task(task_result)
                     # v2 returns the next task; if v1 was used (returns None), immediately
@@ -589,12 +627,17 @@ class AsyncTaskRunner:
                     if task is None and not self._use_update_v2 and not self._shutdown:
                         tasks = await self.__async_batch_poll(1)
                         task = tasks[0] if tasks else None
+                    if task is not None:
+                        self._track_lease(task)
             except Exception as e:
                 logger.error(
                     "Error executing/updating task %s: %s",
                     task.task_id if task else "unknown",
                     traceback.format_exc()
                 )
+            finally:
+                if task is not None:
+                    self._untrack_lease(task.task_id)
 
     async def __async_execute_task(self, task: Task) -> TaskResult:
         """Execute async worker function directly (no threads, no BackgroundEventLoop)."""
@@ -908,6 +951,30 @@ class AsyncTaskRunner:
 
         return None
 
+    # -- Lease extension (heartbeat) delegation to LeaseManager ----------------
+
+    def _track_lease(self, task) -> None:
+        """Start tracking a task for lease extension via the shared LeaseManager."""
+        if not getattr(self.worker, 'lease_extend_enabled', False):
+            return
+        timeout = getattr(task, 'response_timeout_seconds', None) or 0
+        if timeout <= 0:
+            return
+        self._lease_manager.track(
+            task_id=task.task_id,
+            workflow_instance_id=task.workflow_instance_id,
+            response_timeout_seconds=timeout,
+            task_client=self._sync_task_client,
+        )
+        self._tracked_task_ids.add(task.task_id)
+
+    def _untrack_lease(self, task_id: str) -> None:
+        """Stop tracking a task for lease extension."""
+        self._lease_manager.untrack(task_id)
+        self._tracked_task_ids.discard(task_id)
+
+    # --------------------------------------------------------------------------
+
     def __set_worker_properties(self) -> None:
         """
         Resolve worker configuration using hierarchical override (same as TaskRunner).

conductor/client/automator/lease_tracker.py

@@ -0,0 +1,224 @@
+"""Centralized lease extension (heartbeat) management for Conductor task runners.
+
+Architecture:
+    LeaseManager runs a single background daemon thread that periodically checks
+    for tasks needing lease extension heartbeats. Due heartbeats are dispatched
+    to a small fixed ThreadPoolExecutor for parallel, non-blocking API calls.
+
+    This decouples heartbeat work entirely from worker poll loops, preventing
+    heartbeat API calls (and their retries) from blocking task polling.
+
+    Thread-safe: track() and untrack() can be called from any thread or event loop.
+"""
+
+import logging
+import os
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from conductor.client.http.models.task_result import TaskResult
+
+logger = logging.getLogger(__name__)
+
+# Lease extension constants (matches Java SDK)
+LEASE_EXTEND_RETRY_COUNT = 3
+LEASE_EXTEND_DURATION_FACTOR = 0.8
+
+
+@dataclass
+class LeaseInfo:
+    """Tracks when a heartbeat is next due for an in-flight task."""
+    task_id: str
+    workflow_instance_id: str
+    response_timeout_seconds: float
+    last_heartbeat_time: float  # time.monotonic() of last heartbeat (or task start)
+    interval_seconds: float     # 80% of responseTimeoutSeconds
+    task_client: Any = None     # Sync TaskResourceApi for sending heartbeats
+
+
+class LeaseManager:
+    """Centralized lease extension manager for all workers in a process.
+
+    One background daemon thread checks for due heartbeats at a fixed interval.
+    A small ThreadPoolExecutor sends heartbeat API calls in parallel.
+    Poll loops are never blocked by heartbeat work.
+
+    Usage:
+        manager = LeaseManager.get_instance()
+        manager.track(task_id, workflow_id, timeout, task_client)
+        # ... task completes ...
+        manager.untrack(task_id)
+    """
+
+    _instance: Optional['LeaseManager'] = None
+    _instance_lock = threading.Lock()
+    _instance_pid: Optional[int] = None
+
+    @classmethod
+    def get_instance(cls, check_interval: float = 1.0,
+                     max_heartbeat_workers: int = 4) -> 'LeaseManager':
+        """Get or create the process-wide LeaseManager singleton.
+
+        Fork-safe: a new instance is created after fork (threads don't survive fork).
+        """
+        current_pid = os.getpid()
+        if cls._instance is None or cls._instance_pid != current_pid:
+            with cls._instance_lock:
+                if cls._instance is None or cls._instance_pid != current_pid:
+                    cls._instance = cls(
+                        check_interval=check_interval,
+                        max_heartbeat_workers=max_heartbeat_workers,
+                    )
+                    cls._instance_pid = current_pid
+        return cls._instance
+
+    @classmethod
+    def _reset_instance(cls):
+        """Reset the singleton. For testing only."""
+        with cls._instance_lock:
+            if cls._instance is not None:
+                cls._instance.shutdown()
+            cls._instance = None
+            cls._instance_pid = None
+
+    def __init__(self, check_interval: float = 1.0, max_heartbeat_workers: int = 4):
+        self._tracked: Dict[str, LeaseInfo] = {}
+        self._lock = threading.Lock()
+        self._executor = ThreadPoolExecutor(
+            max_workers=max_heartbeat_workers,
+            thread_name_prefix="lease-heartbeat",
+        )
+        self._stop_event = threading.Event()
+        self._check_interval = check_interval
+        self._thread: Optional[threading.Thread] = None
+        self._started = False
+        self._start_lock = threading.Lock()
+
+    def _ensure_started(self) -> None:
+        """Lazily start the background thread on first track() call."""
+        if self._started:
+            return
+        with self._start_lock:
+            if not self._started:
+                self._thread = threading.Thread(
+                    target=self._run, daemon=True, name="lease-manager",
+                )
+                self._thread.start()
+                self._started = True
+                logger.debug(
+                    "LeaseManager started (check_interval=%.1fs)", self._check_interval,
+                )
+
+    def track(self, task_id: str, workflow_instance_id: str,
+              response_timeout_seconds: float, task_client: Any) -> None:
+        """Start tracking a task for lease extension heartbeats.
+
+        Thread-safe. Can be called from any worker thread or event loop.
+
+        Args:
+            task_id: Conductor task ID.
+            workflow_instance_id: Workflow instance this task belongs to.
+            response_timeout_seconds: The task's server-side response timeout.
+            task_client: A **sync** TaskResourceApi for sending heartbeat API calls.
+        """
+        interval = response_timeout_seconds * LEASE_EXTEND_DURATION_FACTOR
+        if interval < 1:
+            logger.debug(
+                "Skipping lease tracking for task %s (interval %.1fs too short)",
+                task_id, interval,
+            )
+            return
+
+        info = LeaseInfo(
+            task_id=task_id,
+            workflow_instance_id=workflow_instance_id,
+            response_timeout_seconds=response_timeout_seconds,
+            last_heartbeat_time=time.monotonic(),
+            interval_seconds=interval,
+            task_client=task_client,
+        )
+        with self._lock:
+            self._tracked[task_id] = info
+        self._ensure_started()
+        logger.debug(
+            "Tracking lease for task %s (timeout=%ss, heartbeat every %ss)",
+            task_id, response_timeout_seconds, interval,
+        )
+
+    def untrack(self, task_id: str) -> None:
+        """Stop tracking a task. Thread-safe."""
+        with self._lock:
+            removed = self._tracked.pop(task_id, None)
+        if removed is not None:
+            logger.debug("Untracked lease for task %s", task_id)
+
+    @property
+    def tracked_count(self) -> int:
+        """Number of currently tracked tasks."""
+        with self._lock:
+            return len(self._tracked)
+
+    # -- Background thread -----------------------------------------------------
+
+    def _run(self) -> None:
+        """Background loop — checks for due heartbeats at fixed intervals."""
+        while not self._stop_event.is_set():
+            try:
+                self._check_and_send()
+            except Exception as e:
+                logger.error("LeaseManager error: %s", e)
+            self._stop_event.wait(self._check_interval)
+
+    def _check_and_send(self) -> None:
+        """Find tasks with due heartbeats and dispatch to the thread pool."""
+        now = time.monotonic()
+        with self._lock:
+            due = [
+                info for info in self._tracked.values()
+                if now - info.last_heartbeat_time >= info.interval_seconds
+            ]
+        for info in due:
+            # Update timestamp immediately to prevent double-dispatch on next tick
+            info.last_heartbeat_time = time.monotonic()
+            self._executor.submit(self._send_heartbeat, info)
+
+    @staticmethod
+    def _send_heartbeat(info: LeaseInfo) -> None:
+        """Send a single lease extension heartbeat with retry.
+
+        Runs in a pool thread — blocking retries only block the pool thread,
+        never a poll loop.
+        """
+        result = TaskResult(
+            task_id=info.task_id,
+            workflow_instance_id=info.workflow_instance_id,
+            extend_lease=True,
+        )
+        for attempt in range(LEASE_EXTEND_RETRY_COUNT):
+            try:
+                info.task_client.update_task(body=result)
+                logger.debug("Extended lease for task %s", info.task_id)
+                return
+            except Exception as e:
+                if attempt < LEASE_EXTEND_RETRY_COUNT - 1:
+                    time.sleep(0.5 * (attempt + 2))
+                else:
+                    logger.error(
+                        "Failed to extend lease for task %s after %d attempts: %s",
+                        info.task_id, LEASE_EXTEND_RETRY_COUNT, e,
+                    )
+
+    # -- Lifecycle -------------------------------------------------------------
+
+    def shutdown(self) -> None:
+        """Stop the background thread and thread pool."""
+        self._stop_event.set()
+        if self._started and self._thread is not None:
+            self._thread.join(timeout=5)
+        self._executor.shutdown(wait=False)
+        with self._lock:
+            self._tracked.clear()
+        logger.debug("LeaseManager shut down")

conductor/client/automator/task_handler.py

@@ -4,6 +4,7 @@ import importlib
 import inspect
 import logging
 import os
+import signal
 import threading
 import time
 from multiprocessing import Process, freeze_support, Queue, set_start_method
@@ -52,6 +53,7 @@ def _run_sync_worker_process(
         event_listeners: Optional[List[Any]],
 ) -> None:
     """Process target: construct TaskRunner after fork/spawn and run forever."""
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
     task_runner = TaskRunner(worker, configuration, metrics_settings, event_listeners)
     task_runner.run()
 
@@ -63,6 +65,7 @@ def _run_async_worker_process(
         event_listeners: Optional[List[Any]],
 ) -> None:
     """Process target: construct AsyncTaskRunner after fork/spawn and run forever."""
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
     async_task_runner = AsyncTaskRunner(worker, configuration, metrics_settings, event_listeners)
     asyncio.run(async_task_runner.run())
 
@@ -70,7 +73,8 @@ def _run_async_worker_process(
 def register_decorated_fn(name: str, poll_interval: int, domain: str, worker_id: str, func,
                          thread_count: int = 1, register_task_def: bool = False,
                          poll_timeout: int = 100, lease_extend_enabled: bool = False, task_def: Optional['TaskDef'] = None,
-                         overwrite_task_def: bool = True, strict_schema: bool = False):
+                         overwrite_task_def: bool = True, strict_schema: bool = False,
+                         register_schema: Optional[bool] = None):
     logger.debug("decorated %s", name)
     _decorated_functions[(name, domain)] = {
         "func": func,
@@ -83,7 +87,8 @@ def register_decorated_fn(name: str, poll_interval: int, domain: str, worker_id:
         "lease_extend_enabled": lease_extend_enabled,
         "task_def": task_def,
         "overwrite_task_def": overwrite_task_def,
-        "strict_schema": strict_schema
+        "strict_schema": strict_schema,
+        "register_schema": register_schema
     }
 
 
@@ -109,7 +114,8 @@ def get_registered_workers() -> List[Worker]:
             paused=False,  # Always default to False, only env vars can set to True
             task_def_template=record.get("task_def"),  # Optional TaskDef configuration
             overwrite_task_def=record.get("overwrite_task_def", True),
-            strict_schema=record.get("strict_schema", False)
+            strict_schema=record.get("strict_schema", False),
+            register_schema=record.get("register_schema") if record.get("register_schema") is not None else False
         )
         workers.append(worker)
     return workers
@@ -146,6 +152,29 @@ class TaskHandler:
         - 10-100x better concurrency for I/O-bound workloads
         - Automatically detected - no configuration needed
 
+    Parameters:
+        scan_for_annotated_workers: When True (default), discovers workers decorated with
+            @worker_task that have already been imported into the current Python process.
+            This does NOT scan the filesystem — it only finds functions registered at
+            import time. Workers defined in modules that have not yet been imported will
+            be silently ignored.
+
+            To load workers from separate files, import them before creating TaskHandler:
+
+                import myapp.workers  # must import first
+                with TaskHandler(configuration=config) as th:
+                    th.start_processes()
+
+            Or use the import_modules parameter to have TaskHandler import them for you:
+
+                with TaskHandler(configuration=config,
+                                 import_modules=['myapp.workers']) as th:
+                    th.start_processes()
+
+        import_modules: List of module dotted-path strings to import before scanning for
+            annotated workers (e.g. ['myapp.workers', 'myapp.other_workers']). Use this
+            instead of manual imports when scan_for_annotated_workers=True.
+
     Usage:
         # Default configuration
         handler = TaskHandler(configuration=config)
@@ -228,9 +257,15 @@ class TaskHandler:
                     'poll_timeout': record.get("poll_timeout", 100),
                     'lease_extend_enabled': record.get("lease_extend_enabled", True),
                     'overwrite_task_def': record.get("overwrite_task_def", True),
-                    'strict_schema': record.get("strict_schema", False)
+                    'strict_schema': record.get("strict_schema", False),
+                    'register_schema': record.get("register_schema")
                 }
 
+                # Apply global Configuration.register_schema as fallback when
+                # the decorator doesn't set it explicitly (None).
+                if code_config['register_schema'] is None and hasattr(configuration, 'register_schema') and configuration.register_schema is not None:
+                    code_config['register_schema'] = configuration.register_schema
+
                 # Resolve configuration with environment variable overrides
                 resolved_config = resolve_worker_config(
                     worker_name=task_def_name,
@@ -249,7 +284,8 @@ class TaskHandler:
                     lease_extend_enabled=resolved_config['lease_extend_enabled'],
                     task_def_template=record.get("task_def"),  # Pass TaskDef configuration
                     overwrite_task_def=resolved_config.get('overwrite_task_def', True),
-                    strict_schema=resolved_config.get('strict_schema', False))
+                    strict_schema=resolved_config.get('strict_schema', False),
+                    register_schema=resolved_config.get('register_schema', False))
                 logger.debug("created worker with name=%s and domain=%s", task_def_name, resolved_config['domain'])
                 workers.append(worker)
 
@@ -607,6 +643,7 @@ def _setup_logging_queue(configuration: Configuration):
 
 # This process performs the centralized logging
 def __logger_process(queue, log_level, logger_format=None):
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
     c_logger = logging.getLogger(
         Configuration.get_logging_formatted_name(
             __name__