pyworkflow-engine 0.4.0__tar.gz → 0.4.1__tar.gz

{pyworkflow_engine-0.4.0 → pyworkflow_engine-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyworkflow-engine
-Version: 0.4.0
+Version: 0.4.1
 Summary: A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow
 Author: PyWorkflow Contributors
 License: MIT
@@ -24,13 +24,13 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: celery<6.0.0,>=5.3.0
 Requires-Dist: cloudpickle>=3.0.0
-Requires-Dist: pydantic<3.0.0,>=2.0.0
+Requires-Dist: pydantic<3.0.0,>=2.14.1
 Requires-Dist: loguru>=0.7.0
 Requires-Dist: click>=8.4.2
 Requires-Dist: inquirerpy>=0.3.4; python_version < "4.0"
 Requires-Dist: httpx>=0.25.0
 Requires-Dist: python-dateutil>=2.8.0
-Requires-Dist: filelock>=3.32.4
+Requires-Dist: filelock>=4.2.4
 Requires-Dist: pyyaml>=6.0.0
 Requires-Dist: croniter>=6.2.4
 Provides-Extra: redis
@@ -38,11 +38,11 @@ Requires-Dist: redis>=8.1.0; extra == "redis"
 Provides-Extra: sqlite
 Requires-Dist: aiosqlite>=0.19.0; extra == "sqlite"
 Provides-Extra: postgres
-Requires-Dist: asyncpg>=0.29.0; extra == "postgres"
+Requires-Dist: asyncpg>=0.32.0; extra == "postgres"
 Provides-Extra: aws
-Requires-Dist: aws-durable-execution-sdk-python>=1.6.0; extra == "aws"
+Requires-Dist: aws-durable-execution-sdk-python>=2.0.1; extra == "aws"
 Provides-Extra: dynamodb
-Requires-Dist: aiobotocore>=3.9.0; extra == "dynamodb"
+Requires-Dist: aiobotocore>=3.9.2; extra == "dynamodb"
 Provides-Extra: cassandra
 Requires-Dist: cassandra-driver>=3.30.1; extra == "cassandra"
 Provides-Extra: mysql
@@ -50,28 +50,28 @@ Requires-Dist: aiomysql>=0.2.0; extra == "mysql"
 Provides-Extra: all
 Requires-Dist: redis>=8.1.0; extra == "all"
 Requires-Dist: aiosqlite>=0.19.0; extra == "all"
-Requires-Dist: asyncpg>=0.29.0; extra == "all"
+Requires-Dist: asyncpg>=0.32.0; extra == "all"
 Requires-Dist: aiomysql>=0.2.0; extra == "all"
-Requires-Dist: aws-durable-execution-sdk-python>=1.6.0; extra == "all"
+Requires-Dist: aws-durable-execution-sdk-python>=2.0.1; extra == "all"
 Requires-Dist: cassandra-driver>=3.30.1; extra == "all"
 Provides-Extra: dev
 Requires-Dist: pytest>=9.1.1; extra == "dev"
 Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
 Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
 Requires-Dist: pytest-celery>=0.0.0; extra == "dev"
-Requires-Dist: pytest-mock>=3.12.0; extra == "dev"
-Requires-Dist: moto[dynamodb]>=5.2.3; extra == "dev"
-Requires-Dist: black>=26.5.1; extra == "dev"
-Requires-Dist: ruff>=0.16.4; extra == "dev"
-Requires-Dist: mypy>=2.3.1; extra == "dev"
-Requires-Dist: pre-commit>=4.6.2; extra == "dev"
+Requires-Dist: pytest-mock>=3.17.0; extra == "dev"
+Requires-Dist: moto[dynamodb]>=5.2.4; extra == "dev"
+Requires-Dist: black>=26.10.1; extra == "dev"
+Requires-Dist: ruff>=0.17.2; extra == "dev"
+Requires-Dist: mypy>=2.4.0; extra == "dev"
+Requires-Dist: pre-commit>=4.7.0; extra == "dev"
 Requires-Dist: types-redis>=4.6.0; extra == "dev"
-Requires-Dist: types-python-dateutil>=2.9.0.20260807; extra == "dev"
-Requires-Dist: types-PyYAML>=6.0.12.20260815; extra == "dev"
-Requires-Dist: flower>=2.1.0; extra == "dev"
+Requires-Dist: types-python-dateutil>=2.9.0.20261021; extra == "dev"
+Requires-Dist: types-PyYAML>=6.0.12.20260906; extra == "dev"
+Requires-Dist: flower>=2.2.0; extra == "dev"
 Requires-Dist: redis>=8.1.0; extra == "dev"
 Requires-Dist: aiosqlite>=0.19.0; extra == "dev"
-Requires-Dist: asyncpg>=0.29.0; extra == "dev"
+Requires-Dist: asyncpg>=0.32.0; extra == "dev"
 Dynamic: license-file
 
 # PyWorkflow

{pyworkflow_engine-0.4.0 → pyworkflow_engine-0.4.1}/docs/guides/configuration.mdx

@@ -392,6 +392,9 @@ Environment variables provide deployment flexibility:
 | `PYWORKFLOW_CELERY_BROKER` | Celery broker URL |
 | `PYWORKFLOW_CELERY_RESULT_BACKEND` | Celery result backend URL |
 | `PYWORKFLOW_DISCOVER` | Modules to import for workflow discovery |
+| `PYWORKFLOW_RESCHEDULE_ON_SIGTERM` | Hand in-flight tasks off to another worker when a worker gets SIGTERM (default on; `0` falls back to Celery's warm shutdown) |
+| `PYWORKFLOW_SOFT_SHUTDOWN_TIMEOUT` | Seconds a SIGTERMed worker lets running tasks finish before terminating them and letting the broker redeliver each once (default `20`). Keep it below your orchestrator's grace period (Kubernetes `terminationGracePeriodSeconds`, 30s by default), or a SIGKILL during the wait leaves tasks unacked until the broker's `visibility_timeout` |
+| `PYWORKFLOW_CANCEL_TASKS_ON_CONNECTION_LOSS` | Terminate running tasks when the worker loses its broker connection (Redis restart, Sentinel failover), because their messages are redelivered to another worker (default on). `0` keeps them running, and the redelivered copy then runs at the same time |
 
 ```bash
 # Example: Production deployment with environment variables

{pyworkflow_engine-0.4.0 → pyworkflow_engine-0.4.1}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.4.0"
+version = "0.4.1"
 description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -33,7 +33,7 @@ dependencies = [
     # Core dependencies
     "celery>=5.3.0,<6.0.0",
     "cloudpickle>=3.0.0",
-    "pydantic>=2.0.0,<3.0.0",
+    "pydantic>=2.14.1,<3.0.0",
     "loguru>=0.7.0",
 
     # CLI dependencies
@@ -43,7 +43,7 @@ dependencies = [
     # Utilities
     "httpx>=0.25.0",
     "python-dateutil>=2.8.0",
-    "filelock>=3.32.4",
+    "filelock>=4.2.4",
     "pyyaml>=6.0.0",
     "croniter>=6.2.4",
 ]
@@ -59,15 +59,15 @@ sqlite = [
 ]
 
 postgres = [
-    "asyncpg>=0.29.0",
+    "asyncpg>=0.32.0",
 ]
 
 # AWS runtime support
 aws = [
-    "aws-durable-execution-sdk-python>=1.6.0",
+    "aws-durable-execution-sdk-python>=2.0.1",
 ]
 dynamodb = [
-    "aiobotocore>=3.9.0",
+    "aiobotocore>=3.9.2",
 ]
 
 cassandra = [
@@ -82,9 +82,9 @@ mysql = [
 all = [
     "redis>=8.1.0",
     "aiosqlite>=0.19.0",
-    "asyncpg>=0.29.0",
+    "asyncpg>=0.32.0",
     "aiomysql>=0.2.0",
-    "aws-durable-execution-sdk-python>=1.6.0",
+    "aws-durable-execution-sdk-python>=2.0.1",
     "cassandra-driver>=3.30.1",
 ]
 
@@ -95,27 +95,27 @@ dev = [
     "pytest-asyncio>=0.21.0",
     "pytest-cov>=4.1.0",
     "pytest-celery>=0.0.0",
-    "pytest-mock>=3.12.0",
-    "moto[dynamodb]>=5.2.3",  # AWS mocking for DynamoDB tests
+    "pytest-mock>=3.17.0",
+    "moto[dynamodb]>=5.2.4",  # AWS mocking for DynamoDB tests
 
     # Code quality
-    "black>=26.5.1",
-    "ruff>=0.16.4",
-    "mypy>=2.3.1",
-    "pre-commit>=4.6.2",
+    "black>=26.10.1",
+    "ruff>=0.17.2",
+    "mypy>=2.4.0",
+    "pre-commit>=4.7.0",
 
     # Type stubs
     "types-redis>=4.6.0",
-    "types-python-dateutil>=2.9.0.20260807",
-    "types-PyYAML>=6.0.12.20260815",
+    "types-python-dateutil>=2.9.0.20261021",
+    "types-PyYAML>=6.0.12.20260906",
 
     # Monitoring (optional)
-    "flower>=2.1.0",  # Celery web UI
+    "flower>=2.2.0",  # Celery web UI
 
     # Include all storage backends for development
     "redis>=8.1.0",
     "aiosqlite>=0.19.0",
-    "asyncpg>=0.29.0",
+    "asyncpg>=0.32.0",
 ]
 
 [project.urls]

{pyworkflow_engine-0.4.0 → pyworkflow_engine-0.4.1}/pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"
 
 # Configuration
 from pyworkflow.config import (

{pyworkflow_engine-0.4.0 → pyworkflow_engine-0.4.1}/pyworkflow/celery/app.py

@@ -12,17 +12,15 @@ warnings in Celery logs. This is a known compatibility issue between Python 3.13
 garbage collector and Celery's saferepr module. It does not affect functionality.
 """
 
-import contextlib
 import os
 from datetime import timedelta
 from typing import Any
 
 from celery import Celery
 from celery.signals import (
-    task_postrun,
-    task_prerun,
     worker_init,
     worker_process_init,
+    worker_ready,
     worker_shutdown,
 )
 from kombu import Exchange, Queue
@@ -288,9 +286,13 @@ def create_celery_app(
         PYWORKFLOW_CELERY_SENTINEL_MASTER: Sentinel master name (used if sentinel_master_name param not provided)
         PYWORKFLOW_WORKER_MAX_MEMORY: Max memory per worker child (KB) (used if worker_max_memory_per_child param not provided)
         PYWORKFLOW_WORKER_MAX_TASKS: Max tasks per worker child (used if worker_max_tasks_per_child param not provided)
-        PYWORKFLOW_RESCHEDULE_ON_SIGTERM: Re-enqueue in-flight tasks immediately when a worker is shut down
+        PYWORKFLOW_RESCHEDULE_ON_SIGTERM: Hand in-flight tasks off to another worker when a worker is shut down
             (SIGTERM / spot reclaim), instead of waiting out the broker visibility_timeout. Default on; set to
-            "0"/"false"/"no" to disable. Implemented as a consumer bootstep; see pyworkflow/celery/reschedule.py.
+            "0"/"false"/"no" to disable. See pyworkflow/celery/reschedule.py.
+        PYWORKFLOW_SOFT_SHUTDOWN_TIMEOUT: Seconds a SIGTERMed worker lets running tasks finish before handing
+            them off (default 20). Must be lower than the orchestrator's grace period.
+        PYWORKFLOW_CANCEL_TASKS_ON_CONNECTION_LOSS: Terminate running tasks when the broker connection is lost,
+            since their messages are redelivered (default on; "0" keeps them running at the risk of duplicates).
 
     Examples:
         # Default configuration (uses env vars if set, otherwise localhost Redis)
@@ -461,6 +463,14 @@ def create_celery_app(
             }
         )
 
+    # SIGTERM handoff: let running tasks finish for a bounded time, then terminate
+    # them and let the broker redeliver each one once. On a broker reconnect,
+    # terminate running tasks whose messages are being redelivered. See reschedule.py.
+    from pyworkflow.celery.reschedule import connection_loss_config, soft_shutdown_config
+
+    config_dict.update(soft_shutdown_config())
+    config_dict.update(connection_loss_config())
+
     # Configure Celery
     app.conf.update(config_dict)
 
@@ -481,14 +491,6 @@ def create_celery_app(
     # to ensure proper initialization AFTER process forking.
     # See on_worker_init() and on_worker_process_init() below.
 
-    # Re-enqueue in-flight tasks on worker shutdown (SIGTERM / spot reclaim) so a
-    # live worker resumes them within seconds instead of waiting out the broker
-    # visibility_timeout. Hooked as a consumer bootstep so it fires exactly as the
-    # consumer is torn down (event-based). See pyworkflow/celery/reschedule.py.
-    from pyworkflow.celery.reschedule import RescheduleConsumerStep
-
-    app.steps["consumer"].add(RescheduleConsumerStep)
-
     # Auto-discover workflows from environment variable or configured modules
     discover_workflows()
 
@@ -573,27 +575,16 @@ def on_worker_shutdown(**kwargs):
         close_worker_loop()
 
 
-@task_prerun.connect
-def on_task_prerun(task_id=None, task=None, args=None, kwargs=None, **_):
+@worker_ready.connect
+def on_worker_ready(sender=None, **_):
     """
-    Track the task now executing in this child process so it can be re-enqueued
-    if the worker is shut down (SIGTERM) before the task finishes.
-
-    Wrapped defensively: tracking must never interfere with task execution.
+    Make SIGTERM hand in-flight tasks off to another worker (soft + cold
+    shutdown) instead of Celery's default warm shutdown. Installed here, after
+    Celery has set up its own signal handlers. See pyworkflow/celery/reschedule.py.
     """
-    from pyworkflow.celery.reschedule import track_task_start
-
-    with contextlib.suppress(Exception):
-        track_task_start(task, task_id, args, kwargs)
-
-
-@task_postrun.connect
-def on_task_postrun(task_id=None, **_):
-    """Clear the in-flight task once it has finished (clean completion)."""
-    from pyworkflow.celery.reschedule import track_task_end
+    from pyworkflow.celery.reschedule import install_sigterm_handoff
 
-    with contextlib.suppress(Exception):
-        track_task_end(task_id)
+    install_sigterm_handoff(getattr(sender, "controller", None))
 
 
 def get_celery_app() -> Celery:

pyworkflow_engine-0.4.1/pyworkflow/celery/reschedule.py

@@ -0,0 +1,178 @@
+"""
+Hand in-flight tasks off to another worker when a worker is shut down (SIGTERM).
+
+When a Celery worker pod is reclaimed (e.g. an AWS spot node drain), it receives
+SIGTERM. Celery's default *warm* shutdown waits for running tasks to finish, and
+if the pod is SIGKILLed first, the task's message (``task_acks_late=True``) stays
+unacked until the Redis ``visibility_timeout`` (an hour by default). The parent
+workflow stays SUSPENDED that whole time, so the user-facing flow silently hangs.
+
+How the handoff works
+---------------------
+SIGTERM runs Celery's *cold* shutdown with a *soft* phase first (Celery >= 5.5):
+
+1. Soft phase: the worker waits ``worker_soft_shutdown_timeout`` seconds, so a
+   task that finishes in time acks normally and nothing is handed off.
+2. Cold phase: Celery cancels the task consumer, terminates the pool processes
+   still running a task, and closes the broker connection. Kombu restores the
+   unacked messages to their queues, and a live worker picks them up within
+   seconds.
+
+Each in-flight message goes back to the queue exactly once, and only after the
+process that was running it is dead.
+
+Why not re-enqueue at the start of shutdown (#733)
+--------------------------------------------------
+The previous design re-published every in-flight task as soon as the consumer
+stopped. Warm shutdown then let the original keep running, and closing the
+consumer connection made kombu restore the original message too, so every
+in-flight task ran three times at once: the original plus two copies. The
+completion checks in the tasks can't prevent that, because nothing is completed
+while all three are still running.
+
+Celery does the same thing when ``REMAP_SIGTERM=SIGQUIT`` is set, but billiard
+reads that variable at import time, which is too early to rely on for every way
+a worker can be launched. So we install the handler ourselves once the worker is
+ready; when ``REMAP_SIGTERM`` is set, Celery already handles SIGTERM and we
+leave it alone. Only the prefork pool is supported: terminating a running task
+requires a separate process.
+
+The timeout must be lower than the orchestrator's grace period (Kubernetes
+``terminationGracePeriodSeconds``, 30s by default). A SIGKILL during the soft
+phase leaves the message unacked until the ``visibility_timeout`` again.
+
+Broker connection loss
+----------------------
+Losing the broker connection (a Redis restart, a Sentinel failover) runs the
+same kombu restore: Celery collects the old connection and the unacked messages
+go back to the queue. By default Celery leaves the running tasks alive, so the
+original and the redelivered copy run at the same time. We enable Celery's
+``worker_cancel_long_running_tasks_on_connection_loss`` (the default from
+Celery 6) so the original is terminated as its message is restored.
+
+Gates:
+- ``PYWORKFLOW_RESCHEDULE_ON_SIGTERM`` (default on; set to ``0``/``false``/``no``
+  to fall back to Celery's plain warm shutdown).
+- ``PYWORKFLOW_SOFT_SHUTDOWN_TIMEOUT`` (seconds, default 20). ``0`` hands tasks
+  off immediately without waiting for them.
+- ``PYWORKFLOW_CANCEL_TASKS_ON_CONNECTION_LOSS`` (default on; ``0`` keeps
+  running tasks alive across a broker reconnect, at the risk of a duplicate).
+"""
+
+import os
+import threading
+from typing import Any
+
+from loguru import logger
+
+DEFAULT_SOFT_SHUTDOWN_TIMEOUT = 20.0
+
+
+def _flag(name: str) -> bool:
+    return os.getenv(name, "1").strip().lower() not in ("0", "false", "no")
+
+
+def _enabled() -> bool:
+    return _flag("PYWORKFLOW_RESCHEDULE_ON_SIGTERM")
+
+
+def soft_shutdown_timeout() -> float:
+    """Seconds a SIGTERMed worker waits for running tasks before handing them off."""
+    raw = os.getenv("PYWORKFLOW_SOFT_SHUTDOWN_TIMEOUT", "").strip()
+    if not raw:
+        return DEFAULT_SOFT_SHUTDOWN_TIMEOUT
+    try:
+        return max(float(raw), 0.0)
+    except ValueError:
+        logger.warning(
+            f"reschedule: invalid PYWORKFLOW_SOFT_SHUTDOWN_TIMEOUT={raw!r}, "
+            f"using {DEFAULT_SOFT_SHUTDOWN_TIMEOUT}s"
+        )
+        return DEFAULT_SOFT_SHUTDOWN_TIMEOUT
+
+
+def soft_shutdown_config() -> dict[str, Any]:
+    """Celery settings for the SIGTERM handoff (empty when disabled)."""
+    if not _enabled():
+        return {}
+    return {"worker_soft_shutdown_timeout": soft_shutdown_timeout()}
+
+
+def connection_loss_config() -> dict[str, Any]:
+    """Celery settings that stop a task before a broker reconnect redelivers it."""
+    return {
+        "worker_cancel_long_running_tasks_on_connection_loss": _flag(
+            "PYWORKFLOW_CANCEL_TASKS_ON_CONNECTION_LOSS"
+        )
+    }
+
+
+def _is_worker_stopping() -> bool:
+    """True once the worker has started shutting down (warm or cold)."""
+    try:
+        from celery.worker import state
+
+        # Identity checks (matching celery.worker.state): the flags are set to the
+        # exit code, which can be 0 (EX_OK) -- `0 == False`.
+        def _set(v: Any) -> bool:
+            return v is not None and v is not False
+
+        return _set(state.should_stop) or _set(state.should_terminate)
+    except Exception:
+        return False
+
+
+def _is_prefork(worker: Any) -> bool:
+    from celery.concurrency.prefork import TaskPool
+
+    return isinstance(getattr(worker, "pool", None), TaskPool)
+
+
+def _on_sigterm(worker: Any) -> None:
+    """Run Celery's soft + cold shutdown, as it does for ``REMAP_SIGTERM=SIGQUIT``."""
+    from celery.apps import worker as celery_worker
+
+    # A second SIGTERM must not start a second soft wait: treat it like a repeated
+    # SIGQUIT (cancel the tasks now, keep waiting for the cold shutdown).
+    celery_worker.install_worker_term_hard_handler(
+        worker, sig="SIGTERM", callback=celery_worker.during_soft_shutdown
+    )
+    logger.warning(
+        "reschedule: SIGTERM received, handing off in-flight tasks after "
+        f"{worker.app.conf.worker_soft_shutdown_timeout}s"
+    )
+    celery_worker.on_cold_shutdown(worker)
+
+
+def install_sigterm_handoff(worker: Any) -> bool:
+    """
+    Make SIGTERM hand in-flight tasks off via Celery's soft + cold shutdown.
+
+    Called once the worker is ready (after Celery installed its own signal
+    handlers). Returns True if the handler was installed. Never raises.
+    """
+    if not _enabled() or worker is None or _is_worker_stopping():
+        return False
+    try:
+        from celery.apps import worker as celery_worker
+
+        if not all(hasattr(celery_worker, n) for n in ("on_cold_shutdown", "during_soft_shutdown")):
+            logger.warning(
+                "reschedule: SIGTERM handoff needs Celery >= 5.5 (soft shutdown), "
+                "keeping Celery's warm shutdown"
+            )
+            return False
+        if getattr(celery_worker, "REMAP_SIGTERM", None):
+            return False  # Celery already maps SIGTERM to the cold shutdown.
+        if not _is_prefork(worker):
+            logger.info("reschedule: SIGTERM handoff requires the prefork pool, skipping")
+            return False
+        if threading.current_thread() is not threading.main_thread():
+            logger.warning("reschedule: not on the main thread, cannot install SIGTERM handler")
+            return False
+
+        celery_worker.install_worker_term_hard_handler(worker, sig="SIGTERM", callback=_on_sigterm)
+        return True
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to install SIGTERM handoff: {exc}")
+        return False