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

pypeline/__init__.py

@@ -0,0 +1 @@
+__version__ = '1.1.0'

pypeline/celery.py

@@ -0,0 +1,270 @@
+""" Configure and instantiate Celery
+"""
+import os
+
+from pypeline.constants import DEFAULT_RETRY_TASK_MAX_TTL, DEFAULT_MAX_RETRY
+
+if os.environ.get('USE_GEVENT', "False").lower() == 'true':
+    from gevent import monkey
+    monkey.patch_all()
+
+import random
+import time
+
+from celery_dyrygent.tasks import register_workflow_processor
+
+import sys
+import logging
+from typing import List
+from celery import Celery
+from pypeline.logging_config import setup_logging
+from pypeline.utils.module_utils import SermosModuleLoader
+from pypeline.utils.task_utils import PipelineGenerator, PipelineResult, \
+    get_service_config_for_worker
+from pypeline.extensions import sermos_config, sermos_client_version
+from pypeline import __version__
+
+logger = logging.getLogger('celery')
+ENABLE_TOOLS = str(os.environ.get('ENABLE_TOOLS', 'false')).lower() == 'true'
+CELERY_TASKS_ACK_LATE = str(os.environ.get('CELERY_TASKS_ACK_LATE', 'false')).lower() == 'true'
+LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO")
+OVERLOAD_ES = os.environ.get('ENV', 'production').lower() == 'production'
+PIPELINE_CHORD_COMPRESSION = os.environ.get('PIPELINE_CHORD_COMPRESSION', None)
+
+setup_logging(app_version=__version__,
+              client_version=sermos_client_version,
+              default_level=LOG_LEVEL,
+              overload_elasticsearch=OVERLOAD_ES,
+              establish_logging_config=True)
+
+
+def pipeline_retry(event: dict):
+    """ Handle pipeline retry and deadletter logic.
+    """
+    access_key = event.get('access_key', None)
+    pipeline_id = event.get('pipeline_id', None)
+    execution_id = event.get('execution_id', None)
+    if pipeline_id is None or execution_id is None:
+        logger.error(f"Unable to retry pipeline {pipeline_id} / "
+                     f"execution {execution_id}.")
+        return False
+
+    # generate_chain() will return `None` if the pipeline has exceeded
+    # max retry count or other erorrs happen.
+    gen = PipelineGenerator(pipeline_id=pipeline_id,
+                            access_key=access_key,
+                            execution_id=execution_id,
+                            queue=event.get('queue', None),
+                            default_task_ttl=event.get('default_task_ttl',
+                                                       None),
+                            add_retry=event.get('add_retry', False),
+                            chain_payload=event.get('chain_payload', None))
+
+    if gen.good_to_go:
+        chain = gen.generate_chain()
+        if chain is not None:
+            # Exponential backoff
+            exponential_backoff = min((3 ** gen.pipeline_wrapper.retry_count) +
+                                      (random.randint(0, 1000) / 1000),
+                                      DEFAULT_RETRY_TASK_MAX_TTL)
+            logger.debug(f"Exponential backoff sleep {exponential_backoff}")
+            time.sleep(exponential_backoff)
+            # Kick it off again.
+            chain.apply_async()
+
+    logger.warning(f"Pipeline retry was invoked for {pipeline_id} "
+                   f"({execution_id})")
+    return True
+
+
+def task_chain_regulator(*args, **kwargs):
+    """ Utility task to ensure celery properly waits between groups in a chain.
+
+        For a chain(), if each element is a group() then celery does not
+        properly adhere to the chain elements occurring sequentially. If you
+        insert a task that is not a group() in between, though, then the
+        chain operates as expected.
+    """
+    return True
+
+
+def pipeline_success(event: dict):
+    """ Utility task to ensure celery properly waits between groups in a chain.
+
+        For a chain(), if each element is a group() then celery does not
+        properly adhere to the chain elements occurring sequentially. If you
+        insert a task that is not a group() in between, though, then the
+        chain operates as expected.
+    """
+    pr = PipelineResult(event['execution_id'])
+    pr.load()
+    pr.save(status='success')
+
+
+class GenerateCeleryTasks(SermosModuleLoader):
+    """ Use the sermos.yaml configuration to turn customer methods into
+        decorated celery tasks that are available for work/pipelines
+    """
+    def __init__(self, config: dict, celery_instance: Celery):
+        super(GenerateCeleryTasks, self).__init__()
+        self.config = config if config else {}
+        self.celery = celery_instance
+
+    def _get_default_tasks(self) -> List[dict]:
+        """ Sermos provides default tasks that all workers should know about.
+        """
+        return [{
+            'handler': 'sermos.celery.pipeline_retry'
+        }, {
+            'handler': 'sermos.celery.task_chain_regulator'
+        }, {
+            'handler': 'sermos.celery.pipeline_success'
+        }]
+
+    def generate(self):
+        """ Loads methods based on sermos config file and decorates them as
+            celery tasks.
+
+            Customer's methods:
+            --------------------------------
+            def demo_task(*args, **kwargs):
+                return True
+
+            Turns into the equivallent of:
+            --------------------------------
+            @celery.task(queue='queue-name')
+            def demo_task(*args, **kwargs):t
+                return True
+        """
+        # Set in k8s deployment as an environment variable when Sermos Cloud
+        # generates the final secrets.yaml file. The name comes from the user's
+        # sermos.yaml file based on serviceConfig[].name. Each 'worker' will
+        # have a single name and each individually registers tasks through its
+        # registeredTasks list. This allows each worker to only attempt
+        # bootstrapping those tasks that are relevant to the worker and not, for
+        # example, attempt to import a package that's not used by this worker
+        service = get_service_config_for_worker(self.config)
+        if not service:
+            return
+        for task in service.get('registeredTasks', []):
+            try:
+                worker_path = task['handler']  # Required, no default
+
+                tmp_handler = self.get_callable(worker_path)
+
+                # Decorate the method as a celery task along with a default
+                # queue if provided in config. Set ChainedTask as the base
+                # which allows chained tasks to pass kwargs correctly.
+                tmp_handler = self.celery.task(tmp_handler)
+            except Exception as e:
+                logger.warning(f"Unable to add a task to celery: {e}")
+        # Sermos provides default tasks that all workers should know about, add
+        # them here.
+        for task in self._get_default_tasks():
+            tmp_handler = self.get_callable(task['handler'])
+            tmp_handler = self.celery.task(tmp_handler)
+
+
+def configure_celery(celery: Celery):
+    """ Configure Sermos-compatible Celery instance. Primarily this means
+    compatibility with Pipelines and Scheduled Tasks through injecting the
+    event kwarg. Also sets prebaked defaults that can be overloaded by user.
+    """
+    REDIS_URL = os.environ.get('REDIS_URL', 'redis://localhost:6379/0')
+    CELERY_BROKER_URL = os.environ.get('CELERY_BROKER_URL', REDIS_URL)
+    CELERY_RESULT_BACKEND = os.environ.get('CELERY_RESULT_BACKEND', REDIS_URL)
+    TaskBase = celery.Task
+
+    class ChainedTask(TaskBase):
+        """ A Celery Task that is used as the _base_ for all dynamically
+        generated tasks (by GenerateCeleryTasks().generate()). This injects
+        `event` into every task's signature, which allows pipelines to pass
+        event information easily through a chain.
+        """
+        abstract = True
+        autoretry_for = (Exception,)
+        max_retries = DEFAULT_MAX_RETRY
+        retry_backoff = True
+        retry_jitter = True
+
+        def __call__(self, *args, **kwargs):
+            """ Allow the return value of one task to update the kwargs of a
+                subsequent task if it's a dictionary. Important to the function
+                of a pipeline to allow event information to flow easily.
+            """
+            # Inject app context
+            if len(args) == 1 and isinstance(args[0], dict):
+                kwargs.update(args[0])
+                args = ()
+
+            # Event holds information used in PipelineRunWrapper and
+            # other areas.
+            if 'event' not in kwargs.keys():
+                kwargs['event'] = {}
+            # This is a special worker from dyrygent that orchestrates our
+            # pipelines.  It provides a patch in fix for celery's poor
+            # implementation of Canvas work-flows
+            if self.__name__ == 'workflow_processor':
+                kwargs.pop('event', None)
+            return super(ChainedTask, self).__call__(*args, **kwargs)
+
+    celery.Task = ChainedTask
+
+    # Configure the broker and tasks
+    celery.conf.broker_url = CELERY_BROKER_URL
+
+    # Use our custom database scheduler for dynamic celery beat updates.
+    celery.conf.beat_scheduler =\
+        'sermos.celery_beat:SermosScheduler'
+
+    # Reasonable defaults, override as necessary
+    celery.conf.worker_redirect_stdouts = True
+    celery.conf.worker_redirect_stdouts_level = LOG_LEVEL
+    celery.conf.worker_hijack_root_logger = False
+
+    if PIPELINE_CHORD_COMPRESSION:
+      celery.conf.task_compression = PIPELINE_CHORD_COMPRESSION
+
+    # NOTE: The broker URL may not be the best result backend. For example,
+    # When using Rabbit as the broker (recommended), you should use Redis
+    # as the result backend, as Rabbit has horrible support as backend.
+    celery.conf.result_backend = CELERY_RESULT_BACKEND
+    celery.conf.task_ignore_result = False  # Must not ignore for Chords
+    celery.conf.task_acks_late = False  # Check per worker
+    celery.conf.result_expires = int(
+        os.environ.get('CELERY_RESULT_EXPIRES', 10800))  # 3 hours by default
+    celery.conf.broker_pool_limit = int(os.environ.get('BROKER_POOL_LIMIT',
+                                                       10))
+    celery.conf.worker_max_tasks_per_child = int(
+        os.environ.get('MAX_TASKS_PER_CHILD', 100))
+    celery.conf.task_soft_time_limit =\
+        int(os.environ.get('TASK_TIMEOUT_SECONDS', 3600))
+    celery.conf.task_time_limit =\
+        int(os.environ.get('TASK_TIMEOUT_SECONDS', 3600)) + 10  # Cleanup buffer
+    celery.conf.task_acks_late = CELERY_TASKS_ACK_LATE
+    celery.conf.task_serializer = 'json'
+    celery.conf.result_serializer = 'json'
+    celery.conf.accept_content = ['json']
+    # Required config options for some brokers we use frequently.
+    transport_options = {}
+    celery.conf.broker_transport_options = transport_options
+
+    # Sermos generally has long-running tasks (relatively speaking), so
+    # limit number of jobs a worker can reserve. This may not be true for
+    # all tasks, so configure this on a per application basis. In the event
+    # mutltiple task kinds exist in an application (short and long), see
+    # http://docs.celeryproject.org/en/latest/userguide/optimizing.html#optimizing-prefetch-limit
+    # for some guidance on combining multiple workers and routing tasks.
+    # TODO make configurable from env
+    celery.conf.worker_prefetch_multiplier = 1
+
+    # Add our application's workers & any other tasks to be made
+    # available
+    register_workflow_processor(celery)
+    try:
+        GenerateCeleryTasks(sermos_config, celery).generate()
+    except Exception as e:
+        logger.error(f"Unable to dynamically generate celery tasks: {e}")
+        sys.exit(1)
+
+    return celery

