scalable-pypeline 2.1.9__py2.py3-none-any.whl → 2.1.11__py2.py3-none-any.whl

pypeline/__init__.py

@@ -1 +1 @@
-__version__ = "2.1.9"
+__version__ = "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:

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"""

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

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:

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()

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.dist-info → scalable_pypeline-2.1.11.dist-info}/METADATA

@@ -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.dist-info → scalable_pypeline-2.1.11.dist-info}/RECORD

@@ -1,9 +1,9 @@
-pypeline/__init__.py,sha256=Gf0zUMmdRLkPmvwnoP7oCpWn7Ep93jlNIlBD6nifugI,22
+pypeline/__init__.py,sha256=zcHMV6HbszEZFZ1gDEpbHxTYg8dJ5fCYdeDd-O6IPP4,23
 pypeline/barrier.py,sha256=oO964l9qOCOibweOHyNivmAvufdXOke9nz2tdgclouo,1172
 pypeline/constants.py,sha256=EGSuLq4KhZ4bxrbtnUgKclELRyya5ipvv0WeybCzNAs,3049
-pypeline/dramatiq.py,sha256=5whdOQjzknLjl9lYNj5-f2jRw5ysRPhQExi5NN4KeOY,12837
+pypeline/dramatiq.py,sha256=nnNxm2akL6hjqfjwQglK76LspZrdrL_2E_rDN5ThwZ8,13432
 pypeline/extensions.py,sha256=BzOTnXhNxap3N7uIUUh_hO6dDwx08Vc_RJDE93_K0Lo,610
-pypeline/pipeline_config_schema.py,sha256=hK2_egtg-YFx_XJDs_NyrOTGKkel7W83X-G0sic52sM,10592
+pypeline/pipeline_config_schema.py,sha256=ApgDqgINFDcmErVP7jJUJgkaVfwNRanj4q12ehWVJUU,11143
 pypeline/pipeline_settings_schema.py,sha256=IBBtWqxrRA0MrUwNZEnGhYJbmag-sus1Ln6VCrUvuyY,20313
 pypeline/pypeline_yaml.py,sha256=Og08sUKwOjq7JYPnkg-NIcGbHravYCkC5Arz22rZEtA,16981
 pypeline/schedule_config_schema.py,sha256=vtZV-5wpGcAiYcXxdBPRkrjsbR6x_9E-1PC2elrKKbE,3611
@@ -14,24 +14,24 @@ pypeline/flask/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuF
 pypeline/flask/api/pipelines.py,sha256=lw1ggsjp_Iha5MhyQGHtVW0akpVJnxIk0hn6NkC3c8s,9314
 pypeline/flask/api/schedules.py,sha256=8PKCMdPucaer8opchNlI5aDssK2UqT79hHpeg5BMtTA,1210
 pypeline/pipelines/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pypeline/pipelines/factory.py,sha256=356v1S0WPvDkd9f0fKk0H9aXVbOqQYSWt47aOl66EKk,3172
+pypeline/pipelines/factory.py,sha256=p2b1XhhMYt_f5h2FAr6Zze1FfWz8SCyO9DEIS_EDDKE,3258
 pypeline/pipelines/composition/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pypeline/pipelines/composition/parallel_pipeline_composition.py,sha256=pTw9Xb9h4JnV4siFc3JStm5lB-i9djUADo3Kh5K3s7g,12976
-pypeline/pipelines/composition/pypeline_composition.py,sha256=UBuDKEfRoIbL-9c-HH2ZTVbzfkwFSlNoFH-AVNqt0QE,7965
+pypeline/pipelines/composition/pypeline_composition.py,sha256=G6hptAN3LaVGm6DE7yhjGs-LDbw7m94-z08pLZCqNXA,8154
 pypeline/pipelines/middleware/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pypeline/pipelines/middleware/get_active_worker_id_middleware.py,sha256=X4ZfRk3L8MD00DTsGHth7oOdy-W7LQV96T8vu5UC42A,755
 pypeline/pipelines/middleware/parallel_pipeline_middleware.py,sha256=kTp6niYoe2nXIiN6EGRfdpxrJyioo0GPxDkfefbGlEk,2821
-pypeline/pipelines/middleware/pypeline_middleware.py,sha256=IXVqzcOlSJ43lsn-i298RkaeygB-PTJjsvdTDtpgfwg,8141
+pypeline/pipelines/middleware/pypeline_middleware.py,sha256=wdqcnya9SrcBs2VOfBOIXY03VC0uadMcS5SD4hGOk8U,8843
 pypeline/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pypeline/utils/config_utils.py,sha256=rAIATyoW7kGETZ_Z2DqiXtGd7bJp5uPfcLtfNPOYsNs,2167
-pypeline/utils/dramatiq_utils.py,sha256=5GDcOvKY-8S8r---wb6Q8QAywhbKVJ-qILjcYNHei8Y,3658
+pypeline/utils/dramatiq_utils.py,sha256=DUdgVywm1182A4i69XzH9EIh1EJ9zAHmJLtOaVSW7pw,3844
 pypeline/utils/module_utils.py,sha256=-yEJIukDCoXnmlZVXB6Dww25tH6GdPE5SoFqv6pfdVU,3682
 pypeline/utils/pipeline_utils.py,sha256=kGP1QwCJikGC5QNRtzRXCDVewyRMpWIqERTNnxGLlSY,4795
 pypeline/utils/schema_utils.py,sha256=Fgl0y9Cuo_TZeEx_S3gaSVnLjn6467LTkjb2ek7Ms98,851
 tests/fixtures/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-scalable_pypeline-2.1.9.dist-info/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
-scalable_pypeline-2.1.9.dist-info/METADATA,sha256=H96u1H3mZdGCb1yfmo52azC_BiVq0rzo2nBYTViHElk,5926
-scalable_pypeline-2.1.9.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-scalable_pypeline-2.1.9.dist-info/entry_points.txt,sha256=uWs10ODfHSBKo2Cx_QaUjPHQTpZ3e77j9VlAdRRmMyg,119
-scalable_pypeline-2.1.9.dist-info/top_level.txt,sha256=C7dpkEOc_-nnsAQb28BfQknjD6XHRyS9ZrvVeoIbV7s,15
-scalable_pypeline-2.1.9.dist-info/RECORD,,
+scalable_pypeline-2.1.11.dist-info/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
+scalable_pypeline-2.1.11.dist-info/METADATA,sha256=HdDTGiEP9ELgzK4H1HrIWem1ztzABQ5SmIv-Y3NPpm8,5927
+scalable_pypeline-2.1.11.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+scalable_pypeline-2.1.11.dist-info/entry_points.txt,sha256=uWs10ODfHSBKo2Cx_QaUjPHQTpZ3e77j9VlAdRRmMyg,119
+scalable_pypeline-2.1.11.dist-info/top_level.txt,sha256=C7dpkEOc_-nnsAQb28BfQknjD6XHRyS9ZrvVeoIbV7s,15
+scalable_pypeline-2.1.11.dist-info/RECORD,,