"""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 time import uuid from collections import deque from datetime import UTC, datetime from functools import cache from typing import Any from sqlalchemy import update 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" PROGRESS_LINES = 12 # rolling window of pipeline output kept on the job PROGRESS_INTERVAL_S = 3 # between writes, unless a new process starts 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=`." ) 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.pipeline_database_url or 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: """Follow the pipeline: record progress while it runs, and the reason if it dies.""" tail: deque[str] = deque(maxlen=PROGRESS_LINES) last_write = 0.0 if proc.stdout is not None: # stderr is merged into stdout async for raw in proc.stdout: line = raw.decode(errors="replace").rstrip() if not line: continue tail.append(line) now = time.monotonic() # A new pipeline step is worth showing at once; otherwise throttle the writes. if line.startswith("[PROCESS") or now - last_write >= PROGRESS_INTERVAL_S: last_write = now await _record_progress(job_id, "\n".join(tail)) returncode = await proc.wait() if returncode == 0: return # the loader marks the job succeeded 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 {returncode}\n" + "\n".join(tail) job.finished_at = datetime.now(UTC) await session.commit() async def _record_progress(job_id: uuid.UUID, text: str) -> None: """Put the latest output on the job, leaving a job that has already finished alone.""" async with SessionLocal() as session: await session.execute( update(Job).where(Job.id == job_id, Job.status == JobStatus.running).values(log=text) ) await session.commit()