dkist-processing-core 4.3.0__py3-none-any.whl

dkist_processing_core/__init__.py

@@ -0,0 +1,13 @@
+"""Package-level setup information."""
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version
+
+from dkist_processing_core.resource_queue import ResourceQueue
+from dkist_processing_core.task import TaskBase
+from dkist_processing_core.workflow import Workflow
+
+try:
+    __version__ = version(distribution_name=__name__)
+except PackageNotFoundError:
+    # package is not installed
+    __version__ = "unknown"

dkist_processing_core/build_utils.py

@@ -0,0 +1,139 @@
+"""Utilities for the build pipeline."""
+import importlib
+from pathlib import Path
+from shutil import rmtree
+from types import ModuleType
+
+from dkist_processing_core import Workflow
+
+
+__all__ = ["validate_workflows", "export_dags", "export_notebook_dockerfile", "export_notebooks"]
+
+
+def validate_workflows(workflow_package: ModuleType, export_path: Path | None = None) -> None:
+    """Validate that workflow engine (airflow) objects are acyclic and that exported workflows compile."""
+    # configure export path.  Clean up after if export path not provided
+    rm_export_path_after_test = not bool(export_path)
+    if export_path is None:
+        export_path = Path("export/")
+    workflows = extract_workflows_from_package(workflow_package)
+    try:
+        _validate_workflows(workflows, export_path)
+    finally:
+        if rm_export_path_after_test:
+            rmtree(export_path)
+
+
+def _validate_workflows(workflows: list[Workflow], export_path: Path) -> None:
+    """Validate workflows by ensuring their exported version compiles as python and that there is at least one node."""
+    for w in workflows:
+        workflow_py = w.export_dag(path=export_path)
+        with workflow_py.open(mode="r") as f:
+            compile(f.read(), filename=f"{workflow_py.stem}.pyc", mode="exec")
+        if len(w.nodes) == 0:
+            raise ValueError(f"Workflow {w.workflow_name} has 0 nodes.")
+
+
+def export_dags(workflow_package: ModuleType, path: str | Path) -> list[Path]:
+    """Export Airflow DAG files."""
+    return [w.export_dag(path=path) for w in extract_workflows_from_package(workflow_package)]
+
+
+def export_notebooks(workflow_package: ModuleType, path: str | Path) -> list[Path]:
+    """Export Jupyter Notebook files."""
+    return [w.export_notebook(path=path) for w in extract_workflows_from_package(workflow_package)]
+
+
+def export_notebook_dockerfile(workflow_package: ModuleType, path: str | Path) -> Path:
+    """Export a dockerfile to containerize notebooks."""
+    path = Path(path)
+    notebook_paths = export_notebooks(workflow_package=workflow_package, path=path)
+    category = extract_category_from_workflows(workflow_package=workflow_package)
+    dockerfile = NotebookDockerfile(notebook_paths=notebook_paths, category=category)
+    dockerfile_path = Path("Dockerfile")
+    dockerfile_path.touch(exist_ok=False)
+    with open(dockerfile_path, mode="w") as f:
+        f.write(dockerfile.contents)
+    return dockerfile_path
+
+
+def extract_category_from_workflows(workflow_package: ModuleType) -> str:
+    """Extract the category from the workflows in the package to provide a unique category for the dockerfile."""
+    workflows = extract_workflows_from_package(workflow_package)
+    categories = {w.category for w in workflows}
+    if len(categories) > 1:
+        raise ValueError(
+            f"Multiple categories found in provided workflows. Categories found: {categories}"
+        )
+    return categories.pop()
+
+
+def extract_workflows_from_package(workflow_package: ModuleType) -> list[Workflow]:
+    """Extract all the Workflow objects from a package."""
+    return extract_objects_from_package_by_type(workflow_package, Workflow)
+
+
+def extract_objects_from_package_by_type(package: ModuleType, object_type: type) -> list:
+    """Extract all objects in public modules of a given type from a package."""
+    modules = parse_unprotected_modules_names_from_package(package)
+    objects = []
+    for module in modules:
+        imported_module = importlib.import_module(f".{module}", package.__name__)
+        objects += [var for var in vars(imported_module).values() if isinstance(var, object_type)]
+    return objects
+
+
+def parse_unprotected_modules_names_from_package(package: ModuleType) -> list[str]:
+    """Parse the names of all modules in a package that are not private i.e. don't begin with an underscore."""
+    package_path = Path(package.__path__[0])
+    return [m.stem for m in package_path.glob("[!_]*.py")]
+
+
+class NotebookDockerfile:
+    """Build a Dockerfile for deployment as a Manual Processing Worker."""
+
+    def __init__(self, notebook_paths: list[Path], category: str):
+        self.notebook_paths = notebook_paths
+        self.validate_notebook_paths_are_relative()
+        self.category = category
+
+    def validate_notebook_paths_are_relative(self):
+        """Validate that the notebook paths are all relative."""
+        return all([not p.is_absolute() for p in self.notebook_paths])
+
+    @property
+    def contents(self) -> str:
+        """Return the Dockerfile body."""
+        return "\n".join(self.preamble + self.setup + self.notebooks + self.command)
+
+    @property
+    def preamble(self) -> list[str]:
+        """Dockerfile preamble lines."""
+        return ["FROM python:3.11", "ENV LANG=C.UTF-8"]
+
+    @property
+    def setup(self) -> list[str]:
+        """Environment setup lines."""
+        return [
+            "COPY . /app",
+            "WORKDIR /app",
+            "RUN python -m pip install -U pip",
+            "RUN pip install notebook",
+            "RUN pip freeze | grep notebook= > constraints.txt",
+            "RUN cat constraints.txt",
+            "RUN python -m pip install -c constraints.txt .",
+        ]
+
+    @property
+    def notebooks(self) -> list[str]:
+        """Generate workflow notebooks and include in Docker container."""
+        return [f"COPY {notebook_path} /notebooks/" for notebook_path in self.notebook_paths]
+
+    @property
+    def command(self) -> list[str]:
+        """Run notebook server on deployment."""
+        port = 8888
+        return [
+            f"EXPOSE {port}",
+            f"CMD jupyter notebook --NotebookApp.allow_root=True --NotebookApp.base_url='/mpw-{self.category}/' --NotebookApp.ip='0.0.0.0' --NotebookApp.port={port} --MappingKernelManager.cull_idle_timeout=300 --notebook-dir=/notebooks --allow-root",
+        ]

