virtool-workflow 0.0.0__py3-none-any.whl

virtool_workflow/runtime/ping.py

@@ -0,0 +1,54 @@
+"""Ping the API to keep the job alive."""
+import asyncio
+from contextlib import asynccontextmanager
+
+from aiohttp import ClientOSError, ServerDisconnectedError
+from structlog import get_logger
+
+from virtool_workflow.api.client import APIClient
+
+logger = get_logger("api")
+
+
+async def _ping_periodically(api: APIClient, job_id: str):
+    retries = 0
+
+    try:
+        while True:
+            if retries > 5:
+                logger.warning("failed to ping server")
+                break
+
+            await asyncio.sleep(0.1)
+
+            try:
+                await api.put_json(f"/jobs/{job_id}/ping", {})
+            except (ClientOSError, ServerDisconnectedError):
+                await asyncio.sleep(0.3)
+                retries += 1
+                continue
+
+            await asyncio.sleep(5)
+    except asyncio.CancelledError:
+        logger.info("stopped pinging server")
+
+
+@asynccontextmanager
+async def ping_periodically(api: APIClient, job_id: str):
+    """Ping the API to keep the job alive.
+
+    While the context manager is open, a task runs that pings the API every 5 seconds.
+    When the context manager is closed, the task is cleanly cancelled.
+
+    The ping request is retried up to 5 times before the task is cancelled.
+
+    :param api: The API client.
+    :param job_id: The ID of the job to ping.
+
+    """
+    task = asyncio.create_task(_ping_periodically(api, job_id))
+
+    yield
+
+    task.cancel()
+    await task

virtool_workflow/runtime/redis.py

@@ -0,0 +1,65 @@
+import asyncio
+from asyncio import CancelledError
+from collections.abc import Callable
+
+from structlog import get_logger
+from virtool.redis import Redis
+
+logger = get_logger("redis")
+
+CANCELLATION_CHANNEL = "channel:cancel"
+
+
+async def get_next_job_with_timeout(
+    list_name: str,
+    redis: Redis,
+    timeout: int | None = None,
+) -> str:
+    """Get the next job ID from a Redis list.
+
+    Raise a  :class:``Timeout`` error if an ID is not found in ``timeout``
+    seconds.
+
+    :param list_name: the name of the list to pop from
+    :param redis: the Redis client
+    :param timeout: seconds to wait before raising :class:``Timeout``
+    :return: the next job ID
+    """
+    logger.info(
+        "Waiting for a job",
+        timeout=f"{timeout if timeout else 'infinity'} seconds",
+    )
+
+    return await asyncio.wait_for(get_next_job(list_name, redis), timeout)
+
+
+async def get_next_job(list_name: str, redis: Redis) -> str:
+    """Get the next job ID from a Redis list.
+
+    :param list_name: the name of the list to pop from
+    :param redis: the Redis client
+    :return: the next job ID
+
+    """
+    if (job_id := await redis.blpop(list_name)) is not None:
+        logger.info("pulled job id from redis", id=job_id)
+        return job_id
+
+    raise ValueError("Unexpected None from job id list")
+
+
+async def wait_for_cancellation(redis: Redis, job_id: str, func: Callable):
+    """Call a function ``func`` when a job matching ``job_id`` is cancelled.
+
+    :param redis: the Redis client
+    :param job_id: the job ID to watch for
+    :param func: the function to call when the job is cancelled
+
+    """
+    try:
+        async for cancelled_job_id in redis.subscribe(CANCELLATION_CHANNEL):
+            if cancelled_job_id == job_id:
+                return func()
+
+    except CancelledError:
+        ...

virtool_workflow/runtime/run.py

