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

pypeline/pypeline_yaml.py

@@ -0,0 +1,468 @@
+""" Definition of the `pypeline.yaml` file.
+
+If using, a basic file may look like::
+    serviceConfig:
+        - name: pypeline-worker
+            registeredTasks:
+                - handler: pypeline_demo_client.workers.demo_worker.demo_worker_task
+                - handler: pypeline_demo_client.workers.demo_worker.demo_model_task
+
+    pipelines:
+        demo-pipeline:
+            name: demo-pipeline
+            description: Demo Pipeline.
+            schemaVersion: 1
+            config:
+                dagAdjacency:
+                    node_a:
+                        - node_b
+                        - node_c
+                metadata:
+                    maxRetry: 3
+                    maxTtl: 60
+                    queue: default-task-queue
+                taskDefinitions:
+                    node_a:
+                        handler: pypeline_demo_client.workers.demo_pipeline.demo_pipeline_node_a
+                    node_b:
+                        handler: pypeline_demo_client.workers.demo_pipeline.demo_pipeline_node_b
+                        queue: node-b-queue
+                    node_c:
+                        handler: pypeline_demo_client.workers.demo_pipeline.demo_pipeline_node_c
+
+    scheduledTasks:
+        demo-model-task:
+            name: Demo Model Task
+            enabled: true
+            config:
+                task: pypeline_demo_client.workers.demo_worker.demo_model_task
+                queue: default-task-queue
+                schedule:
+                    minute: '*'
+                    hour: '*'
+                    dayOfWeek: '*'
+                    dayOfMonth: '*'
+                    monthOfYear: '*'
+            schemaVersion: 1
+
+"""
+import re
+import os
+import logging
+import pkg_resources
+import yaml
+from yaml.loader import SafeLoader
+from marshmallow import Schema, fields, pre_load, EXCLUDE, INCLUDE, validates_schema
+from marshmallow.exceptions import ValidationError
+
+from pypeline.executable_job_config_schema import ExecutableJobSchema
+from pypeline.utils.module_utils import PypelineModuleLoader, normalized_pkg_name
+from pypeline.constants import PYPELINE_YAML_PATH, PYPELINE_CLIENT_PKG_NAME
+from pypeline.pipeline_config_schema import BasePipelineSchema
+from pypeline.schedule_config_schema import BaseScheduleSchema
+
+logger = logging.getLogger(__name__)
+
+
+class InvalidPackagePath(Exception):
+    pass
+
+
+class InvalidPypelineConfig(Exception):
+    pass
+
+
+class MissingPypelineConfig(Exception):
+    pass
+
+
+class ExcludeUnknownSchema(Schema):
+    class Meta:
+        unknown = EXCLUDE
+
+
+class NameSchema(Schema):
+    """Validated name string field."""
+
+    name = fields.String(
+        required=True,
+        description="Name for service or image. Must include "
+        "only alphanumeric characters along with `_` and `-`.",
+        example="my-service-name",
+    )
+
+    @pre_load
+    def validate_characters(self, item, **kwargs):
+        """Ensure name field conforms to allowed characters"""
+        valid_chars = r"^[\w\d\-\_]+$"
+        if not bool(re.match(valid_chars, item["name"])):
+            raise ValueError(
+                f"Invalid name: {item['name']}. Only alphanumeric characters "
+                "allowed along with `-` and `_`."
+            )
+        return item
+
+
+class PypelineRegisteredTaskDetailConfigSchema(Schema):
+    handler = fields.String(
+        required=True,
+        description="Full path to the Method handles work / pipeline tasks.",
+        example="pypeline_customer_client.workers.worker_group.useful_worker",
+    )
+
+    event = fields.Raw(
+        required=False,
+        unknown=INCLUDE,
+        description="Arbitrary user data, passed through `event` arg in task.",
+    )
+
+
+class PypelineCeleryWorkerConfigSchema(Schema):
+    """Attributes for a celery worker.  This worker will run all of the
+    pipelines and scheduled tasks.
+    """
+
+    registeredTasks = fields.List(
+        fields.Nested(PypelineRegisteredTaskDetailConfigSchema, required=True),
+        required=False,
+        _required=True,
+        description="List of task handlers to register for to your Pypeline app.",
+    )
+
+
+class PypelineServiceConfigSchema(
+    ExcludeUnknownSchema, PypelineCeleryWorkerConfigSchema, NameSchema
+):
+    """Base service config object definition for workers."""
+
+    pass
+
+
+class PypelineYamlSchema(ExcludeUnknownSchema):
+    """The primary `pypeline.yaml` file schema. This defines all available
+    properties in a valid Pypeline configuration file.
+    """
+
+    serviceConfig = fields.List(
+        fields.Nested(
+            PypelineServiceConfigSchema,
+            required=True,
+            description="Core service configuration.",
+        ),
+        description="List of workers for Pypeline to manage.",
+        required=True,
+    )
+
+    pipelines = fields.Dict(
+        keys=fields.String(),
+        values=fields.Nested(BasePipelineSchema),
+        description="List of pipelines",
+        required=False,
+    )
+
+    scheduledTasks = fields.Dict(
+        keys=fields.String(),
+        values=fields.Nested(BaseScheduleSchema),
+        description="List of scheduled tasks",
+        required=False,
+    )
+
+    executableJobs = fields.Dict(
+        keys=fields.String(),
+        values=fields.Nested(ExecutableJobSchema),
+        description="List of executable jobs",
+        required=False,
+        allow_none=True,
+    )
+
+    def validate_errors(self, schema: Schema, value: dict):
+        """Run Marshmallow validate() and raise if any errors"""
+        schema = schema()
+        errors = schema.validate(value)
+        if len(errors.keys()) > 0:
+            raise ValidationError(errors)
+
+    @validates_schema
+    def validate_schema(self, data, **kwargs):
+        """Additional validation.
+
+        Nested fields that are not required are not validated by Marshmallow
+        by default. Do a single level down of validation for now.
+
+        imageConfig can provide *either* an install command for Pypeline
+        to use to build the image for customer *or* a Docker repository
+        for Pypeline to pull.
+        """
+        # Vaidate nested
+        key_schema_pairs = (("serviceConfig", PypelineServiceConfigSchema),)
+        for k_s in key_schema_pairs:
+            val = data.get(k_s[0], None)
+            if val is not None:
+                if type(val) == list:
+                    for v in val:
+                        self.validate_errors(k_s[1], v)
+                else:
+                    self.validate_errors(k_s[1], val)
+
+        # Validate the services. We list every service schema field as not
+        # required in order to use them as mixins for a generic service object,
+        # however, they ARE required, so validate here using the custom
+        # metadata property `_required`. Default to value of `required`.
+        for service in data.get("serviceConfig"):
+            schema = PypelineCeleryWorkerConfigSchema
+            for field in schema().fields:
+                try:
+                    if (
+                        schema()
+                        .fields[field]
+                        .metadata.get(
+                            "_required", getattr(schema().fields[field], "required")
+                        )
+                    ):
+                        assert field in service
+                except AssertionError:
+                    raise ValidationError(f"`{field}` missing in worker definition.")
+
+        # Validate unique pipeline ids
+        if "pipelines" in data:
+            pipeline_ids = set()
+            for pipeline_id, pipeline_data in data["pipelines"].items():
+                if pipeline_id in pipeline_ids:
+                    raise ValidationError("All pipeline ids must be unique!")
+                pipeline_ids.add(pipeline_id)
+                schema_version = pipeline_data["schemaVersion"]
+                PipelineSchema = BasePipelineSchema.get_by_version(schema_version)
+                self.validate_errors(PipelineSchema, pipeline_data)
+
+        # Validate unique scheduled tasks names
+        if "scheduledTasks" in data:
+            task_ids = set()
+            for task_id, task_data in data["scheduledTasks"].items():
+                if task_id in task_ids:
+                    raise ValidationError("All schedule ids must be unique!")
+                task_ids.add(task_id)
+                schema_version = task_data["schemaVersion"]
+                TaskSchema = BaseScheduleSchema.get_by_version(schema_version)
+                self.validate_errors(TaskSchema, task_data)
+
+
+class YamlPatternConstructor:
+    """Adds a pattern resolver + constructor to PyYaml.
+
+    Typical/deault usage is for parsing environment variables
+    in a yaml file but this can be used for any pattern you provide.
+
+    See: https://pyyaml.org/wiki/PyYAMLDocumentation
+    """
+
+    def __init__(self, env_var_pattern: str = None, add_constructor: bool = True):
+        self.env_var_pattern = env_var_pattern
+        if self.env_var_pattern is None:
+            # Default pattern is: ${VAR:default}
+            self.env_var_pattern = r"^\$\{(.*)\}$"
+        self.path_matcher = re.compile(self.env_var_pattern)
+
+        if add_constructor:
+            self.add_constructor()
+
+    def _path_constructor(self, loader, node):
+        """Extract the matched value, expand env variable,
+        and replace the match
+
+        TODO: Would need to update this (specifically the parsing) if any
+        pattern other than our default (or a highly compatible variation)
+        is provided.
+        """
+        # Try to match the correct env variable pattern in this node's value
+        # If the value does not match the pattern, return None (which means
+        # this node will not be parsed for ENV variables and instead just
+        # returned as-is).
+        env_var_name = re.match(self.env_var_pattern, node.value)
+        try:
+            env_var_name = env_var_name.group(1)
+        except AttributeError:
+            return None
+
+        # If we get down here, then the 'node.value' matches our specified
+        # pattern, so try to parse. env_var_name is the value inside ${...}.
+        # Split on `:`, which is our delimiter for default values.
+        env_var_name_split = env_var_name.split(":")
+
+        # Attempt to retrieve the environment variable...from the environment
+        env_var = os.environ.get(env_var_name_split[0], None)
+
+        if env_var is None:  # Nothing found in environment
+            # If a default was provided (e.g. VAR:default), return that.
+            # We join anything after first element because the default
+            # value might be a URL or something with a colon in it
+            # which would have 'split' above
+            if len(env_var_name_split) > 1:
+                return ":".join(env_var_name_split[1:])
+            return "unset"  # Return 'unset' if not in environ nor default
+        return env_var
+
+    def add_constructor(self):
+        """Initialize PyYaml with ability to resolve/load environment
+        variables defined in a yaml template when they exist in
+        the environment.
+
+        Add to SafeLoader in addition to standard Loader.
+        """
+        # Add the `!env_var` tag to any scalar (value) that matches the
+        # pattern self.path_matcher. This allows the template to be much more
+        # intuitive vs needing to add !env_var to the beginning of each value
+        yaml.add_implicit_resolver("!env_var", self.path_matcher)
+        yaml.add_implicit_resolver("!env_var", self.path_matcher, Loader=SafeLoader)
+
+        # Add constructor for the tag `!env_var`, which is a function that
+        # converts a node of a YAML representation graph to a native Python
+        # object.
+        yaml.add_constructor("!env_var", self._path_constructor)
+        yaml.add_constructor("!env_var", self._path_constructor, Loader=SafeLoader)
+
+
+def parse_config_file(pypeline_yaml: str):
+    """Parse the `pypeline.yaml` file when it's been loaded.
+
+    Arguments:
+        pypeline_yaml (required): String of loaded pypeline.yaml file.
+    """
+    YamlPatternConstructor()  # Add our env variable parser
+    try:
+        pypeline_yaml_schema = PypelineYamlSchema()
+        # First suss out yaml issues
+        pypeline_config = yaml.safe_load(pypeline_yaml)
+        # Then schema issues
+        pypeline_config = pypeline_yaml_schema.load(pypeline_config)
+    except ValidationError as e:
+        msg = "Invalid Pypeline configuration due to {}".format(e.messages)
+        logger.error(msg)
+        raise InvalidPypelineConfig(msg)
+    except Exception as e:
+        msg = (
+            "Invalid Pypeline configuration, likely due to invalid "
+            "YAML formatting ..."
+        )
+        logger.exception("{} {}".format(msg, e))
+        raise InvalidPypelineConfig(msg)
+    return pypeline_config
+
+
+def _get_pkg_name(pkg_name: str) -> str:
+    """Retrieve the normalized package name."""
+    if pkg_name is None:
+        pkg_name = PYPELINE_CLIENT_PKG_NAME  # From environment
+        if pkg_name is None:
+            return None
+    return normalized_pkg_name(pkg_name)
+
+
+def load_pypeline_config(
+    pkg_name: str = None, pypeline_yaml_filename: str = None, as_dict: bool = True
+):
+    """Load and parse the `pypeline.yaml` file. Issue usable exceptions for
+    known error modes so bootstrapping can handle appropriately.
+
+    Arguments:
+        pkg_name (required): Directory name for your Python
+            package. e.g. my_package_name . If none provided, will check
+            environment for `PYPELINE_CLIENT_PKG_NAME`. If not found,
+            will exit.
+        pypeline_yaml_filename (optional): Relative path to find your
+            `pypeline.yaml` configuration file. Defaults to `pypeline.yaml`
+            which should be found inside your `pkg_name`
+        as_dict (optional): If true (default), return the loaded pypeline
+            configuration as a dictionary. If false, return the loaded
+            string value of the yaml file.
+    """
+    if pypeline_yaml_filename is None:
+        pypeline_yaml_filename = PYPELINE_YAML_PATH
+
+    logger.info(
+        f"Loading `pypeline.yaml` from package `{pkg_name}` "
+        f"and file location `{pypeline_yaml_filename}` ..."
+    )
+    pypeline_config = None
+
+    pkg_name = _get_pkg_name(pkg_name)
+
+    if pkg_name is None:  # Nothing to retrieve at this point
+        logger.warning("Unable to retrieve pypeline.yaml configuration ...")
+        return pypeline_config
+
+    try:
+        pypeline_config_path = pkg_resources.resource_filename(
+            pkg_name, pypeline_yaml_filename
+        )
+    except Exception as e:
+        msg = (
+            "Either pkg_name ({}) or pypeline_yaml_filename ({}) is "
+            "invalid ...".format(pkg_name, pypeline_yaml_filename)
+        )
+        logger.error("{} ... {}".format(msg, e))
+        raise InvalidPackagePath(e)
+
+    try:
+        with open(pypeline_config_path, "r") as f:
+            pypeline_yaml = f.read()
+            pypeline_config = parse_config_file(pypeline_yaml)
+    except InvalidPypelineConfig as e:
+        raise
+    except FileNotFoundError as e:
+        msg = "Pypeline config file could not be found at path {} ...".format(
+            pypeline_config_path
+        )
+        raise MissingPypelineConfig(msg)
+    except Exception as e:
+        raise e
+    if as_dict:
+        return pypeline_config
+    return yaml.safe_dump(pypeline_config)
+
+
+def load_client_config_and_version(
+    pkg_name: str = None, pypeline_yaml_filename: str = None
+):
+    """Load and parse the `pypeline.yaml` file and a client package's version.
+
+        Arguments:
+            pkg_name (required): Directory name for your Python
+                package. e.g. my_package_name . If none provided, will check
+                environment for `PYPELINE_CLIENT_PKG_NAME`. If not found,
+                will exit.
+            pypeline_yaml_filename (optional): Relative path to find your
+                `pypeline.yaml` configuration file. Defaults to `pypeline.yaml`
+                which should be found inside your `pkg_name`
+            as_dict (optional): If true (default), return the loaded pypeline
+                configuration as a dictionary. If false, return the loaded
+                string value of the yaml file.
+
+    For this to work properly, the provided package must be installed in the
+    same environment as this Pypeline package and it must have a `__version__`
+    variable inside its `__init__.py` file, e.g. `__version__ = '0.0.0'`
+    """
+    pypeline_config = None
+    client_version = None
+
+    pkg_name = _get_pkg_name(pkg_name)
+
+    try:
+        loader = PypelineModuleLoader()
+        pkg = loader.get_module(pkg_name + ".__init__")
+        client_version = getattr(pkg, "__version__", "0.0.0")
+        pypeline_config = load_pypeline_config(pkg_name, pypeline_yaml_filename)
+    except MissingPypelineConfig as e:
+        logger.error(e)
+    except InvalidPypelineConfig as e:
+        logger.error(e)
+    except InvalidPackagePath as e:
+        logger.error(e)
+    except Exception as e:
+        logger.error(
+            "Unable to load client's pkg __version__ or "
+            "{} config file for package: {} ... {}".format(
+                pypeline_yaml_filename, pkg_name, e
+            )
+        )
+
+    return pypeline_config, client_version