dkist_processing_core/config.py

@@ -0,0 +1,82 @@
+"""Environment controlled configurations for dkist_processing_core."""
+from dkist_service_configuration import MeshServiceConfigurationBase
+from dkist_service_configuration.settings import MeshService
+from pydantic import Field
+from talus import ConnectionRetryerFactory
+from talus import Exchange
+from talus.models.connection_parameters import ConnectionParameterFactory
+
+
+class DKISTProcessingCoreConfiguration(MeshServiceConfigurationBase):
+    """Environment configurations for dkist_processing_core."""
+
+    isb_username: str = Field(default="guest")
+    isb_password: str = Field(default="guest")
+    isb_exchange: str = Field(default="master.direct.x")
+    isb_queue_type: str = Field(default="classic")
+    elastic_apm_service_name: str = Field(default="dkist-processing-core")
+    elastic_apm_other_options: dict = Field(default_factory=dict)
+    elastic_apm_enabled: bool = False
+    build_version: str = Field(default="dev")
+
+    @property
+    def isb_mesh_service(self) -> MeshService:
+        """Return the mesh service details for the interservice-bus."""
+        return self.service_mesh_detail(
+            service_name="interservice-bus",
+            default_mesh_service=MeshService(mesh_address="localhost", mesh_port=5672),
+        )
+
+    @property
+    def isb_producer_connection_parameters(self) -> ConnectionParameterFactory:
+        """Return the connection parameters for the ISB producer."""
+        return ConnectionParameterFactory(
+            rabbitmq_host=self.isb_mesh_service.host,
+            rabbitmq_port=self.isb_mesh_service.port,
+            rabbitmq_user=self.isb_username,
+            rabbitmq_pass=self.isb_password,
+            connection_name="dkist-processing-core-producer",
+        )
+
+    @property
+    def isb_connection_retryer(self) -> ConnectionRetryerFactory:
+        """Return the connection retryer for the ISB connection."""
+        return ConnectionRetryerFactory(
+            delay_min=1,
+            delay_max=5,
+            backoff=1,
+            jitter_min=1,
+            jitter_max=3,
+            attempts=3,
+        )
+
+    @property
+    def isb_queue_arguments(self) -> dict:
+        """Return the queue arguments for the ISB."""
+        return {
+            "x-queue-type": self.isb_queue_type,
+        }
+
+    @property
+    def isb_publish_exchange(self) -> Exchange:
+        """Return the exchange for the ISB."""
+        return Exchange(name=self.isb_exchange)
+
+    @property
+    def elastic_apm_server_url(self) -> str:
+        """Return the URL for the Elastic APM server."""
+        apm_server = self.service_mesh_detail(service_name="system-monitoring-log-apm")
+        return f"http://{apm_server.host}:{apm_server.port}/"
+
+    @property
+    def apm_config(self) -> dict:
+        """Return the configuration for the Elastic APM."""
+        return {
+            "SERVICE_NAME": self.elastic_apm_service_name,
+            "SERVER_URL": self.elastic_apm_server_url,
+            "ENVIRONMENT": "Workflows",
+            **self.elastic_apm_other_options,
+        }
+
+
+core_configurations = DKISTProcessingCoreConfiguration()