@@ -0,0 +1,276 @@
+import asyncio
+import signal
+import sys
+from asyncio import CancelledError
+from collections.abc import Callable
+from pathlib import Path
+
+from pyfixtures import FixtureScope, runs_in_new_fixture_context
+from structlog import get_logger
+from virtool.jobs.models import JobState
+from virtool.redis import Redis
+
+from virtool_workflow.api.acquire import acquire_job_by_id
+from virtool_workflow.api.client import api_client
+from virtool_workflow.hooks import (
+    cleanup_builtin_status_hooks,
+    on_cancelled,
+    on_error,
+    on_failure,
+    on_finish,
+    on_result,
+    on_step_finish,
+    on_step_start,
+    on_success,
+    on_terminated,
+    on_workflow_start,
+)
+from virtool_workflow.runtime.config import RunConfig
+from virtool_workflow.runtime.discover import (
+    load_builtin_fixtures,
+    load_custom_fixtures,
+    load_workflow_from_file,
+)
+from virtool_workflow.runtime.events import Events
+from virtool_workflow.runtime.path import create_work_path
+from virtool_workflow.runtime.ping import ping_periodically
+from virtool_workflow.runtime.redis import (
+    get_next_job_with_timeout,
+    wait_for_cancellation,
+)
+from virtool_workflow.runtime.sentry import configure_sentry
+from virtool_workflow.utils import configure_logs, get_virtool_workflow_version
+from virtool_workflow.workflow import Workflow
+
+logger = get_logger("runtime")
+
+
+def configure_status_hooks():
+    """Configure built-in job status hooks.
+
+    Push status updates to API when various lifecycle hooks are triggered.
+
+    """
+
+    @on_step_start
+    async def handle_step_start(push_status):
+        await push_status()
+
+    @on_error(once=True)
+    async def handle_error(push_status):
+        await push_status()
+
+    @on_cancelled(once=True)
+    async def handle_cancelled(push_status):
+        await push_status()
+
+    @on_terminated(once=True)
+    async def handle_terminated(push_status):
+        await push_status()
+
+    @on_success(once=True)
+    async def handle_success(push_status):
+        await push_status()
+
+
+async def execute(workflow: Workflow, scope: FixtureScope, events: Events):
+    """Execute a workflow.
+
+    :param workflow: The workflow to execute
+    :param scope: The :class:`FixtureScope` to use for fixture injection
+
+    """
+    await on_workflow_start.trigger(scope)
+
+    scope["_state"] = JobState.RUNNING
+
+    try:
+        for step in workflow.steps:
+            scope["_step"] = step
+
+            bound_step = await scope.bind(step.function)
+
+            await on_step_start.trigger(scope)
+            logger.info("running workflow step", name=step.display_name)
+            await bound_step()
+            await on_step_finish.trigger(scope)
+
+    except CancelledError:
+        logger.info("cancellation or termination interrupted workflow execution")
+
+        if events.cancelled.is_set():
+            logger.info("workflow cancelled")
+
+            scope["_state"] = JobState.CANCELLED
+
+            await asyncio.gather(
+                on_cancelled.trigger(scope),
+                on_failure.trigger(scope),
+            )
+        else:
+            logger.info("workflow terminated")
+
+            scope["_state"] = JobState.TERMINATED
+
+            if not events.terminated.is_set():
+                logger.warning(
+                    "workflow terminated without sigterm. this should not happen.",
+                )
+
+            await asyncio.gather(
+                on_terminated.trigger(scope),
+                on_failure.trigger(scope),
+            )
+
+    except Exception as error:
+        scope["_error"] = error
+        scope["_state"] = JobState.ERROR
+
+        logger.exception(error)
+
+        await asyncio.gather(on_error.trigger(scope), on_failure.trigger(scope))
+
+        if isinstance(error, asyncio.CancelledError):
+            raise
+
+    else:
+        scope["_state"] = JobState.COMPLETE
+        scope["_step"] = None
+
+        if "results" in scope:
+            await on_result.trigger(scope)
+
+        await on_success.trigger(scope)
+
+    finally:
+        await on_finish.trigger(scope)
+
+
+async def run_workflow(
+    config: RunConfig,
+    job_id: str,
+    workflow: Workflow,
+    events: Events,
+):
+    # Configure hooks here so that they can be tested when using `run_workflow`.
+    configure_status_hooks()
+
+    load_builtin_fixtures()
+
+    job = await acquire_job_by_id(config.jobs_api_connection_string, job_id)
+
+    async with (
+        api_client(
+            config.jobs_api_connection_string,
+            job.id,
+            job.key,
+        ) as api,
+        FixtureScope() as scope,
+    ):
+        # These fixtures should not be used directly by the workflow. They are used
+        # by other built-in fixtures.
+        scope["_api"] = api
+        scope["_config"] = config
+        scope["_error"] = None
+        scope["_job"] = job
+        scope["_state"] = JobState.WAITING
+        scope["_step"] = None
+        scope["_workflow"] = workflow
+
+        scope["logger"] = get_logger("workflow")
+        scope["mem"] = config.mem
+        scope["proc"] = config.proc
+        scope["results"] = {}
+
+        async with create_work_path(config) as work_path:
+            scope["work_path"] = work_path
+
+            async with ping_periodically(api, job_id):
+                await execute(workflow, scope, events)
+                cleanup_builtin_status_hooks()
+
+
+@runs_in_new_fixture_context()
+async def start_runtime(
+    dev: bool,
+    jobs_api_connection_string: str,
+    mem: int,
+    proc: int,
+    redis_connection_string: str,
+    redis_list_name: str,
+    sentry_dsn: str,
+    timeout: int,
+    work_path: Path,
+    workflow_loader: Callable[[], Workflow] = load_workflow_from_file,
+):
+    """Start the workflow runtime.
+
+    The runtime loads the workflow and fixtures. It then waits for a job ID to be pushed
+    to the configured Redis list.
+
+    When a job ID is received, the runtime acquires the job from the jobs API and
+    """
+    configure_logs(bool(sentry_dsn))
+
+    logger.info(
+        "found virtool-workflow",
+        version=get_virtool_workflow_version(),
+    )
+
+    workflow = workflow_loader()
+
+    load_builtin_fixtures()
+    load_custom_fixtures()
+
+    configure_sentry(sentry_dsn)
+
+    async with Redis(redis_connection_string) as redis:
+        try:
+            job_id = await get_next_job_with_timeout(redis_list_name, redis, timeout)
+        except TimeoutError:
+            # This happens due to Kubernetes scheduling issues or job cancellations. It
+            # is not an error.
+            logger.warning("timed out while waiting for job id")
+            return
+
+    events = Events()
+
+    run_workflow_task = asyncio.create_task(
+        run_workflow(
+            RunConfig(
+                dev,
+                jobs_api_connection_string,
+                mem,
+                proc,
+                work_path,
+            ),
+            job_id,
+            workflow,
+            events,
+        ),
+    )
+
+    def terminate_workflow(*_):
+        logger.info("received sigterm. terminating workflow.")
+        events.terminated.set()
+        run_workflow_task.cancel()
+
+    signal.signal(signal.SIGTERM, terminate_workflow)
+
+    def cancel_workflow(*_):
+        logger.info("received cancellation signal from redis")
+        events.cancelled.set()
+        run_workflow_task.cancel()
+
+    async with Redis(redis_connection_string) as redis:
+        cancellation_task = asyncio.create_task(
+            wait_for_cancellation(redis, job_id, cancel_workflow),
+        )
+
+        await run_workflow_task
+
+        cancellation_task.cancel()
+        await cancellation_task
+
+    if events.terminated.is_set():
+        sys.exit(124)

