dramatiq-workflow 0.1.0__py3-none-any.whl

dramatiq_workflow/__init__.py

@@ -0,0 +1,13 @@
+from ._base import Workflow
+from ._middleware import WorkflowMiddleware
+from ._models import Chain, Group, Message, WithDelay, WorkflowType
+
+__all__ = [
+    "Chain",
+    "Group",
+    "Message",
+    "WithDelay",
+    "Workflow",
+    "WorkflowMiddleware",
+    "WorkflowType",
+]

dramatiq_workflow/_barrier.py

@@ -0,0 +1,38 @@
+import dramatiq.rate_limits
+
+
+class AtMostOnceBarrier(dramatiq.rate_limits.Barrier):
+    """
+    The AtMostOnceBarrier is a barrier that ensures that it is released at most
+    once.
+
+    We use this because we want to avoid running callbacks in chains multiple
+    times. Running callbacks more than once can have compounding effects
+    especially when groups are involved.
+
+    The downside of this is that we cannot guarantee that the barrier will be
+    released at all. Theoretically a worker could die after releasing the
+    barrier but just before it has a chance to schedule the callbacks.
+    """
+
+    def __init__(self, backend, key, *args, ttl=900000):
+        super().__init__(backend, key, *args, ttl=ttl)
+        self.ran_key = f"{key}_ran"
+
+    def create(self, parties):
+        self.backend.add(self.ran_key, 0, self.ttl)
+        return super().create(parties)
+
+    def wait(self, *args, block=True, timeout=None):
+        if block:
+            # Blocking with an AtMostOnceBarrier is not supported as it could
+            # lead to clients waiting indefinitely if the barrier already
+            # released.
+            raise ValueError("Blocking is not supported by AtMostOnceBarrier")
+
+        released = super().wait(*args, block=False)
+        if released:
+            never_released = self.backend.incr(self.ran_key, 1, 1, self.ttl)
+            return never_released
+
+        return False

dramatiq_workflow/_base.py

