scalable-pypeline 2.1.3__tar.gz → 2.1.31__tar.gz

{scalable-pypeline-2.1.3 → scalable_pypeline-2.1.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
-Name: scalable-pypeline
-Version: 2.1.3
+Name: scalable_pypeline
+Version: 2.1.31
 Summary: PypeLine - Python pipelines for the Real World
 Home-page: https://gitlab.com/bravos2/pypeline
 Author: Bravos Power Corporation

scalable_pypeline-2.1.31/pypeline/__init__.py

@@ -0,0 +1 @@
+__version__ = "2.1.31"

scalable_pypeline-2.1.31/pypeline/barrier.py

@@ -0,0 +1,63 @@
+import time
+import redis
+from redis.sentinel import Sentinel
+from urllib.parse import urlparse
+from pypeline.constants import (
+    REDIS_SENTINEL_MASTER_NAME,
+    DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT,
+    DEFAULT_REDIS_SOCKET_TIMEOUT,
+    DEFAULT_REDIS_RETRY_ON_TIMEOUT,
+    DEFAULT_REDIS_SOCKET_KEEPALIVE,
+    DEFAULT_REDIS_HEALTH_CHECK_INTERVAL,
+)
+
+
+class LockingParallelBarrier:
+    def __init__(self, redis_url, task_key="task_counter", lock_key="task_lock"):
+        # Connect to Redis using the provided URL
+        if REDIS_SENTINEL_MASTER_NAME is not None:
+            parsed_redis_url = urlparse(redis_url)
+            redis_sentinel = Sentinel(
+                sentinels=[(parsed_redis_url.hostname, parsed_redis_url.port)],
+            )
+            self.redis = redis_sentinel.master_for(
+                REDIS_SENTINEL_MASTER_NAME,
+                db=int(parsed_redis_url.path[1]) if parsed_redis_url.path else 0,
+                password=parsed_redis_url.password,
+                socket_connect_timeout=DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT,
+                socket_timeout=DEFAULT_REDIS_SOCKET_TIMEOUT,
+                retry_on_timeout=DEFAULT_REDIS_RETRY_ON_TIMEOUT,
+                socket_keepalive=DEFAULT_REDIS_SOCKET_KEEPALIVE,
+                health_check_interval=DEFAULT_REDIS_HEALTH_CHECK_INTERVAL,
+                decode_responses=True,
+            )
+        else:
+            self.redis = redis.StrictRedis.from_url(redis_url, decode_responses=True)
+        self.task_key = task_key
+        self.lock_key = lock_key
+
+    def acquire_lock(self, timeout=5):
+        """Acquire a lock using Redis."""
+        while True:
+            if self.redis.set(self.lock_key, "locked", nx=True, ex=timeout):
+                return True
+            time.sleep(0.1)
+
+    def release_lock(self):
+        """Release the lock in Redis."""
+        self.redis.delete(self.lock_key)
+
+    def set_task_count(self, count):
+        """Initialize the task counter in Redis."""
+        self.redis.set(self.task_key, count)
+
+    def decrement_task_count(self):
+        """Decrement the task counter in Redis."""
+        return self.redis.decr(self.task_key)
+
+    def task_exists(self):
+        return self.redis.exists(self.task_key)
+
+    def get_task_count(self):
+        """Get the current value of the task counter."""
+        return int(self.redis.get(self.task_key) or 0)

{scalable-pypeline-2.1.3 → scalable_pypeline-2.1.31}/pypeline/constants.py

@@ -1,5 +1,4 @@
-""" Pypeline Constants
-"""
+"""Pypeline Constants"""
 
 import os
 
@@ -15,6 +14,7 @@ DEFAULT_BROKER_CALLABLE = os.environ.get(
 # Pypeline broker connections
 RABBIT_URL = os.environ.get("RABBIT_URL", "amqp://admin:password@127.0.0.1:5672")
 REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
+REDIS_SENTINEL_MASTER_NAME = os.environ.get("REDIS_SENTINEL_MASTER_NAME", None)
 
 # Pypeline task defaults
 PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL = int(
@@ -36,7 +36,20 @@ DEFAULT_BROKER_CONNECTION_ATTEMPTS = int(
 DEFAULT_BROKER_BLOCKED_CONNECTION_TIMEOUT = int(
     os.getenv("DEFAULT_BROKER_BLOCKED_CONNECTION_TIMEOUT", 30)
 )
+DEFAULT_BROKER_HEARTBEAT_TIMEOUT = int(
+    os.getenv("DEFAULT_BROKER_HEARTBEAT_TIMEOUT", 300000)
+)
+DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT = int(
+    os.getenv("DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT", 1)
+)
+DEFAULT_REDIS_SOCKET_TIMEOUT = int(os.getenv("DEFAULT_REDIS_SOCKET_TIMEOUT", 2))
+DEFAULT_REDIS_RETRY_ON_TIMEOUT = bool(os.getenv("DEFAULT_REDIS_RETRY_ON_TIMEOUT", True))
+DEFAULT_REDIS_SOCKET_KEEPALIVE = bool(os.getenv("DEFAULT_REDIS_SOCKET_KEEPALIVE", True))
+DEFAULT_REDIS_HEALTH_CHECK_INTERVAL = int(
+    os.getenv("DEFAULT_REDIS_HEALTH_CHECK_INTERVAL", 30)
+)
 
+MESSAGE_BROKER = os.getenv("MESSAGE_BROKER", "RABBITMQ")
 MS_IN_SECONDS = 1000
 API_PATH_V1 = "/api/v1"
 

{scalable-pypeline-2.1.3 → scalable_pypeline-2.1.31}/pypeline/dramatiq.py

@@ -1,10 +1,14 @@
+import copy
 import typing
 import pika
 import logging
+import os
 
 import click
 from urllib.parse import urlparse
 
+from dramatiq.brokers.redis import RedisBroker
+from redis.sentinel import Sentinel
 from pypeline.extensions import pypeline_config
 from warnings import warn
 from apscheduler.schedulers.blocking import BlockingScheduler
@@ -26,20 +30,29 @@ from flask.cli import with_appcontext
 
 from pypeline.constants import (
     REDIS_URL,
+    REDIS_SENTINEL_MASTER_NAME,
     RABBIT_URL,
     DEFAULT_BROKER_CALLABLE,
     DEFAULT_BROKER_CONNECTION_HEARTBEAT,
     DEFAULT_BROKER_BLOCKED_CONNECTION_TIMEOUT,
     DEFAULT_BROKER_CONNECTION_ATTEMPTS,
+    MESSAGE_BROKER,
+    DEFAULT_BROKER_HEARTBEAT_TIMEOUT,
+    DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT,
+    DEFAULT_REDIS_SOCKET_TIMEOUT,
+    DEFAULT_REDIS_RETRY_ON_TIMEOUT,
+    DEFAULT_REDIS_SOCKET_KEEPALIVE,
+    DEFAULT_REDIS_HEALTH_CHECK_INTERVAL,
+)
+from pypeline.pipelines.middleware.get_active_worker_id_middleware import (
+    GetActiveWorkerIdMiddleware,
 )
 from pypeline.pipelines.middleware.parallel_pipeline_middleware import ParallelPipeline
 from pypeline.pipelines.middleware.pypeline_middleware import PypelineMiddleware
-from pypeline.pipelines.middleware.deduplication_middleware import (
-    DeduplicationMiddleware,
-)
 from pypeline.utils.config_utils import (
     retrieve_latest_schedule_config,
     get_service_config_for_worker,
+    retrieve_executable_job_config,
 )
 from pypeline.utils.dramatiq_utils import (
     guess_code_directory,
@@ -47,83 +60,149 @@ from pypeline.utils.dramatiq_utils import (
     register_lazy_actor,
     LazyActor,
 )
+from pypeline.utils.graceful_shutdown_util import enable_graceful_shutdown
 from pypeline.utils.module_utils import get_callable
-
+from dramatiq.middleware import (
+    Retries,
+    Callbacks,
+    TimeLimit,
+    AgeLimit,
+    ShutdownNotifications,
+    Pipelines,
+)
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
 def configure_default_broker(broker: Broker = None):
-    redis_backend = RedisBackend(url=REDIS_URL)
-    parsed_url = urlparse(RABBIT_URL)
-    credentials = pika.PlainCredentials(parsed_url.username, parsed_url.password)
-    rabbit_broker = (
-        broker
-        if broker is not None
-        else RabbitmqBroker(
-            host=parsed_url.hostname,
-            port=parsed_url.port,
-            credentials=credentials,
-            heartbeat=DEFAULT_BROKER_CONNECTION_HEARTBEAT,
-            connection_attempts=DEFAULT_BROKER_CONNECTION_ATTEMPTS,
-            blocked_connection_timeout=DEFAULT_BROKER_BLOCKED_CONNECTION_TIMEOUT,
+    reworked_defaults=[AgeLimit(), TimeLimit(), ShutdownNotifications(), Callbacks(), Pipelines(), Retries()]
+    redis_client = None
+    if REDIS_SENTINEL_MASTER_NAME is not None:
+        parsed_redis_url = urlparse(REDIS_URL)
+        redis_sentinel = Sentinel(
+            sentinels=[(parsed_redis_url.hostname, parsed_redis_url.port)],
+        )
+        redis_client = redis_sentinel.master_for(
+            REDIS_SENTINEL_MASTER_NAME,
+            db=int(parsed_redis_url.path[1]) if parsed_redis_url.path else 0,
+            password=parsed_redis_url.password,
+            socket_connect_timeout=DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT,
+            socket_timeout=DEFAULT_REDIS_SOCKET_TIMEOUT,
+            retry_on_timeout=DEFAULT_REDIS_RETRY_ON_TIMEOUT,
+            socket_keepalive=DEFAULT_REDIS_SOCKET_KEEPALIVE,
+            health_check_interval=DEFAULT_REDIS_HEALTH_CHECK_INTERVAL,
+        )
+    redis_backend = RedisBackend(client=redis_client, url=REDIS_URL)
+
+    if MESSAGE_BROKER == "RABBITMQ":
+        parsed_url = urlparse(RABBIT_URL)
+        credentials = pika.PlainCredentials(parsed_url.username, parsed_url.password)
+        broker = (
+            broker
+            if broker is not None
+            else RabbitmqBroker(
+                host=parsed_url.hostname,
+                port=parsed_url.port,
+                credentials=credentials,
+                heartbeat=DEFAULT_BROKER_CONNECTION_HEARTBEAT,
+                connection_attempts=DEFAULT_BROKER_CONNECTION_ATTEMPTS,
+                blocked_connection_timeout=DEFAULT_BROKER_BLOCKED_CONNECTION_TIMEOUT,
+                middleware=reworked_defaults
+            )
+        )
+
+    elif MESSAGE_BROKER == "REDIS":
+        broker = (
+            broker
+            if broker is not None
+            else RedisBroker(
+                client=redis_client,
+                url=REDIS_URL,
+                heartbeat_timeout=DEFAULT_BROKER_HEARTBEAT_TIMEOUT,
+                middleware=reworked_defaults
+            )
         )
-    )
-    rabbit_broker.add_middleware(Results(backend=redis_backend))
-    rabbit_broker.add_middleware(ParallelPipeline(redis_url=REDIS_URL))
-    rabbit_broker.add_middleware(PypelineMiddleware(redis_url=REDIS_URL))
-    rabbit_broker.add_middleware(CurrentMessage())
-    register_actors_for_workers(rabbit_broker)
-    rabbit_broker.add_middleware(DeduplicationMiddleware(redis_url=REDIS_URL))
-    set_broker(rabbit_broker)
+
+    broker.add_middleware(Results(backend=redis_backend))
+    broker.add_middleware(ParallelPipeline(redis_url=REDIS_URL))
+    broker.add_middleware(PypelineMiddleware(redis_url=REDIS_URL))
+    broker.add_middleware(GetActiveWorkerIdMiddleware())
+    broker.add_middleware(CurrentMessage())
+    if (
+        os.getenv("RESTRICT_WORKER_SHUTDOWN_WHILE_JOBS_RUNNING", "false").lower()
+        == "true"
+    ):
+        enable_graceful_shutdown(broker=broker, redis_url=REDIS_URL)
+
+    register_actors_for_workers(broker)
+
+    set_broker(broker)
 
 
 def register_actors_for_workers(broker: Broker):
     service = get_service_config_for_worker(pypeline_config)
     scheduled_jobs_config = retrieve_latest_schedule_config()
+    executable_jobs_config = retrieve_executable_job_config()
 
     if not service:
         return
-    for task in service.get("registeredTasks", []):
-        pipeline_meta = None
-        for pipeline_key, pipeline in pypeline_config["pipelines"].items():
-            pipeline_config = pipeline["config"]
+
+    worker_registered_tasks = [
+        task_handler["handler"] for task_handler in service.get("registeredTasks")
+    ]
+
+    # Loop over the pipelines to get metadata and other information about the task for registration
+    for pipeline_key, pipeline in pypeline_config["pipelines"].items():
+        for task, task_handler_meta in pipeline["config"]["taskDefinitions"].items():
             if pipeline["schemaVersion"] == 1:
-                pipeline_tasks = [
-                    t["handler"] for t in pipeline_config["taskDefinitions"].values()
-                ]
+                # Check if any task in this pipeline is registered
+                task_handlers = [task_handler_meta["handler"]]
             elif pipeline["schemaVersion"] == 2:
-                pipeline_tasks = [
-                    handler
-                    for key in pipeline_config["taskDefinitions"]
-                    for handler in pipeline_config["taskDefinitions"][key].get(
-                        "handlers", []
-                    )
-                ]
-            if task["handler"] in pipeline_tasks:
-                pipeline_meta = pipeline_config["metadata"]
-                break
-
-        if pipeline_meta is None:
-            for job in scheduled_jobs_config:
-                config = job["config"]
-                if config["task"] == task["handler"]:
-                    pipeline_meta = {"queue": config.get("queue", "default")}
-
-        if pipeline_meta is None:
-            raise ValueError(
-                f"Registered task {task['handler']} is not defined in a pipeline or scheduled task"
-            )
-
-        try:
-            worker_path = task["handler"]  # Required, no default
-            tmp_handler = get_callable(worker_path)
-            if pipeline_meta and pipeline_meta.get("maxRetry", 0) >= 0:
-                pipeline_meta["store_results"] = True
-                _ = register_lazy_actor(broker, tmp_handler, pipeline_meta)
-        except Exception as e:
-            logger.exception(f"Unable to add a task to dramatiq: {e}")
+                task_handlers = [t for t in task_handler_meta["handlers"]]
+
+            for task_handler in task_handlers:
+                if task_handler in worker_registered_tasks:
+                    server_type = task_handler_meta.get("serverType", None)
+
+                    try:
+                        pipeline_metadata = copy.deepcopy(
+                            pipeline["config"]["metadata"]
+                        )
+                        tmp_handler = get_callable(task_handler)
+                        if pipeline_metadata.get("maxRetry", 0) >= 0:
+                            pipeline_metadata["store_results"] = True
+                            _ = register_lazy_actor(
+                                broker, tmp_handler, pipeline_metadata, server_type
+                            )
+                    except Exception as e:
+                        logger.exception(
+                            f"Unable to add a task {task_handler} to dramatiq: {e}"
+                        )
+    # Loop over the scheduled jobs and create metadata and other information about the task for registration
+    for job in scheduled_jobs_config:
+        config = job["config"]
+        if config["task"] in worker_registered_tasks:
+            pipeline_meta = {"queue": config.get("queue", "default")}
+            try:
+                tmp_handler = get_callable(config["task"])
+                if pipeline_meta and pipeline_meta.get("maxRetry", 0) >= 0:
+                    pipeline_meta["store_results"] = True
+                    _ = register_lazy_actor(broker, tmp_handler, pipeline_meta, None)
+            except Exception as e:
+                logger.exception(f"Unable to add a task to dramatiq: {e}")
+
+    for job in executable_jobs_config or []:
+        config = job["config"]
+        if config["task"] in worker_registered_tasks:
+            pipeline_meta = {"queue": config.get("queue", "default")}
+            try:
+                tmp_handler = get_callable(config["task"])
+                if pipeline_meta and pipeline_meta.get("maxRetry", 0) >= 0:
+                    pipeline_meta["store_results"] = True
+                    _ = register_lazy_actor(broker, tmp_handler, pipeline_meta, None)
+            except Exception as e:
+                logger.exception(f"Unable to add a task to dramatiq: {e}")
 
 
 class Dramatiq:

scalable_pypeline-2.1.31/pypeline/executable_job_config_schema.py

@@ -0,0 +1,35 @@
+from marshmallow import Schema, EXCLUDE, fields
+
+class ExecutableJobConfigSchema(Schema):
+    queue = fields.String(
+        required=True,
+        description="Name of queue on which to place task.",
+        example="my-default-queue",
+    )
+    task = fields.String(
+        required=True,
+        description="Path to task to invoke.",
+        example="my_app.module.method",
+    )
+
+class ExecutableJobSchema(Schema):
+    """Definition of a single schedule entry"""
+    class Meta:
+        unknown = EXCLUDE
+
+    name = fields.String(
+        required=True,
+        description="Name of schedule entry.",
+        example="My Scheduled Task",
+    )
+    schemaVersion = fields.Integer(required=True)
+    config = fields.Dict(required=True)
+    enabled = fields.Boolean(
+        required=True, description="Whether entry is enabled.", example=True
+    )
+    config = fields.Nested(
+        ExecutableJobConfigSchema,
+        required=True,
+        description="Configuration information for this job.",
+    )
+

{scalable-pypeline-2.1.3 → scalable_pypeline-2.1.31}/pypeline/flask/api/pipelines.py

@@ -56,26 +56,6 @@ class InvokePipelineSchema(Schema):
         example={"document_id": "123", "send_alert": True},
         required=False,
     )
-    settings = fields.Raw(
-        description="Payload contains settings for a given pipeline",
-        example={
-            "param1": "Dataset",
-            "param2": 1,
-            "param3": 2,
-        },
-        required=False,
-    )
-
-    task_replacements = fields.Raw(
-        description="A dictionary of task definitions as the key and the value of the index for which handler"
-        " should be executed.  If none provided it will default to the first handler in the list at index position 0.",
-        example={
-            "a": 1,
-            "b": 3,
-        },
-        required=False,
-    )
-
     scenarios = fields.List(
         fields.Nested(PipelineScenarioSchema),
         metadata={"description": "List of scenarios to run for a given pipeline"},
@@ -198,20 +178,16 @@ class PipelineInvoke(MethodView):
         retval = {"pipeline_id": pipeline_id, "status": "starting"}
         try:
             chain_payload = payload.get("chain_payload", {})
-            settings = payload.get("settings", None)
-            task_replacements = payload.get("task_replacements", {})
             scenarios = payload.get("scenarios", [])
             if pipeline_config["schemaVersion"] == 1:
                 pipeline = dag_generator(
                     pipeline_id=pipeline_id,
                     event=chain_payload,
                 )
-            elif pipeline_config["schemaVersion"] == 2 and task_replacements:
+            elif pipeline_config["schemaVersion"] == 2:
                 pipeline = dag_generator(
                     pipeline_id=pipeline_id,
-                    task_replacements=task_replacements,
                     scenarios=scenarios,
-                    settings=settings,
                 )
                 retval["scenarios"] = pipeline.scenarios
             pipeline.run()

scalable_pypeline-2.1.31/pypeline/job_runner.py

@@ -0,0 +1,205 @@
+import os
+import logging
+import argparse
+import threading
+import multiprocessing as mp
+
+# Prefer 'spawn' for user code using multiprocessing
+if mp.get_start_method(allow_none=True) != "spawn":
+    mp.set_start_method("spawn", force=True)
+
+# Avoid staging more than one message; must be set before Dramatiq import path runs
+os.environ.setdefault("dramatiq_queue_prefetch", "1")
+
+from dramatiq import Worker, get_broker, set_broker
+from dramatiq.middleware import Middleware
+
+
+try:
+    # If your project exposes a helper to configure the default broker, use it.
+    from pypeline.dramatiq import configure_default_broker  # adjust import if needed
+
+    broker = configure_default_broker() or get_broker()
+    set_broker(broker)
+except Exception:
+    # Fall back to whatever Dramatiq has as the active broker.
+    import pypeline.dramatiq  # noqa: F401 (ensure module side-effects run)
+
+    broker = get_broker()
+
+
+class OneAndDone(Middleware):
+    """
+    Signals when the first message starts ('got_work') and completes ('done').
+    If stop_on_failure=True, we'll also mark done after the first failure.
+    """
+
+    def __init__(
+        self,
+        got_work: threading.Event,
+        done: threading.Event,
+        *,
+        stop_on_failure: bool = False
+    ):
+        self.got_work = got_work
+        self.done = done
+        self.stop_on_failure = stop_on_failure
+
+    def before_process_message(self, broker, message):
+        # First time we see a message begin processing in this process
+        if not self.got_work.is_set():
+            self.got_work.set()
+
+    def after_process_message(self, broker, message, *, result=None, exception=None):
+        # On success (or also on failure if configured), finish this worker
+        if exception is None or self.stop_on_failure:
+            if not self.done.is_set():
+                self.done.set()
+
+
+def _graceful_stop(worker: Worker, log: logging.Logger):
+    try:
+        log.info("Stopping dramatiq worker...")
+        worker.stop()  # stop consumers; no new messages will start
+        worker.join()
+        log.info("Worker stopped.")
+    except Exception as e:
+        log.exception("Error stopping worker: %s", e)
+
+
+def _close_broker(log: logging.Logger):
+    try:
+        b = get_broker()
+        if b is not None and hasattr(b, "close"):
+            b.close()
+            log.info("Broker closed.")
+    except Exception as e:
+        log.exception("Error closing broker: %s", e)
+
+
+def job_runner(queues, idle_timeout_ms: int = 0, *, stop_on_failure: bool = False):
+    """
+    Start a single-thread Dramatiq worker. Behavior:
+      - Wait up to `idle_timeout_ms` for *a job to start* (time-to-first-job).
+      - Once a job begins, wait indefinitely for it to complete.
+      - After the first successful job completes (or first job, if stop_on_failure=True), stop and exit.
+
+    Args:
+        queues (list[str]): queues to listen to
+        idle_timeout_ms (int): <=0 => wait forever for first job; >0 => exit if no job starts in time
+        stop_on_failure (bool): if True, exit after first job even if it fails
+    """
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
+    log = logging.getLogger("oneshot")
+
+    # Normalize timeout (treat non-positive as "infinite")
+    timeout_ms = (
+        int(idle_timeout_ms) if idle_timeout_ms and int(idle_timeout_ms) > 0 else 0
+    )
+    log.info(
+        "Launching worker with queues=%s, idle_timeout_ms=%s", queues, timeout_ms or "∞"
+    )
+
+    got_work = threading.Event()
+    done = threading.Event()
+    broker.add_middleware(OneAndDone(got_work, done, stop_on_failure=stop_on_failure))
+
+    worker = Worker(
+        broker,
+        worker_threads=1,  # strictly one at a time
+        queues=queues,
+        worker_timeout=1000,  # ms; how often the worker checks for stop
+    )
+
+    worker.start()
+
+    def controller():
+        log.debug("Controller thread started.")
+        try:
+            # Phase 1: Wait for *first job to start*
+            if timeout_ms > 0:
+                started = got_work.wait(timeout_ms / 1000.0)
+                if not started:
+                    log.info(
+                        "Idle timeout reached (%d ms); no jobs started. Stopping worker.",
+                        timeout_ms,
+                    )
+                    return
+            else:
+                got_work.wait()
+
+            log.info("First job started; waiting for it to finish...")
+            # Phase 2: Wait for the first job to complete (no timeout)
+            done.wait()
+            log.info("First job finished; shutting down.")
+        finally:
+            _graceful_stop(worker, log)
+            _close_broker(log)
+            # Hard-exit to ensure K8s Job is marked Succeeded promptly, no lingering threads.
+            os._exit(0)
+
+    t = threading.Thread(target=controller, name="oneshot-controller", daemon=False)
+    t.start()
+    t.join()  # Block until controller completes (which shuts everything down)
+
+
+def _parse_args(argv=None):
+    ap = argparse.ArgumentParser(description="Run a one-shot Dramatiq worker.")
+    ap.add_argument(
+        "-q",
+        "--queue",
+        action="append",
+        default=None,
+        help="Queue to listen to (repeatable). You can also pass a comma-separated list.",
+    )
+    ap.add_argument(
+        "--idle-timeout-ms",
+        type=int,
+        default=int(os.getenv("IDLE_TIMEOUT_MS", "0")),
+        help="Exit if no job starts within this time (<=0 = wait forever).",
+    )
+    ap.add_argument(
+        "--stop-on-failure",
+        action="store_true",
+        help="Exit after the first job even if it fails.",
+    )
+    return ap.parse_args(argv)
+
+
+def main(argv=None):
+    args = _parse_args(argv)
+
+    # Build queue list from flags or env, support comma-separated entries.
+    raw_entries = (
+        args.queue if args.queue else [os.getenv("JOB_QUEUE", "pipeline-queue")]
+    )
+    queues = []
+    for entry in raw_entries:
+        queues.extend([q.strip() for q in str(entry).split(",") if q and q.strip()])
+
+    if not queues:
+        raise SystemExit("No queues provided. Use -q ... or set JOB_QUEUE.")
+
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
+    log = logging.getLogger("oneshot")
+
+    pid = os.getpid()
+    ppid = os.getppid()
+    log.info(
+        "Starting one-shot worker PID=%s, Parent PID=%s, queues=%s, idle_timeout_ms=%s, stop_on_failure=%s",
+        pid,
+        ppid,
+        queues,
+        args.idle_timeout_ms if args.idle_timeout_ms > 0 else "∞",
+        args.stop_on_failure,
+    )
+
+    job_runner(
+        queues,
+        idle_timeout_ms=args.idle_timeout_ms,
+        stop_on_failure=args.stop_on_failure,
+    )
+
+
+if __name__ == "__main__":
+    main()