rarelens/api/app/services/events.py

"""Start the annotation pipeline for a job.

Serverless (CLOUDRUN_JOB set): execute a Cloud Run job that runs the Nextflow driver, which
submits the pipeline tasks to Google Batch. Nothing runs, or is billed, between annotations.
Kubernetes (PUBSUB_TOPIC set): publish to Pub/Sub; an Argo Events sensor submits the annotate-vcf
WorkflowTemplate, whose exit handler marks the job failed if the workflow does not succeed.
Local: run Nextflow as a child process and watch it. The docker-compose API image has neither
Nextflow nor Docker, so there the job fails at once with instructions to run it from the host.
"""
import asyncio
import json
import logging
import os
import shutil
import uuid
from datetime import UTC, datetime
from functools import cache
from typing import Any

from app.config import settings
from app.db import SessionLocal
from app.models import Job, JobStatus

logger = logging.getLogger(__name__)

PIPELINE_ENTRYPOINT = "/pipeline/main.nf"
LOG_TAIL_CHARS = 4000
PUBLISH_TIMEOUT_S = 30
_watchers: set[asyncio.Task[None]] = set()


class LaunchError(Exception):
    """The pipeline could not be started; the message is stored on the job for the user."""


async def launch(job_id: uuid.UUID, vcf_uri: str, assembly: str) -> str:
    """Start the pipeline and return a reference to the run (stored as jobs.workflow_ref)."""
    if settings.cloudrun_job:
        return await _run_cloud_run_job(job_id, vcf_uri, assembly)
    if settings.pubsub_topic:
        message = {"job_id": str(job_id), "vcf_uri": vcf_uri, "assembly": assembly}
        return await _publish(settings.pubsub_topic, message)
    return await _run_local(job_id, vcf_uri, assembly)


async def drain() -> None:
    """Wait for local pipeline watchers (tests, graceful shutdown)."""
    await asyncio.gather(*_watchers)


@cache
def _jobs_client() -> Any:
    from google.cloud import run_v2

    return run_v2.JobsClient()


async def _run_cloud_run_job(job_id: uuid.UUID, vcf_uri: str, assembly: str) -> str:
    if not settings.gcp_project:
        raise LaunchError("CLOUDRUN_JOB is set but GCP_PROJECT is not")
    name = (
        f"projects/{settings.gcp_project}/locations/{settings.gcp_region}"
        f"/jobs/{settings.cloudrun_job}"
    )
    args = [
        "run", PIPELINE_ENTRYPOINT, "-profile", "gcp",
        "--vcf", vcf_uri, "--job_id", str(job_id), "--assembly", assembly,
    ]
    request = {"name": name, "overrides": {"container_overrides": [{"args": args}]}}
    try:
        operation = await asyncio.to_thread(_jobs_client().run_job, request=request)
    except Exception as e:
        logger.exception("starting Cloud Run job %s failed", name)
        raise LaunchError(f"could not start {settings.cloudrun_job}: {e}") from e
    name_or_job = getattr(getattr(operation, "metadata", None), "name", "")
    execution = name_or_job or settings.cloudrun_job or "started"
    return f"cloudrun:{execution.rsplit('/', 1)[-1]}"


@cache
def _publisher() -> Any:
    from google.cloud import pubsub_v1

    return pubsub_v1.PublisherClient()


async def _publish(topic_name: str, message: dict[str, str]) -> str:
    if not settings.gcp_project:
        raise LaunchError("PUBSUB_TOPIC is set but GCP_PROJECT is not")
    publisher = _publisher()
    topic = publisher.topic_path(settings.gcp_project, topic_name)
    try:
        future = publisher.publish(topic, json.dumps(message).encode())
        message_id = await asyncio.to_thread(future.result, timeout=PUBLISH_TIMEOUT_S)
    except Exception as e:
        logger.exception("publishing to %s failed", topic)
        raise LaunchError(f"could not publish to {topic}: {e}") from e
    return f"pubsub:{message_id}"


async def _run_local(job_id: uuid.UUID, vcf_uri: str, assembly: str) -> str:
    nextflow = shutil.which("nextflow")
    if nextflow is None:
        raise LaunchError(
            "Nextflow is not installed where the API runs. Start the pipeline from the host with "
            f"`make annotate JOB={job_id} VCF=<path to the VCF>`."
        )
    cmd = [
        nextflow, "run", str(settings.pipeline_dir / "main.nf"),
        "-profile", settings.nextflow_profile,
        "--vcf", vcf_uri, "--job_id", str(job_id), "--assembly", assembly,
    ]
    # The loader reads DATABASE_URL from its environment; keep it off the command line.
    env = {**os.environ, "DATABASE_URL": settings.database_url}
    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            cwd=settings.pipeline_dir,
            env=env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
    except OSError as e:
        raise LaunchError(f"could not start Nextflow: {e}") from e
    task = asyncio.create_task(_watch(job_id, proc))
    _watchers.add(task)
    task.add_done_callback(_watchers.discard)
    return f"nextflow:pid-{proc.pid}"


async def _watch(job_id: uuid.UUID, proc: Any) -> None:
    out, err = await proc.communicate()
    if proc.returncode == 0:
        return  # the loader marks the job succeeded
    tail = ((out or b"") + (err or b"")).decode(errors="replace")[-LOG_TAIL_CHARS:]
    async with SessionLocal() as session:
        job = await session.get(Job, job_id)
        if job is not None and job.status != JobStatus.succeeded:
            job.status = JobStatus.failed
            job.log = f"nextflow exited with code {proc.returncode}\n{tail}"
            job.finished_at = datetime.now(UTC)
            await session.commit()