pypeline/celery_beat.py

@@ -0,0 +1,254 @@
+""" Custom Sermos Scheduler and Celery Entry classes used for dynamic beat.
+"""
+import datetime
+import os
+import logging
+from rhodb.redis_conf import RedisConnector
+from celery.beat import Scheduler, ScheduleEntry
+from celery import current_app
+from celery.utils.time import is_naive
+from celery.schedules import schedule as c_schedule, crontab as c_crontab
+from pypeline.utils.config_utils import retrieve_latest_schedule_config, \
+    update_schedule_config
+from pypeline.constants import CONFIG_REFRESH_RATE, SCHEDULE_DATE_FORMAT, \
+    USING_SERMOS_CLOUD
+
+logger = logging.getLogger(__name__)
+redis_conn = RedisConnector().get_connection()
+
+
+def convert_to_datetime(
+        datetime_str: str,
+        datetime_format: str = SCHEDULE_DATE_FORMAT) -> datetime.datetime:
+    """ Accept a string in the standard format and return a datetime object
+    """
+    return datetime.datetime.strptime(datetime_str, datetime_format)
+
+
+def instantiate_celery_schedule(schedule_entry: dict) -> c_schedule:
+    """ From a schedule entry and the full schedule from Sermos, create a
+        celery `schedule` object.
+    """
+    scheduleType = schedule_entry['config']['scheduleType']
+
+    if scheduleType == 'interval':
+        # Create a timedelta object
+        period = schedule_entry['config']['schedule']['period']
+        every = schedule_entry['config']['schedule']['every']
+        the_delta = datetime.timedelta(**{period: every})
+        # Instantiate the celery schedule object
+        return c_schedule(run_every=the_delta)
+
+    if scheduleType == 'crontab':
+        return c_crontab(
+            minute=schedule_entry['config']['schedule']['minute'],
+            hour=schedule_entry['config']['schedule']['hour'],
+            day_of_week=schedule_entry['config']['schedule']['dayOfWeek'],
+            day_of_month=schedule_entry['config']['schedule']['dayOfMonth'],
+            month_of_year=schedule_entry['config']['schedule']['monthOfYear'])
+
+    raise ValueError(f"Unsupported scheduleType ({scheduleType} ...")
+
+
+class SermosEntry(ScheduleEntry):
+    """ Create a beat entry with additional functionality for Sermos scheduler.
+
+        https://docs.celeryproject.org/en/latest/userguide/periodic-tasks.html
+    """
+    def __init__(self, schedule_entry: dict = None, **kwargs):
+        schedule_entry = schedule_entry if schedule_entry else {}
+        if schedule_entry:
+            # This event is being instantiated directly with the Sermos
+            # schedule entry
+            celery_schedule = instantiate_celery_schedule(schedule_entry)
+
+            # celery.beat.ScheduleEntry expects these keys in a dictionary
+            # called `options`. See
+            # https://docs.celeryproject.org/en/stable/userguide/calling.html
+            # In the case of Sermos, we require the queue in the
+            # ScheduleEntrySchema, others are all optional.
+            options = dict()
+            optional_keys = ('queue', 'exchange', 'routing_key', 'expires')
+            for key in optional_keys:
+                value = schedule_entry['config'].get(key, None)
+                if value is not None:
+                    options[key] = value
+
+            last_run_at = schedule_entry.get('lastRunAt')
+            if last_run_at is None:
+                last_run_at = current_app.now()
+                schedule_entry['lastRunAt'] = last_run_at
+            if isinstance(schedule_entry['lastRunAt'], str):
+                last_run_at = convert_to_datetime(
+                    schedule_entry['lastRunAt'])
+
+            # Verify times are accurate
+            orig = last_run_at
+            if not is_naive(last_run_at):
+                last_run_at = last_run_at.replace(tzinfo=None)
+            assert orig.hour == last_run_at.hour  # timezone sanity
+
+            if USING_SERMOS_CLOUD:
+                # We need to keep track of the id because this used to send
+                # updates to sermos cloud.  The name can't be concatenated
+                # with the sermos id or else it will be created as duplicate
+                # celery beat task.
+                name = schedule_entry['name']
+                self.sermos_id = schedule_entry['id']
+            else:
+                name = schedule_entry['name']
+
+            super().__init__(app=current_app._get_current_object(),
+                             name=name,
+                             task=schedule_entry['config']['task'],
+                             args=schedule_entry.get('args', None),
+                             kwargs=schedule_entry.get('kwargs', None),
+                             options=options,
+                             schedule=celery_schedule,
+                             last_run_at=last_run_at,
+                             total_run_count=schedule_entry.get(
+                                 'totalRunCount', 0))
+        else:
+            # This is a task issued directly by celery's scheduler so won't
+            # have the schedule_entry argument. Still not entirely clear why
+            # this is seen.  Pop the id before initializing the super class.
+            # Add it back after so we can keep sermos up to date w/ config.
+            if USING_SERMOS_CLOUD:
+                sermos_id = kwargs.pop('sermos_id')
+                super().__init__(**kwargs)
+                self.sermos_id = sermos_id
+            else:
+                super().__init__(**kwargs)
+
+        # Ensure all events have 'event' key - this is populated by ChainedTask
+        if 'event' not in self.kwargs.keys():
+            self.kwargs['event'] = {}
+
+
+class SermosScheduler(Scheduler):
+    """ Sermos' implementation of a Celery Scheduler. Leverages a Sermos
+    configuration server to provide the up-to-date schedule and provides to
+    this scheduler for in-memory tracking.
+    """
+    Entry = SermosEntry
+    _last_refresh = None  # Internal time keeper for Sermos syncing
+    _refresh_rate = CONFIG_REFRESH_RATE * 1000000  # Turn to microseconds
+    _schedule = None  # Holds latest Celery schedule with only enabled tasks
+    _schedule_full = None  # Holds latest schedule, regardless of enabled
+    _initial_read = True  # Set to False upon initial bootstrapping
+
+    def __init__(self, *args, **kwargs):
+        logger.info("Initializing SermosScheduler ...")
+        # This step ensures the latest schedule is pulled from Sermos/cache
+        # and bootstraps the local time checker we use.
+        self.set_under_schedule()
+        self._last_refresh = datetime.datetime.utcnow()
+
+        # Default 60 second max interval here so our schedule is always
+        # forced to be up to date.
+        max_interval = int(
+            os.environ.get('CELERY_BEAT_SYNC_MAX_INTERVAL',
+                           CONFIG_REFRESH_RATE))
+        kwargs['max_interval'] = max_interval
+
+        kwargs['schedule'] = self._schedule
+        Scheduler.__init__(self, *args, **kwargs)
+
+    def set_under_schedule(self):
+        """ Parse the latest schedule config and set self._schedule with parsed
+        schedule including only those that are enabled.
+        """
+        s = {}
+        s_full = []
+        s_full_orig = [s.copy() for s in self._schedule_full
+                       ] if self._schedule_full else []
+        latest_schedule = retrieve_latest_schedule_config()
+        for sched in latest_schedule:
+            s_full.append(sched)  # Append to full list regardless of enabled
+            if sched['enabled']:
+                s[sched['name']] = SermosEntry(sched)
+        self._schedule = s
+        self._schedule_full = s_full
+
+        # Report if schedule changed
+        if self._schedule_full != s_full_orig:
+            logger.info("SermosScheduler: Schedule updated ...")
+            logger.info(f"SermosScheduler: {self._schedule}")
+
+    def get_current_sermos_schedule(self):
+        """ Unpack Celery's current representation of the schedule into Sermos
+        format. This is used to send updates back to Sermos related to dynamic
+        properties such as last_run_at and total_run_count.
+        """
+
+        sched = {'schedules': []}
+        for entry_name, entry in self.schedule.items():
+            sched['schedules'].append({
+                'id': entry.sermos_id,
+                'lastRunAt': entry.last_run_at.isoformat(),
+                'totalRunCount': entry.total_run_count
+            })
+
+        return sched
+
+    def setup_schedule(self):
+        self.install_default_entries(self.data)
+        # Overload default behavior and instead bootstrap with our _schedule
+        # instead of app.conf.beat_schedule.
+        self.merge_inplace(self._schedule)
+
+    def should_refresh(self):
+        """ Determine if enough time has elapsed to perform a schedule refresh.
+
+        We turn everything into microseconds so we don't spam external services
+        intra-second as most of the time, more than one task exists in the
+        schedule and therefore we need to check the scheduler's `schedule`
+        on each task very rapidly when issuing tasks.
+        """
+        now = datetime.datetime.utcnow()
+        microseconds_since_last_refresh = float(
+            str((now - self._last_refresh).seconds) + "." +
+            str((now - self._last_refresh).microseconds)) * 1000000
+        res = bool(microseconds_since_last_refresh > self._refresh_rate)
+        if res is True:
+            self._last_refresh = now - datetime.timedelta(milliseconds=1)
+        return res
+
+    def sync(self):
+        """ Sync local schedule with Sermos and update Celery's representation
+        TODO check this vis-a-vis local vs cloud
+        """
+        if self.schedule and USING_SERMOS_CLOUD:
+            update_schedule_config(self.get_current_sermos_schedule())
+        self.set_under_schedule()  # Internal representation
+        self.merge_inplace(self._schedule)  # Celery representation
+
+    def get_schedule(self):
+        """ Overload default Scheduler get_schedule method to check for updates
+
+        Note: Celery uses a property function, e.g.:
+        https://www.tutorialsteacher.com/python/property-function
+        for getting/setting the schedule internally. We only override the
+        get_schedule method here.
+        """
+        update = False
+        if self._initial_read:
+            logger.info('SermosScheduler: Initial read ...')
+            update = True
+            self._initial_read = False
+        elif self.should_refresh():
+            logger.info('SermosScheduler: Refreshing schedule ...')
+            update = True
+
+        if update:
+            self.sync()
+
+        return self._schedule
+
+    def set_schedule(self, schedule):
+        """ Redefine Celery set_schedule method
+        """
+        self.data = schedule
+
+    # Redefine Celery schedule property()
+    schedule = property(get_schedule, set_schedule)

