scalable-pypeline 2.1.9__tar.gz → 2.1.11__tar.gz

{scalable-pypeline-2.1.9/scalable_pypeline.egg-info → scalable-pypeline-2.1.11}/PKG-INFO

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

scalable-pypeline-2.1.11/pypeline/__init__.py

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

{scalable-pypeline-2.1.9 → scalable-pypeline-2.1.11}/pypeline/dramatiq.py

@@ -1,3 +1,4 @@
+import copy
 import typing
 import pika
 import logging
@@ -101,45 +102,50 @@ def register_actors_for_workers(broker: Broker):
 
     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}")
 
 
 class Dramatiq:

{scalable-pypeline-2.1.9 → scalable-pypeline-2.1.11}/pypeline/pipeline_config_schema.py

@@ -4,6 +4,7 @@
 import yaml
 from marshmallow import Schema, fields, EXCLUDE, validates_schema
 from marshmallow.exceptions import ValidationError
+from marshmallow.validate import OneOf
 
 from pypeline.pipeline_settings_schema import PipelineSettingsSchema
 
@@ -116,6 +117,14 @@ class TaskDefinitionsSchemaV1(ExcludeUnknownSchema):
         example="custom-queue-name",
     )
 
+    serverType = fields.String(
+        required=False,
+        validate=OneOf(
+            ["xs", "s", "m", "l", "xl", "xxl", "xxxl"],
+            error="Invalid serverType. Available options: 'xs', 's', 'm', 'l', 'xl', 'xxl', 'xxxl'.",
+        ),
+    )
+
 
 class TaskDefinitionsSchemaV2(ExcludeUnknownSchema):
     """Schema for a single task's configuration"""
@@ -140,6 +149,14 @@ class TaskDefinitionsSchemaV2(ExcludeUnknownSchema):
         example="custom-queue-name",
     )
 
+    serverType = fields.String(
+        required=False,
+        validate=OneOf(
+            ["xs", "s", "m", "l", "xl", "xxl", "xxxl"],
+            error="Invalid serverType. Available options: 'xs', 's', 'm', 'l', 'xl', 'xxl', 'xxxl'.",
+        ),
+    )
+
 
 class PipelineConfigSchemaBase(Schema):
     """Overall pipeline configuration schema"""

{scalable-pypeline-2.1.9 → scalable-pypeline-2.1.11}/pypeline/pipelines/composition/pypeline_composition.py

@@ -19,7 +19,6 @@ class Pypeline:
         pipeline: dict,
         scenarios: dict = {},
         broker=None,
-        execution_id=None,
     ):
         # Construct initial properties
         self.pipeline = pipeline
@@ -78,10 +77,12 @@ class Pypeline:
                 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
@@ -111,10 +112,12 @@ class Pypeline:
         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

{scalable-pypeline-2.1.9 → scalable-pypeline-2.1.11}/pypeline/pipelines/factory.py

@@ -63,9 +63,10 @@ def dag_generator(
         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"]
+                broker, tmp_handler, pipeline_config["metadata"], server_type
             )
             registered_actors[task] = lazy_actor
             if args and not kwargs:

{scalable-pypeline-2.1.9 → scalable-pypeline-2.1.11}/pypeline/pipelines/middleware/pypeline_middleware.py

@@ -31,16 +31,14 @@ class PypelineMiddleware(Middleware):
         task_name = message.options["task_name"]
         task_key = f"{execution_id}-{task_name}"
 
-        # Signal to other jobs that current job is finished
+        # Signal to other jobs that current task is finished
         locking_parallel_barrier = LockingParallelBarrier(
             self.redis_url,
             task_key=task_key,
             lock_key=f"{message.options['base_case_execution_id']}-lock",
         )
         try:
-            locking_parallel_barrier.acquire_lock(
-                timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
-            )
+            locking_parallel_barrier.acquire_lock(timeout=10)
             _ = locking_parallel_barrier.decrement_task_count()
         finally:
             locking_parallel_barrier.release_lock()
@@ -50,16 +48,16 @@ class PypelineMiddleware(Middleware):
 
         messages = []
         for child in children_tasks:
-            child_ancestors = sorted(nx.ancestors(graph, child))
+            child_ancestors = sorted(graph.predecessors(child))
 
             ancestor_tasks_complete = True
 
             for ancestor in child_ancestors:
