dr-queues 0.1.0__py3-none-any.whl

dr_queues/__init__.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+__version__ = "0.1.0"
+
+from dr_queues.analysis.filter import filter_run_event_dicts, filter_run_events
+from dr_queues.events import (
+    AmqpEventSink,
+    CompositeEventSink,
+    EventKind,
+    EventSink,
+    MemoryEventSink,
+    MongoEventSink,
+    PipelineEvent,
+)
+from dr_queues.manifest import (
+    RunManifest,
+    RunStageManifest,
+    format_worker_commands,
+    load_run_manifest,
+    manifest_path,
+    parse_workers_arg,
+    write_run_manifest,
+)
+from dr_queues.pipeline import (
+    JobEnvelope,
+    TerminalTap,
+    WorkerPool,
+    seed_jobs,
+)
+from dr_queues.pipeline.runner import (
+    run_in_process,
+    seed_manifest_jobs,
+    setup_run_queues,
+    spawn_all_stage_workers,
+    spawn_stage_worker_process,
+)
+from dr_queues.workflow import (
+    HandlerRegistry,
+    Pipeline,
+    PipelineDefinition,
+    PipelineLane,
+    PipelineStep,
+)
+
+__all__ = [
+    "AmqpEventSink",
+    "CompositeEventSink",
+    "EventKind",
+    "EventSink",
+    "HandlerRegistry",
+    "JobEnvelope",
+    "MemoryEventSink",
+    "MongoEventSink",
+    "Pipeline",
+    "PipelineDefinition",
+    "PipelineEvent",
+    "PipelineLane",
+    "PipelineStep",
+    "RunManifest",
+    "RunStageManifest",
+    "TerminalTap",
+    "WorkerPool",
+    "__version__",
+    "filter_run_event_dicts",
+    "filter_run_events",
+    "format_worker_commands",
+    "load_run_manifest",
+    "manifest_path",
+    "parse_workers_arg",
+    "run_in_process",
+    "seed_jobs",
+    "seed_manifest_jobs",
+    "setup_run_queues",
+    "spawn_all_stage_workers",
+    "spawn_stage_worker_process",
+    "write_run_manifest",
+]

dr_queues/amqp/__init__.py

@@ -0,0 +1,18 @@
+from dr_queues.amqp.connection import (
+    ChannelSession,
+    PikaDeliveryMode,
+    delivery_tag,
+    open_connection,
+    publish_job,
+)
+from dr_queues.amqp.queues import StageQueues, build_stage_queues
+
+__all__ = [
+    "ChannelSession",
+    "PikaDeliveryMode",
+    "StageQueues",
+    "build_stage_queues",
+    "delivery_tag",
+    "open_connection",
+    "publish_job",
+]

