pyworkflow-engine 0.3.2__tar.gz → 0.3.6__tar.gz

{pyworkflow_engine-0.3.2 → pyworkflow_engine-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyworkflow-engine
-Version: 0.3.2
+Version: 0.3.6
 Summary: A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow
 Author: PyWorkflow Contributors
 License: MIT
@@ -26,11 +26,11 @@ 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: loguru>=0.7.0
-Requires-Dist: click>=8.3.3
+Requires-Dist: click>=8.4.1
 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.12.0
+Requires-Dist: filelock>=3.29.3
 Requires-Dist: pyyaml>=6.0.0
 Requires-Dist: croniter>=2.0.0
 Provides-Extra: redis
@@ -42,7 +42,7 @@ Requires-Dist: asyncpg>=0.29.0; extra == "postgres"
 Provides-Extra: aws
 Requires-Dist: aws-durable-execution-sdk-python>=0.1.0; extra == "aws"
 Provides-Extra: dynamodb
-Requires-Dist: aiobotocore>=3.6.0; extra == "dynamodb"
+Requires-Dist: aiobotocore>=3.7.0; extra == "dynamodb"
 Provides-Extra: cassandra
 Requires-Dist: cassandra-driver>=3.29.0; extra == "cassandra"
 Provides-Extra: mysql
@@ -60,10 +60,10 @@ 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.0; extra == "dev"
-Requires-Dist: black>=23.0.0; extra == "dev"
-Requires-Dist: ruff>=0.15.12; extra == "dev"
-Requires-Dist: mypy>=1.7.0; extra == "dev"
+Requires-Dist: moto[dynamodb]>=5.2.2; extra == "dev"
+Requires-Dist: black>=26.5.1; extra == "dev"
+Requires-Dist: ruff>=0.15.17; extra == "dev"
+Requires-Dist: mypy>=2.1.0; extra == "dev"
 Requires-Dist: pre-commit>=4.6.0; extra == "dev"
 Requires-Dist: types-redis>=4.6.0; extra == "dev"
 Requires-Dist: types-python-dateutil>=2.8.0; extra == "dev"

{pyworkflow_engine-0.3.2 → pyworkflow_engine-0.3.6}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.3.2"
+version = "0.3.6"
 description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -37,13 +37,13 @@ dependencies = [
     "loguru>=0.7.0",
 
     # CLI dependencies
-    "click>=8.3.3",
+    "click>=8.4.1",
     "inquirerpy>=0.3.4; python_version < '4.0'",
 
     # Utilities
     "httpx>=0.25.0",
     "python-dateutil>=2.8.0",
-    "filelock>=3.12.0",
+    "filelock>=3.29.3",
     "pyyaml>=6.0.0",
     "croniter>=2.0.0",
 ]
@@ -67,7 +67,7 @@ aws = [
     "aws-durable-execution-sdk-python>=0.1.0",
 ]
 dynamodb = [
-    "aiobotocore>=3.6.0",
+    "aiobotocore>=3.7.0",
 ]
 
 cassandra = [
@@ -96,12 +96,12 @@ dev = [
     "pytest-cov>=4.1.0",
     "pytest-celery>=0.0.0",
     "pytest-mock>=3.12.0",
-    "moto[dynamodb]>=5.2.0",  # AWS mocking for DynamoDB tests
+    "moto[dynamodb]>=5.2.2",  # AWS mocking for DynamoDB tests
 
     # Code quality
-    "black>=23.0.0",
-    "ruff>=0.15.12",
-    "mypy>=1.7.0",
+    "black>=26.5.1",
+    "ruff>=0.15.17",
+    "mypy>=2.1.0",
     "pre-commit>=4.6.0",
 
     # Type stubs

{pyworkflow_engine-0.3.2 → pyworkflow_engine-0.3.6}/pyworkflow/__init__.py

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

{pyworkflow_engine-0.3.2 → pyworkflow_engine-0.3.6}/pyworkflow/celery/app.py

@@ -12,12 +12,19 @@ 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 worker_init, worker_process_init, worker_shutdown
+from celery.signals import (
+    task_postrun,
+    task_prerun,
+    worker_init,
+    worker_process_init,
+    worker_shutdown,
+)
 from kombu import Exchange, Queue
 
 from pyworkflow.observability.logging import configure_logging
@@ -281,6 +288,9 @@ 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
+            (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.
 
     Examples:
         # Default configuration (uses env vars if set, otherwise localhost Redis)
@@ -471,6 +481,14 @@ 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()
 
@@ -555,6 +573,29 @@ def on_worker_shutdown(**kwargs):
         close_worker_loop()
 
 
+@task_prerun.connect
+def on_task_prerun(task_id=None, task=None, args=None, kwargs=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.
+    """
+    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
+
+    with contextlib.suppress(Exception):
+        track_task_end(task_id)
+
+
 def get_celery_app() -> Celery:
     """
     Get the global Celery application instance.

pyworkflow_engine-0.3.6/pyworkflow/celery/reschedule.py

@@ -0,0 +1,339 @@
+"""
+Reschedule in-flight tasks when a worker is shut down (SIGTERM).
+
+When a Celery worker pod is reclaimed (e.g. an AWS spot node drain), it receives
+SIGTERM and warm-shuts-down. Any step task that was executing is *not* promptly
+redelivered: with ``task_acks_late=True`` the broker only re-queues it after the
+Redis ``visibility_timeout`` (an hour by default), and the parent workflow stays
+SUSPENDED that whole time -- so the user-facing flow silently hangs with no output.
+
+This module makes the worker re-enqueue its in-flight task(s) *immediately* on
+shutdown so a live worker resumes them within seconds instead of ~1h.
+
+Why a shared Redis registry (not process-local state)
+-----------------------------------------------------
+Tasks execute in forked **child** processes, but the shutdown hook runs in the
+**main** process. The main process can't see a child's in-memory task args, so
+children publish what they are running to a per-pod Redis hash
+(``pyworkflow:inflight:<pod>``) on ``task_prerun`` and remove it on
+``task_postrun``. On shutdown the main process reads that hash and re-enqueues
+every entry.
+
+Why a consumer bootstep (event-based, no sleep)
+-----------------------------------------------
+The re-enqueue must happen *after* this worker's broker consumer is cancelled --
+otherwise the worker's own still-active consumer grabs the message straight back
+into ``unacked`` (stranded until the visibility_timeout). There is no Celery
+*signal* for "consumer cancelled" (``worker_shutting_down`` fires before it;
+``worker_shutdown`` fires after the pool has drained the long task -- far too
+late). So we hook the consumer blueprint directly with a bootstep that
+``requires`` the ``Tasks`` step: its ``stop()`` runs as the consumer is torn
+down (early in shutdown, before the pool drain). There we cancel the task
+consumer ourselves -- guaranteeing this worker stops fetching -- and only then
+re-enqueue, so the message lands in the *ready* queue and another worker's
+blocking BRPOP picks it up immediately. Event-based, deterministic, no sleep.
+
+Safety against double execution
+-------------------------------
+``execute_step_task`` short-circuits on an existing ``STEP_COMPLETED`` event (and
+on a terminal run status), and ``_record_step_completion_and_resume`` re-checks
+before recording, so at most one completion is ever recorded no matter how many
+copies run. Only ``SingletonWorkflowTask`` tasks that declare ``unique_on`` are
+rescheduled -- those are idempotent by run_id/step_id. Tasks without a uniqueness
+key are left alone.
+
+Gate: ``PYWORKFLOW_RESCHEDULE_ON_SIGTERM`` (default on; set to ``0``/``false``/``no``
+to disable and fall back to the old visibility_timeout redelivery behaviour).
+"""
+
+import json
+import os
+import socket
+import threading
+from typing import Any
+
+from celery import bootsteps
+from loguru import logger
+
+# Registry hash self-expires so a hard-killed pod (no postrun) cannot leak entries.
+_REGISTRY_TTL_SECONDS = 6 * 3600
+_REGISTRY_KEY_PREFIX = "pyworkflow:inflight:"
+
+_backend_lock = threading.Lock()
+_backend: Any = None
+_backend_resolved = False
+
+
+def _enabled() -> bool:
+    return os.getenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", "1").strip().lower() not in (
+        "0",
+        "false",
+        "no",
+    )
+
+
+def _pod_id() -> str:
+    """Stable per-pod id shared by the main process and its forked children."""
+    return os.environ.get("HOSTNAME") or socket.gethostname()
+
+
+def _registry_key() -> str:
+    return f"{_REGISTRY_KEY_PREFIX}{_pod_id()}"
+
+
+def _get_backend() -> Any:
+    """
+    Cached sentinel-aware Redis backend, reusing the singleton lock configuration
+    (same broker/sentinel the task locks already use). Returns None when no Redis
+    backend is configured, in which case rescheduling is a no-op.
+    """
+    global _backend, _backend_resolved
+    if _backend_resolved:
+        return _backend
+    with _backend_lock:
+        if _backend_resolved:
+            return _backend
+        try:
+            from pyworkflow.celery.app import celery_app
+            from pyworkflow.celery.singleton import RedisLockBackend, SingletonConfig
+
+            cfg = SingletonConfig(celery_app)
+            url = cfg.backend_url
+            if url:
+                _backend = RedisLockBackend(
+                    url,
+                    is_sentinel=cfg.is_sentinel,
+                    sentinel_master=cfg.sentinel_master,
+                )
+        except Exception as exc:
+            logger.warning(f"reschedule: could not initialise redis backend: {exc}")
+            _backend = None
+        _backend_resolved = True
+    return _backend
+
+
+def _resolve_task(name: str) -> Any:
+    """Look up a registered task instance by name (used to re-enqueue)."""
+    from pyworkflow.celery.app import celery_app
+
+    return celery_app.tasks.get(name)
+
+
+def _resolve_queue(task: Any) -> str | None:
+    """Best-effort queue the task was delivered on, so the re-enqueue lands there."""
+    try:
+        delivery_info = getattr(task.request, "delivery_info", None) or {}
+        routing_key = delivery_info.get("routing_key")
+        if routing_key:
+            return routing_key
+    except Exception:
+        pass
+    return getattr(task, "queue", None)
+
+
+def track_task_start(
+    task: Any,
+    task_id: str | None,
+    args: Any,
+    kwargs: Any,
+) -> None:
+    """Publish a task as in-flight (per pod) so it can be re-enqueued on SIGTERM."""
+    if not _enabled() or task is None or task_id is None:
+        return
+
+    # Import here to avoid import cycles at module load.
+    from pyworkflow.celery.singleton import SingletonWorkflowTask
+
+    # Only re-enqueue tasks keyed by a uniqueness field (run_id/step_id). These are
+    # idempotent; blindly re-enqueuing a non-unique task could start duplicate runs.
+    if not isinstance(task, SingletonWorkflowTask) or not task.unique_on:
+        return
+
+    backend = _get_backend()
+    if backend is None:
+        return
+
+    descriptor = json.dumps(
+        {
+            "name": task.name,
+            "args": list(args or ()),
+            "kwargs": dict(kwargs or {}),
+            "queue": _resolve_queue(task),
+        },
+        default=str,
+    )
+    key = _registry_key()
+    try:
+        backend._execute_with_refresh(lambda: backend.redis.hset(key, task_id, descriptor))
+        backend._execute_with_refresh(lambda: backend.redis.expire(key, _REGISTRY_TTL_SECONDS))
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to record in-flight task {task_id}: {exc}")
+
+
+def track_task_end(task_id: str | None) -> None:
+    """Remove a task from the in-flight registry once it has finished."""
+    if not _enabled() or task_id is None:
+        return
+    backend = _get_backend()
+    if backend is None:
+        return
+    try:
+        backend._execute_with_refresh(lambda: backend.redis.hdel(_registry_key(), task_id))
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to clear in-flight task {task_id}: {exc}")
+
+
+def reschedule_inflight_on_shutdown() -> None:
+    """
+    Re-enqueue every task still recorded as in-flight for this pod. Assumes the
+    caller has already cancelled this worker's task consumer (see
+    :class:`RescheduleConsumerStep`) so the fresh copy is not re-grabbed. Never
+    raises -- shutdown must not be blocked.
+    """
+    if not _enabled():
+        return
+    backend = _get_backend()
+    if backend is None:
+        return
+
+    key = _registry_key()
+    try:
+        entries = backend._execute_with_refresh(lambda: backend.redis.hgetall(key)) or {}
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to read in-flight registry: {exc}")
+        return
+
+    if not entries:
+        return
+
+    for task_id, raw in entries.items():
+        try:
+            descriptor = json.loads(raw)
+            task = _resolve_task(descriptor["name"])
+            if task is None:
+                logger.warning(
+                    f"reschedule: unknown task {descriptor.get('name')!r}, "
+                    f"cannot re-enqueue {task_id}"
+                )
+                continue
+
+            args = descriptor.get("args") or []
+            kwargs = descriptor.get("kwargs") or {}
+            queue = descriptor.get("queue")
+
+            # Release the singleton lock held by the dying task so the re-enqueued
+            # copy is not rejected as a duplicate. release_lock() regenerates the
+            # exact lock key from the same args/kwargs the running task holds.
+            try:
+                task.release_lock(task_args=list(args), task_kwargs=kwargs)
+            except Exception as exc:
+                logger.warning(
+                    f"reschedule: failed to release lock for {task.name} "
+                    f"(old_task_id={task_id}): {exc}"
+                )
+
+            options = {"queue": queue} if queue else {}
+            task.apply_async(args=list(args), kwargs=kwargs, **options)
+            logger.warning(
+                f"reschedule: re-enqueued in-flight task on shutdown "
+                f"task={task.name} old_task_id={task_id} queue={queue}"
+            )
+        except Exception as exc:
+            logger.opt(exception=True).error(
+                f"reschedule: failed to re-enqueue task {task_id}: {exc}"
+            )
+
+    # Drop the registry for this pod; a fresh worker starts with a clean slate.
+    try:
+        backend._execute_with_refresh(lambda: backend.redis.delete(key))
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to clear in-flight registry: {exc}")
+
+
+def _is_worker_stopping() -> bool:
+    """
+    True only when the worker is actually shutting down (SIGTERM warm or cold),
+    not on a transient consumer restart (e.g. broker reconnect), which also tears
+    the consumer down. Celery's SIGTERM handler sets these flags before the
+    consumer is stopped.
+    """
+    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`, so membership/equality
+        # against False would wrongly treat a clean exitcode-0 shutdown as "not
+        # stopping" and skip the reschedule.
+        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 on_consumer_stop(consumer: Any) -> None:
+    """
+    Invoked from the consumer bootstep as the consumer is torn down. Cancels this
+    worker's task consumer (so it stops fetching) and re-enqueues in-flight work.
+    """
+    if not _enabled() or not _is_worker_stopping():
+        return
+
+    # Cancel our own task consumer first: once we stop BRPOPing, the re-enqueued
+    # message stays in the ready list for another worker instead of being grabbed
+    # back here into `unacked`. Tasks.stop() will cancel again later (idempotent).
+    task_consumer = getattr(consumer, "task_consumer", None)
+    if task_consumer is not None:
+        try:
+            task_consumer.cancel()
+        except Exception as exc:
+            logger.warning(f"reschedule: failed to cancel task consumer: {exc}")
+
+    # Cancelling is not enough on the Redis transport: a BRPOP this worker already
+    # issued is accepted by the server and will still pop the *next* message pushed
+    # to those queues -- including the copy we are about to re-enqueue. The dying
+    # worker then exits before recording it in `unacked`, so the message is lost
+    # entirely (neither in the ready list nor restorable via visibility_timeout).
+    # Drop the consumer's broker connection so any in-flight BRPOP is aborted
+    # server-side *before* we re-publish; the re-enqueue itself uses the producer
+    # pool (a separate connection), so it is unaffected.
+    _abort_consumer_fetching(consumer)
+
+    reschedule_inflight_on_shutdown()
+
+
+def _abort_consumer_fetching(consumer: Any) -> None:
+    """Close the consumer's broker connection to abort any in-flight BRPOP."""
+    conn = getattr(consumer, "connection", None)
+    if conn is None:
+        return
+    try:
+        # collect() tears down channels + the underlying socket, aborting the
+        # pending BRPOP. Never raises into shutdown.
+        conn.collect()
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to drop consumer connection: {exc}")
+
+
+class RescheduleConsumerStep(bootsteps.StartStopStep):
+    """
+    Consumer bootstep that re-enqueues this worker's in-flight tasks the moment
+    the consumer is torn down on shutdown -- event-based, no polling/sleep.
+
+    ``requires`` the task consumer (``Tasks``) so that (a) ``c.task_consumer``
+    exists when our ``stop()`` runs and (b) we run *before* ``Tasks.stop()``,
+    letting us cancel the consumer ourselves and re-enqueue while it's guaranteed
+    not fetching.
+    """
+
+    requires = ("celery.worker.consumer:Tasks",)
+
+    def start(self, c: Any) -> None:  # noqa: D401 - nothing to start
+        pass
+
+    def stop(self, c: Any) -> None:
+        try:
+            on_consumer_stop(c)
+        except Exception as exc:  # never block consumer teardown
+            logger.warning(f"reschedule: consumer-stop hook failed: {exc}")

{pyworkflow_engine-0.3.2 → pyworkflow_engine-0.3.6}/pyworkflow/cli/utils/interactive.py

@@ -76,15 +76,12 @@ def select(
     try:
         # Convert string list to Choice objects
         choice_objects: list[Choice] = []
-        if choices and isinstance(choices[0], str):
-            for c in choices:
-                if isinstance(c, str):
-                    choice_objects.append(Choice(value=c, name=c))
-        else:
-            # Dict format: {"name": "...", "value": "..."}
-            for c in choices:
-                if isinstance(c, dict):
-                    choice_objects.append(Choice(value=c.get("value", c["name"]), name=c["name"]))
+        for c in choices:
+            if isinstance(c, str):
+                choice_objects.append(Choice(value=c, name=c))
+            else:
+                # Dict format: {"name": "...", "value": "..."}
+                choice_objects.append(Choice(value=c.get("value", c["name"]), name=c["name"]))
 
         return inquirer.select(
             message=message,
@@ -222,22 +219,19 @@ def multiselect(
     try:
         # Convert string list to Choice objects
         choice_objects: list[Choice] = []
-        if choices and isinstance(choices[0], str):
-            for c in choices:
-                if isinstance(c, str):
-                    choice_objects.append(Choice(value=c, name=c, enabled=(c in (default or []))))
-        else:
-            # Dict format: {"name": "...", "value": "..."}
-            for c in choices:
-                if isinstance(c, dict):
-                    value = c.get("value", c["name"])
-                    choice_objects.append(
-                        Choice(
-                            value=value,
-                            name=c["name"],
-                            enabled=(value in (default or [])),
-                        )
+        for c in choices:
+            if isinstance(c, str):
+                choice_objects.append(Choice(value=c, name=c, enabled=(c in (default or []))))
+            else:
+                # Dict format: {"name": "...", "value": "..."}
+                value = c.get("value", c["name"])
+                choice_objects.append(
+                    Choice(
+                        value=value,
+                        name=c["name"],
+                        enabled=(value in (default or [])),
                     )
+                )
 
         return inquirer.checkbox(
             message=message,

{pyworkflow_engine-0.3.2 → pyworkflow_engine-0.3.6}/pyworkflow/storage/citus.py

@@ -308,6 +308,14 @@ class CitusStorageBackend(PostgresStorageBackend):
             await conn.execute(
                 "CREATE INDEX IF NOT EXISTS idx_runs_created_at ON workflow_runs(created_at DESC)"
             )
+            # Composite index backing keyset pagination + ORDER BY created_at DESC, run_id DESC
+            # for the unfiltered Runs list (issue #482). The summary column projection
+            # (RUN_SUMMARY_COLUMNS) already drops the heavy payload columns from the scatter-gather;
+            # no INCLUDE covering columns, to avoid duplicating the summary columns into the index.
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_runs_created_at_run_id "
+                "ON workflow_runs(created_at DESC, run_id DESC)"
+            )
             # Non-unique index: Citus cannot enforce global uniqueness on non-distribution columns
             await conn.execute(
                 "CREATE INDEX IF NOT EXISTS idx_runs_idempotency_key ON workflow_runs(idempotency_key) WHERE idempotency_key IS NOT NULL"