-                task_key = f"{execution_id}-{ancestor}"
+                ancestor_task_key = f"{execution_id}-{ancestor}"
 
                 locking_parallel_barrier = LockingParallelBarrier(
                     self.redis_url,
-                    task_key=task_key,
+                    task_key=ancestor_task_key,
                     lock_key=f"{message.options['base_case_execution_id']}-lock",
                 )
                 try:
@@ -74,12 +72,17 @@ class PypelineMiddleware(Middleware):
                 finally:
                     locking_parallel_barrier.release_lock()
 
-                if not remaining_tasks:
-                    task_key = f"{message.options['base_case_execution_id']}-{ancestor}"
+                # If the lock didn't exist for the current tasks execution id then it would indicate
+                # that this is the start of a new scenario.  Therefore we need to find the ancestor
+                # that is executed in the base case execution id and make sure it has completed
+                if remaining_tasks is None:
+                    ancestor_task_key = (
+                        f"{message.options['base_case_execution_id']}-{ancestor}"
+                    )
 
                     locking_parallel_barrier = LockingParallelBarrier(
                         self.redis_url,
-                        task_key=task_key,
+                        task_key=ancestor_task_key,
                         lock_key=f"{message.options['base_case_execution_id']}-lock",
                     )
                     try:
@@ -95,7 +98,7 @@ class PypelineMiddleware(Middleware):
                             )
                     finally:
                         locking_parallel_barrier.release_lock()
-                if remaining_tasks >= 1:
+                if remaining_tasks is None or remaining_tasks >= 1:
                     ancestor_tasks_complete = False
                     break
 
@@ -103,6 +106,7 @@ class PypelineMiddleware(Middleware):
             if not ancestor_tasks_complete:
                 break
 
+            # Handle situation where base case kicks off new scenario
             if (
                 message.options["base_case_execution_id"]
                 == message.options["execution_id"]
@@ -124,11 +128,13 @@ class PypelineMiddleware(Middleware):
                         handler = task_definitions[child]["handlers"][
                             task_replacements.get(child, 0)
                         ]
+                        server_type = task_definitions[child].get("serverType", None)
 
                         lazy_actor = register_lazy_actor(
                             broker,
                             get_callable(handler),
                             pipeline_config["metadata"],
+                            server_type,
                         )
                         scenario_message = lazy_actor.message()
                         scenario_message.options["pipeline"] = pipeline
@@ -154,6 +160,8 @@ class PypelineMiddleware(Middleware):
                                 scenario["execution_id"]
                             )
                         messages.append(scenario_message)
+
+            # Kick off child task for current scenario
             task_key = f"{execution_id}-{child}"
             locking_parallel_barrier = LockingParallelBarrier(
                 self.redis_url,
@@ -164,10 +172,12 @@ class PypelineMiddleware(Middleware):
             handler = task_definitions[child]["handlers"][
                 task_replacements.get(child, 0)
             ]
+            server_type = task_definitions[child].get("serverType", None)
             lazy_actor = register_lazy_actor(
                 broker,
                 get_callable(handler),
                 pipeline_config["metadata"],
+                server_type,
             )
 
             child_message = lazy_actor.message()

{scalable-pypeline-2.1.9 → scalable-pypeline-2.1.11}/pypeline/utils/dramatiq_utils.py

@@ -45,9 +45,13 @@ def register_lazy_actor(
     broker: Broker,
     fn: Optional[Callable[P, Union[Awaitable[R], R]]] = None,
     pipeline_meta: typing.Dict = {},
+    server_type: Optional[str] = None,
     **kwargs,
 ) -> typing.Type["LazyActor"]:
-    kwargs["queue_name"] = pipeline_meta.get("queue", "default")
+    if server_type:
+        kwargs["queue_name"] = server_type + "-" + pipeline_meta.get("queue", "default")
+    else:
+        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"] = (
@@ -110,7 +114,7 @@ class LazyActor(object):
 
     def register(self, broker):
         self.actor = register_actor(
-            actor_name=f"{self.fn.__module__}.{self.fn.__name__}",
+            actor_name=f"{self.fn.__module__}.{self.fn.__name__}-{self.kw['queue_name']}",
             broker=broker,
             **self.kw,
         )(ensure_return_value(default_value=True)(self.fn))

{scalable-pypeline-2.1.9 → scalable-pypeline-2.1.11/scalable_pypeline.egg-info}/PKG-INFO

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

scalable-pypeline-2.1.9/pypeline/__init__.py

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