dkist_processing_core/failure_callback.py

@@ -0,0 +1,96 @@
+"""Define the failure callback functionality."""
+import logging
+from contextlib import contextmanager
+from typing import Callable
+from typing import Type
+
+from talus import Binding
+from talus import DurableProducer
+from talus import MessageBodyBase
+from talus import PublishMessageBase
+from talus import Queue
+
+from dkist_processing_core.config import core_configurations
+
+logger = logging.getLogger(__name__)
+
+
+# Recipe run failure message Definition
+class RecipeRunFailureMessageBody(MessageBodyBase):
+    """Schema for the recipe run failure message body."""
+
+    workflowName: str
+    workflowVersion: str
+    taskName: str
+    dagRunId: str | None = None
+    logUrl: str | None = None
+
+
+class RecipeRunFailureMessage(PublishMessageBase):
+    """Recipe run failure message including the message body and other publication information."""
+
+    message_body_cls: Type[RecipeRunFailureMessageBody] = RecipeRunFailureMessageBody
+    default_routing_key: str = "recipe.run.failure.m"
+
+
+@contextmanager
+def recipe_run_failure_message_producer_factory() -> DurableProducer:
+    """Create message producer for recipe run failure messages."""
+    # Configure the queue the messages should be routed to
+    recipe_run_failure_queue = Queue(
+        name="recipe.run.failure.q", arguments=core_configurations.isb_queue_arguments
+    )
+    # Configure the exchange and queue bindings for publishing
+    bindings = [Binding(queue=recipe_run_failure_queue, message=RecipeRunFailureMessage)]
+    try:
+        with DurableProducer(
+            queue_bindings=bindings,
+            publish_exchange=core_configurations.isb_publish_exchange,
+            connection_parameters=core_configurations.isb_producer_connection_parameters,
+            connection_retryer=core_configurations.isb_connection_retryer,
+        ) as producer:
+            yield producer
+    finally:
+        pass
+
+
+def parse_dag_run_id_from_context(context: dict) -> str | None:
+    """Find dag run id."""
+    return context.get("run_id", None)
+
+
+def parse_log_url_from_context(context: dict) -> str | None:
+    """Given an airflow context, find the URL of the logs created by the task."""
+    ti = context.get("task_instance", object)
+    try:
+        return ti.log_url
+    except AttributeError:
+        pass
+
+
+def chat_ops_notification(
+    context: dict,
+    workflow_name: str,
+    workflow_version: str,
+    task_name: str,
+    producer_factory: Callable[[], DurableProducer] = recipe_run_failure_message_producer_factory,
+) -> RecipeRunFailureMessage:
+    """Publish message with information regarding a task failure for publication to a chat service."""
+    dag_run_id = parse_dag_run_id_from_context(context)
+    log_url = parse_log_url_from_context(context)
+    body = RecipeRunFailureMessageBody(
+        workflowName=workflow_name,
+        workflowVersion=workflow_version,
+        taskName=task_name,
+        logUrl=log_url,
+        dagRunId=dag_run_id,
+    )
+    message = RecipeRunFailureMessage(body)
+
+    try:
+        with producer_factory() as producer:
+            logger.warning(f"Publishing failure callback message: {message=}")
+            producer.publish(message)
+            return message
+    except Exception as e:  # pragma: no cover
+        logger.error(f"Error raised executing failure callback: {e=}")  # pragma: no cover

