scalable-pypeline 1.2.3__py2.py3-none-any.whl → 2.0.2__py2.py3-none-any.whl

pypeline/dramatiq.py

@@ -0,0 +1,470 @@
+import importlib
+import os.path
+import sys
+import typing
+import logging
+import click
+from pypeline.extensions import pypeline_config
+from warnings import warn
+from functools import wraps
+from apscheduler.schedulers.blocking import BlockingScheduler
+from apscheduler.triggers.cron import CronTrigger
+from typing import Awaitable, Callable, Optional, Union, TYPE_CHECKING, TypeVar
+from dramatiq import Broker, Middleware, actor as register_actor, set_broker, get_broker
+from dramatiq.brokers.rabbitmq import RabbitmqBroker
+from dramatiq.cli import (
+    CPUS,
+    HAS_WATCHDOG,
+    main as dramatiq_worker,
+    make_argument_parser as dramatiq_argument_parser,
+    import_object,
+)
+from dramatiq.middleware import default_middleware
+from dramatiq.results import Results
+from dramatiq.results.backends.redis import RedisBackend
+from flask import current_app, Flask
+from flask.cli import with_appcontext
+
+from pypeline.constants import (
+    REDIS_URL,
+    RABBIT_URL,
+    DEFAULT_BROKER_CALLABLE,
+    MS_IN_SECONDS,
+    DEFAULT_TASK_TTL,
+    DEFAULT_RESULT_TTL,
+    DEFAULT_TASK_MAX_RETRY,
+    DEFAULT_TASK_MIN_BACKOFF,
+    DEFAULT_TASK_MAX_BACKOFF,
+)
+from pypeline.middleware import ParallelPipeline
+from pypeline.utils.config_utils import (
+    retrieve_latest_schedule_config,
+    get_service_config_for_worker,
+)
+
+if TYPE_CHECKING:
+    from typing_extensions import ParamSpec
+
+    P = ParamSpec("P")
+else:
+    P = TypeVar("P")
+R = TypeVar("R")
+
+logger = logging.getLogger(__name__)
+
+
+def configure_default_broker(broker: Broker = None):
+    redis_backend = RedisBackend(url=REDIS_URL)
+    rabbit_broker = broker if broker is not None else RabbitmqBroker(url=RABBIT_URL)
+    rabbit_broker.add_middleware(Results(backend=redis_backend))
+    rabbit_broker.add_middleware(ParallelPipeline(redis_url=REDIS_URL))
+    register_actors_for_workers(rabbit_broker)
+    set_broker(rabbit_broker)
+
+
+def guess_code_directory(broker):
+    actor = next(iter(broker.actors.values()))
+    modname, *_ = actor.fn.__module__.partition(".")
+    mod = sys.modules[modname]
+    return os.path.dirname(mod.__file__)
+
+
+def get_module(resource_dot_path: str):
+    """Retrieve the module based on a 'resource dot path'.
+    e.g. package.subdir.feature_file.MyCallable
+    """
+    module_path = ".".join(resource_dot_path.split(".")[:-1])
+    module = importlib.import_module(module_path)
+    return module
+
+
+def get_callable_name(resource_dot_path: str) -> str:
+    """Retrieve the callable based on config string.
+    e.g. package.subdir.feature_file.MyCallable
+    """
+    callable_name = resource_dot_path.split(".")[-1]
+    return callable_name
+
+
+def get_callable(resource_dot_path: str) -> Callable:
+    """Retrieve the actual handler class based on config string.
+    e.g. package.subdir.feature_file.MyCallable
+    """
+    module = get_module(resource_dot_path)
+    callable_name = get_callable_name(resource_dot_path)
+    return getattr(module, callable_name)
+
+
+def register_lazy_actor(
+    broker: Broker,
+    fn: Optional[Callable[P, Union[Awaitable[R], R]]] = None,
+    pipeline_meta: typing.Dict = {},
+    **kwargs,
+) -> typing.Type["LazyActor"]:
+    kwargs["queue_name"] = pipeline_meta.get("queue", "default")
+    kwargs["max_retries"] = pipeline_meta.get("maxRetry", DEFAULT_TASK_MAX_RETRY)
+    # Convert from seconds to milliseconds
+    kwargs["min_backoff"] = (
+        pipeline_meta.get("retryBackoff", DEFAULT_TASK_MIN_BACKOFF) * MS_IN_SECONDS
+    )
+    kwargs["max_backoff"] = (
+        pipeline_meta.get("retryBackoffMax", DEFAULT_TASK_MAX_BACKOFF) * MS_IN_SECONDS
+    )
+    kwargs["time_limit"] = pipeline_meta.get("maxTtl", DEFAULT_TASK_TTL) * MS_IN_SECONDS
+    # Always store results for registered pipeline actors
+    kwargs["store_results"] = pipeline_meta.get("store_results", False)
+    if kwargs["store_results"]:
+        kwargs["result_ttl"] = (
+            pipeline_meta.get("result_ttl", DEFAULT_RESULT_TTL) * MS_IN_SECONDS
+        )
+    lazy_actor: LazyActor = LazyActor(fn, kwargs)
+    lazy_actor.register(broker)
+    return lazy_actor
+
+
+def register_actors_for_workers(broker: Broker):
+    service = get_service_config_for_worker(pypeline_config)
+    scheduled_jobs_config = retrieve_latest_schedule_config()
+
+    if not service:
+        return
+    for task in service.get("registeredTasks", []):
+        print(task)
+        pipeline_meta = None
+        for pipeline_key, pipeline in pypeline_config["pipelines"].items():
+            pipeline_config = pipeline["config"]
+            pipeline_tasks = [
+                t["handler"] for t in pipeline_config["taskDefinitions"].values()
+            ]
+            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:
+            print(e)
+            print(f"Unable to add a task to dramatiq: {e}")
+
+
+class Dramatiq:
+    """Flask extension bridging Dramatiq broker and Flask app.
+
+    Dramatiq API is eager. Broker initialisation precede actor declaration.
+    This breaks application factory pattern and other way to initialize
+    configuration after import.
+
+    This class enables lazy initialization of Dramatiq. Actual Dramatiq broker
+    is instanciated only once Flask app is created.
+
+    .. automethod:: actor
+    .. automethod:: init_app
+    """
+
+    def __init__(
+        self,
+        app: Flask = None,
+        name: str = "dramatiq",
+        config_prefix: str = None,
+        middleware: typing.List[Middleware] = None,
+    ):
+        """
+        :app: Flask application if created. See :meth:`init_app`.
+
+        :param broker_configuration_callable_module: In order to work in fork and spawn mode
+            we need to configure our broker using a callable function.  Default is specified as
+            "pypeline.flask_dramatiq:configure_default_broker".  This allows the user to
+            override if necessary.
+
+        :param name: Unique identifier for multi-broker app.
+
+        :param config_prefix: Flask configuration option prefix for this
+            broker. By default, it is derived from ``name`` parameter,
+            capitalized.
+
+        :param middleware: List of Dramatiq middleware instances to override
+             Dramatiq defaults.
+
+        Flask-Dramatiq always prepend a custom middleware to the middleware
+        stack that setup Flask context. This way, every middleware can use
+        Flask app context.
+
+        """
+        self.actors = []
+        self.app = None
+        self.config_prefix = config_prefix or name.upper() + "_BROKER"
+        self.name = name
+        self.broker = None
+        if middleware is None:
+            middleware = [m() for m in default_middleware]
+        self.middleware = middleware
+        if app:
+            self.init_app(app)
+
+    def __repr__(self) -> str:
+        return "<%s %s>" % (self.__class__.__name__, self.name)
+
+    def init_app(self, app: Flask):
+        """Initialize extension for one Flask application
+
+        This method triggers Dramatiq broker instantiation and effective actor
+        registration.
+
+        """
+        if self.app is not None:
+            warn(
+                "%s is used by more than one flask application. "
+                "Actor's context may be set incorrectly." % (self,),
+                stacklevel=2,
+            )
+        self.app = app
+        app.extensions["dramatiq-" + self.name] = self
+
+        module_name, broker_or_callable = import_object(DEFAULT_BROKER_CALLABLE)
+
+        # Callable function is expected to setBroker()
+        if callable(broker_or_callable):
+            print(f"Configuring broker via {DEFAULT_BROKER_CALLABLE}")
+            broker_or_callable()
+        else:
+            raise TypeError("DEFAULT_BROKER_CALLABLE must point to a callable function")
+        self.broker = get_broker()
+        for actor in self.actors:
+            actor.register(broker=self.broker)
+
+    def actor(self, fn=None, **kw):
+        """Register a callable as Dramatiq actor.
+
+        This decorator lazily register a callable as a Dramatiq actor. The
+        actor can't be called before :meth:`init_app` is called.
+
+        :param kw: Keywords argument passed to :func:`dramatiq.actor`.
+
+        """
+        # Substitute dramatiq.actor decorator to return a lazy wrapper. This
+        # allows to register actors in extension before the broker is
+        # effectively configured by init_app.
+
+        def decorator(fn):
+            lazy_actor = LazyActor(self, fn, kw)
+            self.actors.append(lazy_actor)
+            if self.app:
+                lazy_actor.register(self.broker)
+            return lazy_actor
+
+        if fn:
+            return decorator(fn)
+        return decorator
+
+
+def format_actor(actor):
+    return "%s@%s" % (actor.actor_name, actor.queue_name)
+
+
+def ensure_return_value(default_value=None):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Call the original function
+            result = func(*args, **kwargs)
+            # Check if the function has returned a value
+            if result is None:
+                # Return the default value if the function returned None
+                return default_value
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+class LazyActor(object):
+    # Intermediate object that register actor on broker an call.
+
+    def __init__(self, fn, kw):
+        self.fn = fn
+        self.kw = kw
+        self.actor = None
+
+    def __call__(self, *a, **kw):
+        return self.fn(*a, **kw)
+
+    def __repr__(self):
+        return "<%s %s.%s>" % (
+            self.__class__.__name__,
+            self.fn.__module__,
+            self.fn.__name__,
+        )
+
+    def __getattr__(self, name):
+        if not self.actor:
+            raise AttributeError(name)
+        return getattr(self.actor, name)
+
+    def register(self, broker):
+        self.actor = register_actor(broker=broker, **self.kw)(
+            ensure_return_value(default_value=True)(self.fn)
+        )
+
+    # Next is regular actor API.
+    def send(self, *a, **kw):
+        return self.actor.send(*a, **kw)
+
+    def message(self, *a, **kw):
+        return self.actor.message(*a, **kw)
+
+    def send_with_options(self, *a, **kw):
+        return self.actor.send_with_options(*a, **kw)
+
+
+def list_managed_actors(broker, queues):
+    queues = set(queues)
+    all_actors = broker.actors.values()
+    if not queues:
+        return all_actors
+    else:
+        return [a for a in all_actors if a.queue_name in queues]
+
+
+@click.command("cron-scheduler")
+def cron_scheduler():  # pragma: no cover
+    # Configure our broker that we will schedule registered tasks for
+    scheduler = BlockingScheduler()
+    configure_default_broker()
+    broker = get_broker()
+    jobs = retrieve_latest_schedule_config()
+
+    for job in jobs:
+        if job["enabled"]:
+            config = job["config"]
+            worker_path = config["task"]
+            tmp_handler = get_callable(worker_path)
+            pipeline_meta = {"queue": config.get("queue", "default")}
+            actor = register_lazy_actor(broker, tmp_handler, pipeline_meta)
+            schedule = config["schedule"]
+            scheduler.add_job(
+                actor.send,
+                CronTrigger.from_crontab(
+                    f"{schedule['minute']} {schedule['hour']} {schedule['dayOfMonth']} {schedule['monthOfYear']} {schedule['dayOfWeek']}"
+                ),
+            )
+
+    try:
+        scheduler.start()
+    except KeyboardInterrupt:
+        scheduler.shutdown()
+
+
+@click.command("pypeline-worker")
+@click.argument("broker_name", default="dramatiq")
+@click.option(
+    "-v", "--verbose", default=0, count=True, help="turn on verbose log output"
+)
+@click.option(
+    "-p",
+    "--processes",
+    default=CPUS,
+    metavar="PROCESSES",
+    show_default=True,
+    help="the number of worker processes to run",
+)
+@click.option(
+    "-t",
+    "--threads",
+    default=8,
+    metavar="THREADS",
+    show_default=True,
+    help="the number of worker treads per processes",
+)
+@click.option(
+    "-Q",
+    "--queues",
+    type=str,
+    default=None,
+    metavar="QUEUES",
+    show_default=True,
+    help="listen to a subset of queues, comma separated",
+)
+@click.option(
+    "--use-spawn",
+    type=bool,
+    default=False,
+    metavar="USE_SPAWN",
+    show_default=True,
+    help="start processes by spawning (default: fork on unix, spawn on windows)",
+)
+@with_appcontext
+def pypeline_worker(
+    verbose, processes, threads, queues, broker_name, use_spawn
+):  # pragma: no cover
+    """Run dramatiq workers.
+
+    Setup Dramatiq with broker and task modules from Flask app.
+
+    \b
+    examples:
+      # Run dramatiq with 1 thread per process.
+      $ flask worker --threads 1
+
+    \b
+      # Listen only to the "foo" and "bar" queues.
+      $ flask worker -Q foo,bar
+
+    \b
+      # Consuming from a specific broker
+      $ flask worker mybroker
+    """
+    # Plugin for flask.commands entrypoint.
+    #
+    # Wraps dramatiq worker CLI in a Flask command. This is private API of
+    # dramatiq.
+
+    parser = dramatiq_argument_parser()
+
+    # Set worker broker globally.
+    needle = "dramatiq-" + broker_name
+    broker = current_app.extensions[needle].broker
+    set_broker(broker)
+
+    command = [
+        "--processes",
+        str(processes),
+        "--threads",
+        str(threads),
+        # Fall back to flask_dramatiq global broker
+        DEFAULT_BROKER_CALLABLE,
+    ]
+
+    if use_spawn:
+        command += ["--use-spawn"]
+
+    if current_app.config["DEBUG"]:
+        verbose = max(1, verbose)
+        if HAS_WATCHDOG:
+            command += ["--watch", guess_code_directory(broker)]
+
+    queues = queues.split(",") if queues else []
+    if queues:
+        command += ["--queues"] + queues
+    command += verbose * ["-v"]
+    args = parser.parse_args(command)
+    current_app.logger.info("Able to execute the following actors:")
+    for actor in list_managed_actors(broker, queues):
+        current_app.logger.info("    %s.", format_actor(actor))
+
+    dramatiq_worker(args)

