scalable-pypeline 2.0.9__py2.py3-none-any.whl → 2.1.0__py2.py3-none-any.whl

pypeline/__init__.py

@@ -1 +1 @@
-__version__ = "2.0.9"
+__version__ = "2.1.0"

pypeline/barrier.py

@@ -29,6 +29,9 @@ class LockingParallelBarrier:
         """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)

pypeline/dramatiq.py

@@ -1,6 +1,3 @@
-import importlib
-import os.path
-import sys
 import typing
 import pika
 import logging
@@ -10,11 +7,9 @@ from urllib.parse import urlparse
 
 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 import Broker, Middleware, set_broker, get_broker
 from dramatiq.brokers.rabbitmq import RabbitmqBroker
 from dramatiq.cli import (
     CPUS,
@@ -33,29 +28,24 @@ 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,
     DEFAULT_BROKER_CONNECTION_HEARTBEAT,
     DEFAULT_BROKER_BLOCKED_CONNECTION_TIMEOUT,
     DEFAULT_BROKER_CONNECTION_ATTEMPTS,
 )
-from pypeline.middleware import ParallelPipeline
+from pypeline.pipelines.middleware.parallel_pipeline_middleware import ParallelPipeline
+from pypeline.pipelines.middleware.pypeline_middleware import PypelineMiddleware
 from pypeline.utils.config_utils import (
     retrieve_latest_schedule_config,
     get_service_config_for_worker,
 )
+from pypeline.utils.dramatiq_utils import (
+    guess_code_directory,
+    list_managed_actors,
+    register_lazy_actor,
+    LazyActor,
+)
+from pypeline.utils.module_utils import get_callable
 
-if TYPE_CHECKING:
-    from typing_extensions import ParamSpec
-
-    P = ParamSpec("P")
-else:
-    P = TypeVar("P")
-R = TypeVar("R")
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -79,71 +69,12 @@ def configure_default_broker(broker: Broker = None):
     )
     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)
     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()
@@ -154,9 +85,18 @@ def register_actors_for_workers(broker: Broker):
         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 pipeline["schemaVersion"] == 1:
+                pipeline_tasks = [
+                    t["handler"] for t in pipeline_config["taskDefinitions"].values()
+                ]
+            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
@@ -292,77 +232,6 @@ class Dramatiq:
         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(
-            actor_name=f"{self.fn.__module__}.{self.fn.__name__}",
-            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
@@ -464,6 +333,9 @@ def pypeline_worker(
     # Wraps dramatiq worker CLI in a Flask command. This is private API of
     # dramatiq.
 
+    def format_actor(actor):
+        return "%s@%s" % (actor.actor_name, actor.queue_name)
+
     parser = dramatiq_argument_parser()
 
     # Set worker broker globally.

pypeline/flask/api/pipelines.py

@@ -3,6 +3,7 @@
 
 import importlib.metadata
 import logging
+from http import HTTPStatus
 
 from flask import jsonify
 from flask.views import MethodView
@@ -11,12 +12,17 @@ from marshmallow import Schema, fields
 from marshmallow.exceptions import ValidationError
 from webargs.flaskparser import abort
 from packaging import version
-from pypeline.composition import PipelineResult
+from pypeline.pipelines.composition.parallel_pipeline_composition import PipelineResult
 from pypeline.constants import API_DOC_RESPONSES, API_DOC_PARAMS, API_PATH_V1
 from pypeline.flask.decorators import require_accesskey
 from pypeline.pipeline_config_schema import BasePipelineSchema, PipelineSchemaV1
+from pypeline.pipeline_settings_schema import (
+    MissingSettingsException,
+    PipelineScenarioSchema,
+)
+from pypeline.pipelines.factory import dag_generator
 from pypeline.utils.config_utils import retrieve_latest_pipeline_config
-from pypeline.utils.pipeline_utils import dag_generator
+from pypeline.utils.schema_utils import get_clean_validation_messages
 
 logger = logging.getLogger(__name__)
 bp = Blueprint("pipelines", __name__, url_prefix=API_PATH_V1 + "/pipelines")
@@ -50,6 +56,31 @@ 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"},
+        required=False,
+    )
 
 
 class InvokePipelineResponseSchema(Schema):
@@ -166,12 +197,37 @@ class PipelineInvoke(MethodView):
 
         retval = {"pipeline_id": pipeline_id, "status": "starting"}
         try:
-            payload = payload["chain_payload"] if "chain_payload" in payload else {}
-            pipeline = dag_generator(pipeline_id=pipeline_id, event=payload)
+            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:
+                pipeline = dag_generator(
+                    pipeline_id=pipeline_id,
+                    task_replacements=task_replacements,
+                    scenarios=scenarios,
+                    settings=settings,
+                )
+                retval["scenarios"] = pipeline.scenarios
             pipeline.run()
             pipeline_result = PipelineResult(pipeline.execution_id)
             pipeline_result.create_result_entry(pipeline.to_json())
             retval["execution_id"] = pipeline.execution_id
+        except MissingSettingsException:
+            abort(
+                HTTPStatus.BAD_REQUEST,
+                message="Missing required settings in the request.",
+            )
+        except ValidationError as ve:
+            abort(
+                HTTPStatus.BAD_REQUEST,
+                message=get_clean_validation_messages(ve),
+            )
         except Exception as e:
             msg = "Failed to invoke pipeline ... {}".format(pipeline_id)
             logger.error(msg)

pypeline/flask/api/schedules.py

@@ -1,16 +1,14 @@
 """ API Endpoints for Scheduled Tasks
 """
-import os
+
 import logging
 from flask import jsonify, request
 from flask.views import MethodView
-from marshmallow import Schema, fields
 from flask_smorest import Blueprint
 from flask import abort
 from marshmallow.exceptions import ValidationError
 from pypeline.constants import API_DOC_RESPONSES, API_DOC_PARAMS, API_PATH_V1
 from pypeline.utils.config_utils import retrieve_latest_schedule_config
-from pypeline.schedule_config_schema import BaseScheduleSchema
 from pypeline.flask.decorators import require_accesskey
 
 logger = logging.getLogger(__name__)

pypeline/pipeline_config_schema.py

@@ -1,9 +1,12 @@
 """ Schemas for Pipelines
 """
+
 import yaml
 from marshmallow import Schema, fields, EXCLUDE, validates_schema
 from marshmallow.exceptions import ValidationError
 
+from pypeline.pipeline_settings_schema import PipelineSettingsSchema
+
 
 class ExcludeUnknownSchema(Schema):
     """Remove unknown keys from loaded dictionary"""
@@ -81,8 +84,17 @@ class MetadataSchema(Schema):
         example=600,
     )
 
+    groupName = fields.String(
+        required=False,
+        metadata={
+            "description": "If two pipelines logically belong to a group the user can identify that two.  "
+            "Imagine pipeline_a and pipeline_b both process data for images.  "
+            'Logically we could give them a mutual group name of "Image Processing Pipelines"'
+        },
+    )
+
 
-class TaskDefinitionsSchema(ExcludeUnknownSchema):
+class TaskDefinitionsSchemaV1(ExcludeUnknownSchema):
     """Schema for a single task's configuration"""
 
     handler = fields.String(
@@ -105,7 +117,31 @@ class TaskDefinitionsSchema(ExcludeUnknownSchema):
     )
 
 
