<commit_before># -*- coding: utf-8 -*-
"""Mock Processing Block Controller Celery task.

export CELERY_BROKER=redis://localhost:6379/1
export CELERY_BACKEND=redis://localhost:6379/2
celery -A mock_processing_block_controller.tasks worker -l info
"""
import json
import os
import time
from copy import deepcopy

import jinja2
import yaml
from celery import Celery

from sip_config_db.scheduling import ProcessingBlock, ProcessingBlockList
from sip_config_db.scheduling.workflow_stage import WorkflowStage
from sip_docker_swarm import DockerClient, __version__ as sip_swarm_api_version
from sip_logging import init_logger
from .release import LOG, __version__

BROKER = os.getenv('CELERY_BROKER_URL', 'redis://localhost:6379/1')
BACKEND = os.getenv('CELERY_RESULT_BACKEND', 'redis://localhost:6379/2')
APP = Celery(broker=BROKER, backend=BACKEND)


@APP.task
def echo(value: str):
    """Test Celery task to echo a string.

    Args:
        value (str): string to return.

    Returns:
        str, the input argument string.

    """
    return value


@APP.task
def version():
    """Celery task which returns the returns the PBC version.

    Returns:
        str, the PBC version

    """
    return __version__


def _start_workflow_stages(pb: ProcessingBlock, pb_id: str,
                           workflow_stage_dict: dict,
                           workflow_stage: WorkflowStage,
                           docker: DockerClient):
    """Start a workflow stage by starting a number of docker services.

    This function first assesses if the specified workflow stage can be
    started based on its dependencies. If this is found to be the case,
    the workflow stage is stared by first resolving and template arguments
    in the workflow stage configuration, and then using the Docker Swarm Client
    API to start workflow stage services. As part of this, the
    workflow_stage_dict data structure is updated accordingly.

    TODO(BMo) This function will need refactoring at some point as part \
    of an update to the way workflow state metadata is stored in the
    configuration database. Currently the stage_data dictionary
    is a bit of a hack for a badly specified Configuration Database
    backed WorkflowStage object.

    This function is used by `execute_processing_block`.

    Args:
        pb (ProcessingBlock): Configuration database Processing Block data
            object
        pb_id (str): Processing Block identifier
        workflow_stage_dict (dict): Workflow stage metadata structure
        workflow_stage (WorkflowStage): Workflow state configuration database
            data object.
        docker (DockerClient): Docker Swarm Client object.

    """
    # FIXME(BMo) replace pb_id argument, get this from the pb instead!
    stage_data = workflow_stage_dict[workflow_stage.id]
    stage_data['start'] = False

    # Determine if the stage can be started.
    if stage_data['status'] == 'none':
        if not workflow_stage.dependencies:
            stage_data['start'] = True
        else:
            dependency_status = []
            for dependency in workflow_stage.dependencies:
                dependency_status.append(
                    workflow_stage_dict[dependency['value']][
                        'status'] == 'complete')
                # ii += 1
            stage_data['start'] = all(dependency_status)

    # Start the workflow stage.
    if stage_data['start']:
        # Configure EE (set up templates)
        LOG.info('-- Starting workflow stage: %s --', workflow_stage.id)
        LOG.info('Configuring EE templates.')
        args_template = jinja2.Template(workflow_stage.args_template)
        stage_params = pb.workflow_parameters[workflow_stage.id]
        template_params = {**workflow_stage.config, **stage_params}
        args = args_template.render(stage=template_params)
        LOG.info('Resolving workflow script arguments.')

        args = json.dumps(json.loads(args))
        compose_template = jinja2.Template(
            workflow_stage.compose_template)
        compose_str = compose_template.render(stage=dict(args=args))

        # Prefix service names with the PB id
        compose_dict = yaml.load(compose_str)
        service_names = compose_dict['services'].keys()
        new_service_names = [
            '{}_{}_{}'.format(pb_id, pb.workflow_id, name)
            for name in service_names]
        for new, old in zip(new_service_names, service_names):
            compose_dict['services'][new] = \
                compose_dict['services'].pop(old)
        compose_str = yaml.dump(compose_dict)

        # Run the compose file
        service_ids = docker.create_services(compose_str)
        LOG.info('Staring workflow containers:')
        for service_id in service_ids:
            service_name = docker.get_service_name(service_id)
            LOG.info("  %s, %s ", service_name, service_id)
            stage_data['services'][service_id] = {}
            LOG.info('Created Services: %s', service_ids)
            stage_data['services'][service_id] = dict(
                name=docker.get_service_name(service_id),
                status='running',
                complete=False
            )
        stage_data["status"] = 'running'