@@ -0,0 +1,185 @@
+import logging
+import time
+from uuid import uuid4
+
+import dramatiq
+import dramatiq.rate_limits
+
+from ._constants import CALLBACK_BARRIER_TTL, OPTION_KEY_CALLBACKS
+from ._helpers import workflow_with_completion_callbacks
+from ._middleware import WorkflowMiddleware, workflow_noop
+from ._models import Barrier, Chain, CompletionCallbacks, Group, Message, WithDelay, WorkflowType
+from ._serialize import serialize_callbacks, serialize_workflow
+
+logger = logging.getLogger(__name__)
+
+
+class Workflow:
+    """
+    A workflow allows running tasks in parallel and in sequence. It is a way to
+    define a workflow of tasks, a combination of chains and groups in any
+    order and nested as needed.
+
+    Example:
+
+    Let's assume we want a workflow that looks like this:
+
+                 ╭────────╮  ╭────────╮
+                 │ Task 2 │  │ Task 5 │
+              ╭──┼●      ●┼──┼●      ●┼╮
+    ╭────────╮│  ╰────────╯  ╰────────╯│  ╭────────╮
+    │ Task 1 ││  ╭────────╮            │  │ Task 8 │
+    │       ●┼╯  │ Task 3 │            ╰──┼●       │
+    │       ●┼───┼●      ●┼───────────────┼●       │
+    │       ●┼╮  ╰────────╯             ╭─┼●       │
+    ╰────────╯│  ╭────────╮   ╭────────╮│╭┼●       │
+              │  │ Task 4 │   │ Task 6 │││╰────────╯
+              ╰──┼●      ●┼───┼●      ●┼╯│
+                 │       ●┼╮  ╰────────╯ │
+                 ╰────────╯│             │
+                           │  ╭────────╮ │
+                           │  │ Task 7 │ │
+                           ╰──┼●      ●┼─╯
+                              ╰────────╯
+
+    We can define this workflow as follows:
+
+    ```python
+    from dramatiq_workflow import Workflow, Chain, Group
+
+    workflow = Workflow(
+        Chain(
+            task1.message(),
+            Group(
+                Chain(
+                    task2.message(),
+                    task5.message(),
+                ),
+                task3.message(),
+                Chain(
+                    task4.message(),
+                    Group(
+                        task6.message(),
+                        task7.message(),
+                    ),
+                ),
+            ),
+            task8.message(),
+        ),
+    )
+    workflow.run() # Schedules the workflow to run in the background
+    ```
+
+    In this example, the execution would look like this*:
+    1. Task 1 runs
+    2. Task 2, 3, and 4 run in parallel once Task 1 finishes
+    3. Task 5 runs once Task 2 finishes
+    4. Task 6 and 7 run in parallel once Task 4 finishes
+    5. Task 8 runs once Task 5, 6, and 7 finish
+
+    * This is a simplified example. The actual execution order may vary because
+    tasks that can run in parallel (i.e. in a Group) are not guaranteed to run
+    in the order they are defined in the workflow.
+    """
+
+    def __init__(
+        self,
+        workflow: WorkflowType,
+        broker: dramatiq.Broker | None = None,
+    ):
+        self.workflow = workflow
+        self.broker = broker or dramatiq.get_broker()
+
+        self._delay = None
+        self._completion_callbacks = []
+
+        while isinstance(self.workflow, WithDelay):
+            self._delay = (self._delay or 0) + self.workflow.delay
+            self.workflow = self.workflow.task
+
+    def run(self):
+        current = self.workflow
+        completion_callbacks = self._completion_callbacks.copy()
+
+        if isinstance(current, Message):
+            current = self.__augment_message(current, completion_callbacks)
+            self.broker.enqueue(current, delay=self._delay)
+            return
+
+        if isinstance(current, Chain):
+            tasks = current.tasks[:]
+            if not tasks:
+                self.__schedule_noop(completion_callbacks)
+                return
+
+            task = tasks.pop(0)
+            if tasks:
+                completion_id = self.__create_barrier(1)
+                completion_callbacks.append((completion_id, Chain(*tasks), False))
+            self.__workflow_with_completion_callbacks(task, completion_callbacks).run()
+            return
+
+        if isinstance(current, Group):
+            tasks = current.tasks[:]
+            if not tasks:
+                self.__schedule_noop(completion_callbacks)
+                return
+
+            completion_id = self.__create_barrier(len(tasks))
+            completion_callbacks.append((completion_id, None, True))
+            for task in tasks:
+                self.__workflow_with_completion_callbacks(task, completion_callbacks).run()
+            return
+
+        raise TypeError(f"Unsupported workflow type: {type(current)}")
+
+    def __workflow_with_completion_callbacks(self, task, completion_callbacks) -> "Workflow":
+        return workflow_with_completion_callbacks(
+            task,
+            self.broker,
+            completion_callbacks,
+            delay=self._delay,
+        )
+
+    def __schedule_noop(self, completion_callbacks: CompletionCallbacks):
+        noop_message = workflow_noop.message()
+        noop_message = self.__augment_message(noop_message, completion_callbacks)
+        self.broker.enqueue(noop_message, delay=self._delay)
+
+    def __augment_message(self, message: Message, completion_callbacks: CompletionCallbacks) -> Message:
+        return message.copy(
+            # We reset the message timestamp to better represent the time the
+            # message was actually enqueued.  This is to avoid tripping the max_age
+            # check in the broker.
+            message_timestamp=time.time() * 1000,
+            options={OPTION_KEY_CALLBACKS: serialize_callbacks(completion_callbacks)},
+        )
+
+    @property
+    def __rate_limiter_backend(self):
+        if not hasattr(self, "__cached_rate_limiter_backend"):
+            for middleware in self.broker.middleware:
+                if isinstance(middleware, WorkflowMiddleware):
+                    self.__cached_rate_limiter_backend = middleware.rate_limiter_backend
+                    break
+            else:
+                raise RuntimeError(
+                    "WorkflowMiddleware middleware not found! Did you forget "
+                    "to set it up? It is required if you want to use "
+                    "workflows."
+                )
+        return self.__cached_rate_limiter_backend
+
+    def __create_barrier(self, count: int):
+        if count == 1:
+            # No need to create a distributed barrier if there is only one task
+            return None
+
+        completion_uuid = str(uuid4())
+        completion_barrier = Barrier(self.__rate_limiter_backend, completion_uuid, ttl=CALLBACK_BARRIER_TTL)
+        completion_barrier.create(count)
+        logger.debug("Barrier created: %s (%d tasks)", completion_uuid, count)
+        return completion_uuid
+
+    def __str__(self):
+        return f"Workflow({serialize_workflow(self.workflow)})"

