scalable-pypeline 2.0.10__py2.py3-none-any.whl → 2.1.1__py2.py3-none-any.whl

pypeline/pipelines/composition/pypeline_composition.py

@@ -0,0 +1,188 @@
+import json
+import typing
+from copy import copy
+from uuid import uuid4
+
+import networkx as nx
+from dramatiq import get_broker
+
+from pypeline.barrier import LockingParallelBarrier
+from pypeline.constants import REDIS_URL, PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+from pypeline.utils.dramatiq_utils import register_lazy_actor
+from pypeline.utils.module_utils import get_callable
+from pypeline.utils.pipeline_utils import get_execution_graph
+
+
+class Pypeline:
+    def __init__(
+        self,
+        pipeline: dict,
+        pipeline_settings: dict = None,
+        task_replacements: dict = {},
+        scenarios: dict = {},
+        broker=None,
+        execution_id=None,
+    ):
+        # Construct initial properties
+        self.pipeline = pipeline
+        self.broker = broker or get_broker()
+        self.execution_id = execution_id or str(uuid4())
+        self._starting_messages = []
+        self.scenarios = scenarios
+        self.pipeline_settings = pipeline_settings
+        self.task_replacements = task_replacements
+
+        # Get pipeline dag graph and find first task
+        pipeline_config = pipeline["config"]
+        self.graph = get_execution_graph(pipeline_config)
+        self.number_of_tasks = len(self.graph.nodes)
+        task_definitions = pipeline_config["taskDefinitions"]
+        first_task = list(pipeline_config["dagAdjacency"].keys())[0]
+
+        # Process the scenarios one by one
+        for scenario in self.scenarios:
+            tasks_in_reruns = scenario["taskReruns"]
+
+            # Find any tasks that have replacements for this scenario
+            tasks_in_replacements = list(scenario["taskReplacements"].keys())
+
+            distinct_scenario_tasks = list(set(tasks_in_reruns + tasks_in_replacements))
+            tasks_to_be_rerun_in_scenario = distinct_scenario_tasks
+
+            tasks_to_be_rerun_in_scenario = list(
+                set(
+                    task
+                    for task in distinct_scenario_tasks
+                    for task in nx.descendants(self.graph, task)
+                )
+                | set(tasks_to_be_rerun_in_scenario)
+            )
+
+            self.number_of_tasks = self.number_of_tasks + len(
+                tasks_to_be_rerun_in_scenario
+            )
+            scenario["tasksToRunInScenario"] = tasks_to_be_rerun_in_scenario
+            scenario["execution_id"] = scenario.get("execution_id", None) or str(
+                uuid4()
+            )
+
+            # Check if any of the scenarios need to be kicked off now
+            if first_task in tasks_to_be_rerun_in_scenario:
+                handler = task_definitions[first_task]["handlers"][
+                    scenario["taskReplacements"].get(first_task, 0)
+                ]
+                lazy_actor = register_lazy_actor(
+                    self.broker,
+                    get_callable(handler),
+                    pipeline_config["metadata"],
+                )
+                message = lazy_actor.message()
+                message.options["pipeline"] = pipeline
+                message.options["task_replacements"] = self.task_replacements
+                message.options["execution_id"] = scenario["execution_id"]
+                message.options["task_name"] = first_task
+                message.options["root_execution_id"] = self.execution_id
+                if self.pipeline_settings:
+                    message.kwargs["settings"] = copy(self.pipeline_settings)
+                    message.kwargs["settings"]["execution_id"] = scenario[
+                        "execution_id"
+                    ]
+                self._starting_messages.append(message)
+
+        for m in self._starting_messages:
+            m.options["scenarios"] = self.scenarios
+
+        handler = task_definitions[first_task]["handlers"][
+            self.task_replacements.get(first_task, 0)
+        ]
+        lazy_actor = register_lazy_actor(
+            self.broker,
+            get_callable(handler),
+            pipeline_config["metadata"],
+        )
+        message = lazy_actor.message()
+        message.options["pipeline"] = pipeline
+        message.options["task_replacements"] = self.task_replacements
+        message.options["execution_id"] = self.execution_id
+        message.options["task_name"] = first_task
+        message.options["scenarios"] = self.scenarios
+        message.options["root_execution_id"] = self.execution_id
+
+        if self.pipeline_settings:
+            message.kwargs["settings"] = copy(self.pipeline_settings)
+            message.kwargs["settings"]["execution_id"] = self.execution_id
+
+        self._starting_messages.append(message)
+
+    def run(self, *, delay=None):
+        for message in self._starting_messages:
+            task_key = (
+                f"{message.options['execution_id']}-{message.options['task_name']}"
+            )
+            locking_parallel_barrier = LockingParallelBarrier(
+                REDIS_URL, task_key=task_key, lock_key=f"{self.execution_id}-lock"
+            )
+            locking_parallel_barrier.set_task_count(1)
+            self.broker.enqueue(message, delay=delay)
+
+        return self
+
+    def __len__(self):
+        return self.number_of_tasks
+
+    def completed(self):
+        redis_task_keys = [
+            f"{self.execution_id}-{node}" for node in list(self.graph.nodes)
+        ]
+        redis_lock_key = f"{self.execution_id}-lock"
+        for scenario in self.scenarios:
+            scenario_task_keys = [
+                f"{scenario['execution_id']}-{task}"
+                for task in scenario["tasksToRunInScenario"]
+            ]
+            redis_task_keys = redis_task_keys + scenario_task_keys
+
+        for task_key in redis_task_keys:
+            locking_parallel_barrier = LockingParallelBarrier(
+                REDIS_URL, task_key=task_key, lock_key=redis_lock_key
+            )
+            try:
+                locking_parallel_barrier.acquire_lock(
+                    timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+                )
+                task_complete = True
+                if locking_parallel_barrier.task_exists():
+                    remaining_tasks = locking_parallel_barrier.get_task_count()
+                    if remaining_tasks >= 1:
+                        task_complete = False
+                else:
+                    task_complete = False
+            finally:
+                locking_parallel_barrier.release_lock()
+            if not task_complete:
+                return task_complete
+
+        return True
+
+    def to_json(self) -> str:
+        return json.dumps(
+            {
+                "pipeline": self.pipeline,
+                "pipeline_settings": self.pipeline_settings,
+                "task_replacements": self.task_replacements,
+                "scenarios": self.scenarios,
+                "execution_id": self.execution_id,
+            }
+        )
+
+    @classmethod
+    def from_json(cls, json_data: str) -> typing.Type["Pypeline"]:
+        data = json.loads(json_data)
+
+        return cls(
+            data["pipeline"],
+            pipeline_settings=data["pipeline_settings"],
+            task_replacements=data["task_replacements"],
+            scenarios=data["scenarios"],
+            execution_id=data["execution_id"],
+        )

