Source code for swh.scheduler.cli.task_type

# Copyright (C) 2016-2019  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information

from importlib import import_module
import logging
from typing import Mapping

# WARNING: do not import unnecessary things here to keep cli startup time under
# control
import click
from pkg_resources import iter_entry_points

from . import cli

logger = logging.getLogger(__name__)


DEFAULT_TASK_TYPE = {
    "full": {  # for tasks like 'list_xxx_full()'
        "default_interval": "90 days",
        "min_interval": "90 days",
        "max_interval": "90 days",
        "backoff_factor": 1,
    },
    "*": {  # value if not suffix matches
        "default_interval": "1 day",
        "min_interval": "1 day",
        "max_interval": "1 day",
        "backoff_factor": 1,
    },
}


PLUGIN_WORKER_DESCRIPTIONS = {
    entry_point.name: entry_point for entry_point in iter_entry_points("swh.workers")
}


@cli.group("task-type")
@click.pass_context
def task_type(ctx):
    """Manipulate task types."""
    scheduler = ctx.obj["scheduler"]
    if not scheduler:
        raise ValueError("Scheduler class (local/remote) must be instantiated")


@task_type.command("list")
@click.option("--verbose", "-v", is_flag=True, default=False, help="Verbose mode")
@click.option(
    "--task_type",
    "-t",
    multiple=True,
    default=None,
    help="List task types of given type",
)
@click.option(
    "--task_name",
    "-n",
    multiple=True,
    default=None,
    help="List task types of given backend task name",
)
@click.pass_context
def list_task_types(ctx, verbose, task_type, task_name):
    click.echo("Known task types:")
    if verbose:
        tmpl = (
            click.style("{type}: ", bold=True)
            + """{backend_name}
  {description}
  interval: {default_interval} [{min_interval}, {max_interval}]
  backoff_factor: {backoff_factor}
  max_queue_length: {max_queue_length}
  num_retries: {num_retries}
  retry_delay: {retry_delay}
"""
        )
    else:
        tmpl = "{type}:\n  {description}"
    for tasktype in sorted(
        ctx.obj["scheduler"].get_task_types(), key=lambda x: x["type"]
    ):
        if task_type and tasktype["type"] not in task_type:
            continue
        if task_name and tasktype["backend_name"] not in task_name:
            continue
        click.echo(tmpl.format(**tasktype))


@task_type.command("register")
@click.option(
    "--plugins",
    "-p",
    "plugins",
    multiple=True,
    default=("all",),
    type=click.Choice(["all"] + list(PLUGIN_WORKER_DESCRIPTIONS)),
    help="Registers task-types for provided plugins. " "Defaults to all",
)
@click.pass_context
def register_task_types(ctx, plugins):
    """Register missing task-type entries in the scheduler.

    According to declared tasks in each loaded worker (e.g. lister, loader,
    ...) plugins.

    """
    import celery.app.task

    scheduler = ctx.obj["scheduler"]

    if plugins == ("all",):
        plugins = list(PLUGIN_WORKER_DESCRIPTIONS)

    for plugin in plugins:
        entrypoint = PLUGIN_WORKER_DESCRIPTIONS[plugin]
        logger.info("Loading entrypoint for plugin %s", plugin)
        registry_entry = entrypoint.load()()

        for task_module in registry_entry["task_modules"]:
            mod = import_module(task_module)
            for task_name in (x for x in dir(mod) if not x.startswith("_")):
                logger.debug("Loading task name %s", task_name)
                taskobj = getattr(mod, task_name)
                if isinstance(taskobj, celery.app.task.Task):
                    tt_name = task_name.replace("_", "-")
                    task_cfg = registry_entry.get("task_types", {}).get(tt_name, {})
                    ensure_task_type(task_module, tt_name, taskobj, task_cfg, scheduler)


[docs]def ensure_task_type( task_module: str, task_type: str, swhtask, task_config: Mapping, scheduler ): """Ensure a given task-type (for the task_module) exists in the scheduler. Args: task_module: task module we are currently checking for task type consistency task_type: the type of the task to check/insert (correspond to the 'type' field in the db) swhtask (SWHTask): the SWHTask instance the task-type correspond to task_config: a dict with specific/overloaded values for the task-type to be created scheduler: the scheduler object used to access the scheduler db """ for suffix, defaults in DEFAULT_TASK_TYPE.items(): if task_type.endswith("-" + suffix): task_type_dict = defaults.copy() break else: task_type_dict = DEFAULT_TASK_TYPE["*"].copy() task_type_dict["type"] = task_type task_type_dict["backend_name"] = swhtask.name if swhtask.__doc__: task_type_dict["description"] = swhtask.__doc__.splitlines()[0] task_type_dict.update(task_config) current_task_type = scheduler.get_task_type(task_type) if current_task_type: # Ensure the existing task_type is consistent in the scheduler if current_task_type["backend_name"] != task_type_dict["backend_name"]: logger.warning( "Existing task type %s for module %s has a " "different backend name than current " "code version provides (%s vs. %s)", task_type, task_module, current_task_type["backend_name"], task_type_dict["backend_name"], ) else: logger.info("Create task type %s in scheduler", task_type) logger.debug(" %s", task_type_dict) scheduler.create_task_type(task_type_dict)
@task_type.command("add") @click.argument("type", required=True) @click.argument("task-name", required=True) @click.argument("description", required=True) @click.option( "--default-interval", "-i", default="90 days", help='Default interval ("90 days" by default)', ) @click.option( "--min-interval", default=None, help="Minimum interval (default interval if not set)", ) @click.option( "--max-interval", "-i", default=None, help="Maximal interval (default interval if not set)", ) @click.option("--backoff-factor", "-f", type=float, default=1, help="Backoff factor") @click.pass_context def add_task_type( ctx, type, task_name, description, default_interval, min_interval, max_interval, backoff_factor, ): """Create a new task type """ task_type = dict( type=type, backend_name=task_name, description=description, default_interval=default_interval, min_interval=min_interval, max_interval=max_interval, backoff_factor=backoff_factor, max_queue_length=None, num_retries=None, retry_delay=None, ) ctx.obj["scheduler"].create_task_type(task_type) click.echo("OK")