dramatiq_workflow/_constants.py

@@ -0,0 +1,2 @@
+CALLBACK_BARRIER_TTL = 86_400_000
+OPTION_KEY_CALLBACKS = "workflow_completion_callbacks"

dramatiq_workflow/_helpers.py

@@ -0,0 +1,18 @@
+import dramatiq
+
+from ._models import CompletionCallbacks, WorkflowType
+
+
+def workflow_with_completion_callbacks(
+    workflow: WorkflowType,
+    broker: dramatiq.Broker,
+    completion_callbacks: CompletionCallbacks,
+    delay: int | None = None,
+):
+    from ._base import Workflow
+
+    w = Workflow(workflow, broker)
+    w._completion_callbacks = completion_callbacks
+    if delay is not None:
+        w._delay = (w._delay or 0) + delay
+    return w

dramatiq_workflow/_middleware.py

@@ -0,0 +1,67 @@
+import logging
+
+import dramatiq
+import dramatiq.rate_limits
+
+from ._constants import OPTION_KEY_CALLBACKS
+from ._helpers import workflow_with_completion_callbacks
+from ._models import Barrier
+from ._serialize import unserialize_callbacks, unserialize_workflow
+
+logger = logging.getLogger(__name__)
+
+
+class WorkflowMiddleware(dramatiq.Middleware):
+    def __init__(self, rate_limiter_backend: dramatiq.rate_limits.RateLimiterBackend):
+        self.rate_limiter_backend = rate_limiter_backend
+
+    def after_process_boot(self, broker: dramatiq.Broker):
+        broker.declare_actor(workflow_noop)
+
+    def after_process_message(
+        self, broker: dramatiq.Broker, message: dramatiq.broker.MessageProxy, *, result=None, exception=None
+    ):
+        if exception is not None:
+            # TODO: Add a way to handle exceptions in the workflow?
+            return
+
+        if message.failed:
+            return
+
+        completion_callbacks: list[dict] | None = message.options.get(OPTION_KEY_CALLBACKS)
+        if completion_callbacks is None:
+            return
+
+        # Go through the completion callbacks backwards until we hit the first non-completed barrier
+        while len(completion_callbacks) > 0:
+            completion_id, remaining_workflow, propagate = completion_callbacks[-1]
+            if completion_id is not None:
+                barrier = Barrier(self.rate_limiter_backend, completion_id)
+                if not barrier.wait(block=False):
+                    logger.debug("Barrier not completed: %s", completion_id)
+                    break
+
+            logger.debug("Barrier completed: %s", completion_id)
+            completion_callbacks.pop()
+            if remaining_workflow is None:
+                continue
+            workflow_with_completion_callbacks(
+                unserialize_workflow(remaining_workflow),
+                broker,
+                # TODO: This is somewhat inefficient because we're unserializing all callbacks
+                # even though we are just going to serialize them again.
+                unserialize_callbacks(completion_callbacks),
+            ).run()
+
+            if not propagate:
+                break
+
+
+@dramatiq.actor
+def workflow_noop():
+    """
+    This task does nothing. It is used as a placeholder when a Group or a Chain
+    has no tasks. This allows us to schedule the execution of the completion
+    callbacks via the WorkflowMiddleware.
+    """
+    pass

dramatiq_workflow/_models.py