virtool_workflow/runtime/run_subprocess.py

@@ -0,0 +1,168 @@
+"""Code for running and managing subprocesses."""
+
+import asyncio
+from asyncio.subprocess import Process
+from collections.abc import Callable, Coroutine
+from pathlib import Path
+from typing import Protocol
+
+from pyfixtures import fixture
+from structlog import get_logger
+from virtool.utils import timestamp
+
+from virtool_workflow.errors import SubprocessFailedError
+
+logger = get_logger("subprocess")
+
+
+class LineOutputHandler(Protocol):
+    async def __call__(self, line: bytes):
+        """Handle input from stdin, or stderr, line by line.
+
+        :param line: A line of output from the stream.
+        """
+        raise NotImplementedError
+
+
+class RunSubprocess(Protocol):
+    async def __call__(
+        self,
+        command: list[str],
+        cwd: str | Path | None = None,
+        env: dict | None = None,
+        stderr_handler: LineOutputHandler | None = None,
+        stdout_handler: LineOutputHandler | None = None,
+    ) -> Process:
+        """Run a shell command in a subprocess.
+
+        :param command: A shell command
+        :param stdout_handler: A function to handle stdout output line by line
+        :param stderr_handler: A function to handle stderr output line by line
+        :param env: environment variables which should be available to the subprocess
+        :param cwd: The current working directory
+        :raise SubprocessFailed: The subprocess has exited with a non-zero exit code
+        :return: An :class:`.Process` instance
+        """
+        raise NotImplementedError
+
+
+async def watch_pipe(
+    stream: asyncio.StreamReader,
+    handler: LineOutputHandler,
+):
+    """Watch the stdout or stderr stream and pass lines to the `handler` callback function.
+
+    :param stream: a stdout or stderr file object
+    :param handler: a handler coroutine for output lines
+
+    """
+    while True:
+        line = await stream.readline()
+
+        if not line:
+            return
+
+        await handler(line)
+
+
+def stderr_logger(line: bytes):
+    """Log a line of stderr output and try to decode it as UTF-8.
+
+    If the line is not decodable, log it as a string.
+
+    :param line: a line of stderr output
+    """
+    line = line.rstrip()
+
+    try:
+        logger.info("stderr", line=line.decode())
+    except UnicodeDecodeError:
+        logger.info("stderr", line=line)
+
+
+async def _run_subprocess(
+    command: list[str],
+    stdout_handler: LineOutputHandler | None = None,
+    stderr_handler: Callable[[str], Coroutine] | None = None,
+    env: dict | None = None,
+    cwd: str | None = None,
+) -> asyncio.subprocess.Process:
+    """An implementation of :class:`RunSubprocess` using `asyncio.subprocess`."""
+    log = logger.bind()
+    log.info("running subprocess", command=command)
+
+    stdout = asyncio.subprocess.PIPE if stdout_handler else asyncio.subprocess.DEVNULL
+
+    if stderr_handler:
+
+        async def _stderr_handler(line):
+            stderr_logger(line)
+            await stderr_handler(line)
+
+    else:
+
+        async def _stderr_handler(line):
+            stderr_logger(line)
+
+    process = await asyncio.create_subprocess_exec(
+        *(str(arg) for arg in command),
+        cwd=cwd,
+        env=env,
+        limit=1024 * 1024 * 128,
+        stderr=asyncio.subprocess.PIPE,
+        stdout=stdout,
+    )
+
+    log.info("started subprocess", pid=process.pid, timestamp=timestamp().isoformat())
+
+    aws = [watch_pipe(process.stderr, _stderr_handler)]
+
+    if stdout_handler:
+        aws.append(watch_pipe(process.stdout, stdout_handler))
+
+    watcher_future = asyncio.gather(*aws)
+
+    try:
+        await watcher_future
+    except asyncio.CancelledError:
+        logger.info("terminating subprocess")
+
+        process.terminate()
+
+        # Have to do this in Python 3.10 to avoid Event loop closed error.
+        # https://github.com/python/cpython/issues/88050
+        try:
+            process._transport.close()
+        except AttributeError:
+            pass
+
+        await process.wait()
+        logger.info("subprocess exited", code=process.returncode)
+
+        await watcher_future
+
+        return process
+
+    await process.wait()
+
+    # Exit code 15 indicates that the process was terminated. This is expected
+    # when the workflow fails for some other reason, hence not an exception
+    if process.returncode not in [0, 15, -15]:
+        raise SubprocessFailedError(
+            f"{command[0]} failed with exit code {process.returncode}\n"
+            f"arguments: {command}\n",
+        )
+
+    log.info(
+        "subprocess finished",
+        return_code=process.returncode,
+        timestamp=timestamp().isoformat(),
+    )
+
+    return process
+
+
+@fixture(protocol=RunSubprocess)
+def run_subprocess() -> RunSubprocess:
+    """Fixture to run subprocesses and handle stdin and stderr output line-by-line."""
+    return _run_subprocess

