ergon-framework-python 0.3.0__tar.gz → 0.3.1__tar.gz

{ergon_framework_python-0.3.0/src/ergon_framework_python.egg-info → ergon_framework_python-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ergon-framework-python
-Version: 0.3.0
+Version: 0.3.1
 Summary: Ergon internal task-oriented project bootstrapper
 Author-email: Ergondata Technologies <anza.vossos@protonmail.com>
 License-Expression: MIT

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ergon-framework-python"
-version = "0.3.0"
+version = "0.3.1"
 description = "Ergon internal task-oriented project bootstrapper"
 readme = "README.md"
 license = "MIT"

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/src/ergon/__init__.py

@@ -10,4 +10,4 @@ __all__ = [
     "manager",
 ]
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/src/ergon/connector/ergon_platform/channels/connector.py

@@ -208,6 +208,10 @@ class ErgonPlatformChannelsConnector(Connector):
         if requeue and transaction.id:
             self._seen_event_ids.discard(transaction.id)
 
+    def release_transaction(self, transaction: Transaction) -> None:
+        """Give back a fetched but unstarted inbox event (worker shutdown): requeued with no delay."""
+        self.nack_transaction(transaction, requeue=True, delay_seconds=0)
+
     def download_attachments(
         self,
         transaction: Transaction,

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/src/ergon/connector/ergon_platform/workflows/connector.py

@@ -159,6 +159,13 @@ class ErgonPlatformConnector(Connector):
             **fields,
         )
 
+    def release_transaction(self, transaction: Transaction) -> None:
+        """Give back a transaction that was fetched but never started (worker shutdown).
+
+        A plain release: no failure reason, no cooldown, so it is picked up again right away.
+        """
+        self.release_item(transaction.id)
+
     def nack_transaction(
         self,
         transaction: Transaction,

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/src/ergon/connector/nylas/connector.py

@@ -119,5 +119,10 @@ class NylasConnector(Connector):
             return
         logger.debug("nack_transaction is a no-op for Nylas; message %s remains available for refetch", transaction.id)
 
+    def release_transaction(self, transaction: Transaction) -> None:
+        """Give back a fetched but unstarted message (worker shutdown): unchanged, refetchable."""
+        self.service.reset_pagination()
+        logger.debug("Nylas message %s released unchanged", transaction.id)
+
     def close(self) -> None:
         self.service.close()

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/src/ergon/connector/outlook/connector.py

@@ -194,5 +194,10 @@ class OutlookGraphConnector(Connector):
             self.service.reset_pagination()
             logger.debug("Outlook message %s was made available for refetch", transaction.id)
 
+    def release_transaction(self, transaction: Transaction) -> None:
+        """Give back a fetched but unstarted message (worker shutdown): unchanged, refetchable."""
+        self.service.reset_pagination()
+        logger.debug("Outlook message %s released unchanged", transaction.id)
+
     def close(self) -> None:
         self.service.close()

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/src/ergon/task/helpers.py

@@ -309,12 +309,17 @@ def run_fn(
     return run()
 
 
+_STOP_POLL_SECONDS = 0.25
+
+
 def multithread_execute(
     submissions: Iterable[Callable[[], futures.Future]],
     *,
     concurrency: int,
     limit: Optional[int] = None,
     timeout: Optional[float] = None,
+    stop: Optional[Callable[[], bool]] = None,
+    on_stop: Optional[Callable[[], None]] = None,
 ) -> int:
     """
     Execute submitted callables with bounded concurrency and refill.
@@ -322,6 +327,8 @@ def multithread_execute(
     - submissions yields callables that RETURN a Future when called
     - this function manages in-flight futures only
     - no domain semantics, no retries, no tracing
+    - ``stop`` is polled while futures are in flight; once it returns True nothing more is
+      submitted, ``on_stop`` runs once right away, and the in-flight futures are still awaited
 
     Returns:
         Number of completed submissions
@@ -329,6 +336,10 @@ def multithread_execute(
 
     in_flight: dict[futures.Future, float] = {}
     completed = 0
+    stopping = False
+    wait_timeout = timeout
+    if stop is not None:
+        wait_timeout = _STOP_POLL_SECONDS if timeout is None else min(timeout, _STOP_POLL_SECONDS)
 
     submit_iter = iter(submissions)
 
@@ -363,7 +374,7 @@ def multithread_execute(
         done, _ = futures.wait(
             in_flight.keys(),
             return_when=futures.FIRST_COMPLETED,
-            timeout=timeout,
+            timeout=wait_timeout,
         )
 
         # ----------------------------
@@ -391,6 +402,14 @@ def multithread_execute(
                     in_flight.pop(fut)
                     completed += 1
 
+        if not stopping and stop is not None and stop():
+            stopping = True
+            logger.info("Stop requested; waiting for %d in-flight submission(s), submitting no more", len(in_flight))
+            if on_stop is not None:
+                on_stop()
+        if stopping:
+            continue
+
         # ============================================================
         # REFILL
         # ============================================================

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/src/ergon/task/mixins/consumer.py

@@ -10,7 +10,7 @@ from async_timeout import timeout as timeout_after
 from opentelemetry import context as otel_context
 
 from ... import connector, telemetry
-from .. import base, exceptions, helpers, liveness, policies, utils
+from .. import base, exceptions, helpers, liveness, policies, shutdown, utils
 from . import metrics as mixin_metrics
 from . import producer
 
@@ -304,6 +304,38 @@ class ConsumerMixin(ABC):
             return next(iter(self.connectors.values()))
         raise ValueError("Multiple connectors configured; specify one in policy")
 
+    # =====================================================================
+    # SHUTDOWN HAND-BACK
+    # =====================================================================
+    def _hand_back_unstarted(self, conn, transactions: List[connector.Transaction]) -> None:
+        """Give back fetched transactions that were never started because the process is stopping.
+
+        ``release_transaction`` (no verdict) when the connector has it, else
+        ``nack_transaction(requeue=True)``. Never raises: the stop must go on, and a
+        transaction that cannot be handed back returns when its claim or lease expires.
+        """
+        release = getattr(conn, "release_transaction", None)
+        nack = getattr(conn, "nack_transaction", None)
+        for tr in transactions:
+            try:
+                if callable(release):
+                    release(tr)
+                elif callable(nack):
+                    nack(tr, requeue=True)
+                else:
+                    logger.warning(
+                        "Connector %s cannot hand back transaction %s; it returns when its claim or lease expires",
+                        conn.__class__.__name__,
+                        tr.id,
+                    )
+                    continue
+                logger.info("Handed back unstarted transaction %s", tr.id)
+            except Exception:
+                logger.exception(
+                    "Could not hand back transaction %s; it returns when its claim or lease expires",
+                    tr.id,
+                )
+
     # =====================================================================
     # PUBLIC CONSUME LOOP
     # =====================================================================
@@ -338,6 +370,10 @@ class ConsumerMixin(ABC):
                 )
 
             while True:
+                if shutdown.is_requested():
+                    logger.info("Shutdown requested; consume loop stops before the next fetch")
+                    break
+
                 batch_number += 1
 
                 # -------------------------
@@ -345,6 +381,9 @@ class ConsumerMixin(ABC):
                 # -------------------------
                 logger.debug(f"Fetching transactions batch with fetch policy: {policy.fetch.model_dump_json(indent=2)}")
                 success, result = self._handle_fetch(conn, policy.fetch)
+                if not success and shutdown.is_requested():
+                    logger.warning(f"Fetch failed while stopping; nothing was fetched, stopping → {result}")
+                    break
                 if not success:
                     logger.error(f"Fetch failed → {result}")
                     executor.shutdown(wait=False)
@@ -354,6 +393,14 @@ class ConsumerMixin(ABC):
 
                 transactions = result
 
+                if transactions and shutdown.is_requested():
+                    logger.info(
+                        "Shutdown requested during fetch; handing back %d unstarted transaction(s)",
+                        len(transactions),
+                    )
+                    self._hand_back_unstarted(conn, list(transactions))
+                    break
+
                 # -------------------------
                 # EMPTY QUEUE HANDLING
                 # -------------------------
@@ -375,6 +422,7 @@ class ConsumerMixin(ABC):
                         policy.fetch.empty.backoff_multiplier,
                         policy.fetch.empty.backoff_cap,
                         empty_count,
+                        interruptible=True,
                     )
                     empty_count += 1
                     continue
@@ -417,11 +465,24 @@ class ConsumerMixin(ABC):
                         "streaming": policy.loop.streaming,
                     },
                 ):
+                    not_started = list(transactions)
 
                     def submissions():
-                        for tr in transactions:
+                        while not_started and not shutdown.is_requested():
+                            tr = not_started.pop(0)
                             yield lambda tr=tr: submit_start_processing(tr, policy)
 
+                    def hand_back_not_started():
+                        remaining = list(not_started)
+                        not_started.clear()
+                        if remaining:
+                            logger.info(
+                                "Shutdown requested; handing back %d unstarted transaction(s) of this batch, "
+                                "in-flight ones will finish",
+                                len(remaining),
+                            )
+                            self._hand_back_unstarted(conn, remaining)
+
                     logger.debug(
                         f"Submitting {len(transactions)} transactions for processing "
                         f"with concurrency policy: {policy.loop.concurrency.model_dump_json(indent=2)}."
@@ -432,10 +493,17 @@ class ConsumerMixin(ABC):
                         concurrency=policy.loop.concurrency.value,
                         limit=policy.loop.limit,
                         timeout=policy.transaction_runtime.timeout,
+                        stop=shutdown.is_requested,
+                        on_stop=hand_back_not_started,
                     )
 
                 processed += count
 
+                if shutdown.is_requested():
+                    hand_back_not_started()
+                    logger.info("Shutdown requested; consume loop stops after the in-flight transactions")
+                    break
+
                 if policy.loop.limit and processed >= policy.loop.limit:
                     break
 
@@ -446,6 +514,7 @@ class ConsumerMixin(ABC):
                         multiplier=policy.fetch.batch.interval.backoff_multiplier,
                         cap=policy.fetch.batch.interval.backoff_cap,
                         attempt=0,
+                        interruptible=True,
                     )
 
             executor.shutdown()

{ergon_framework_python-0.3.0 → ergon_framework_python-0.3.1}/src/ergon/task/runner.py

@@ -1,6 +1,8 @@
 import asyncio
 import inspect
 import logging as stdlib_logging
+import multiprocessing
+import multiprocessing.connection
 import os
 import signal
 import threading
@@ -8,13 +10,13 @@ import time
 import traceback
 import uuid
 from collections.abc import Callable
-from concurrent.futures import ProcessPoolExecutor
 from datetime import datetime
 from enum import IntEnum
 from typing import Any, Literal
 
 from ..connector import Transaction
 from ..telemetry import logging, metrics, tracing
+from . import shutdown
 from .base import (
     BaseAsyncTask,
     BaseTask,
@@ -38,63 +40,24 @@ class ExitCode(IntEnum):
 
 
 # =============================================================
-# SHUTDOWN STATE (PROCESS-LOCAL)
+# SHUTDOWN STATE (PROCESS-LOCAL, see ``shutdown``)
 # =============================================================
 
-_shutdown_event = threading.Event()
-_shutdown_signal: int | None = None
-_shutdown_callbacks: set[Callable[[], None]] = set()
-_shutdown_callbacks_lock = threading.Lock()
 _supervisor_logger = stdlib_logging.getLogger(__name__)
 
-
-def _signal_handler(signum, frame):
-    global _shutdown_signal
-    _shutdown_signal = signum
-    _shutdown_event.set()
-    with _shutdown_callbacks_lock:
-        callbacks = tuple(_shutdown_callbacks)
-    for callback in callbacks:
-        try:
-            callback()
-        except Exception:
-            _supervisor_logger.exception("Async shutdown callback failed")
-
-
-def _install_signal_handlers():
-    signal.signal(signal.SIGINT, _signal_handler)
-    signal.signal(signal.SIGTERM, _signal_handler)
-
-
-def _reset_shutdown_state() -> None:
-    global _shutdown_signal
-    _shutdown_event.clear()
-    _shutdown_signal = None
-    with _shutdown_callbacks_lock:
-        _shutdown_callbacks.clear()
-
-
-def _register_shutdown_callback(callback: Callable[[], None]) -> None:
-    with _shutdown_callbacks_lock:
-        _shutdown_callbacks.add(callback)
-        shutdown_requested = _shutdown_event.is_set()
-    if shutdown_requested:
-        callback()
-
-
-def _unregister_shutdown_callback(callback: Callable[[], None]) -> None:
-    with _shutdown_callbacks_lock:
-        _shutdown_callbacks.discard(callback)
-
-
-def is_shutdown_requested() -> bool:
-    return _shutdown_event.is_set()
+_signal_handler = shutdown._handle_signal
+_install_signal_handlers = shutdown.install_handlers
+_reset_shutdown_state = shutdown.reset
+_register_shutdown_callback = shutdown.register_callback
+_unregister_shutdown_callback = shutdown.unregister_callback
+is_shutdown_requested = shutdown.is_requested
 
 
 def get_shutdown_exit_code() -> ExitCode:
-    if _shutdown_signal == signal.SIGINT:
+    requested = shutdown.requested_signal()
+    if requested == signal.SIGINT:
         return ExitCode.SIGINT
-    if _shutdown_signal == signal.SIGTERM:
+    if requested == signal.SIGTERM:
         return ExitCode.SIGTERM
     return ExitCode.ERROR
 
@@ -497,6 +460,103 @@ def __run_task_sync(
                     instance.exit()
 
 
+def _run_worker_sync(config: TaskConfig, mode: Literal["task", "transaction"] = "task", *args, **kwargs):
+    """Worker-process entrypoint: cooperative handlers, no escalation, keeps an early stop request."""
+    shutdown.prepare_worker()
+    if shutdown.is_requested():
+        _supervisor_logger.warning("Stop requested before worker %s started; not running it", kwargs.get("worker_id"))
+        return None
+    return __run_task_sync(config, mode, *args, **kwargs)
+
+
+def _worker_signal_relays(worker_pids: Callable[[], list[int]]) -> tuple[Callable[[], None], Callable[[], None]]:
+    """(forward the parent's signal to the workers, SIGKILL the workers).
+
+    ``worker_pids`` is read on every call, so workers started after registration count.
+    ``kill_workers`` runs right before the process exits, so it does not log.
+    """
+
+    def _send(signum: int) -> list[tuple[int, OSError]]:
+        failures = []
+        for pid in worker_pids():
+            try:
+                os.kill(pid, signum)
+            except ProcessLookupError:
+                continue
+            except OSError as exc:
+                failures.append((pid, exc))
+        return failures
+
+    def forward_shutdown() -> None:
+        for pid, exc in _send(shutdown.requested_signal() or signal.SIGTERM):
+            _supervisor_logger.error("Could not signal worker process %s: %s", pid, exc)
+
+    def kill_workers() -> None:
+        _send(signal.SIGKILL)
+
+    return forward_shutdown, kill_workers
+
+
+def _start_worker(config: TaskConfig, mode: Literal["task", "transaction"], args: tuple, worker_kwargs: dict):
+    process = multiprocessing.Process(
+        target=_run_worker_sync,
+        args=(config, mode, *args),
+        kwargs=worker_kwargs,
+        name=f"{config.name}-worker-{worker_kwargs['worker_id']}",
+    )
+    process.start()
+    return process
+
+
+def _run_workers(config: TaskConfig, mode: Literal["task", "transaction"], args: tuple, kwargs: dict) -> bool:
+    """Run ``max_workers`` worker processes to completion; True when one of them failed.
+
+    Plain processes, not a pool: a worker whose task returns simply exits. A worker that
+    fails or dies stops its siblings cooperatively, as a broken pool used to.
+    """
+    processes: list[multiprocessing.Process] = []
+    has_error = False
+
+    def alive_pids() -> list[int]:
+        return [p.pid for p in processes if p.pid is not None and p.exitcode is None]
+
+    forward_shutdown, kill_workers = _worker_signal_relays(alive_pids)
+    shutdown.register_callback(forward_shutdown)
+    shutdown.register_escalation_callback(kill_workers)
+    try:
+        # Workers inherit the blocked mask, so a signal that arrives while they start
+        # stays pending until ``prepare_worker`` installed their handlers.
+        with shutdown.blocked_signals():
+            for worker_id in range(config.max_workers):
+                if is_shutdown_requested():
+                    break
+                worker_kwargs = {**kwargs, "worker_id": worker_id, "total_workers": config.max_workers}
+                processes.append(_start_worker(config, mode, args, worker_kwargs))
+        if is_shutdown_requested():
+            forward_shutdown()
+
+        running = list(processes)
+        while running:
+            multiprocessing.connection.wait([p.sentinel for p in running], timeout=1.0)
+            for p in [p for p in running if p.exitcode is not None]:
+                running.remove(p)
+                p.join()
+                if p.exitcode == 0:
+                    continue
+                has_error = True
+                _supervisor_logger.error("Worker process %s (pid %s) exited with %s", p.name, p.pid, p.exitcode)
+                for sibling in running:
+                    if sibling.pid is not None:
+                        try:
+                            os.kill(sibling.pid, signal.SIGTERM)
+                        except ProcessLookupError:
+                            pass
+    finally:
+        shutdown.unregister_callback(forward_shutdown)
+        shutdown.unregister_escalation_callback(kill_workers)
+    return has_error
+
+
 # =============================================================
 # PUBLIC API — RUNNER
 # =============================================================
@@ -515,9 +575,10 @@ def run_task(
     Returns POSIX-compatible exit code.
     """
 
-    _reset_shutdown_state()
-    _install_signal_handlers()
     is_async = issubclass(config.task, BaseAsyncTask)  # type: ignore[arg-type]
+    _reset_shutdown_state()
+    # Async: every signal re-cancels and the liveness supervisor owns the hard deadline.
+    _install_signal_handlers(escalate=not is_async)
 
     # ---------------------------------------------------------
     # SINGLE PROCESS
@@ -552,24 +613,7 @@ def run_task(
     if is_async:
         raise RuntimeError("Async tasks cannot be executed with multiple processes. Use debug=True or max_workers=1.")
 
-    has_error = False
-
-    with ProcessPoolExecutor(max_workers=config.max_workers) as executor:
-        futures = []
-        for worker_id in range(config.max_workers):
-            worker_kwargs = {
-                **kwargs,
-                "worker_id": worker_id,
-                "total_workers": config.max_workers,
-            }
-            futures.append(executor.submit(__run_task_sync, config, mode, *args, **worker_kwargs))
-
-        for f in futures:
-            try:
-                f.result()
-            except Exception:
-                traceback.print_exc()
-                has_error = True
+    has_error = _run_workers(config, mode, args, kwargs)
 
     if is_shutdown_requested():
         return int(get_shutdown_exit_code())