pypeline/cli/config_server.py

@@ -0,0 +1,48 @@
+""" Command Line Utilities for starting the local configuration server
+"""
+import logging
+import click
+from pypeline.lib.config_server import api, set_api_config
+
+logger = logging.getLogger(__name__)
+
+
+@click.group()
+def config_server():
+    """ Deployment command group.
+    """
+
+
+@config_server.command()
+@click.option('--base-dir', required=False, default=None)
+@click.option('--pipelines-yaml', required=False, default=None)
+@click.option('--schedules-json', required=False, default=None)
+@click.option('--port', required=False, default=8000)
+def local_config_api(base_dir: str = None,
+                     pipelines_yaml: str = None,
+                     schedules_json: str = None,
+                     port: int = 8000):
+    """ Start a local configuration API server for development.
+
+    This will use the provided pipelines.yaml and schedules.json file to
+    mock the API endpoints available to managed Sermos Deployments. To use
+    for development, make sure to set the DEFAULT_BASE_URL in your application's
+    environment (e.g. DEFAULT_BASE_URL=http://localhost:8000/api/v1/)
+
+    Arguments::
+
+        base-dir (optional): Directory name where your development config
+            files reside. Defaults to `dev`.
+
+        pipelines-yaml (optional): Path to find your `pipelines.yaml`
+            configuration file. Defaults to `pipelines.yaml`
+
+        schedules-json (optional): Path to find your `schedules.json`
+            configuration file. Defaults to `schedules.json`
+    """
+    click.echo("Starting Local Sermos Configuration Server ...")
+    set_api_config(base_dir=base_dir,
+                   pipelines_yaml=pipelines_yaml,
+                   schedules_json=schedules_json)
+
+    api.run(port=port, host='0.0.0.0')

pypeline/cli/core.py

@@ -0,0 +1,32 @@
+""" Primary CLI group entrypoint
+"""
+import logging
+import click
+from pypeline.logging_config import setup_logging
+
+setup_logging(default_level='INFO', establish_logging_config=False)
+logger = logging.getLogger(__name__)
+
+# Not all CLI tools will be functional / available depending on which extras
+# are installed. For example, the config server won't work if the `workers`
+# extra isn't available, which installs celery and networkx.
+collection = []
+
+warning_msg = "{} CLI tools are not available. This is most "\
+               "likely due to a missing import. Verify you have the correct "\
+               "Sermos extras installed."
+try:
+    from pypeline.cli.deploy import deployment
+    collection.append(deployment)
+except ImportError as e:
+    logger.warning(warning_msg.format("Deployment"))
+    logger.warning(f"{e}")
+
+try:
+    from pypeline.cli.config_server import config_server
+    collection.append(config_server)
+except ImportError as e:
+    logger.warning(warning_msg.format("Configuration Server"))
+    logger.warning(f"{e}")
+
+sermos = click.CommandCollection(sources=collection)