pypeline/extensions.py

@@ -5,12 +5,13 @@ import logging
 logger = logging.getLogger(__name__)
 
 try:
-    # Client packages *should* provide a `sermos.yaml` file. This
-    # loads the configuration file with the provided name ofthe client
-    # package (e.g. sermos_demo_client)
-    from pypeline.sermos_yaml import load_client_config_and_version
-    sermos_config, sermos_client_version = load_client_config_and_version()
+    # Client packages *should* provide a `pypeline.yaml` file. This
+    # loads the configuration file with the provided name of the client
+    # package (e.g. pypeline_demo)
+    from pypeline.pypeline_yaml import load_client_config_and_version
+
+    pypeline_config, pypeline_client_version = load_client_config_and_version()
 except Exception as e:
-    sermos_config = None
-    sermos_client_version = None
-    logger.warning("Unable to load client Sermos config ... {}".format(e))
+    pypeline_config = None
+    pypeline_client_version = None
+    logger.warning("Unable to load client Pypeline config ... {}".format(e))

pypeline/flask/__init__.py

@@ -8,11 +8,9 @@ try:
     from flask_smorest import Blueprint, Api
     from flask import abort
 except Exception as e:
-    logger.error("Unable to import Web services (Blueprint, API, abort)"
-                 f" ... {e}")
+    logger.error("Unable to import Web services (Blueprint, API, abort)" f" ... {e}")
 
 try:
-    from pypeline.flask.flask_sermos import FlaskSermos
+    from pypeline.flask.flask_pypeline import FlaskPypeline
 except Exception as e:
-    logger.exception("Unable to import Sermos services (FlaskSermos)"
-                     f" ... {e}")
+    logger.exception("Unable to import Sermos services (FlaskPypeline)" f" ... {e}")