pypeline/pipelines/factory.py

@@ -0,0 +1,107 @@
+import typing
+from dramatiq import get_broker, Message
+from pypeline.pipelines.composition.parallel_pipeline_composition import (
+    parallel_pipeline,
+)
+from pypeline.dramatiq import LazyActor
+from pypeline.utils.dramatiq_utils import register_lazy_actor
+from pypeline.pipeline_settings_schema import (
+    MissingSettingsException,
+    create_pipeline_settings_schema,
+    PipelineScenarioSchema,
+)
+from pypeline.pipelines.composition.pypeline_composition import Pypeline
+from pypeline.utils.config_utils import retrieve_latest_pipeline_config
+from pypeline.utils.module_utils import get_callable
+from pypeline.utils.pipeline_utils import (
+    get_execution_graph,
+    topological_sort_with_parallelism,
+)
+
+
+def dag_generator(
+    pipeline_id: str,
+    task_replacements: dict = {},
+    scenarios: typing.List[typing.Dict] = [],
+    *args,
+    **kwargs
+) -> typing.Union[parallel_pipeline, Pypeline]:
+    """Generates a pipeline dag from a pre-defined pipeline yaml
+
+    :param pipeline_id: Id of the pipeline to generate
+    :param task_replacements: A dictionary of task names and handler index to run. E.g. {"a": 1} would run the handler
+        in the second index position.
+    :param scenarios:
+    :param args:
+    :param kwargs:
+    :return: Returns a parallel_pipeline object which can be run
+    """
+    pipeline = retrieve_latest_pipeline_config(pipeline_id=pipeline_id)
+
+    pipeline_config = pipeline["config"]
+    broker = get_broker()
+    broker.actors.clear()
+
+    if pipeline["schemaVersion"] == 2:
+        # If the pipeline_config expects settings ensure we have them
+        if (
+            "settings" in pipeline_config
+            and len(pipeline_config["settings"]["required"]) > 0
+            and "settings" not in kwargs
+        ):
+            raise MissingSettingsException()
+
+        # If we're here we expect to  have settings.  Pop them out of kwargs to validate
+        inputted_settings = kwargs.pop("settings", {})
+        if "settings" in pipeline_config:
+            supplied_pipeline_settings_schema = create_pipeline_settings_schema(
+                pipeline_config["settings"]
+            )
+
+            # Validate scenarios settings to make sure they look okay
+            validated_scenarios = PipelineScenarioSchema(many=True).load(scenarios)
+
+            for scenario in validated_scenarios:
+                supplied_pipeline_settings_schema.load(scenario["settings"])
+
+            validated_settings = supplied_pipeline_settings_schema.load(
+                inputted_settings
+            )
+            p = Pypeline(
+                pipeline,
+                pipeline_settings=validated_settings,
+                task_replacements=task_replacements,
+                scenarios=scenarios,
+                broker=broker,
+            )
+        else:
+            p = Pypeline(pipeline, task_replacements=task_replacements, broker=broker)
+        return p
+    graph = get_execution_graph(pipeline_config)
+    optimal_execution_graph = topological_sort_with_parallelism(graph.copy())
+    registered_actors: typing.Dict[str, LazyActor] = {}
+
+    messages: typing.List[typing.List[Message]] = []
+
+    task_definitions = pipeline_config["taskDefinitions"]
+    for task_group in optimal_execution_graph:
+        message_group = []
+        for task in task_group:
+            module_path = task_definitions[task]["handler"]
+            tmp_handler = get_callable(module_path)
+            lazy_actor = register_lazy_actor(
+                broker, tmp_handler, pipeline_config["metadata"]
+            )
+            registered_actors[task] = lazy_actor
+            if args and not kwargs:
+                message_group.append(registered_actors[task].message(*args))
+            elif kwargs and not args:
+                message_group.append(registered_actors[task].message(**kwargs))
+            elif args and kwargs:
+                message_group.append(registered_actors[task].message(*args, **kwargs))
+            else:
+                message_group.append(registered_actors[task].message())
+        messages.append(message_group)
+    p = parallel_pipeline(messages)
+
+    return p