pypeline/schedule_config_schema.py

@@ -0,0 +1,125 @@
+""" Schemas for Schedule Configuration
+"""
+import re
+
+# from celery.schedules import crontab_parser
+from croniter import croniter
+from marshmallow.exceptions import ValidationError
+from marshmallow import Schema, fields, EXCLUDE, pre_load, validates_schema
+
+
+class ExcludeUnknownSchema(Schema):
+    """Remove unknown keys from loaded dictionary"""
+
+    class Meta:
+        unknown = EXCLUDE
+
+
+class CrontabScheduleSchema(Schema):
+    minute = fields.String(required=True)
+    hour = fields.String(required=True)
+    dayOfWeek = fields.String(required=True)
+    dayOfMonth = fields.String(required=True)
+    monthOfYear = fields.String(required=True)
+
+    @validates_schema
+    def validate_values(self, data, **kwargs):
+        if (
+            data["minute"] is None
+            or data["hour"] is None
+            or data["dayOfWeek"] is None
+            or data["dayOfMonth"] is None
+            or data["monthOfYear"] is None
+        ):
+            raise ValidationError("Empty crontab value")
+
+        test_cron_expression = (
+            f"{data['minute']} {data['hour']} {data['dayOfMonth']} "
+            f"{data['monthOfYear']} {data['dayOfWeek']}"
+        )
+
+        if not croniter.is_valid(test_cron_expression):
+            return ValidationError("Invalid crontab value")
+
+
+class Schedule(fields.Dict):
+    def _serialize(self, value, attr, obj, **kwargs):
+        return value
+
+    def _deserialize(self, value, attr, data, **kwargs):
+        schema = CrontabScheduleSchema()
+        return schema.load(value)
+
+
+class ScheduleConfigSchemaV1(ExcludeUnknownSchema):
+    """Definition of a single schedule entry"""
+
+    queue = fields.String(
+        required=True,
+        description="Name of queue on which to place task.",
+        example="my-default-queue",
+    )
+    task = fields.String(
+        required=True,
+        description="Path to task to invoke.",
+        example="my_app.module.method",
+    )
+
+    schedule = Schedule(required=True)
+
+
+class BaseScheduleSchema(ExcludeUnknownSchema):
+    __schema_version__ = 0
+
+    name = fields.String(
+        required=True,
+        description="Name of schedule entry.",
+        example="My Scheduled Task",
+    )
+    schemaVersion = fields.Integer(required=True)
+    config = fields.Dict(required=True)
+    enabled = fields.Boolean(
+        required=True, description="Whether entry is enabled.", example=True
+    )
+
+    @classmethod
+    def get_by_version(cls, version):
+        for subclass in cls.__subclasses__():
+            if subclass.__schema_version__ == version:
+                return subclass
+
+        return None
+
+    @classmethod
+    def get_latest(cls):
+        max_version = 0
+        max_class = None
+        for subclass in cls.__subclasses__():
+            if subclass.__schema_version__ > max_version:
+                max_version = max_version
+                max_class = subclass
+
+        return max_class
+
+    @validates_schema
+    def validate_scheduled_tasks(self, data, **kwargs):
+        schema_version = data["schemaVersion"]
+        TaskSchema = BaseScheduleSchema.get_by_version(schema_version)
+        schema = TaskSchema()
+        schema.load(data)
+
+
+class ScheduleSchemaV1(BaseScheduleSchema):
+    __schema_version__ = 1
+
+    config = fields.Nested(
+        ScheduleConfigSchemaV1,
+        required=True,
+        description="Configuration information for this schedule.",
+    )
+
+    def validate_scheduled_tasks(self, data, **kwargs):
+        # We need to add this function to avoid infinite recursion since
+        # the BaseScheduleSchema class above uses the same method for
+        # validation
+        pass