def _update_workflow_stages(stage_data: dict, workflow_stage: WorkflowStage,
                            docker: DockerClient):
    """Check and update the status of a workflow stage.

    This function checks and updates the status of a workflow stage
    specified by the parameters in the specified stage_data dictionary.

    If the workflow stage is not marked as complete, this function will
    check with the Docker Swarm API on the status of Docker services
    defined for the stage. If **all** services are found to be complete
    (based on their service state being reported as 'shutdown',
    the workflow stage is marked complete.

    This function is used by `execute_processing_block`.

    TODO(BMo) This function will need refactoring at some point as part \
        of an update to the way workflow state metadata is stored in the
        configuration database. Currently the stage_data dictionary
        is a bit of a hack for a badly specified Configuration Database
        backed WorkflowStage object.

    Args:
        stage_data (dict): Dictionary holding workflow stage metadata.
        workflow_stage (WorkflowStage): Workflow stage data object.
        docker (DockerClient): Docker Swarm Client object.

    """
    service_status_complete = []

    # FIXME(BMo) is not "complete" -> is "running"
    if stage_data["status"] != "complete":
        for service_id, service_dict in stage_data['services'].items():
            service_state = docker.get_service_state(service_id)
            if service_state == 'shutdown':
                docker.delete_service(service_id)
            service_dict['status'] = service_state
            service_dict['complete'] = (service_state == 'shutdown')
            service_status_complete.append(service_dict['complete'])
            if all(service_status_complete):
                LOG.info('Workflow stage service %s complete!',
                         workflow_stage.id)
                stage_data['status'] = "complete"


def _abort_workflow(pb: ProcessingBlock, workflow_stage_dict: dict,
                    docker: DockerClient):
    """Abort the workflow.

    TODO(BMo): This function currently does nothing as the abort flag \
        is hardcoded to False!

    This function is used by `execute_processing_block`.<commit_msg>:memo: Updated docstrings (apparently \'s are not allowed!)<commit_after># -*- coding: utf-8 -*-
"""Mock Processing Block Controller Celery task.

