scalable-pypeline 2.1.31__py2.py3-none-any.whl

pypeline/pipelines/composition/parallel_pipeline_composition.py

@@ -0,0 +1,375 @@
+from __future__ import annotations
+
+import copy
+import json
+import time
+import typing
+from uuid import uuid4
+from urllib.parse import urlparse
+
+from dramatiq.broker import get_broker
+from dramatiq.results import ResultMissing
+from db_medley.redis_conf import RedisConnector
+from redis.exceptions import RedisError
+from redis.sentinel import Sentinel
+from pypeline.constants import (
+    REDIS_URL,
+    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,
+)
+from pypeline.barrier import LockingParallelBarrier
+from pypeline.constants import DEFAULT_RESULT_TTL
+from pypeline.dramatiq import REDIS_URL
+
+from dramatiq.message import Message
+
+
+class parallel_pipeline:
+    """Chain actors together, passing the result of one actor to the
+    next one in line.
+
+    Parameters:
+      children(typing.List[typing.List[Message]]): A sequence of messages or
+        pipelines.  Child pipelines are flattened into the resulting
+        pipeline.
+      broker(Broker): The broker to run the pipeline on.  Defaults to
+        the current global broker.
+    """
+
+    messages: list[Message]
+
+    def __init__(self, messages: typing.List[typing.List[Message]], broker=None):
+        self.broker = broker or get_broker()
+        self.messages = messages
+        self.execution_id = str(uuid4())
+        execution_graph = []
+
+        for message_group in self.messages:
+            sub_execution_group = []
+            group_completion_uuid = str(uuid4())
+            for m in message_group:
+                m.kwargs["event"]["execution_id"] = self.execution_id
+                m.options["group_completion_uuid"] = group_completion_uuid
+                message_dict = copy.deepcopy(m.asdict())
+                sub_execution_group.append(message_dict)
+            # Last item in the group is the id of the group to be executed
+            execution_graph.append(sub_execution_group)
+
+        self.execution_graph = execution_graph
+
+        for m in self.messages[0]:
+            m.options["execution_graph"] = execution_graph
+
+    def __len__(self):
+        """Returns the length of the parallel_pipeline."""
+        count = 0
+        for message_group in self.messages:
+            count = count + len(message_group)
+
+        return count
+
+    def __str__(self):  # pragma: no cover
+        """Return a string representation of the parallel_pipeline.
+
+        This representation shows the order of execution for each group of messages.
+        """
+        result = []
+
+        for i, message_group in enumerate(self.messages):
+            group_str = f"Group {i + 1}: [\n"
+            for j, message in enumerate(message_group):
+                message_str = f"  Message {j + 1}: {message.actor_name}\n"
+                group_str += message_str
+            group_str += "]\n"
+            result.append(group_str)
+
+        return "".join(result)
+
+    @property
+    def completed(self):
+        return self.completed_count == len(self)
+
+    @property
+    def completed_count(self):
+        count = 0
+
+        for message_group in self.messages:
+            for message in message_group:
+                try:
+                    message.get_result()
+                    count = count + 1
+                except ResultMissing:
+                    pass
+        return count
+
+    def run(self, *, delay=None):
+        """Run this parallel_pipeline.
+
+        Parameters:
+          delay(int): The minimum amount of time, in milliseconds, the
+            parallel_pipeline should be delayed by. If both parallel_pipeline's delay and
+            first message's delay are provided, the bigger value will be
+            used.
+
+        Returns:
+          parallel_pipeline: Itself.
+        """
+        starting_group = self.messages[0]
+
+        completion_uuid = starting_group[0].options["group_completion_uuid"]
+        locking_parallel_barrier = LockingParallelBarrier(
+            REDIS_URL, task_key=completion_uuid, lock_key=f"{completion_uuid}-lock"
+        )
+        locking_parallel_barrier.set_task_count(len(starting_group))
+
+        for m in starting_group:
+            self.broker.enqueue(m, delay=delay)
+
+        return self
+
+    def get_result(self, *, block=False, timeout=None):
+        """Get the result of this pipeline.
+
+        Pipeline results are represented by the result of the last
+        message in the chain.
+
+        Parameters:
+          block(bool): Whether or not to block until a result is set.
+          timeout(int): The maximum amount of time, in ms, to wait for
+            a result when block is True.  Defaults to 10 seconds.
+
+        Raises:
+          ResultMissing: When block is False and the result isn't set.
+          ResultTimeout: When waiting for a result times out.
+
+        Returns:
+          object: The result.
+        """
+        last_message = self.messages[-1][-1]
+
+        backend = self.broker.get_results_backend()
+        return last_message.get_result(backend=backend, block=block, timeout=timeout)
+
+    def get_results(self, *, block=False, timeout=None):
+        """Get the results of each job in the pipeline.
+
+        Parameters:
+          block(bool): Whether or not to block until a result is set.
+          timeout(int): The maximum amount of time, in ms, to wait for
+            a result when block is True.  Defaults to 10 seconds.
+
+        Raises:
+          ResultMissing: When block is False and the result isn't set.
+          ResultTimeout: When waiting for a result times out.
+
+        Returns:
+          A result generator.
+        """
+        deadline = None
+        if timeout:
+            deadline = time.monotonic() + timeout / 1000
+
+        for message_group in self.messages:
+            for message in message_group:
+                if deadline:
+                    timeout = max(0, int((deadline - time.monotonic()) * 1000))
+
+                backend = self.broker.get_results_backend()
+                yield {
+                    message.actor_name: message.get_result(
+                        backend=backend, block=block, timeout=timeout
+                    )
+                }
+
+    def to_json(self) -> str:
+        """Convert the execution graph to a JSON string representation.
+
+        This method serializes the execution graph of the pipeline into a JSON string.
+        This serialized form can be used to save the pipeline state or share it across different systems,
+        enabling the retrieval of a pipeline "run" for obtaining its results at a later time.
+
+        :return: A JSON string representing the execution graph.
+        :rtype: str
+        """
+        return json.dumps(self.execution_graph)
+
+    @classmethod
+    def from_json(cls, json_data: str) -> parallel_pipeline:
+        """Create a ParallelPipeline object from a JSON string representation of the execution graph.
+
+        This class method deserializes a JSON string into a list of messages, each representing
+        a task or operation in the pipeline. The method reconstructs the execution graph using
+        the `dramatiq.message.Message` objects and returns an instance of the `parallel_pipeline` class.
+
+        :param json_data: A JSON string containing the serialized execution graph.
+        :type json_data: str
+        :return: An instance of `parallel_pipeline` reconstructed from the JSON data.
+        :rtype: parallel_pipeline
+        """
+        execution_graph = json.loads(json_data)
+
+        messages = []
+
+        for message_group in execution_graph:
+            temp_group = []
+            for message in message_group:
+                temp_group.append(Message(**message))
+            messages.append(temp_group)
+
+        return cls(messages)
+
+
+class PipelineResult:
+    """
+    A class to manage and retrieve the results of a parallel pipeline execution.
+
+    The `PipelineResult` class provides methods for creating a result entry in a Redis database,
+    loading pipeline data from Redis, and retrieving the status and results of the pipeline execution.
+
+    Attributes:
+        pipeline (parallel_pipeline): The pipeline object representing the execution graph.
+        execution_id (str): A unique identifier for the execution of the pipeline.
+        redis_key (str): The key used to store and retrieve pipeline data from Redis.
+        redis_conn: A Redis connection object used to interact with the Redis database.
+        result_ttl (int): Time-to-live (TTL) for the result entry in Redis, in seconds.
+    """
+
+    def __init__(self, execution_id: str, result_ttl: int = DEFAULT_RESULT_TTL):
+        """
+        Initialize a PipelineResult object with an execution ID and optional result TTL.
+
+        :param execution_id: A unique identifier for the pipeline execution.
+        :type execution_id: str
+        :param result_ttl: The time-to-live (TTL) for the result entry in Redis. Defaults to DEFAULT_RESULT_TTL.
+        :type result_ttl: int
+        """
+        self.pipeline: parallel_pipeline = None
+        self.execution_id = execution_id
+        self.redis_key = f"{execution_id}-results-key"
+        self.result_ttl = result_ttl
+
+        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_conn = 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,
+            )
+        else:
+            self.redis_conn = RedisConnector().get_connection()
+
+    def create_result_entry(self, pipeline_json_str: str):
+        """
+        Store the serialized pipeline data in Redis with a specified TTL.
+
+        This method saves the JSON string representation of the pipeline in the Redis database
+        using the execution ID as the key. The entry is stored with a time-to-live (TTL) defined by `result_ttl`.
+
+        :param pipeline_json_str: A JSON string representing the pipeline execution graph.
+        :type pipeline_json_str: str
+        :raises ValueError: If the provided pipeline data is None or an empty string.
+        :raises RedisError: If there is an issue connecting to Redis or setting the value.
+        """
+        if not pipeline_json_str:
+            raise ValueError("No pipeline data passed to create result store")
+
+        try:
+            self.redis_conn.setex(self.redis_key, self.result_ttl, pipeline_json_str)
+        except RedisError as e:
+            raise RuntimeError(f"Failed to store pipeline data in Redis: {e}")
+
+    def load(self):
+        """
+        Load the pipeline data from Redis and reconstruct the pipeline object.
+
+        This method retrieves the JSON string stored in Redis and deserializes it
+        into a `parallel_pipeline` object, enabling access to the pipeline's execution details.
+
+        :raises RedisError: If there is an issue connecting to Redis or retrieving the data.
+        """
+        try:
+            pipeline_data = self.redis_conn.get(self.redis_key)
+            if pipeline_data:
+                self.pipeline = parallel_pipeline.from_json(pipeline_data)
+            else:
+                self.pipeline = None
+        except RedisError as e:
+            raise RuntimeError(f"Failed to load pipeline data from Redis: {e}")
+
+    @property
+    def status(self) -> str:
+        """
+        Get the current status of the pipeline execution.
+
+        This property checks the completion status of the pipeline and returns its current state.
+
+        :return: The status of the pipeline execution, which can be "complete", "pending", or "unavailable".
+        :rtype: str
+        """
+        if not self.pipeline:
+            return "unavailable"
+        return "complete" if self.pipeline.completed else "pending"
+
+    def get_results(self) -> dict:
+        """
+        Retrieve all results from the pipeline execution with unique actor identifiers.
+
+        This method aggregates results from the pipeline and ensures that each actor's result
+        has a unique identifier by appending a numeric suffix to duplicate actor names.
+
+        :return: A dictionary containing all results from the pipeline execution, keyed by unique actor identifiers.
+        :rtype: dict
+        """
+        if not self.pipeline:
+            return {}
+
+        results = {}
+        for result in self.pipeline.get_results():
+            for actor, res in result.items():
+                unique_actor = self._get_unique_actor_name(actor, results)
+                results[unique_actor] = res
+        return results
+
+    def get_result(self):
+        """
+        Retrieve a single result from the pipeline execution.
+
+        This method returns the result of a single execution step from the pipeline, if available.
+
+        :return: The result of a single execution step from the pipeline, or None if no pipeline is loaded.
+        """
+        if self.pipeline:
+            return self.pipeline.get_result()
+
+    def _get_unique_actor_name(self, actor: str, results: dict) -> str:
+        """
+        Generate a unique actor name by appending a numeric suffix if necessary.
+
+        :param actor: The base name of the actor.
+        :type actor: str
+        :param results: The current dictionary of results to check for uniqueness.
+        :type results: dict
+        :return: A unique actor name.
+        :rtype: str
+        """
+        if actor not in results:
+            return actor
+
+        suffix = 0
+        new_actor = f"{actor}-{suffix}"
+        while new_actor in results:
+            suffix += 1
+            new_actor = f"{actor}-{suffix}"
+        return new_actor