-class PipelineConfigSchemaV1(Schema):
+class TaskDefinitionsSchemaV2(ExcludeUnknownSchema):
+    """Schema for a single task's configuration"""
+
+    handlers = fields.List(
+        fields.String(
+            required=True,
+            description="Path to the worker task definition",
+            example="client.workers.my_task",
+        )
+    )
+    maxTtl = fields.Integer(
+        required=False,
+        description="Max TTL for a task in seconds.",
+        default=60,
+        example=60,
+    )
+
+    queue = fields.String(
+        required=False,
+        description="Non-default queue for this task.",
+        example="custom-queue-name",
+    )
+
+
+class PipelineConfigSchemaBase(Schema):
     """Overall pipeline configuration schema"""
 
     metadata = fields.Nested(
@@ -128,6 +164,31 @@ class PipelineConfigSchemaV1(Schema):
         required=True,
         description="The DAG Adjacency definition.",
     )
+
+
+class PipelineConfigSchemaV1(PipelineConfigSchemaBase):
+    """Overall pipeline configuration schema"""
+
+    taskDefinitions = fields.Dict(
+        keys=fields.String(
+            required=True,
+            description="Task's node name. *Must* match related key in dagAdjacency.",
+            example="node_a",
+        ),
+        values=fields.Nested(
+            TaskDefinitionsSchemaV1,
+            required=True,
+            description="Definition of each task in the pipeline.",
+            example={"handler": "abc.task", "maxRetry": 1},
+        ),
+        required=True,
+        description="Configuration for each node defined in DAG.",
+    )
+
+
+class PipelineConfigSchemaV2(PipelineConfigSchemaBase):
+    """Overall pipeline configuration schema"""
+
     taskDefinitions = fields.Dict(
         keys=fields.String(
             required=True,
@@ -135,7 +196,7 @@ class PipelineConfigSchemaV1(Schema):
             example="node_a",
         ),
         values=fields.Nested(
-            TaskDefinitionsSchema,
+            TaskDefinitionsSchemaV2,
             required=True,
             description="Definition of each task in the pipeline.",
             example={"handler": "abc.task", "maxRetry": 1},
@@ -144,6 +205,14 @@ class PipelineConfigSchemaV1(Schema):
         description="Configuration for each node defined in DAG.",
     )
 
+    settings = fields.Nested(
+        PipelineSettingsSchema,
+        required=False,
+        metadata={
+            "description": "Settings schema to validate the actual settings being passed through to the pipelines."
+        },
+    )
+
 
 class BasePipelineSchema(ExcludeUnknownSchema):
     __schema_version__ = None
@@ -185,6 +254,25 @@ class BasePipelineSchema(ExcludeUnknownSchema):
         schema.load(data)
 
 
+class PipelineSchemaV2(BasePipelineSchema):
+    __schema_version__ = 2
+
+    class Meta:
+        unknown = EXCLUDE
+
+    config = fields.Nested(
+        PipelineConfigSchemaV2,
+        required=True,
+        description="Metadata and configuration information for this pipeline.",
+    )
+
+    def validate_pipeline(self, data, **kwargs):
+        # We need to add this function to avoid infinite recursion since
+        # the BasePipelineSchema class above uses the same method for
+        # validation
+        pass
+
+
 class PipelineSchemaV1(BasePipelineSchema):
     __schema_version__ = 1