dr_queues/amqp/connection.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import os
+from enum import IntEnum
+from functools import lru_cache
+from typing import Any
+
+import pika
+import pika.adapters.blocking_connection as pika_blocking
+from pydantic import BaseModel, ConfigDict
+
+DEFAULT_AMQP_URL = "amqp://guest:guest@localhost:5672/"
+
+type PikaBlockingChannel = pika_blocking.BlockingChannel
+type PikaBasicProperties = pika.spec.BasicProperties
+type PikaDeliveryTag = int
+type PikaDeliveryMethod = pika.spec.Basic.Deliver
+type PikaGetOkMethod = pika.spec.Basic.GetOk
+
+
+class PikaDeliveryMode(IntEnum):
+    TRANSIENT = pika.spec.TRANSIENT_DELIVERY_MODE
+    PERSISTENT = pika.spec.PERSISTENT_DELIVERY_MODE
+
+
+class ReceivedMessage(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    method: PikaGetOkMethod | PikaDeliveryMethod | None
+    body: bytes = rb""
+    properties: PikaBasicProperties | None
+
+    @classmethod
+    def from_get_tuple(
+        cls,
+        method: Any,
+        properties: Any,
+        body: Any,
+    ) -> ReceivedMessage:
+        return ReceivedMessage(
+            method=method,
+            body=body or b"",
+            properties=properties,
+        )
+
+    @property
+    def has_messages(self) -> bool:
+        return self.method is not None
+
+    @property
+    def payload(self) -> tuple[bytes, PikaBasicProperties | None]:
+        return (self.body, self.properties)
+
+
+def delivery_tag(
+    method: PikaDeliveryMethod | PikaGetOkMethod,
+) -> PikaDeliveryTag:
+    tag = method.delivery_tag
+    if tag is None:
+        msg = "Missing delivery tag on message."
+        raise RuntimeError(msg)
+    return tag
+
+
+def amqp_url() -> str:
+    return os.environ.get("AMQP_URL", DEFAULT_AMQP_URL)
+
+
+@lru_cache(maxsize=1)
+def _parameters() -> pika.URLParameters:
+    return pika.URLParameters(amqp_url())
+
+
+def open_connection() -> pika_blocking.BlockingConnection:
+    return pika.BlockingConnection(_parameters())
+
+
+def make_delivery_props(
+    delivery_mode: PikaDeliveryMode,
+) -> pika.BasicProperties:
+    return pika.BasicProperties(delivery_mode=delivery_mode)
+
+
+def publish_job(
+    channel: PikaBlockingChannel,
+    queue_name: str,
+    body: bytes,
+    *,
+    exchange: str = "",
+    delivery_mode: PikaDeliveryMode = PikaDeliveryMode.PERSISTENT,
+    properties: pika.spec.BasicProperties | None = None,
+) -> None:
+    props = properties or make_delivery_props(delivery_mode=delivery_mode)
+    if not channel.is_open:
+        raise RuntimeError("Channel is closed")
+    channel.basic_publish(
+        exchange=exchange,
+        routing_key=queue_name,
+        body=body,
+        properties=props,
+    )
+
+
+class ChannelSession:
+    def __init__(
+        self,
+        connection: pika.BlockingConnection,
+        channel: PikaBlockingChannel,
+        delivery_mode: PikaDeliveryMode = PikaDeliveryMode.PERSISTENT,
+    ):
+        self.connection = connection
+        self.channel = channel
+        self._delivery_mode = delivery_mode
+
+    @classmethod
+    def open_session(
+        cls,
+        *,
+        delivery_mode: PikaDeliveryMode = PikaDeliveryMode.PERSISTENT,
+    ) -> ChannelSession:
+        connection = open_connection()
+        return cls(
+            connection=connection,
+            channel=connection.channel(),
+            delivery_mode=delivery_mode,
+        )
+
+    @classmethod
+    def ensure_channel(
+        cls,
+        *,
+        channel: PikaBlockingChannel | None = None,
+        delivery_mode: PikaDeliveryMode = PikaDeliveryMode.PERSISTENT,
+    ) -> tuple[ChannelSession | None, PikaBlockingChannel]:
+        session = None
+        if channel is None:
+            session = cls.open_session(delivery_mode=delivery_mode)
+            channel = session.channel
+        return session, channel
+
+    @classmethod
+    def declare_durable_queue(
+        cls,
+        *,
+        queue_name: str,
+        channel: PikaBlockingChannel | None = None,
+        delivery_mode: PikaDeliveryMode = PikaDeliveryMode.PERSISTENT,
+    ) -> None:
+        session = None
+        if channel is None:
+            session = cls.open_session(delivery_mode=delivery_mode)
+            channel = session.channel
+
+        try:
+            channel.queue_declare(queue=queue_name, durable=True)
+        finally:
+            if session is not None:
+                session.close()
+
+    @classmethod
+    def ensure_durable_queue(
+        cls,
+        *,
+        queue_name: str,
+        channel: PikaBlockingChannel | None = None,
+        delivery_mode: PikaDeliveryMode = PikaDeliveryMode.PERSISTENT,
+    ) -> None:
+        ChannelSession.declare_durable_queue(
+            queue_name=queue_name,
+            channel=channel,
+            delivery_mode=delivery_mode,
+        )
+
+    @property
+    def channel_props(self) -> pika.BasicProperties:
+        return make_delivery_props(delivery_mode=self._delivery_mode)
+
+    def publish_job(
+        self,
+        queue_name: str,
+        body: bytes,
+    ) -> None:
+        publish_job(
+            channel=self.channel,
+            queue_name=queue_name,
+            body=body,
+            properties=self.channel_props,
+        )
+
+    def close(self) -> None:
+        if self.channel.is_open:
+            self.channel.close()
+        if self.connection.is_open:
+            self.connection.close()

dr_queues/amqp/queues.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from dr_queues.amqp.connection import (
+    ChannelSession,
+    PikaBlockingChannel,
+    PikaDeliveryMode,
+)
+
+
+class StageQueues(BaseModel):
+    prefix: str
+    delivery_mode: PikaDeliveryMode
+    pending_name: str
+    completed_name: str
+
+    @classmethod
+    def get_queue_name(cls, prefix: str, role: str) -> str:
+        return f"{prefix}.{role}"
+
+    def declare_queues(
+        self,
+        *,
+        channel: PikaBlockingChannel | None = None,
+    ) -> None:
+        build_queue_session, channel = ChannelSession.ensure_channel(
+            channel=channel,
+            delivery_mode=self.delivery_mode,
+        )
+        try:
+            ChannelSession.declare_durable_queue(
+                queue_name=self.pending_name,
+                channel=channel,
+                delivery_mode=self.delivery_mode,
+            )
+            if self.completed_name != self.pending_name:
+                ChannelSession.declare_durable_queue(
+                    queue_name=self.completed_name,
+                    channel=channel,
+                    delivery_mode=self.delivery_mode,
+                )
+        finally:
+            if build_queue_session is not None:
+                build_queue_session.close()
+
+
+def build_stage_queues(
+    *,
+    prefix: str,
+    pending: str | None = None,
+    completed: str | None = None,
+    delivery_mode: PikaDeliveryMode = PikaDeliveryMode.PERSISTENT,
+) -> StageQueues:
+    pending_name = pending or StageQueues.get_queue_name(prefix, "pending")
+    completed_name = completed or StageQueues.get_queue_name(
+        prefix, "completed"
+    )
+    stage_queues = StageQueues(
+        prefix=prefix,
+        delivery_mode=delivery_mode,
+        pending_name=pending_name,
+        completed_name=completed_name,
+    )
+    stage_queues.declare_queues()
+    return stage_queues

dr_queues/analysis/__init__.py

@@ -0,0 +1,3 @@
+from dr_queues.analysis.filter import filter_run_event_dicts, filter_run_events
+
+__all__ = ["filter_run_event_dicts", "filter_run_events"]

dr_queues/analysis/filter.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import Any
+
+from dr_queues.events.schema import PipelineEvent
+
+
+def filter_run_events(
+    events: list[PipelineEvent],
+    run_id: str,
+) -> list[PipelineEvent]:
+    return [event for event in events if event.run_id == run_id]
+
+
+def filter_run_event_dicts(
+    events: list[dict[str, Any]],
+    run_id: str,
+) -> list[dict[str, Any]]:
+    return [event for event in events if event.get("run_id") == run_id]

dr_queues/cli/__init__.py

@@ -0,0 +1,6 @@
+from dr_queues.cli.commands import (
+    STAGE_WORKER_ENTRYPOINT,
+    stage_worker_command_prefix,
+)
+
+__all__ = ["STAGE_WORKER_ENTRYPOINT", "stage_worker_command_prefix"]

dr_queues/cli/commands.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import shutil
+import sys
+
+STAGE_WORKER_ENTRYPOINT = "dr-queues-stage-worker"
+
+
+def stage_worker_command_prefix() -> list[str]:
+    """Resolve the stage-worker CLI for installed or editable runs."""
+    executable = shutil.which(STAGE_WORKER_ENTRYPOINT)
+    if executable is not None:
+        return [executable]
+    return [sys.executable, "-m", "dr_queues.cli.stage_worker"]

dr_queues/cli/demo.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import importlib
+import shutil
+import sys
+from uuid import uuid4
+
+import typer
+
+from dr_queues import (
+    AmqpEventSink,
+    CompositeEventSink,
+    EventKind,
+    MongoEventSink,
+    Pipeline,
+    PipelineDefinition,
+    PipelineLane,
+    PipelineStep,
+    filter_run_events,
+    manifest_path,
+    parse_workers_arg,
+    run_in_process,
+    seed_manifest_jobs,
+    setup_run_queues,
+)
+
+DEFAULT_WORKERS = "slow=4,transform=4,finalize=2"
+HANDLERS_MODULE = "dr_queues.demo_handlers"
+
+app = typer.Typer(add_completion=False)
+
+
+def _load_registry(module_path: str):
+    module = importlib.import_module(module_path)
+    if hasattr(module, "registry"):
+        return module.registry
+    msg = f"Module {module_path!r} has no registry attribute."
+    raise typer.BadParameter(msg)
+
+
+def _build_pipeline(
+    *,
+    lanes: int,
+    handlers_module: str,
+) -> Pipeline:
+    registry = _load_registry(handlers_module)
+    definition = PipelineDefinition(
+        id="demo_pipeline",
+        lanes=[PipelineLane(id=f"lane-{index}") for index in range(lanes)],
+        steps=[
+            PipelineStep(name="slow", handler_key="sleep_ms"),
+            PipelineStep(name="transform", handler_key="add_prefix"),
+            PipelineStep(name="finalize", handler_key="record_artifact"),
+        ],
+    )
+    return Pipeline(definition, registry)
+
+
+def _build_event_sink(sink: str):
+    if sink == "mongo":
+        return MongoEventSink()
+    if sink == "amqp":
+        return AmqpEventSink()
+    if sink == "both":
+        return CompositeEventSink(
+            [MongoEventSink(), AmqpEventSink()],
+        )
+    msg = f"Unknown sink {sink!r}; expected mongo, amqp, or both."
+    raise typer.BadParameter(msg)
+
+
+@app.command()
+def main(
+    repeats: int = typer.Option(2, "--repeats"),
+    lanes: int = typer.Option(2, "--lanes"),
+    workers: str = typer.Option(DEFAULT_WORKERS, "--workers"),
+    run_id: str | None = typer.Option(None, "--run-id"),
+    sink: str = typer.Option("mongo", "--sink"),
+    handlers_module: str = typer.Option(
+        HANDLERS_MODULE,
+        "--handlers-module",
+    ),
+    completion_timeout: float = typer.Option(120.0, "--completion-timeout"),
+) -> None:
+    resolved_run_id = run_id or f"demo-{uuid4().hex[:8]}"
+    pipeline = _build_pipeline(lanes=lanes, handlers_module=handlers_module)
+    workers_by_stage = parse_workers_arg(
+        workers,
+        pipeline.step_names(),
+        default=2,
+    )
+    expected = pipeline.expected_job_count(repeats)
+    event_sink = _build_event_sink(sink)
+
+    manifest = setup_run_queues(
+        pipeline=pipeline,
+        run_id=resolved_run_id,
+        workers_by_stage=workers_by_stage,
+        expected_jobs=expected,
+    )
+    jobs = pipeline.make_seed_jobs(run_id=resolved_run_id, repeats=repeats)
+    seed_manifest_jobs(manifest, jobs)
+
+    typer.echo(f"run_id={resolved_run_id}")
+    typer.echo(f"manifest={manifest_path(resolved_run_id)}")
+    typer.echo(f"expected_jobs={expected} sink={sink}")
+
+    run_in_process(
+        manifest=manifest,
+        pipeline=pipeline,
+        workers_by_stage=workers_by_stage,
+        event_sink=event_sink,
+        completion_timeout=completion_timeout,
+    )
+
+    events = filter_run_events(
+        event_sink.read_by_run_id(resolved_run_id),
+        resolved_run_id,
+    )
+    event_sink.close()
+    terminals = [
+        event for event in events if event.event == EventKind.TERMINAL
+    ]
+    typer.echo(f"events={len(events)} terminals={len(terminals)}")
+    if terminals:
+        sample = terminals[0].payload
+        typer.echo(
+            "sample_terminal="
+            f"lane={sample.get('lane')} "
+            f"transform={sample.get('step_outputs', {}).get('transform')}",
+        )
+
+    if len(terminals) != expected:
+        typer.echo("Terminal count mismatch.", err=True)
+        raise typer.Exit(code=1)
+
+
+def run() -> None:
+    app()
+
+
+if __name__ == "__main__":
+    run()

dr_queues/cli/stage_worker.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import importlib
+import os
+import signal
+import sys
+import time
+from pathlib import Path
+
+import typer
+
+from dr_queues.events.mongo import MongoEventSink
+from dr_queues.manifest import (
+    load_run_manifest,
+    manifest_path,
+    read_pid,
+    remove_pid,
+    stage_pid_path,
+    write_pid,
+)
+from dr_queues.pipeline.workers import WorkerPool
+from dr_queues.workflow.pipeline import Pipeline
+from dr_queues.workflow.registry import HandlerRegistry
+
+app = typer.Typer(add_completion=False)
+
+
+def _stop_pid(pid: int, *, timeout: float = 30.0) -> None:
+    try:
+        os.kill(pid, signal.SIGTERM)
+    except ProcessLookupError:
+        return
+
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        try:
+            os.kill(pid, 0)
+        except ProcessLookupError:
+            return
+        time.sleep(0.2)
+
+
+def _load_registry(module_path: str) -> HandlerRegistry:
+    module = importlib.import_module(module_path)
+    registry = getattr(module, "registry", None)
+    if registry is None:
+        msg = f"Module {module_path!r} has no registry attribute."
+        raise typer.BadParameter(msg)
+    return registry
+
+
+@app.command()
+def main(
+    run_id: str | None = typer.Option(None, "--run-id"),
+    stage: str = typer.Option(..., "--stage"),
+    workers: int = typer.Option(..., "--workers"),
+    manifest: Path | None = typer.Option(None, "--manifest"),
+    handlers_module: str = typer.Option(
+        "dr_queues.demo_handlers",
+        "--handlers-module",
+    ),
+    replace: bool = typer.Option(False, "--replace"),
+) -> None:
+    resolved_manifest_path = manifest or (
+        manifest_path(run_id) if run_id else None
+    )
+    if resolved_manifest_path is None or not resolved_manifest_path.exists():
+        typer.echo(
+            "Manifest not found. Pass --manifest or --run-id.",
+            err=True,
+        )
+        raise typer.Exit(code=1)
+
+    run_manifest = load_run_manifest(resolved_manifest_path)
+    stage_entry = next(
+        (item for item in run_manifest.stages if item.name == stage),
+        None,
+    )
+    if stage_entry is None:
+        typer.echo(f"Unknown stage {stage!r} in manifest.", err=True)
+        raise typer.Exit(code=1)
+
+    pid_path = stage_pid_path(run_manifest.run_id, stage)
+    if replace:
+        existing_pid = read_pid(pid_path)
+        if existing_pid is not None:
+            typer.echo(f"Stopping existing worker pid={existing_pid}...")
+            _stop_pid(existing_pid)
+            remove_pid(pid_path)
+
+    registry = _load_registry(handlers_module)
+    pipeline = Pipeline(run_manifest.pipeline_definition, registry)
+    handler = pipeline.make_handler(stage_entry.step_index)
+    event_sink = MongoEventSink()
+    pool = WorkerPool(
+        input_queue=stage_entry.input_queue,
+        output_queue=stage_entry.output_queue,
+        handler=handler,
+        event_sink=event_sink,
+        workers=workers,
+        stage_name=stage_entry.name,
+    )
+
+    write_pid(pid_path, os.getpid())
+    typer.echo(
+        f"stage={stage} workers={workers} input={stage_entry.input_queue}",
+    )
+
+    def _shutdown(_signum: int, _frame: object) -> None:
+        typer.echo(f"Stopping stage {stage}...")
+        pool.stop()
+
+    signal.signal(signal.SIGTERM, _shutdown)
+    signal.signal(signal.SIGINT, _shutdown)
+
+    pool.start()
+    try:
+        while not pool._stop.is_set():
+            time.sleep(0.5)
+    finally:
+        pool.stop()
+        pool.join(timeout=5)
+        event_sink.close()
+        remove_pid(pid_path)
+
+
+def run() -> None:
+    app()
+
+
+if __name__ == "__main__":
+    run()