@@ -0,0 +1,47 @@
+import dramatiq
+import dramatiq.rate_limits
+
+from ._barrier import AtMostOnceBarrier
+
+
+class Chain:
+    def __init__(self, *tasks: "WorkflowType"):
+        self.tasks = list(tasks)
+
+    def __str__(self):
+        return f"Chain({self.tasks})"
+
+    def __eq__(self, other):
+        return isinstance(other, Chain) and self.tasks == other.tasks
+
+
+class Group:
+    def __init__(self, *tasks: "WorkflowType"):
+        self.tasks = list(tasks)
+
+    def __str__(self):
+        return f"Group({self.tasks})"
+
+    def __eq__(self, other):
+        return isinstance(other, Group) and self.tasks == other.tasks
+
+
+class WithDelay:
+    def __init__(self, task: "WorkflowType", delay: int):
+        self.task = task
+        self.delay = delay
+
+    def __str__(self):
+        return f"WithDelay({self.task}, {self.delay})"
+
+    def __eq__(self, other):
+        return isinstance(other, WithDelay) and self.task == other.task and self.delay == other.delay
+
+
+Barrier = AtMostOnceBarrier
+Message = dramatiq.Message
+WorkflowType = Message | Chain | Group | WithDelay
+
+# CompletionCallback is a tuple containing: ID, WorkflowType, propagate
+CompletionCallback = tuple[str, WorkflowType | None, bool]
+CompletionCallbacks = list[CompletionCallback]

dramatiq_workflow/_serialize.py

@@ -0,0 +1,79 @@
+import typing
+
+from ._models import Chain, CompletionCallbacks, Group, Message, WithDelay, WorkflowType
+
+
+def serialize_workflow(workflow: WorkflowType | None) -> typing.Any:
+    """
+    Return a serialized version of the workflow that can be JSON-encoded.
+    """
+    if workflow is None:
+        return None
+
+    if isinstance(workflow, Message):
+        return {
+            "__type__": "message",
+            **workflow.asdict(),
+        }
+    if isinstance(workflow, Chain):
+        return {
+            "__type__": "chain",
+            "children": [serialize_workflow(task) for task in workflow.tasks],
+        }
+    if isinstance(workflow, Group):
+        return {
+            "__type__": "group",
+            "children": [serialize_workflow(task) for task in workflow.tasks],
+        }
+    if isinstance(workflow, WithDelay):
+        return {
+            "__type__": "with_delay",
+            "delay": workflow.delay,
+            "task": serialize_workflow(workflow.task),
+        }
+
+    # NOTE: This line is unreachable by design but I don't know how to tell
+    # Pyright that.
+    raise TypeError(f"Unsupported workflow type: {type(workflow)}")
+
+
+def unserialize_workflow(workflow: typing.Any) -> WorkflowType:
+    """
+    Return an unserialized version of the workflow that can be used to create
+    a Workflow instance.
+    """
+    if workflow is None:
+        raise ValueError("Cannot unserialize None")
+
+    if not isinstance(workflow, dict):
+        raise TypeError(f"Unsupported data type: {type(workflow)}")
+
+    workflow_type = workflow.pop("__type__")
+    match workflow_type:
+        case "message":
+            return Message(**workflow)
+        case "chain":
+            return Chain(*[unserialize_workflow(task) for task in workflow["children"]])
+        case "group":
+            return Group(*[unserialize_workflow(task) for task in workflow["children"]])
+        case "with_delay":
+            return WithDelay(
+                task=unserialize_workflow(workflow["task"]),
+                delay=workflow["delay"],
+            )
+
+    raise TypeError(f"Unsupported workflow type: {workflow_type}")
+
+
+def unserialize_workflow_or_none(workflow: typing.Any) -> WorkflowType | None:
+    if workflow is None:
+        return None
+    return unserialize_workflow(workflow)
+
+
+def serialize_callbacks(callbacks: CompletionCallbacks) -> list[tuple[str, typing.Any, bool]]:
+    return [(id, serialize_workflow(g), propagate) for id, g, propagate in callbacks]
+
+
+def unserialize_callbacks(callbacks: list[dict]) -> CompletionCallbacks:
+    return [(id, unserialize_workflow_or_none(g), propagate) for id, g, propagate in callbacks]

dramatiq_workflow/tests/test_workflow.py