virtool_workflow/runtime/sentry.py

@@ -0,0 +1,28 @@
+import logging
+
+import sentry_sdk
+from sentry_sdk.integrations.logging import LoggingIntegration
+from structlog import get_logger
+
+from virtool_workflow.utils import get_virtool_workflow_version
+
+logger = get_logger("runtime")
+
+
+def configure_sentry(dsn: str):
+    """Initialize Sentry for log aggregation."""
+    if dsn:
+        logger.info("initializing sentry", dsn=f"{dsn[:15]}...")
+
+        sentry_sdk.init(
+            dsn=dsn,
+            integrations=[
+                LoggingIntegration(
+                    event_level=logging.WARNING,
+                ),
+            ],
+            release=get_virtool_workflow_version(),
+            traces_sample_rate=0.2,
+        )
+    else:
+        logger.info("sentry disabled because no dsn was provided")

virtool_workflow/utils.py

@@ -0,0 +1,90 @@
+import asyncio
+import logging
+import sys
+import tarfile
+from collections.abc import Callable
+from functools import wraps
+from importlib import metadata
+from inspect import iscoroutinefunction
+from pathlib import Path
+
+import structlog
+from structlog.processors import LogfmtRenderer
+from structlog_sentry import SentryProcessor
+
+
+def coerce_to_coroutine_function(func: Callable):
+    """Wrap a non-async function in an async function."""
+    if iscoroutinefunction(func):
+        return func
+
+    @wraps(func)
+    async def _func(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    return _func
+
+
+def configure_logs(use_sentry: bool):
+    logging.basicConfig(
+        format="%(message)s",
+        stream=sys.stdout,
+        level=logging.INFO,
+    )
+
+    def normalize_log_level(logger, method_name, event_dict):
+        """Map EXCEPTION level to ERROR since logging module doesn't have EXCEPTION."""
+        if event_dict.get("level") == "EXCEPTION":
+            event_dict["level"] = "ERROR"
+        return event_dict
+
+    processors = [
+        structlog.stdlib.filter_by_level,
+        structlog.stdlib.add_logger_name,
+        structlog.stdlib.add_log_level,
+        normalize_log_level,
+        structlog.stdlib.PositionalArgumentsFormatter(),
+        structlog.processors.TimeStamper(fmt="%Y-%m-%dT%H:%M:%SZ"),
+        structlog.processors.StackInfoRenderer(),
+        structlog.processors.UnicodeDecoder(),
+    ]
+
+    if use_sentry:
+        processors.append(
+            SentryProcessor(event_level=logging.WARNING, level=logging.INFO),
+        )
+
+    processors.append(
+        LogfmtRenderer(
+            key_order=["timestamp", "level", "logger", "event"],
+        ),
+    )
+
+    structlog.configure(
+        cache_logger_on_first_use=True,
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        processors=processors,
+        wrapper_class=structlog.stdlib.BoundLogger,
+    )
+
+
+def get_virtool_workflow_version() -> str:
+    """Get the version of the installed virtool-workflow package."""
+    try:
+        return metadata.version("virtool-workflow")
+    except metadata.PackageNotFoundError:
+        return "0.0.0"
+
+
+async def make_directory(path: Path):
+    await asyncio.to_thread(path.mkdir, exist_ok=True, parents=True)
+
+
+def untar(path: Path, target_path: Path):
+    with tarfile.open(path, "r:gz") as tar:
+        tar.extractall(target_path)
+
+
+def move_all_model_files(source_path: Path, target_path: Path):
+    for file in source_path.iterdir():
+        file.rename(target_path / file.name)