export CELERY_BROKER=redis://localhost:6379/1
export CELERY_BACKEND=redis://localhost:6379/2
celery -A mock_processing_block_controller.tasks worker -l info
"""
import json
import os
import time
from copy import deepcopy

import jinja2
import yaml
from celery import Celery

from sip_config_db.scheduling import ProcessingBlock, ProcessingBlockList
from sip_config_db.scheduling.workflow_stage import WorkflowStage
from sip_docker_swarm import DockerClient, __version__ as sip_swarm_api_version
from sip_logging import init_logger
from .release import LOG, __version__

BROKER = os.getenv('CELERY_BROKER_URL', 'redis://localhost:6379/1')
BACKEND = os.getenv('CELERY_RESULT_BACKEND', 'redis://localhost:6379/2')
APP = Celery(broker=BROKER, backend=BACKEND)


@APP.task
def echo(value: str):
    """Test Celery task to echo a string.

    Args:
        value (str): string to return.

    Returns:
        str, the input argument string.

    """
    return value


@APP.task
def version():
    """Celery task which returns the returns the PBC version.

    Returns:
        str, the PBC version

    """
    return __version__


def _start_workflow_stages(pb: ProcessingBlock, pb_id: str,
                           workflow_stage_dict: dict,
                           workflow_stage: WorkflowStage,
                           docker: DockerClient):
    """Start a workflow stage by starting a number of docker services.

    This function first assesses if the specified workflow stage can be
    started based on its dependencies. If this is found to be the case,
    the workflow stage is stared by first resolving and template arguments
    in the workflow stage configuration, and then using the Docker Swarm Client
    API to start workflow stage services. As part of this, the
    workflow_stage_dict data structure is updated accordingly.

    TODO(BMo) This function will need refactoring at some point as part
    of an update to the way workflow state metadata is stored in the
    configuration database. Currently the stage_data dictionary
    is a bit of a hack for a badly specified Configuration Database
    backed WorkflowStage object.

    This function is used by `execute_processing_block`.

    Args:
        pb (ProcessingBlock): Configuration database Processing Block data
            object
        pb_id (str): Processing Block identifier
        workflow_stage_dict (dict): Workflow stage metadata structure
        workflow_stage (WorkflowStage): Workflow state configuration database
            data object.
        docker (DockerClient): Docker Swarm Client object.

    """
    # FIXME(BMo) replace pb_id argument, get this from the pb instead!
    stage_data = workflow_stage_dict[workflow_stage.id]
    stage_data['start'] = False

    # Determine if the stage can be started.
    if stage_data['status'] == 'none':
        if not workflow_stage.dependencies:
            stage_data['start'] = True
        else:
            dependency_status = []
            for dependency in workflow_stage.dependencies:
                dependency_status.append(
                    workflow_stage_dict[dependency['value']][
                        'status'] == 'complete')
                # ii += 1
            stage_data['start'] = all(dependency_status)

    # Start the workflow stage.
    if stage_data['start']:
        # Configure EE (set up templates)
        LOG.info('-- Starting workflow stage: %s --', workflow_stage.id)
        LOG.info('Configuring EE templates.')
        args_template = jinja2.Template(workflow_stage.args_template)
        stage_params = pb.workflow_parameters[workflow_stage.id]
        template_params = {**workflow_stage.config, **stage_params}
        args = args_template.render(stage=template_params)
        LOG.info('Resolving workflow script arguments.')

        args = json.dumps(json.loads(args))
        compose_template = jinja2.Template(
            workflow_stage.compose_template)
        compose_str = compose_template.render(stage=dict(args=args))

        # Prefix service names with the PB id
        compose_dict = yaml.load(compose_str)
        service_names = compose_dict['services'].keys()
        new_service_names = [
            '{}_{}_{}'.format(pb_id, pb.workflow_id, name)
            for name in service_names]
        for new, old in zip(new_service_names, service_names):
            compose_dict['services'][new] = \
                compose_dict['services'].pop(old)
        compose_str = yaml.dump(compose_dict)

        # Run the compose file
        service_ids = docker.create_services(compose_str)
        LOG.info('Staring workflow containers:')
        for service_id in service_ids:
            service_name = docker.get_service_name(service_id)
            LOG.info("  %s, %s ", service_name, service_id)
            stage_data['services'][service_id] = {}
            LOG.info('Created Services: %s', service_ids)
            stage_data['services'][service_id] = dict(
                name=docker.get_service_name(service_id),
                status='running',
                complete=False
            )
        stage_data["status"] = 'running'


def _update_workflow_stages(stage_data: dict, workflow_stage: WorkflowStage,
                            docker: DockerClient):
    """Check and update the status of a workflow stage.

    This function checks and updates the status of a workflow stage
    specified by the parameters in the specified stage_data dictionary.

    If the workflow stage is not marked as complete, this function will
    check with the Docker Swarm API on the status of Docker services
    defined for the stage. If **all** services are found to be complete
    (based on their service state being reported as 'shutdown',
    the workflow stage is marked complete.

    This function is used by `execute_processing_block`.

    TODO(BMo) This function will need refactoring at some point as part
        of an update to the way workflow state metadata is stored in the
        configuration database. Currently the stage_data dictionary
        is a bit of a hack for a badly specified Configuration Database
        backed WorkflowStage object.

    Args:
        stage_data (dict): Dictionary holding workflow stage metadata.
        workflow_stage (WorkflowStage): Workflow stage data object.
        docker (DockerClient): Docker Swarm Client object.

    """
    service_status_complete = []

    # FIXME(BMo) is not "complete" -> is "running"
    if stage_data["status"] != "complete":
        for service_id, service_dict in stage_data['services'].items():
            service_state = docker.get_service_state(service_id)
            if service_state == 'shutdown':
                docker.delete_service(service_id)
            service_dict['status'] = service_state
            service_dict['complete'] = (service_state == 'shutdown')
            service_status_complete.append(service_dict['complete'])
            if all(service_status_complete):
                LOG.info('Workflow stage service %s complete!',
                         workflow_stage.id)
                stage_data['status'] = "complete"


def _abort_workflow(pb: ProcessingBlock, workflow_stage_dict: dict,
                    docker: DockerClient):
    """Abort the workflow.

    TODO(BMo): This function currently does nothing as the abort flag
        is hardcoded to False!

    This function is used by `execute_processing_block`.
{'fraction_non_alphanumeric': 0.0768621236133122, 'fraction_numerical': 0.003301637612255679, 'mean_word_length': 2.8024102435350238, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '138743', 'n_tokens_mistral': 4167, 'n_tokens_neox': 3941, 'n_words': 1316}