@@ -0,0 +1,383 @@
+import unittest
+from unittest import mock
+
+import dramatiq
+
+from .. import Chain, Group, WithDelay, Workflow, WorkflowMiddleware
+from .._serialize import serialize_workflow, unserialize_workflow
+
+
+class WorkflowTests(unittest.TestCase):
+    def setUp(self):
+        self.broker = mock.MagicMock(middleware=[WorkflowMiddleware(mock.MagicMock())])
+        self.task = mock.MagicMock()
+        self.task.message.side_effect = lambda *args, **kwargs: self.__make_message(
+            self.__generate_id(), *args, **kwargs
+        )
+        self.message_ids = []
+        self.message_timestamp = 1717526084640
+
+    def __generate_id(self):
+        current_id = len(self.message_ids) + 1
+        self.message_ids.append(current_id)
+        return current_id
+
+    def __make_message(self, message_id, *args, message_options={}, message_timestamp=None, **kwargs):
+        return dramatiq.Message(
+            message_id=message_id,
+            message_timestamp=message_timestamp or self.message_timestamp,
+            queue_name="default",
+            actor_name="task",
+            args=args,
+            kwargs=kwargs,
+            options=message_options,
+        )
+
+    @mock.patch("dramatiq_workflow._base.time.time")
+    def test_simple_workflow(self, time_mock):
+        time_mock.return_value = 1717526000.12
+        updated_timestamp = time_mock.return_value * 1000
+        workflow = Workflow(
+            Chain(
+                self.task.message(),
+                Group(
+                    Chain(
+                        self.task.message(),
+                        self.task.message(),
+                    ),
+                    self.task.message(),
+                    Chain(
+                        self.task.message(),
+                        Group(
+                            self.task.message(),
+                            self.task.message(),
+                        ),
+                    ),
+                ),
+                Chain(),
+                self.task.message(),
+            ),
+            broker=self.broker,
+        )
+        workflow.run()
+
+        time_mock.assert_called_once()
+
+        self.broker.enqueue.assert_called_once_with(
+            self.__make_message(
+                1,
+                message_timestamp=updated_timestamp,
+                message_options={
+                    "workflow_completion_callbacks": [
+                        (
+                            None,
+                            {
+                                "__type__": "chain",
+                                "children": [
+                                    {
+                                        "__type__": "group",
+                                        "children": [
+                                            {
+                                                "__type__": "chain",
+                                                "children": [
+                                                    {
+                                                        "__type__": "message",
+                                                        "queue_name": "default",
+                                                        "actor_name": "task",
+                                                        "args": (),
+                                                        "kwargs": {},
+                                                        "options": {},
+                                                        "message_id": 2,
+                                                        "message_timestamp": self.message_timestamp,
+                                                    },
+                                                    {
+                                                        "__type__": "message",
+                                                        "queue_name": "default",
+                                                        "actor_name": "task",
+                                                        "args": (),
+                                                        "kwargs": {},
+                                                        "options": {},
+                                                        "message_id": 3,
+                                                        "message_timestamp": self.message_timestamp,
+                                                    },
+                                                ],
+                                            },
+                                            {
+                                                "__type__": "message",
+                                                "queue_name": "default",
+                                                "actor_name": "task",
+                                                "args": (),
+                                                "kwargs": {},
+                                                "options": {},
+                                                "message_id": 4,
+                                                "message_timestamp": self.message_timestamp,
+                                            },
+                                            {
+                                                "__type__": "chain",
+                                                "children": [
+                                                    {
+                                                        "__type__": "message",
+                                                        "queue_name": "default",
+                                                        "actor_name": "task",
+                                                        "args": (),
+                                                        "kwargs": {},
+                                                        "options": {},
+                                                        "message_id": 5,
+                                                        "message_timestamp": self.message_timestamp,
+                                                    },
+                                                    {
+                                                        "__type__": "group",
+                                                        "children": [
+                                                            {
+                                                                "__type__": "message",
+                                                                "queue_name": "default",
+                                                                "actor_name": "task",
+                                                                "args": (),
+                                                                "kwargs": {},
+                                                                "options": {},
+                                                                "message_id": 6,
+                                                                "message_timestamp": self.message_timestamp,
+                                                            },
+                                                            {
+                                                                "__type__": "message",
+                                                                "queue_name": "default",
+                                                                "actor_name": "task",
+                                                                "args": (),
+                                                                "kwargs": {},
+                                                                "options": {},
+                                                                "message_id": 7,
+                                                                "message_timestamp": self.message_timestamp,
+                                                            },
+                                                        ],
+                                                    },
+                                                ],
+                                            },
+                                        ],
+                                    },
+                                    {"__type__": "chain", "children": []},
+                                    {
+                                        "__type__": "message",
+                                        "queue_name": "default",
+                                        "actor_name": "task",
+                                        "args": (),
+                                        "kwargs": {},
+                                        "options": {},
+                                        "message_id": 8,
+                                        "message_timestamp": self.message_timestamp,
+                                    },
+                                ],
+                            },
+                            False,
+                        )
+                    ]
+                },
+            ),
+            delay=None,
+        )
+
+    @mock.patch("dramatiq_workflow._base.time.time")
+    def test_noop_workflow(self, time_mock):
+        time_mock.return_value = 1717526000.12
+        workflow = Workflow(Chain(), broker=self.broker)
+        workflow.run()
+
+        self.broker.enqueue.assert_called_once()
+
+    def test_unsupported_workflow(self):
+        with self.assertRaises(TypeError):
+            Workflow(1).run()  # type: ignore
+
+    @mock.patch("dramatiq_workflow._base.time.time")
+    def test_chain_with_delay(self, time_mock):
+        time_mock.return_value = 1717526000.12
+        updated_timestamp = time_mock.return_value * 1000
+        workflow = Workflow(
+            WithDelay(
+                Chain(self.task.message(), self.task.message()),
+                delay=10,
+            ),
+            broker=self.broker,
+        )
+        workflow.run()
+
+        self.broker.enqueue.assert_called_once_with(
+            self.__make_message(
+                1,
+                message_timestamp=updated_timestamp,
+                message_options={
+                    "workflow_completion_callbacks": [
+                        (
+                            None,
+                            {
+                                "__type__": "chain",
+                                "children": [
+                                    {
+                                        "__type__": "message",
+                                        "queue_name": "default",
+                                        "actor_name": "task",
+                                        "args": (),
+                                        "kwargs": {},
+                                        "options": {},
+                                        "message_id": 2,
+                                        "message_timestamp": self.message_timestamp,
+                                    },
+                                ],
+                            },
+                            False,
+                        )
+                    ]
+                },
+            ),
+            delay=10,
+        )
+
+    @mock.patch("dramatiq_workflow._base.time.time")
+    @mock.patch("dramatiq_workflow._base.Barrier")
+    def test_group_with_delay(self, barrier_mock, time_mock):
+        time_mock.return_value = 1717526000.12
+        updated_timestamp = time_mock.return_value * 1000
+        workflow = Workflow(
+            WithDelay(
+                Group(self.task.message(), self.task.message()),
+                delay=10,
+            ),
+            broker=self.broker,
+        )
+        workflow.run()
+        self.broker.enqueue.assert_has_calls(
+            [
+                mock.call(
+                    self.__make_message(
+                        1,
+                        message_timestamp=updated_timestamp,
+                        message_options={
+                            "workflow_completion_callbacks": [
+                                (
+                                    mock.ANY,  # Accept any callback ID
+                                    None,  # No more workflow to run after this
+                                    True,
+                                )
+                            ]
+                        },
+                    ),
+                    delay=10,
+                ),
+                mock.call(
+                    self.__make_message(
+                        2,
+                        message_timestamp=updated_timestamp,
+                        message_options={
+                            "workflow_completion_callbacks": [
+                                (
+                                    mock.ANY,
+                                    None,
+                                    True,
+                                )
+                            ]
+                        },
+                    ),
+                    delay=10,
+                ),
+            ],
+            any_order=True,
+        )
+
+    def test_serialize_unserialize(self):
+        workflow = Workflow(
+            Chain(
+                self.task.message(),
+                WithDelay(
+                    Group(
+                        self.task.message(),
+                        self.task.message(),
+                    ),
+                    delay=10,
+                ),
+            ),
+        )
+
+        serialized = serialize_workflow(workflow.workflow)
+        unserialized = unserialize_workflow(serialized)
+        self.assertEqual(workflow.workflow, unserialized)
+
+    @mock.patch("dramatiq_workflow._base.time.time")
+    @mock.patch("dramatiq_workflow._base.Barrier")
+    def test_additive_delays(self, barrier_mock, time_mock):
+        time_mock.return_value = 1717526000.12
+        updated_timestamp = time_mock.return_value * 1000
+        workflow = Workflow(
+            WithDelay(
+                Group(
+                    self.task.message(),
+                    WithDelay(
+                        self.task.message(),
+                        delay=10,
+                    ),
+                ),
+                delay=10,
+            ),
+            broker=self.broker,
+        )
+        workflow.run()
+        self.broker.enqueue.assert_has_calls(
+            [
+                mock.call(
+                    self.__make_message(
+                        1,
+                        message_timestamp=updated_timestamp,
+                        message_options={
+                            "workflow_completion_callbacks": [
+                                (
+                                    mock.ANY,  # Accept any callback ID
+                                    None,  # No more workflow to run after this
+                                    True,
+                                )
+                            ]
+                        },
+                    ),
+                    delay=10,
+                ),
+                mock.call(
+                    self.__make_message(
+                        2,
+                        message_timestamp=updated_timestamp,
+                        message_options={
+                            "workflow_completion_callbacks": [
+                                (
+                                    mock.ANY,
+                                    None,
+                                    True,
+                                )
+                            ]
+                        },
+                    ),
+                    delay=20,
+                ),
+            ],
+            any_order=True,
+        )
+
+    @mock.patch("dramatiq_workflow._base.time.time")
+    def test_nested_delays(self, time_mock):
+        time_mock.return_value = 1717526000.12
+        updated_timestamp = time_mock.return_value * 1000
+        workflow = Workflow(
+            WithDelay(
+                WithDelay(
+                    self.task.message(),
+                    delay=10,
+                ),
+                delay=10,
+            ),
+            broker=self.broker,
+        )
+
+        workflow.run()
+        self.broker.enqueue.assert_called_once_with(
+            self.__make_message(
+                1,
+                message_timestamp=updated_timestamp,
+                message_options={"workflow_completion_callbacks": []},
+            ),
+            delay=20,
+        )