pypeline/pipelines/composition/pypeline_composition.py

@@ -0,0 +1,215 @@
+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,
+        scenarios: dict = {},
+        broker=None,
+    ):
+        # Construct initial properties
+        self.pipeline = pipeline
+        self.broker = broker or get_broker()
+        self._starting_messages = []
+        self.scenarios = scenarios
+
+        # 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]
+
+        base_case_execution_id = None
+
+        # Process the scenarios one by one
+        for scenario in self.scenarios:
+            # The first scenario is the base case and always runs
+            if self.scenarios.index(scenario) == 0:
+                base_case_execution_id = scenario.get("execution_id", None) or str(
+                    uuid4()
+                )
+                scenario["execution_id"] = base_case_execution_id
+                scenario["base_case_execution_id"] = base_case_execution_id
+                scenario["tasksToRunInScenario"] = list(self.graph.nodes)
+                continue
+            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["base_case_execution_id"] = base_case_execution_id
+            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)
+                ]
+                server_type = task_definitions[first_task].get("serverType", None)
+                lazy_actor = register_lazy_actor(
+                    self.broker,
+                    get_callable(handler),
+                    pipeline_config["metadata"],
+                    server_type,
+                )
+                message = lazy_actor.message()
+                message.options["pipeline"] = pipeline
+                if pipeline_config["metadata"].get("maxRetry", None) is not None:
+                    message.options["max_retries"] = pipeline_config["metadata"][
+                        "maxRetry"
+                    ]
+                message.options["task_replacements"] = copy(
+                    scenario["taskReplacements"]
+                )
+                message.options["execution_id"] = scenario["execution_id"]
+                message.options["task_name"] = first_task
+                message.options["base_case_execution_id"] = base_case_execution_id
+                if scenario["settings"]:
+                    message.kwargs["settings"] = copy(scenario["settings"])
+                    message.kwargs["settings"]["execution_id"] = scenario[
+                        "execution_id"
+                    ]
+                    message.kwargs["settings"][
+                        "base_case_execution_id"
+                    ] = base_case_execution_id
+                self._starting_messages.append(message)
+
+        for m in self._starting_messages:
+            m.options["scenarios"] = self.scenarios
+
+        # Run the first task of the first scenario no matter what
+        first_scenario_task_replacements = scenarios[0]["taskReplacements"]
+        first_scenario_settings = scenarios[0].get("settings", None)
+
+        handler = task_definitions[first_task]["handlers"][
+            first_scenario_task_replacements.get(first_task, 0)
+        ]
+        server_type = task_definitions[first_task].get("serverType", None)
+        lazy_actor = register_lazy_actor(
+            self.broker,
+            get_callable(handler),
+            pipeline_config["metadata"],
+            server_type,
+        )
+        message = lazy_actor.message()
+        message.options["pipeline"] = pipeline
+        if pipeline_config["metadata"].get("maxRetry", None) is not None:
+            message.options["max_retries"] = pipeline_config["metadata"]["maxRetry"]
+        message.options["task_replacements"] = first_scenario_task_replacements
+        message.options["execution_id"] = base_case_execution_id
+        message.options["task_name"] = first_task
+        message.options["scenarios"] = self.scenarios
+        message.options["base_case_execution_id"] = base_case_execution_id
+
+        if first_scenario_settings:
+            message.kwargs["settings"] = copy(first_scenario_settings)
+            message.kwargs["settings"]["execution_id"] = base_case_execution_id
+            message.kwargs["settings"][
+                "base_case_execution_id"
+            ] = base_case_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"{message.options['base_case_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):
+        locks = []
+
+        for scenario in self.scenarios:
+            locks.append(
+                {
+                    "scenario_task_keys": [
+                        f"{scenario['execution_id']}-{task}"
+                        for task in scenario["tasksToRunInScenario"]
+                    ],
+                    "redis_lock_key": f"{scenario['base_case_execution_id']}-lock",
+                }
+            )
+
+        for lock in locks:
+            for task_key in lock["scenario_task_keys"]:
+                locking_parallel_barrier = LockingParallelBarrier(
+                    REDIS_URL, task_key=task_key, lock_key=lock["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,
+                "scenarios": self.scenarios,
+            }
+        )
+
+    @classmethod
+    def from_json(cls, json_data: str) -> typing.Type["Pypeline"]:
+        data = json.loads(json_data)
+
+        return cls(
+            data["pipeline"],
+            scenarios=data["scenarios"],
+        )

pypeline/pipelines/factory.py

@@ -0,0 +1,86 @@
+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, 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:
+        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"])
+
+        p = Pypeline(pipeline, scenarios=scenarios, 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"]
+            server_type = task_definitions[task].get("serverType", None)
+            tmp_handler = get_callable(module_path)
+            lazy_actor = register_lazy_actor(
+                broker, tmp_handler, pipeline_config["metadata"], server_type
+            )
+            registered_actors[task] = lazy_actor
+            if args and not kwargs:
+                msg = registered_actors[task].message(*args)
+            elif kwargs and not args:
+                msg = registered_actors[task].message(**kwargs)
+            elif args and kwargs:
+                msg = registered_actors[task].message(*args, **kwargs)
+            else:
+                msg = registered_actors[task].message()
+            msg.options["task_ttl"] = pipeline_config["metadata"]["maxTtl"]
+            message_group.append(msg)
+
+        messages.append(message_group)
+    p = parallel_pipeline(messages)
+
+    return p