pypeline/pipelines/middleware/pypeline_middleware.py

@@ -0,0 +1,188 @@
+from copy import copy
+
+import networkx as nx
+from dramatiq import Middleware
+
+from pypeline.barrier import LockingParallelBarrier
+from pypeline.constants import PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+from pypeline.utils.module_utils import get_callable
+from pypeline.utils.pipeline_utils import get_execution_graph
+from pypeline.utils.dramatiq_utils import register_lazy_actor
+
+
+class PypelineMiddleware(Middleware):
+
+    def __init__(self, redis_url):
+        self.redis_url = redis_url
+
+    def after_process_message(self, broker, message, *, result=None, exception=None):
+
+        if exception is not None:
+            return
+
+        if "pipeline" not in message.options:
+            return
+
+        pipeline = message.options["pipeline"]
+        pipeline_config = pipeline["config"]
+        task_replacements = message.options["task_replacements"]
+        execution_id = message.options["execution_id"]
+        task_definitions = pipeline_config["taskDefinitions"]
+        task_name = message.options["task_name"]
+        task_key = f"{execution_id}-{task_name}"
+
+        # Signal to other jobs that current job is finished
+        locking_parallel_barrier = LockingParallelBarrier(
+            self.redis_url,
+            task_key=task_key,
+            lock_key=f"{message.options['root_execution_id']}-lock",
+        )
+        try:
+            locking_parallel_barrier.acquire_lock(
+                timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+            )
+            _ = locking_parallel_barrier.decrement_task_count()
+        finally:
+            locking_parallel_barrier.release_lock()
+
+        graph = get_execution_graph(pipeline_config)
+        children_tasks = pipeline_config["dagAdjacency"].get(task_name, [])
+
+        messages = []
+        for child in children_tasks:
+            child_ancestors = sorted(nx.ancestors(graph, child))
+
+            ancestor_tasks_complete = True
+
+            for ancestor in child_ancestors:
+                task_key = f"{execution_id}-{ancestor}"
+
+                locking_parallel_barrier = LockingParallelBarrier(
+                    self.redis_url,
+                    task_key=task_key,
+                    lock_key=f"{message.options['root_execution_id']}-lock",
+                )
+                try:
+                    locking_parallel_barrier.acquire_lock(
+                        timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+                    )
+
+                    if locking_parallel_barrier.task_exists():
+                        remaining_tasks = locking_parallel_barrier.get_task_count()
+                    else:
+                        remaining_tasks = None
+                finally:
+                    locking_parallel_barrier.release_lock()
+
+                if not remaining_tasks:
+                    task_key = f"{message.options['root_execution_id']}-{ancestor}"
+
+                    locking_parallel_barrier = LockingParallelBarrier(
+                        self.redis_url,
+                        task_key=task_key,
+                        lock_key=f"{message.options['root_execution_id']}-lock",
+                    )
+                    try:
+                        locking_parallel_barrier.acquire_lock(
+                            timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+                        )
+
+                        if locking_parallel_barrier.task_exists():
+                            remaining_tasks = locking_parallel_barrier.get_task_count()
+                        else:
+                            raise Exception(
+                                f"Parent task {ancestor} barrier lock not found"
+                            )
+                    finally:
+                        locking_parallel_barrier.release_lock()
+                if remaining_tasks >= 1:
+                    ancestor_tasks_complete = False
+                    break
+
+            # If the child's ancestor tasks aren't complete move onto the next child to check
+            if not ancestor_tasks_complete:
+                break
+
+            if message.options["root_execution_id"] == message.options["execution_id"]:
+                for scenario in message.options["scenarios"]:
+                    child_predecessors = list(graph.predecessors(child))
+                    if (
+                        child in scenario["tasksToRunInScenario"]
+                        and task_name in child_predecessors
+                        and task_name not in scenario["tasksToRunInScenario"]
+                    ):
+                        task_key = f"{scenario['execution_id']}-{child}"
+                        locking_parallel_barrier = LockingParallelBarrier(
+                            self.redis_url,
+                            task_key=task_key,
+                            lock_key=f"{message.options['root_execution_id']}-lock",
+                        )
+                        locking_parallel_barrier.set_task_count(1)
+                        handler = task_definitions[child]["handlers"][
+                            task_replacements.get(child, 0)
+                        ]
+
+                        lazy_actor = register_lazy_actor(
+                            broker,
+                            get_callable(handler),
+                            pipeline_config["metadata"],
+                        )
+                        scenario_message = lazy_actor.message()
+                        scenario_message.options["pipeline"] = pipeline
+                        scenario_message.options["task_replacements"] = (
+                            task_replacements
+                        )
+                        scenario_message.options["execution_id"] = scenario[
+                            "execution_id"
+                        ]
+
+                        scenario_message.options["task_name"] = child
+                        scenario_message.options["root_execution_id"] = message.options[
+                            "root_execution_id"
+                        ]
+                        scenario_message.options["scenarios"] = message.options[
+                            "scenarios"
+                        ]
+                        if "settings" in message.kwargs:
+                            scenario_message.kwargs["settings"] = copy(
+                                message.kwargs["settings"]
+                            )
+                            scenario_message.kwargs["settings"]["execution_id"] = (
+                                scenario["execution_id"]
+                            )
+                        messages.append(scenario_message)
+            task_key = f"{execution_id}-{child}"
+            locking_parallel_barrier = LockingParallelBarrier(
+                self.redis_url,
+                task_key=task_key,
+                lock_key=f"{message.options['root_execution_id']}-lock",
+            )
+            locking_parallel_barrier.set_task_count(1)
+            handler = task_definitions[child]["handlers"][
+                task_replacements.get(child, 0)
+            ]
+            lazy_actor = register_lazy_actor(
+                broker,
+                get_callable(handler),
+                pipeline_config["metadata"],
+            )
+
+            child_message = lazy_actor.message()
+            child_message.options["pipeline"] = pipeline
+            child_message.options["task_replacements"] = task_replacements
+            child_message.options["execution_id"] = execution_id
+            child_message.options["task_name"] = child
+            child_message.options["root_execution_id"] = message.options[
+                "root_execution_id"
+            ]
+            child_message.options["scenarios"] = message.options["scenarios"]
+            if "settings" in message.kwargs:
+                child_message.kwargs["settings"] = message.kwargs["settings"]
+                child_message.kwargs["settings"]["execution_id"] = message.options[
+                    "execution_id"
+                ]
+
+            messages.append(child_message)
+
+        for new_message in messages:
+            broker.enqueue(new_message)