dkist_processing_core/node.py

@@ -0,0 +1,169 @@
+"""Abstraction layer to construct a workflow node using and airflow operator."""
+from collections.abc import Iterable
+from typing import Type
+
+from airflow.operators.bash import BashOperator
+
+from dkist_processing_core.resource_queue import ResourceQueue
+from dkist_processing_core.task import TaskBase
+
+
+task_type_hint = Type[TaskBase]
+upstreams_type_hint = list[task_type_hint] | task_type_hint | None
+
+
+class Node:
+    """Abstraction to instantiate a Task in a Workflow graph for target execution environments."""
+
+    def __init__(
+        self,
+        workflow_name: str,
+        workflow_version: str,
+        workflow_package: str,
+        task: task_type_hint,
+        resource_queue: ResourceQueue,
+        upstreams: upstreams_type_hint = None,
+        pip_extras: list[str] | None = None,
+    ):
+        """Node setup."""
+        # Task type checking
+        upstreams = upstreams or []
+        if not isinstance(upstreams, Iterable):
+            upstreams = [
+                upstreams,
+            ]
+        if not all([issubclass(t, TaskBase) for t in [task] + upstreams]):
+            raise TypeError(
+                "Only task classes inheriting from "
+                "dkist_processing_core.TaskBase can be added to a workflow"
+            )
+
+        self.workflow_name = workflow_name
+        self.workflow_version = workflow_version
+        self.task = task
+        self.workflow_package = workflow_package
+        self.upstreams = upstreams
+        self.resource_queue = resource_queue
+        self.pip_extras = pip_extras
+
+    @property
+    def operator(self) -> BashOperator:
+        """Native engine node."""
+        from datetime import timedelta
+        from dkist_processing_core.failure_callback import chat_ops_notification
+        from functools import partial
+
+        return eval(self.operator_definition)
+
+    @property
+    def notebook_cell(self) -> str:
+        """Render the node as python code for a notebook cell."""
+        lines = [
+            f"from {self.task.__module__} import {self.task.__name__}",
+            f"with {self.task.__name__}(recipe_run_id=recipe_run_id, workflow_name='{self.workflow_name}', workflow_version='{self.workflow_version}') as t:\n    #t.is_task_manual = True\n    t()\n    #t.rollback()",
+        ]
+        return "\n".join(lines)
+
+    @property
+    def operator_definition(self) -> str:
+        """Airflow style command to define a bash operator."""
+        return f"""BashOperator(
+    task_id='{self.task.__name__}',
+    bash_command='''{self.bash_script}''',
+    retries={self.task.retries},
+    retry_delay=timedelta(seconds={self.task.retry_delay_seconds}),
+    on_failure_callback=partial(
+        chat_ops_notification,
+        workflow_name='{self.workflow_name}',
+        workflow_version='{self.workflow_version}',
+        task_name='{self.task.__name__}'
+    ),
+    owner="DKIST Data Center",
+    queue="{self.resource_queue.value}",
+    output_processor=str,
+)
+"""
+
+    @property
+    def dependencies(self) -> list[tuple[str, str]]:
+        """List of upstream, downstream task name tuples."""
+        return [(upstream.__name__, self.task.__name__) for upstream in self.upstreams]
+
+    @property
+    def bash_script(self) -> str:
+        """Format bash script for the BashOperator."""
+        command = f"""{self.install_command}
+{self.run_command}"""
+        return self.bash_template(command)
+
+    @staticmethod
+    def bash_template(command: str) -> str:
+        """Return the bash script with a template wrapped command."""
+        return f"""#!/bin/bash
+echo Working Directory
+pwd
+echo Worker Identification
+echo NOMAD_ALLOC_ID
+echo $NOMAD_ALLOC_ID
+echo NOMAD_GROUP_NAME
+echo $NOMAD_GROUP_NAME
+echo NOMAD_HOST_ADDR_worker
+echo $NOMAD_HOST_ADDR_worker
+echo NOMAD_ALLOC_NAME
+echo $NOMAD_ALLOC_NAME
+echo Host Python Environment i.e. system-site-packages
+python3 -m pip install --upgrade --user pip
+pip list
+echo Creating Virtual Environment
+python3 -m venv --system-site-packages .task_venv
+echo Activate Environment
+. .task_venv/bin/activate
+echo Python Interpreter Location
+which python
+echo Run Main Command
+{command}
+export exit_code=$?
+echo Deactivate Environment
+deactivate
+echo Remove Virtual Environment
+rm -rf .task_venv
+echo Exit with code from main command: $exit_code
+exit $exit_code"""
+
+    @property
+    def formatted_pip_extras(self) -> str:
+        """Format pip extras for the installation command."""
+        if self.pip_extras:
+            extra_requirements = ",".join(self.pip_extras)
+            return f"'[{extra_requirements}]'"
+        return ""
+
+    @property
+    def install_command(self) -> str:
+        """Format the installation command for the bash script."""
+        repo_name = self.workflow_package.split(".")[0].replace("_", "-")
+        version = self.workflow_version
+        extras = self.formatted_pip_extras
+        return f"""python -m pip install --upgrade pip
+python -m pip install {repo_name}{extras}=={version}"""
+
+    @property
+    def run_command(self) -> str:
+        """Return the python bash command to execute the task."""
+        return f'python -c "{self.python}"'
+
+    @property
+    def python(self) -> str:
+        """Return the python code to execute the task."""
+        return f"""from {self.task.__module__} import {self.task.__name__}
+with {self.task.__name__}(recipe_run_id={{{{dag_run.conf['recipe_run_id']}}}}, workflow_name='{self.workflow_name}', workflow_version='{self.workflow_version}') as task:
+    task()
+"""
+
+    def __repr__(self):
+        """Render node instantiation as a string."""
+        return f"Node(workflow_name={self.workflow_name}, workflow_version={self.workflow_version}, workflow_package={self.workflow_package}, task={self.task!r}, upstreams={self.upstreams}, queue={self.resource_queue!r})"
+
+    def __str__(self):
+        """Render node instance as a string."""
+        return repr(self)

dkist_processing_core/resource_queue.py

@@ -0,0 +1,9 @@
+"""Resource queue names for Workflow Node specification which specify different resource needs."""
+from enum import StrEnum
+
+
+class ResourceQueue(StrEnum):
+    """Supported queue names."""
+
+    DEFAULT: str = "default"
+    HIGH_MEMORY: str = "high_memory"