dramatiq_workflow-0.1.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Parnassus Labs, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

dramatiq_workflow-0.1.0.dist-info/METADATA

@@ -0,0 +1,111 @@
+Metadata-Version: 2.1
+Name: dramatiq-workflow
+Version: 0.1.0
+Summary: A library for running workflows (chains and groups of tasks) using the Python background task processing library dramatiq.
+Home-page: https://github.com/Outset-AI/dramatiq-workflow
+Author: Outset
+Author-email: engineering@outset.ai
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: dramatiq>=1.10.0
+
+# dramatiq-workflow
+
+`dramatiq-workflow` allows running workflows (chains and groups of tasks) using
+the Python background task processing library [dramatiq](https://dramatiq.io/).
+
+A workflow allows running tasks in parallel and in sequence. It is a way to
+define a workflow of tasks, a combination of chains and groups in any order and
+nested as needed.
+
+## Features
+
+- Define workflows with tasks running in parallel and in sequence.
+- Nest chains and groups of tasks to create complex workflows.
+- Schedules workflows to run in the background using dramatiq.
+
+**Note:** `dramatiq-workflow` does not support passing the results from one task
+to the next one in a chain. We recommend using a database to store intermediate
+results if needed.
+
+## Installation
+
+You can install `dramatiq-workflow` from PyPI:
+
+```sh
+pip install dramatiq-workflow
+```
+
+## Example
+
+Let's assume we want a workflow that looks like this:
+
+```text
+             ╭────────╮  ╭────────╮
+             │ Task 2 │  │ Task 5 │
+          ╭──┼●      ●┼──┼●      ●┼╮
+╭────────╮│  ╰────────╯  ╰────────╯│  ╭────────╮
+│ Task 1 ││  ╭────────╮            │  │ Task 8 │
+│       ●┼╯  │ Task 3 │            ╰──┼●       │
+│       ●┼───┼●      ●┼───────────────┼●       │
+│       ●┼╮  ╰────────╯             ╭─┼●       │
+╰────────╯│  ╭────────╮   ╭────────╮│╭┼●       │
+          │  │ Task 4 │   │ Task 6 │││╰────────╯
+          ╰──┼●      ●┼───┼●      ●┼╯│
+             │       ●┼╮  ╰────────╯ │
+             ╰────────╯│             │
+                       │  ╭────────╮ │
+                       │  │ Task 7 │ │
+                       ╰──┼●      ●┼─╯
+                          ╰────────╯
+```
+
+We can define this workflow as follows:
+
+```python
+from dramatiq_workflow import Workflow, Chain, Group
+
+workflow = Workflow(
+    Chain(
+        task1.message(),
+        Group(
+            Chain(
+                task2.message(),
+                task5.message(),
+            ),
+            task3.message(),
+            Chain(
+                task4.message(),
+                Group(
+                    task6.message(),
+                    task7.message(),
+                ),
+            ),
+        ),
+        task8.message(),
+    ),
+)
+workflow.run()  # Schedules the workflow to run in the background
+```
+
+### Execution Order
+
+In this example, the execution would look like this:
+
+1. Task 1 runs
+2. Task 2, 3, and 4 run in parallel once Task 1 finishes
+3. Task 5 runs once Task 2 finishes
+4. Task 6 and 7 run in parallel once Task 4 finishes
+5. Task 8 runs once Task 5, 6, and 7 finish
+
+*This is a simplified example. The actual execution order may vary because
+tasks that can run in parallel (i.e., in a Group) are not guaranteed to run in
+the order they are defined in the workflow.*
+
+## License
+
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

dramatiq_workflow-0.1.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+dramatiq_workflow/__init__.py,sha256=Dz9OGbZyGefD04-5afcESdjGZqv3mYtf3KEdjkTUOFU,275
+dramatiq_workflow/_barrier.py,sha256=fprBihVkXPdmK-W-KWikFuDw5sAtHgjbkeBXV27cZy4,1414
+dramatiq_workflow/_base.py,sha256=Y5tcG-Ke5eqxce-izmbC3he2Kxd8Ye0fxG4RRgtQIJo,7329
+dramatiq_workflow/_constants.py,sha256=rQ08oG6jiZ0-cvFn7H7nU7hxH53r-9zqjBwGK5dZ7zc,89
+dramatiq_workflow/_helpers.py,sha256=6JXHJuWx5oy3bSSAVg-RuU0CR2sOHdUsOcfDqce4Hkw,452
+dramatiq_workflow/_middleware.py,sha256=i2Iq74qsw7mcnim25RLj5fdTAwdZ5j8_25ZVu6aGIfY,2434
+dramatiq_workflow/_models.py,sha256=01loxf-aOJvadmz7t-VAOTOCxFeaB-XfktsqMLXa53I,1228
+dramatiq_workflow/_serialize.py,sha256=xe0Vk1MVB8BLQGn-Ae2_WPwjmJ8TtB9qoKkK6hSRekg,2645
+dramatiq_workflow/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dramatiq_workflow/tests/test_workflow.py,sha256=qj3LIKD0DWKIqaDlkZzSt4WMhN9gLxGZwltd3fi_5zA,16171
+dramatiq_workflow-0.1.0.dist-info/LICENSE,sha256=4_cXQK7MP9oD26jid8lsdQGc3asMOKyIWbZ_-qeXfzg,1077
+dramatiq_workflow-0.1.0.dist-info/METADATA,sha256=RbYpuMEYko9xI4OlG-czS-WXtHi4hD6bsebhYoZcDIc,3984
+dramatiq_workflow-0.1.0.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+dramatiq_workflow-0.1.0.dist-info/top_level.txt,sha256=y5C0rpDQONUprkF7q1mHKVx2ep-obyNGvEP4JDIrbNE,18
+dramatiq_workflow-0.1.0.dist-info/RECORD,,

dramatiq_workflow-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.6.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

dramatiq_workflow-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+dramatiq_workflow