pypeline/utils/dramatiq_utils.py

@@ -0,0 +1,126 @@
+import os.path
+import sys
+import typing
+from typing import Optional, Callable, Union, Awaitable
+from functools import wraps
+from typing import TYPE_CHECKING, TypeVar
+from dramatiq import Broker, actor as register_actor
+
+from pypeline.constants import (
+    DEFAULT_TASK_MAX_RETRY,
+    DEFAULT_TASK_MIN_BACKOFF,
+    MS_IN_SECONDS,
+    DEFAULT_TASK_MAX_BACKOFF,
+    DEFAULT_TASK_TTL,
+    DEFAULT_RESULT_TTL,
+)
+
+if TYPE_CHECKING:
+    from typing_extensions import ParamSpec
+
+    P = ParamSpec("P")
+else:
+    P = TypeVar("P")
+
+R = TypeVar("R")
+
+
+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 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]
+
+
+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", True)
+    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 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)

pypeline/utils/module_utils.py

@@ -1,11 +1,10 @@
 """ Utilities for loading modules/callables based on strings.
 """
-import os
+
 import re
 import logging
 import importlib
 from typing import Callable
-from pypeline.constants import API_ACCESS_KEY, PYPELINE_CLIENT_PKG_NAME
 
 logger = logging.getLogger(__name__)
 
@@ -81,3 +80,29 @@ def match_prefix_suffix(string: str, prefix_p: str, suffix_p: str) -> bool:
     if match_prefix(string, prefix_p) and match_suffix(string, suffix_p):
         return True
     return False
+
+
+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)