pypeline/utils/config_utils.py

@@ -0,0 +1,81 @@
+import logging
+from typing import Union
+
+from pypeline.constants import WORKER_NAME
+from pypeline.pypeline_yaml import load_pypeline_config
+
+logger = logging.getLogger(__name__)
+
+
+def retrieve_latest_pipeline_config(
+    pipeline_id: Union[str, None] = None
+) -> Union[dict, list]:
+    """Retrieve the 'latest' pipeline configuration for a given pipeline."""
+    pypeline_config = load_pypeline_config()
+    if "pipelines" in pypeline_config:
+        pipelines = []
+        found_pipeline = None
+        for p_id, config in pypeline_config["pipelines"].items():
+            if pipeline_id == p_id:
+                found_pipeline = config
+                break
+            pipelines.append(config)
+
+        if pipeline_id:
+            if found_pipeline:
+                return found_pipeline
+            raise ValueError(f"Invalid pipeline {pipeline_id}")
+
+        return pipelines
+    return None
+
+
+def retrieve_latest_schedule_config():
+    """Retrieve the 'latest' scheduled tasks configuration."""
+    pypeline_config = load_pypeline_config()
+    if "scheduledTasks" in pypeline_config:
+        tasks = []
+        for task_id, config in pypeline_config["scheduledTasks"].items():
+            tasks.append(config)
+        return tasks
+    return None
+
+
+def retrieve_executable_job_config():
+    pypeline_config = load_pypeline_config()
+
+    if not pypeline_config:
+        return None
+    if "executableJobs" in pypeline_config:
+        tasks = []
+        for task_id, config in pypeline_config["executableJobs"].items():
+            tasks.append(config)
+        return tasks
+    return None
+
+
+def get_service_config_for_worker(
+    pypeline_config: dict, worker_name: str = None
+) -> Union[dict, None]:
+    """For the current WORKER_NAME (which must be present in the environment
+    of this worker instance for a valid deployment), return the worker's
+    serviceConfig object.
+    """
+    if pypeline_config is None:
+        raise ValueError("Pypeline config was not provided")
+    if worker_name is None:
+        worker_name = WORKER_NAME
+    if worker_name is None:
+        return None
+
+    service_config = pypeline_config.get("serviceConfig", [])
+    for service in service_config:
+        if service["name"] == worker_name:
+            return service
+
+    raise ValueError(
+        "Could not find a service config for worker "
+        f"`{worker_name}`. Make sure you have added the service in"
+        f" your pypeline.yaml with `name: {worker_name}` and "
+        "`type: celery-worker`."
+    )